Riemannian Flow Matching for Brain Connectivity Matrices via Pullback Geometry

📄 arXiv: 2505.18193v2 📥 PDF

作者: Antoine Collas, Ce Ju, Nicolas Salvy, Bertrand Thirion

分类: cs.LG, eess.SP, stat.ML

发布日期: 2025-05-20 (更新: 2025-10-23)

备注: Accepted at NeurIPS 2025

🔗 代码/项目: GITHUB


💡 一句话要点

提出DiffeoCFM,通过拉回几何实现脑连接矩阵的黎曼流匹配生成模型。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 脑连接矩阵 黎曼流形 条件流匹配 拉回几何 生成模型 fMRI EEG

📋 核心要点

  1. 现有方法在黎曼流形上进行生成建模时,需要重新定义核心操作,导致计算效率低下,限制了其在大规模脑连接矩阵生成中的应用。
  2. DiffeoCFM利用拉回度量,将黎曼流形上的条件流匹配问题转化为欧几里得空间上的标准条件流匹配问题,从而避免了复杂的黎曼计算。
  3. 实验表明,DiffeoCFM在多个大规模fMRI和EEG数据集上实现了快速训练和最先进的性能,同时保证了生成矩阵的流形约束。

📝 摘要(中文)

生成逼真的脑连接矩阵对于分析大脑组织的群体异质性、理解疾病以及在具有挑战性的分类问题中扩充数据至关重要。功能连接矩阵位于受约束的空间中,例如可以建模为黎曼流形的对称正定矩阵或相关矩阵集。然而,使用黎曼工具通常需要重新定义核心操作(测地线、范数、积分),这使得生成建模的计算效率低下。本文提出DiffeoCFM,它通过利用欧几里得空间上全局微分同胚所诱导的拉回度量,在矩阵流形上实现条件流匹配(CFM)。我们证明了使用这种度量的黎曼CFM等价于在数据转换后应用标准CFM。这种等价性允许高效的向量场学习,以及使用标准ODE求解器的快速采样。我们用两种不同的设置实例化DiffeoCFM:协方差矩阵的矩阵对数和相关矩阵的归一化Cholesky分解。我们在三个大规模fMRI数据集(ADNI、ABIDE、OASIS-3)和两个EEG运动想象数据集(BNCI2014-002和BNCI2015-001)上评估了DiffeoCFM。它能够快速训练并实现最先进的性能,同时保持流形约束。代码:https://github.com/antoinecollas/DiffeoCFM

🔬 方法详解

问题定义:论文旨在解决脑连接矩阵生成的问题,特别是功能连接矩阵位于黎曼流形上,直接应用现有生成模型计算复杂度高,效率低下的问题。现有方法需要重新定义黎曼流形上的测地线、范数和积分等操作,导致计算成本显著增加。

核心思路:论文的核心思路是利用拉回几何,通过全局微分同胚将黎曼流形上的问题映射到欧几里得空间,从而可以使用标准的条件流匹配(CFM)方法。这种方法避免了在黎曼流形上直接进行复杂计算,提高了效率。

技术框架:DiffeoCFM的整体框架包括以下几个主要步骤:1) 使用全局微分同胚将脑连接矩阵从黎曼流形映射到欧几里得空间。2) 在欧几里得空间上应用标准的条件流匹配(CFM)方法学习向量场。3) 使用ODE求解器进行快速采样,生成新的脑连接矩阵。4) 将生成的矩阵通过微分同胚的逆映射回黎曼流形,保证流形约束。

关键创新:论文的关键创新在于将拉回几何与条件流匹配相结合,提出了一种在矩阵流形上进行高效生成建模的方法。通过这种方式,可以将复杂的黎曼计算转化为标准的欧几里得空间计算,大大提高了计算效率。此外,该方法能够保证生成矩阵的流形约束,确保生成结果的有效性。

关键设计:DiffeoCFM的关键设计包括:1) 选择合适的全局微分同胚,例如矩阵对数(用于协方差矩阵)和归一化Cholesky分解(用于相关矩阵)。2) 使用标准的神经网络结构(如MLP或CNN)来学习欧几里得空间上的向量场。3) 使用标准的ODE求解器(如Runge-Kutta方法)进行采样。4) 损失函数采用标准的CFM损失函数,无需针对黎曼流形进行修改。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DiffeoCFM在三个大规模fMRI数据集(ADNI、ABIDE、OASIS-3)和两个EEG运动想象数据集(BNCI2014-002和BNCI2015-001)上进行了评估,实验结果表明DiffeoCFM能够快速训练并实现最先进的性能,同时保持流形约束。具体性能数据未知,但论文强调了其在训练速度和生成质量上的优势。

🎯 应用场景

DiffeoCFM在神经科学领域具有广泛的应用前景,可用于生成逼真的脑连接矩阵,从而分析大脑组织的群体异质性,理解疾病机制,并扩充数据以提高分类问题的性能。该方法还可以应用于其他领域,例如金融风险建模、社交网络分析等,只要数据可以表示为矩阵流形。

📄 摘要(原文)

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