Riemannian Flow Matching for Brain Connectivity Matrices via Pullback Geometry
This work addresses a domain-specific problem for neuroscience and medical imaging by providing an efficient generative model for brain connectivity data, though it is incremental as it builds on existing flow matching techniques.
The paper tackles the problem of generating realistic brain connectivity matrices, which lie in constrained Riemannian manifolds, by proposing DiffeoCFM, a method that enables efficient conditional flow matching via pullback metrics. It achieves state-of-the-art performance on large-scale fMRI and EEG datasets with over 4600 scans and 30000 trials, while preserving manifold constraints and allowing fast training and sampling.
Generating realistic brain connectivity matrices is key to analyzing population heterogeneity in brain organization, understanding disease, and augmenting data in challenging classification problems. Functional connectivity matrices lie in constrained spaces, such as the set of symmetric positive definite or correlation matrices, that can be modeled as Riemannian manifolds. However, using Riemannian tools typically requires redefining core operations (geodesics, norms, integration), making generative modeling computationally inefficient. In this work, we propose DiffeoCFM, an approach that enables conditional flow matching (CFM) on matrix manifolds by exploiting pullback metrics induced by global diffeomorphisms on Euclidean spaces. We show that Riemannian CFM with such metrics is equivalent to applying standard CFM after data transformation. This equivalence allows efficient vector field learning, and fast sampling with standard ODE solvers. We instantiate DiffeoCFM with two different settings: the matrix logarithm for covariance matrices and the normalized Cholesky decomposition for correlation matrices. We evaluate DiffeoCFM on three large-scale fMRI datasets with more than 4600 scans from 2800 subjects (ADNI, ABIDE, OASIS-3) and two EEG motor imagery datasets with over 30000 trials from 26 subjects (BNCI2014-002 and BNCI2015-001). It enables fast training and achieves state-of-the-art performance, all while preserving manifold constraints. Code: https://github.com/antoinecollas/DiffeoCFM