Learning fMRI dictionaries across individual geometries via optimal transport

Dictionary learning is a powerful tool for creating interpretable representations from complex data. When applied to functional magnetic resonance imaging (fMRI) data, the resulting patterns of brain activity can be used for various downstream tasks, such as brain activity decoding or population-level analysis. However, the variability in brain geometry across individuals poses a fundamental challenge. This is usually addressed by projecting each individual brain geometry onto a common template, which removes subject-specific information. In this work, we introduce a novel approach to dictionary learning on fMRI data that explicitly accounts for this variability. We use the optimal transport-based Fused Gromov-Wasserstein (FGW) distance to compare graphs with different geometries and features. To address the challenge of computing multiple FGW distances for large graphs such as those derived from fMRI data, we use amortized optimization to train a neural network that approximates the optimal transport plans. We also learn graph dictionary atoms that depend on the FGW trade-off parameter, which controls the balance between feature alignment and structural consistency. Numerical experiments on the HCP dataset demonstrate that the proposed approach captures meaningful information about both contrasts and subjects. It outperforms standard fMRI dictionary learning and is robust to transfer across graph scales.