Riemannian Flow Matching for Brain Connectivity Matrices via Pullback Geometry

We thank the reviewer for their detailed feedback and for highlighting the strengths of our work. We are pleased that the use of pullback metrics, the theoretical justification provided in the supplementary material, and the relevance of the considered applications and datasets were appreciated.

We now address the reviewer’s comments and questions, recalling each point in their own words before providing our response.

Choice of baselines

Important baselines are missing, making it hard to evaluate real benefits of the proposed method w.r.t. existing works. The authors do not use the methods in [1,2] for comparison with the proposed method. [...] Could the authors elaborate on their choice of baselines ?

Adding baseline SPD-CFM [1]. DiffeoCFM (proposed) is, to our knowledge, the first generative model that handles both SPD and correlation matrices while respecting manifold structure and being computationaly efficient. We acknowledge that [1] includes an SPD-manifold implementation, which we initially overlooked, and thank the reviewer for pointing it out. We adapted the official implementation of SPD-CFM to use the same conditional prior as DiffeoCFM and evaluated on BNCI 2014‑002 by holding out each of the first two sessions in turn for testing while training on all remaining sessions. The results are shown in the next table; BNCI 2015‑001 yielded similar results.

Quality metric: $\alpha$,$\beta$-F1

Classification Accuracy Score (CAS) metrics: ROC-AUC, F1

Results interpretation. SPD-CFM [1] performs slightly worse than DiffeoCFM. Moreover, SPD-CFM is 6x slower, which is expected since it relies on geodesic computation and its derivative, the Riemannian norm, and a Riemannian ODE solver. In contrast, DiffeoCFM first maps data to a Euclidean space via the diffeomorphism $\phi$, trains a standard Euclidean CFM [3] (which is very fast), and then maps generated data back to the manifold via $\phi^{-1}$. We will complete Table 1 of the manuscript with this additional baseline.

SPD-DDPM [4]. The Riemannian diffusion framework of [2] is instantiated on the SPD manifold in SPD-DDPM [4]. We have already benchmarked this model in Appendices G–H. Figure 3 (Appendix H) shows that SPD-DDPM is over 1000 × slower than DiffeoCFM due to the heavy computational cost of Riemannian manifold operations, particularly manifold sampling.

Choice of baselines. The other baselines were chosen for their complementary properties:

• DiffeoGauss fits a multivariate Gaussian distribution per class, which—despite its simplicity—performs surprisingly well, suggesting that simple yet effective priors can be valuable for neuroscience applications.
• TriangCFM and TriangDDPM are straightforward parametrizations that learn the lower-triangular parts with a Euclidean CFM; they are as fast as DiffeoCFM but, lacking diffeomorphic parametrizations, show lower precision/recall.
• The Real Data baseline compares real-vs-real samples with different splits, providing an upper bound on attainable performance. On EEG, DiffeoCFM reaches this bound, whereas on fMRI, there remains a gap to close.

This setup exposes key challenges in neuroscience and motivates the use of DiffeoCFM as a simple yet high-performing method.

[1] Chen et al., Flow Matching on General Geometries, ICLR 2024

[2] De Bortoli et al., Riemannian score-based generative modelling, NeurIPS 2022

[3] Lipman et al., Flow Matching for Generative Modeling, ICLR 2023

[4] Li et al., Spd-ddpm: Denoising diffusion probabilistic models in the symmetric positive definite space, AAAI 2024

Choice of evaluation metrics

State-of-the-art works for generative modeling on Riemannian manifolds such as [1,2] use NLL (Negative log-likelihood) scores and NFE (number of function evaluations) as evaluation metrics. [...] Currently, those metrics are not considered in the experiments. Could the authors elaborate on this choice of evaluation metrics ?

NFE. First, all models in our experiments solve the ODE with the same fixed-step midpoint solver ($dt = 0.001$), yielding an identical number of function evaluations (NFE). We therefore did not report NFE.

NLL. Then, regarding negative log-likelihood (NLL) versus sample-level metrics $\alpha$‑precision, $\beta$‑recall, and $\alpha,\beta$‑F1 [5]:

• NLL feasibility: Its computation requires a diffeomorphism between the prior and the learned distribution so that the change-of-variables log-det Jacobian is well defined. Two important baselines—TriangCFM and TriangDDPM—generate samples in ambient space and project them to the SPD (or correlation) manifold. The projection breaks injectivity, so the log-det Jacobian (and hence NLL) is undefined.
• Interpretability: The sample-level metrics compare generated to real data on a 0–1 scale, with an empirical upper bound from real-vs-real comparisons. They are therefore directly interpretable: in Table 1, DiffeoCFM reaches this bound on EEG and approaches it on fMRI. Hence they highlight generative-model challenges in Neuroscience, whereas absolute NLL values are harder to interpret.
• Conditional generation: Unlike [1, 2], our work targets conditional generation—essential in neuroscience, where data must be generated given patient-level variables (here, diagnosis). To assess conditional fidelity, we use the Classification Accuracy Score (CAS) [6], which trains a state-of-the-art classifier on generated data and evaluates it on real data, directly testing preservation of the conditional structure.

[5] Alaa et al., How faithful is your synthetic data? sample-level metrics for evaluating and auditing generative models, ICML 2022

[6] Ravuri et al., Classification accuracy score for conditional generative models, NeurIPS 2019.

Additional reference: generation on Lie groups

The work in [7] is closely related to the present work but is missing in the references. The proposed generative method for SPD manifolds can be directly derived from [7], since SPD manifolds with Log-Euclidean metrics have Lie group structure. I think it should be cited in the related works and eventually be used in comparison with the proposed method.

Contributions of [7]. The authors extend the reparameterisation trick to Lie groups and demonstrate how normalising flows can leverage this approach to learn probability densities on such spaces.

Manuscript update. Since both [7] and DiffeoCFM rely on reparameterisation tricks to estimate probability densities, [7] can be viewed as an early precursor to our approach. We have therefore added the following to the manuscript:

Added in Section 3.3 – Related Work
In [7], the authors leverage reparameterisation with normalising flows to learn probability densities on Lie groups (non-Euclidean spaces), and thus can be seen as an early precursor to our approach.

How we differ. Our contribution differs in both method and scale. We adopt the newer flow-matching framework and show that Riemannian Flow Matching [7] with pull-back metrics reduces to our proposed DiffeoCFM. The latter performs all operations in Euclidean space, yielding greater numerical stability and improved training and sampling efficiency. Moreover, while [7] presents low-dimensional examples, our experiments reach 741-dimensional fMRI connectivity (39 × 39 matrices), a particularly challenging regime.

[7] Luca Falorsi et al., Reparameterizing Distributions on Lie Groups, AISTATS 2019

We hope these results clarify our choices and demonstrate DiffeoCFM’s advantages. Thank you again for the constructive feedback.

Thank you for the detailed answer.

So the statement "The projection breaks injectivity, so the log-det Jacobian (and hence NLL) is undefined." is not true for all of your methods if I got it right ?

Concerning the vector field architecture, how many layers did you use ? For the solver, it seems that the authors of SPD-CFM used a different solver from the one you used (see Appendix H, page 22 of that paper).

We thank the reviewer for their thoughtful and positive feedback. We are glad that the paper’s technical soundness, clarity, and originality were appreciated, that the pullback approach was recognized as mathematically grounded and computationally efficient, and that its relevance to neuroimaging was noted. We also appreciate the acknowledgment of our clear exposition and the significance of DiffeoCFM for generative modeling on manifolds.

We now address the reviewer’s comments and questions, recalling each point in their own words before providing our response.

Experimental details in Appendices

Key implementation and experimental details—e.g., hyperparameters and projection analyses—are buried in the appendices (J, L); a concise summary in the main text would greatly improve self-containment.

We moved from the appendices key implementation details in the main text, complementing Appendices J and L, to improve readability without sacrificing completeness.

First, we added the following information for the metrics:

Added in Section 4.1 Metrics:
We follow the CAS protocol, training a classifier on generated samples and evaluating it on real test data. Specifically, we assess classification utility using a logistic regression [...].

Then, we detailed TriangDDPM and TriangCFM and the deep learning setup:

Added in Section 4.3 Baselines:
Since $\phi$ is not a diffeomorphism, generated matrices that do not lie on the manifold are projected back onto it. In particular, generated matrices are not necessarily positive definite, so we apply a projection [...].

TriangDDPM, TriangCFM and DiffeoCFM employ a two-layer MLP with $512$ hidden units, trained using AdamW with a learning rate of $10^{-3}$ and batch size $64$. Training runs for $200$ epochs on fMRI and $2000$ epochs on EEG. Time integration follows the midpoint method with a step size of $10^{-3}$.

Differentials $\mathrm{D} \phi(x)$, $(\mathrm{D} \phi)^{-1}(x)$: usage, cost, stability, backpropagation

Algorithms 1 and 2 omit pseudocode for computing crucial terms such as Dϕ(x) and Dϕ⁻¹(z), which are necessary for understanding Eq. (7) and ensuring reproducibility.

The vector field transformation in Eq. (7) relies on computing the differential of the diffeomorphism and its inverse. What are the computational costs and numerical stability issues, particularly for high-dimensional Corr matrices? Have the authors observed any instability in Dϕ(Σ) when Σ is near singular?

During back-propagation through Dϕ, did you encounter exploding/vanishing gradients?

Implementing Riemannian CFM [1] with a pullback metric requires $\mathrm{D}\phi(x)$, $(\mathrm{D}\phi)^{-1}(x)$, and backpropagation through them—leading to numerical instability, gradient explosion/vanishing, and high computational cost, which demand careful implementation and tuning.

This motivated DiffeoCFM, which, as shown in Propositions 1–3, is mathematically equivalent to Riemannian CFM with the pullback metric induced by a diffeomorphism $\phi$. In practice, DiffeoCFM proceeds by:

1. apply $\phi$ to data on $\mathcal{M}$;
2. train and sample a classical CFM in the Euclidean space $E$;
3. map generated data back to $\mathcal{M}$ via $\phi^{-1}$.

These steps, detailed in Algorithms 1–2, avoid computing $\mathrm{D}\phi(x)$ or any Riemannian operations—only $\phi$ and $\phi^{-1}$ are needed, both easy to implement. It should be noted that Eq (7) is used solely to prove Propositions 1, 2, and 3.

[1] Chen et al., Flow Matching on General Geometries, ICLR 2024

[2] Lipman et al., Flow Matching for Generative Modeling, ICLR 2023

Step size and network depth

No ablation is reported on ODE step size h or network depth, so the robustness–efficiency trade-off of DIFFEOCFM remains unclear.

We studied the robustness–efficiency trade‑off by varying network depth and ODE step size on the fMRI ADNI dataset, reporting $\alpha,\beta$‑F1 (mean ± std), training, and sampling times over 5‑fold cross‑validation.

Results show that performance saturates with 1 hidden layer, while deeper networks increase computation without improving quality. For the ODE step size, $10^{-3}$ offers a good trade‑off between sample quality and sampling time. This confirms that DiffeoCFM maintains high quality across a broad range of settings, with a favourable robustness–efficiency balance.

Extension to manifold without global diffeomorphisms

How would the method extend to manifolds without known diffeomorphic embeddings (e.g., the Stiefel manifold)? Could approximate embeddings or learned bijections (e.g., neural ODEs) be considered?

This work targets brain‑connectivity data—SPD and correlation matrices—where known global diffeomorphisms $\phi$ enable fast, stable training via Euclidean CFM [2] on embedded data.

As noted in the conclusion, DiffeoCFM requires such a diffeomorphism and does not apply to manifolds like the Stiefel. Even general algorithms [1] fail on manifolds without closed‑form geodesics. Two possible extensions are:

1. Approximate geodesic operators using retractions/inverses, as in $R$/$RL$-barycentres on Stiefel and Grassmann manifolds [3, 4].
2. Ambient‑space methods that progressively map back to the manifold without explicit retractions, e.g., landing algorithms for the orthogonal group [5].

Exploring these directions is promising but beyond our present scope.

[3] Kaneko, et al., Empirical arithmetic averaging over the compact Stiefel manifold IEEE Transactions on Signal Processing, 2012

[4] Bouchard, et al., Beyond $R$-barycenters: an effective averaging method on Stiefel and Grassmann manifolds IEEE Signal Processing Letters, 2025

[5] Ablin, et al., Fast and accurate optimization on the orthogonal manifold without retraction, ICML 2022

Quantitative evaluation of spatial filters

The CSP maps in Figure 2(b) are presented with colour bars and topographic projections. It would be helpful to also report quantitative similarity metrics (e.g., structural similarity index) between real and synthetic filters for more rigorous comparison.

Our study includes such a task‑level test: Table 1 reports CAS [6] when training on real or synthetic data and testing on held‑out real sets. For EEG, real vs synthetic scores match within the standard deviation, showing CSP filters from both have the same discriminative power.

Figure 2 visually confirms this result: in both $\alpha$- (8–12 Hz) and $\beta$- (13–30 Hz) bands, filters from synthetic data reproduce the contralateral C3/C4 patterns of real data during right-hand vs. feet imagery [7,8], as quantified in Table 1.

To clarify this point, we have added the following paragraph in the paper:

Added in Section 5.2 Neurophysiological Plausibility Study: We visualize CSP spatial filters in the $\alpha$ (8–12 Hz) and $\beta$ (13–30 Hz) bands that distinguish imagined right-hand from feet movements. The filters trained on real and on generated data concentrate on the same contralateral sensorimotor regions, mirroring the close CAS scores between Real data and DiffeoCFM in Table 1. This confirms that the generative model preserves the physiologically relevant decision boundaries.

Applying SSIM to EEG topographies can be misleading: (1) topographies are computed at subject/group level, making SSIM across trials impossible to compute; (2) single‑trial maps contain high non‑neural noise, so pixel‑wise metrics risk capturing artifacts rather than $\alpha$/$\beta$ ERD/ERS patterns. This motivated using CAS [6] and CSP filter inspection.

[6] Ravuri et al., Classification accuracy score for conditional generative models, NeurIPS 2019.

[7] Pfurtscheller et al., Event-related EEG/MEG synchronization and desynchronization: basic principles, Clinical Neurophysiology, 1999.

[8] Pfurtscheller et al., Mu rhythm (de)synchronization and EEG single-trial classification of different motor imagery tasks, NeuroImage, 2006.

Manifold-aware α-precision and β-recall

α-precision and β-recall are geometry-agnostic. Could authors provide a discussion of a manifold-aware version (e.g., geodesic-based recall) in future work?

This remark echoes our conclusion: these metrics [9], standard in generative modeling, assess generated–real pairwise correlation match via the Frobenius norm $||\Sigma_{\text{real}} - \Sigma_{\text{gen}}||_F^2$, a long‑standing choice in covariance analysis.

A Riemannian distance could reveal other aspects—for instance, the affine‑invariant metric compares eigenvalues in log space, making multiplicative variance changes (e.g., fixed dB shifts in EEG power) linear in distance. While promising, this would change the meaning of $\alpha$‑precision and $\beta$‑recall and is beyond our scope.

[9] Alaa et al., How faithful is your synthetic data? sample-level metrics for evaluating and auditing generative models, ICML 2022

We thank the reviewer for their thoughtful and encouraging feedback. We are pleased that they found the addressed problem both interesting and valuable, and that the proposed use of the pullback metric to enable Euclidean optimization was recognized as a key strength of our approach. We also appreciate their remarks on the clarity and completeness of the manuscript, as well as their positive assessment of the experimental results on real data.

We now address the reviewer’s comment, recalling the point in their own words before providing our response.

Estimation of Connectivity Matrices

Since the generation is on connectivity where the connectivity matrices are obtained from brain signal data (such as fMRI), the way of defining connectivity could possibly influence the "ground-truth" thus introduce bias. It would be great if authors can discuss more on it.

Oracle Approximating Shrinkage (OAS) [1] and Ledoit–Wolf [2] are the state-of-the-art shrinkage estimators in recent fMRI and EEG studies [3–5], where they consistently deliver the best classification performance. Accordingly, we estimate covariance matrices with OAS and obtain correlation matrices by z-scoring the time series before applying OAS. Indeed, we chose OAS to maximise the Classification Accuracy Score (CAS) of the “Real data” upper bound in Table 1. CAS is the accuracy of a classifier trained on generated data (conditioned on disease) and tested on real data; a high upper bound gives the strongest baseline for evaluating our generative model.

We agree that the estimator itself may influence the ground-truth and, hence, our results. Because our contribution lies in generative modelling rather than connectivity estimation, we fixed a community-standard pipeline to keep this factor constant. To clarify the limitation, we added the following to the manuscript:

Added in Section 6 Conclusions, Limitations, and Future Work:
A full study of how alternative connectivity definitions—e.g., partial correlation or graphical-Lasso precision—affect generation quality remains an open problem that needs to be addressed, as the choice of definition directly determines the ground truth used in the experiments.

[1] Chen et al., Shrinkage algorithms for MMSE covariance estimation, IEEE Transactions on Signal Processing 2010

[2] Ledoit et al., A Well-Conditioned Estimator for Large-Dimensional Covariance Matrices, Journal of Multivariate Analysis 2004

[3] Kamalaker et al., Benchmarking functional connectome-based predictive models for resting-state fMRI, NeuroImage 2019

[4] Pervaiz, et al., Optimising network modelling methods for fMRI, NeuroImage 2020

[5] Sabbagh et al., Predictive regression modeling with MEG/EEG: from source power to signals and cognitive states, NeuroImage, 2022

We thank the reviewer for their positive and encouraging feedback. We are glad that the motivation, writing, and integration of pullback geometry with conditional flow matching in our proposed method, DiffeoCFM, were appreciated. We are also pleased that the reviewer found our state-of-the-art results and the biological plausibility analysis convincing, as these are central to the impact we aim to achieve in neuroscience applications.

We now address the reviewer’s comments and questions, recalling each point in their own words before providing our response.

Impact of the choice of diffeomorphism and vector field

The impact of diffeomorphism choice (log-Euclidean vs. affine-invariant, other Cholesky normalizations) or vector-field capacity is not discussed.

The choice of a two-layer 512-unit MLP is not fully justified. Could the authors try transformer-style vector fields or neural ODEs, and also explore alternative pullbacks?

First of all, the number of available diffeomorphisms for SPD and correlation manifolds is limited. We selected the normalized Cholesky map for correlation matrices and the logarithm map for SPD matrices because of their simplicity and ease of implementation. Then, identifying alternative mappings and systematically comparing their properties remains an open problem and we believe is beyond the scope of this work. We have added the following note about the choice of the two diffeomorphisms in Section 3.2:

Added in Section 3.2: We selected these two diffeomorphisms for their simplicity and ease of implementation. However, other choices are possible. For correlation matrices, alternative parameterizations are discussed in [1]. For SPD matrices, one can use the Cholesky factor with a logarithm applied to the diagonal, leading to the log-Cholesky metric [2], which also defines a global diffeomorphism.

Regarding the vector field, it is implemented as a two-layer MLP with 512 hidden units. This architectural choice is consistent with prior work [3, 4, 5], and in our experience, the results remained stable as long as the network was sufficiently large. In our setting, the pullback yields an unstructured vector representation, lacking spatial or sequential structure. As such, transformer architectures do not offer additional meaningful inductive bias while increasing model complexity. Similarly, replacing the MLP with a neural ODE backbone would duplicate the existing continuous-time formulation without a clear benefit.

[1] Thanwerdas, Riemannian and stratified geometries of covariance and correlation matrices, Thesis 2022.

[2] Lin et al., Riemannian geometry of symmetric positive definite matrices via Cholesky decomposition, SIAM Journal on Matrix Analysis and Applications 2019.

[3] Kotelnikov et al., TabDDPM: Modelling Tabular Data with Diffusion Models, ICML 2023.

[4] Shi et al., Diffusion Schrödinger Bridge Matching, NeurIPS 2023.

[5] Albergo et al., Building Normalizing Flows with Stochastic Interpolants, ICLR 2023.

Comparisons with other true manifold baselines

Comparisons with other true manifold baselines such as SPD-DDPM in terms of quality and runtime would be helpful.

SPD-DDPM A comparison of training time of DiffeoCFM with SPD-DDPM is provided in Appendices G–H. First, we note that SPD-DDPM is designed specifically for SPD matrices (not correlation matrices) and operates entirely on the SPD manifold. As a result, it relies on expensive Riemannian training and sampling (Riemannian optimization of SPDNet, exp/log map, Riemannian sampling solver, ...), making it impractical even for medium-sized datasets (e.g., 100 samples of dimension 30). Indeed, as shown in Figure 3 (Appendix H), it is over 1000× slower than DiffeoCFM, which shares the same computational footprint as CFM, a method explicitly designed for scalability and widely adopted in computer vision. Due to this computational burden, we were unable to assess the sample quality of SPD-DDPM in our experimental setting.

SPD-CFM We compare DiffeoCFM against SPD-CFM [6], which will be added to Table 1. We adapted the official implementation of SPD-CFM to use the same conditional prior as DiffeoCFM and evaluated on BNCI 2014‑002 by holding out each of the first two sessions in turn for testing while training on all remaining sessions. The results are shown in the next table; BNCI 2015‑001 yielded similar results.

Quality metrics = $\alpha$-Precision, $\beta$-Recall, $\alpha$,$\beta$-F1 Classification Accuracy Score (CAS) metrics: ROC-AUC, F1

Results interpretation SPD-CFM [6] performs slightly worse than DiffeoCFM. Moreover, SPD-CFM is approximately 6x slower, which is expected since it relies on geodesic computation and its derivative, the Riemannian norm, and a Riemannian ODE solver. In contrast, DiffeoCFM first maps data to a Euclidean space via the diffeomorphism $\phi$, trains a standard Euclidean CFM (which is very fast), and then maps generated data back to the manifold via $\phi^{-1}$. We will complete Table 1 of the manuscript with this additional baseline.

[6] Chen et al., Flow Matching on General Geometries, ICLR 2024

Additional literature

Additional literature could be discussed, such as Jo et al. and Kapusniak et al., which also generate data on manifolds through learned or mixture-based interpolants.

We have added a brief discussion of [7] and [8] in the Section 3.3 Related-work:

Added in Section 3.3: References [7] and [8] explore more general settings by learning bridge matches on arbitrary manifolds or data-driven Riemannian metrics, whereas our approach focuses on Riemannian geometries defined via known pullback diffeomorphisms.

Key distinctions are summarized below:

• Reference [7] addresses a more general setting, learning bridge matching on manifolds where Riemannian logarithms may be undefined or costly to compute. They rely on the spectral decomposition of the Laplace-Beltrami operator to define distances and then reduce the simulation budget during training. It should be noted that they still require the Riemannian norm in their loss function during training and the geodesic (Riemannian exponential mapping) during sampling. In contrast, our work focuses on pullback manifolds such as SPD and correlation spaces, where known diffeomorphisms enable fast and scalable training and sampling without complex Riemannian operators.
• Reference [8] replaces straight interpolants with curved ones derived from a learned Riemannian metric. While complementary, this line of work is orthogonal to our focus on leveraging existing pullback maps (and their associated Riemannian metrics) for efficient training and generation.

[7] Jo et al., Generative Modeling on Manifolds Through Mixture of Riemannian Diffusion Processes, ICML 2024.

[8] Kapusniak et al., Metric Flow Matching for Smooth Interpolations on the Data Manifold, NeurIPS 2024

Scalability to high-dimensional manifolds

Discussion on scalability to high-dimensional manifolds would be helpful.

First, learning generative models when the ambient dimension $d$ approaches the sample size $n$ is statistically hard. For example, Theorem 3.1 of [9] shows that the estimation error of neural networks trained with diffusion decays no faster than $n^{-1/d}$, which is the curse of dimensionality. In our fMRI experiments we adopted the MSDL atlas with 39 regions, giving a manifold dimension of $39 \times 38 / 2 = 741$, which remains below the number of available samples (e.g. 1900 ADNI scans). Using finer atlases would push us into the $n<d$ regime where theory predicts over-fitting and degraded generative quality. Designing priors or low-rank parameterisations to tame this regime is an important avenue for future work.

We have inserted the following clarification into the Section 6:

Added in Section 6 Conclusions, Limitations, and Future Works:
Learning brain connectivities at higher parcellation granularity is challenging: the manifold dimension grows quadratically with the number of regions, while the number of required samples increases exponentially with dimension [9]. Addressing this high-dimensional regime, e.g. through latent representations, remains an important direction for future work.

Second, from a computational standpoint, DiffeoCFM inherits the scalability of CFM, which was originally developed for high-dimensional data and large sample sizes in vision tasks. In contrast, manifold-based methods like SPD-DDPM scale poorly in practice due to their reliance on costly geometric operations.

[9] Oko, et al. Diffusion models are minimax optimal distribution estimators. ICML 2023