Riemannian Flow Matching for Brain Connectivity Matrices via Pullback Geometry

TL;DR

This paper introduces DiffeoCFM, a novel method for efficiently generating realistic brain connectivity matrices on Riemannian manifolds by leveraging pullback metrics, enabling fast training and state-of-the-art results.

Contribution

It proposes a new approach that simplifies Riemannian generative modeling by transforming data to Euclidean space, allowing standard flow matching techniques to be applied efficiently.

Findings

Achieves state-of-the-art performance on large-scale brain imaging datasets.

Enables fast training and sampling while preserving manifold constraints.

Demonstrates effectiveness on both fMRI and EEG datasets.

Abstract

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