🤖 machine learning

Learning fMRI activations dictionaries across individual geometries via optimal transport

This paper proposes a novel dictionary learning framework for fMRI data that leverages the Fused Gromov-Wasserstein distance and amortized neural optimization to capture individual brain geometry variability without relying on template-based projection, thereby preserving subject-specific information for downstream analysis.

Original authors: Sonia Mazelet, Rémi Flamary, Bertrand Thirion

Published 2026-05-21

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Sonia Mazelet, Rémi Flamary, Bertrand Thirion

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: Mapping the Brain's "Fingerprint"

Imagine you are trying to create a dictionary of brain activities. You want to find the basic "building blocks" (atoms) that make up how our brains work when we do things like feel fear, move a hand, or solve a math problem.

The Problem:
Every human brain is shaped differently. Some are wrinkled like a walnut, others are smoother; some are bigger, some smaller.

The Old Way: To compare brains, scientists used to flatten everyone's brain onto a single, standard "cookie cutter" template. It's like taking a unique, hand-knitted sweater and forcing it into a machine that squishes it into a perfect, flat square so everyone looks the same.
The Catch: When you squish the sweater, you lose the unique texture and shape of the original knitter's work. You lose the "individual fingerprint" of the person's brain.

The New Solution (AGDL):
The authors propose a new method called Amortized Graph Dictionary Learning (AGDL). Instead of squishing everyone into a cookie cutter, they treat every brain as its own unique, 3D landscape. They want to find the building blocks of brain activity while respecting the unique shape of each person's brain.

How It Works: The "Universal Translator"

To compare two different brain maps without squishing them, the researchers use a mathematical tool called Optimal Transport.

The Analogy:
Imagine you have two different cities.

City A has a grid layout with wide avenues.
City B has winding, narrow streets following a river.
You want to move furniture (brain activity) from City A to City B.

If you just try to match "Street 1" to "Street 1," it won't work because the streets are shaped differently. You need a Universal Translator that figures out the most efficient way to move the furniture from a specific house in City A to the best matching house in City B, even if the streets look totally different.

In this paper, that "Universal Translator" is a neural network trained to predict these matches instantly.

The Three Big Innovations

1. The Speed Trick (Amortized Optimization)

Calculating the perfect match between two complex brain maps is like trying to solve a massive jigsaw puzzle where every piece changes shape. Doing this for thousands of people would take a supercomputer years.

The Fix: The researchers trained a "smart assistant" (a neural network) to look at the puzzle and guess the solution almost instantly. Instead of solving the math problem from scratch every time, the assistant just says, "I've seen this before; here is the best match." This makes the process fast enough to handle huge datasets.

2. The "Dial" for Shape vs. Content (The $\alpha$ Parameter)

When comparing brains, you have to decide what matters more:

Feature Alignment: "Does this spot light up when I feel fear?" (Content)
Structural Alignment: "Is this spot in the same physical fold of the brain?" (Shape)

The researchers built a "dial" (parameter $\alpha$ ) that lets you adjust the balance.

Turn the dial to 0: You only care about the content (the activity), ignoring the shape.
Turn the dial to 1: You care deeply about the shape and structure.
The Magic: They didn't just train one dictionary for one setting. They trained a dictionary that changes based on where you turn the dial. You can instantly see what the brain looks like if you focus on shape, or if you focus on activity, without having to retrain the whole system.

3. Learning from the "Native" Shape

Instead of forcing brains onto a standard template, this method learns the dictionary atoms directly from the "native" (real, individual) shapes of the brains.

The Result: The dictionary atoms are like a set of universal Lego bricks. Even though every person's brain (the house they build) is shaped differently, the bricks (the dictionary) can still fit together perfectly to describe the activity, preserving the unique "fingerprint" of each individual.

What Did They Find?

The researchers tested this on a massive dataset of 1,200 people (the HCP dataset) doing various tasks like gambling, moving hands, or listening to stories.

It Works: The method successfully created a dictionary that could describe brain activity for different tasks.
It Captures Individuality: When they tried to guess which person was in the scan based on the brain activity, their method was much better than the old "cookie cutter" method. This proves that keeping the unique brain shape helps identify the person.
The "Sweet Spot": They found that the best results came from a mix of caring about both the shape and the activity. If you only cared about one or the other, the results weren't as good.
Smarter Models: They tried two ways to build the dictionary: a simple linear one and a more complex one (MLP). The complex one worked better, suggesting that brain activity is too complicated for simple rules; it needs a flexible, expressive model to capture it.

Summary

This paper introduces a way to study brain activity that respects the fact that every brain is shaped differently. By using a "smart assistant" to quickly match brain maps and a flexible "dial" to balance shape vs. content, they created a dictionary that captures both the universal tasks we do (like moving a hand) and the unique, individual fingerprints of our brains.

Technical Summary: Learning fMRI Activations Dictionaries Across Individual Geometries via Optimal Transport

Problem Statement

Functional magnetic resonance imaging (fMRI) data is high-dimensional, noisy, and inherently subject-specific. A central challenge in multi-subject fMRI analysis is the substantial variability in brain anatomy and functional organization across individuals. Standard approaches address this by projecting all subjects onto a common anatomical template (e.g., fsaverage). While practical, this registration process discards subject-specific structural information, relies on imperfect correspondences, and can distort individual brain organization.

Existing dictionary learning methods for fMRI assume signals are registered to a common geometry. While Graph Dictionary Learning (GDL) using the Fused Gromov-Wasserstein (FGW) distance offers a framework to compare graphs with different geometries, its application to large-scale neuroimaging data is currently infeasible. The computational cost of solving the FGW optimization problem scales cubically (or quartically depending on the loss variant) with the number of nodes, restricting previous GDL applications to small graphs with only tens of nodes, whereas fMRI graphs typically contain thousands of nodes.

Methodology

The authors propose Amortized Graph Dictionary Learning (AGDL), a novel framework designed to learn interpretable dictionary representations directly from individual subject geometries without projecting them to a common template. The method integrates three core components:

1. Amortized Optimization for FGW Approximation

To overcome the prohibitive computational cost of exact FGW distance calculation, the authors employ amortized optimization. They utilize a pre-trained neural network based on the Unsupervised Learning of Optimal Transport (ULOT) architecture.

Training: The ULOT model is trained on pairs of graphs sampled from the dataset to predict the optimal transport plan $P^\alpha_\theta(G_1, G_2)$ that minimizes the FGW loss. The training uses the Fused Unbalanced Gromov-Wasserstein (FUGW) loss as a relaxation to encourage marginal constraints without requiring ground truth correspondences.
Inference: Once trained, the model predicts transport plans via a simple forward pass with $O(n_1 n_2)$ complexity, enabling the application of FGW-based methods to large graphs (approx. 1000 nodes).

2. Dictionary Learning with Fixed Atom Geometry

The method learns a dictionary of atoms $(F^D_k, C^D_k)$ where the structural component $C$ (representing the graph geometry) is fixed to a common geometry (e.g., the fsaverage template) to ensure interpretability and comparability across subjects. The feature component $F^D_k$ is learned.

Reconstruction: Each subject's graph $G_i$ is reconstructed as a linear combination of dictionary atoms: $(\sum \omega_k F^D_k, \sum \omega_k C^D_k)$ .
Loss Function: The reconstruction error is computed using the FGW distance between the subject's native graph and the reconstructed graph, utilizing the transport plan predicted by the amortized ULOT model. This allows the model to account for individual geometric variability through the transport plan while maintaining a common dictionary structure.

3. $\alpha$ -Conditioned Dictionary Learning

The FGW distance includes a trade-off parameter $\alpha \in [0, 1]$ balancing feature alignment (Wasserstein term) and structural alignment (Gromov-Wasserstein term). Instead of retraining the dictionary for every $\alpha$ , the authors propose learning dictionaries that explicitly depend on $\alpha$ :

Linear Interpolation: A simple model where the dictionary atoms are a linear interpolation between two endpoint atoms: $F^D = \alpha F^D_0 + (1-\alpha)F^D_1$ .
MLP-based Model: A more expressive model using a Multi-Layer Perceptron (MLP) that takes a soft-binned encoding of $\alpha$ as input to output the dictionary atoms. This allows the dictionary to smoothly evolve with $\alpha$ , capturing different levels of geometric variability without computational overhead for downstream parameter selection.

Key Contributions

Scalable FGW Dictionary Learning: The paper introduces AGDL, the first method to apply FGW-based dictionary learning to large-scale fMRI data (thousands of nodes) by leveraging amortized optimization to approximate transport plans.
Native Geometry Learning: Unlike standard pipelines that project data to a common template, AGDL learns representations directly from individual subject geometries, preserving subject-specific structural information.
Parameter-Dependent Dictionaries: The method introduces a mechanism to learn dictionaries conditioned on the FGW trade-off parameter $\alpha$ , enabling the exploration of different balances between feature and structural alignment without retraining.
Empirical Validation on HCP: The approach is validated on the Human Connectome Project (HCP) dataset, demonstrating its ability to learn meaningful representations for both task-based and resting-state data.

Experimental Results

Experiments were conducted on the HCP dataset involving 1,200 subjects across 7 cognitive tasks. The dataset included both native individual geometries and the common fsaverage geometry.

Reconstruction Quality: Visualizations of reconstructed contrast maps (e.g., "story - math") showed that AGDL captures main activation patterns while smoothing noise, effectively bridging the gap between native and common geometries.
Contrast Classification: An SVM classifier trained on the learned embeddings achieved peak accuracy at $\alpha = 0.22$ . While the baseline (common geometry) slightly outperformed AGDL in this specific task, the authors attribute this to the baseline canceling out individual geometric "noise." The MLP-based AGDL model outperformed the linear interpolation model, suggesting higher expressivity is required.
Subject Classification: In a more challenging task of classifying subjects based on their embeddings, AGDL significantly outperformed the common geometry baseline. The best performance was achieved at $\alpha = 0.78$ , indicating that capturing structural information is crucial for distinguishing individuals. t-SNE visualizations confirmed that AGDL embeddings cluster by subject identity, whereas the baseline did not.
Atom Evolution: Visualizations of dictionary atoms across varying $\alpha$ values demonstrated that increasing $\alpha$ (emphasizing structural alignment) leads to sharper, more localized activations corresponding to specific contrasts.

Significance and Claims

The paper claims that AGDL represents a significant step forward in fMRI analysis by moving beyond the standard practice of projecting all subjects to a common geometry. By explicitly modeling individual geometries through optimal transport, the method preserves subject-specific information that is typically lost in registration.

The authors position this work as the first scalable method for learning joint representations of graphs with different geometries while retaining information about individual geometries. They emphasize that while the method currently uses a fixed geometry for dictionary atoms to maintain interpretability, it successfully captures subject-specific variability through the transport plans. The results suggest that accounting for individual geometry is essential for tasks requiring subject discrimination, and that the proposed amortized approach makes such analysis computationally feasible for large-scale neuroimaging datasets.

The authors modestly note that this is a first attempt at heterogeneous geometry dictionary learning, with future work planned to explore joint hemispheric learning, relaxing the fixed geometry assumption for atoms, and fully amortizing the reconstruction coefficients for near real-time processing.