Riemannian Flow Matching for Brain Connectivity Matrices via Pullback Geometry

TL;DR

This paper introduces DiffeoCFM, a pullback-geometry-based conditional flow matching framework for generating brain connectivity matrices that lie on non-Euclidean manifolds (SPD and correlation matrices). By mapping manifold data to Euclidean space via global diffeomorphisms (log for SPD, normalized Cholesky for Corr), the method trains standard Euclidean CFM and maps samples back, guaranteeing manifold-constraint outputs while avoiding costly Riemannian operations. The approach is theoretically shown to be equivalent to Riemannian CFM on pullback manifolds and is demonstrated on large-scale fMRI and EEG datasets, achieving state-of-the-art quality and classification performance with strong neurophysiological plausibility. The work significantly reduces computational overhead for geometry-aware generative modeling in neuroscience, enabling scalable and reliable synthesis of brain connectivity data for disease analysis and brain-computer interfaces.

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

Figures (20)

• Figure 1: Overview of DiffeoCFM.DiffeoCFM is a principled framework for deep generative modeling on matrix manifolds. It reformulates Riemannian conditional flow matching (CFM) on a pullback manifold $(\mathcal{M}, \phi^* g_E)$ as conventional CFM in Euclidean space$E$, via a global diffeomorphism $\phi \colon \mathcal{M} \to E$. The reformulation preserves geometry in two ways: (i) the learned Euclidean vector field $u_\theta^E$ satisfies \ref{['eq:pullback_vector_field']}, ensuring that training $u_\theta^E$ is equivalent to training $u_\theta^\mathcal{M}$; (ii) the flow trajectories obey $\phi(x(t)) = z(t)$, so that integrating in $E$ and pulling back via $\phi^{-1}$ yields the same samples as integrating directly on $\mathcal{M}$. This allows both training and sampling to be carried out efficiently in $E$, while remaining equivalent to operating on $\mathcal{M}$. On the left, fMRI correlation and EEG spatial covariance matrices lie on $\mathcal{M} = {\text{Corr}_d}$ and $\mathcal{M} = {\mathbb{S}_{d}^{++}}$, respectively. These matrices are mapped to $E$ through $\phi$, a time-dependent vector field $u_\theta^E$ is trained in $E$, and integration is performed in $E$ before mapping back via $\phi^{-1}$ to yield connectivity manifold constrained matrices.
• Figure 2: Neurophysiological Plausibility Study of DiffeoCFM.
• Figure 3: Trade-off between generative performance and computational cost for fMRI (top) and EEG (bottom) data. The figure plots the Average F1 Score against Training Time (left) and Sampling Time (right). The Average F1 Score is the mean of the quality metric ($\alpha,\beta$-F1) and the CAS F1-score from \ref{['tab:results']}. Each point marks the mean performance across all splits and datasets for a given modality, with error bars and shaded regions indicating the standard deviation. The dashed gray line represents the Real Data baseline, which serves as an empirical upper bound. Time is shown on a logarithmic scale.
• Figure 4: Training efficiency on ${\mathbb{S}_{d}^{++}}$ (simulated data). Log-scale training time over $1000$ epochs for generating SPD matrices. We vary either the number of samples $n$ or the matrix dimension $d$, keeping the other fixed. DiffeoCFM leverages the diffeomorphism $\phi_{\mathbb{S}_{d}^{++}}$\ref{['eq:diffeo_spd']} and is at least $1000\times$ faster than SPD-DDPMli2024spd. Importantly, DiffeoCFM can also generate natively correlation matrices (${\text{Corr}_d}$), and it performs with same training time as on ${\mathbb{S}_{d}^{++}}$.
• Figure 5: Comparison of fMRI matrices before and after projection

Theorems & Definitions (9)

• Proposition 1: Riemannian CFM loss function on pullback manifolds
• Proposition 2: Equivalence of ODE solutions
• Proposition 3: Equivalence of Runge--Kutta iterates
• Proposition 1: Riemannian CFM loss function on pullback manifolds
• proof
• Proposition 2: Equivalence of ODE solutions
• proof
• Proposition 3: Equivalence of Runge--Kutta iterates
• proof