MAGNET: Manifold-Aware Graph Diffusion Network for Connectome Generation¶

会议: NeurIPS 2026
arXiv: 2610.17014
代码: https://github.com/protyayofficial/MAGNET
领域: 图学习
关键词: 脑连接组生成, 图扩散网络, 流形感知, 归一化Cholesky坐标, 条件支撑集恢复

一句话总结¶

针对粗粒度临床标签导致脑功能连接组生成时出现亚群模式坍缩与失真问题，MAGNET在归一化Cholesky正定相关流形图表上构建基于脑区Token的图扩散架构，并引入类别锚定高斯混合子型、平摊结构桥与相关性保留加噪三大机制，在严格保真数学有效性的同时实现了临床异质性与拓扑支撑集的精准恢复。

研究背景与动机¶

静息态功能性脑连接组（functional connectome）通过脑区间的成对相互作用矩阵刻画全脑神经连接，是神经影像学生物标记物挖掘、临床队列分析与个性化病程预测的核心基石。然而在临床实际中，高保真功能连接组的生成面临着苛刻的几何约束与双重挑战：一方面，生成的矩阵必须严格满足数学有效性，即属于对称半正定且对角线恒为1的有效相关矩阵流形 $\mathcal{C}_d$；另一方面，模型必须能够精确落入对应临床诊断类别的真实子流形支撑集（class-conditional support）。过去的研究主要聚焦于几何有效性，例如利用黎曼流形匹配（Riemannian flow matching）或归一化 Cholesky 坐标变换（normalized-Cholesky pullback）在构造上消除非正定矩阵，但这仅保证了“数学上是一个相关矩阵”，却远不足以恢复临床真实的病理分布。

在真实神经影像数据集中，临床标签通常极其粗粒度（例如简单的二元标签“阿尔茨海默病 vs 健康对照”或“自闭症 vs 对照”），完全忽视了疾病内部极其丰富的亚群变异（subtypes/stages）与跨队列协变量偏移。当条件生成模型以这种粗粒度标签为先验时，模型被迫在相互正交的潜变量子空间中取均值，退化成一种几何上有效但在临床上失真、模态坍缩的虚假“平均连接组”。此外，流形图表本身的非线性导致各向同性高斯加噪并不中立，等方差噪声会不成比例地破坏关键坐标；而在少步逆向采样阶段，局部的带噪状态又缺乏感知全脑全局拓扑的能力，常规分类器无关引导（CFG）极易引发过特化与多样性崩溃。

面对这一几何有效性与临床保真度脱节的核心矛盾，本文认为几何有效性只是连通图生成的必要条件而非充分条件，核心瓶颈在于流形之上的条件支撑集恢复（conditional support recovery）。核心 idea：在无约束变换保证相关流形严格有效的 normalized-Cholesky 坐标系中，将展开向量重构为脑区级（ROI）图元并结合关系归纳偏置，通过类别锚定混合子型、平摊全局结构桥以及相关性自适应加噪三重协同机制，直接攻克粗标签模糊性、全局拓扑感知缺失与坐标各向同性破坏难题。

方法详解¶

整体框架¶

MAGNET 将功能连接组的生成建立在归一化 Cholesky 图表（normalized-Cholesky chart）之上。给定具有 $d$ 个脑区的全脑相关矩阵 $C \in \mathcal{C}_d$，通过 Cholesky 分解 $C = L L^\top$ 与对角归一化 $\widetilde{L} = \operatorname{diag}(L)^{-1} L$，提取严格下三角坐标 $z = \operatorname{vech}_0(\widetilde{L}) \in \mathbb{R}^m$（其中自由度 $m = d(d-1)/2$），并在标准化坐标 $x = (z - \mu_z)/\sigma_z$ 上构建扩散过程。解码时重建下三角矩阵并通过对角重归一化 $\widehat{C} = \operatorname{diag}(G)^{-1/2} G \operatorname{diag}(G)^{-1/2}$（其中 $G = \widehat{L} \widehat{L}^\top$），天然从代数构造上确保任何解码输出均为对称、正定且对角线为1的合法相关矩阵。

为赋予模型捕捉全脑复杂网络交互的关系归纳偏置（relational inductive bias），MAGNET 摒弃了将流形坐标展平为无结构一维向量的常规做法，将带噪坐标 $x_t$ 重塑对称且对角线为0的脑区交互邻接矩阵 $A(x_t) \in \mathbb{R}^{d \times d}$，矩阵的每一行对应一个 ROI 脑区的局部连接轮廓。在此基础上，网络将脑区拓扑表征、可学习解剖位置编码及扩散时间步、类别子型、结构桥先验融合为脑区元（ROI tokens），经由预归一化（Pre-norm）Transformer 架构执行全图交互去噪，最后投影并对称化回图表坐标空间。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["带噪流形坐标 xt<br/>重构为 ROI 交互矩阵 A(xt)"] --> Stage1["类别锚定条件化<br/>EM拟合GMM分解组内亚型 s"]
    Stage1 --> Stage2["平摊结构桥<br/>从带噪状态预估紧凑图谱先验 bt"]
    Stage2 --> Stage3["相关性保留加噪<br/>条件差分自适应屏蔽关键坐标方差"]
    Stage3 --> Core["ROI 图去噪 Transformer<br/>结合解剖位置编码与动态门控条件"]
    Core --> Stage4["正交化饱和引导采样<br/>DDIM 逆向去噪与残差正交 CFG"]
    Stage4 --> Out["归一化 Cholesky 解码<br/>保真输出有效正定相关流形 C_hat"]

关键设计¶

1. 类别锚定条件化：用高斯混合子型分解类内潜在结构异质性 传统方法直接使用粗粒度二元类别标签 $y$ 作为引导，这迫使去噪网络在面临跨受试者病理差异时只能学习多峰分布的数学均值，导致临床子结构的塌缩。MAGNET 在训练阶段引入无监督的高斯混合模型（GMM），在标准化图表坐标空间上针对每个粗类别 $y$ 单独通过期望最大化（EM）算法拟合对角协方差高斯混合分布 $p(x \mid y) \approx \sum_{s=1}^{K_y} \pi_{y,s} p(x \mid y, s)$，并通过贝叶斯信息准则（BIC）自动确定混合分支数 $K_y$。在训练时，样本根据后验概率被赋予硬子型分配 $s_i = \arg\max_s p(s \mid x_i, y_i)$；在采样生成时，先依先验分布采样 $s \sim p(s \mid y)$，再将其嵌入向量 $e_s$ 注入去噪网络。从理论上看，该设计直接降低了贝叶斯去噪风险： $$\mathcal{R}_Y^\star - \mathcal{R}_{Y,S}^\star = \mathbb{E}\left[\left\| \mathbb{E}[X_0 \mid X_t, Y, S] - \mathbb{E}[X_0 \mid X_t, Y] \right\|_2^2\right] \ge 0$$ 这一风险差严格非负，且当且仅当子型变量 $S$ 捕捉到了粗标签 $Y$ 所掩盖的类内流形子结构时取得严格正增益，从而在源头上阻断了模式坍缩。

2. 平摊结构桥：在少步逆向采样中提供轻量级全局谱先验 尽管子型划分减轻了标签模糊性，但在仅有极少逆向步数（如 6 步 DDIM）的推理预算下，局部的带噪状态 $x_t$ 难以在短时间内感知并恢复全脑大尺度宏观拓扑（如谱矩、网络秩、条件数等全局图特征）。MAGNET 提出平摊结构桥（Amortized Structural Bridge）机制，首先定义一个代理结构特征提取器 $\varphi(C) \in \mathbb{R}^{16}$（覆盖特征值矩、图谱秩、条件数及非对角连通性统计量），并为每个子型 $s$ 预先拟合包含紧凑代表性原子的集合 $\mathcal{A}_s = \{a_{s,1}, \dots, a_{s,K_s}\} \subset \mathbb{R}^{16}$。教师目标 $b^\star(x_0, s, t)$ 由这组原子的距离加权凸组合给出。由于推理时无法获取纯净状态 $x_0$，MAGNET 构建了一个轻量化平摊预估网络 $B_\psi(x_t, t, s) = f_\psi(\varphi_{\mathrm{cheap}}(x_t), e_t, e_s)$，直接从当前带噪状态中平摊拟合教师谱结构。训练时引入平方蒸馏目标： $$L_{\mathrm{bridge}} = \left\| B_\psi(x_t, t, s) - b^\star(x_0, s, t) \right\|_2^2$$ 在逆向采样阶段，$B_\psi$ 完全取代耗时的纯净态迭代估计，通过学习得到的门控单元 $\xi(\widetilde{e}_b) \odot \widetilde{e}_b$ 注入到每个脑区 Token 中，以极低的计算开销跨越了局部扰动与宏观图拓扑之间的信息代沟，将采样速度相对教师桥提升了 25–32 倍。

3. 相关性保留加噪：通过条件去噪差分自适应保护关键流形坐标 归一化 Cholesky 坐标变换是非线性的，在此图表上施加标准的各向同性高斯白噪声会导致扰动在流形上产生非中立的畸变，特别是在大步长的少步采样中，关键坐标的剧烈退化会导致生成轨迹不可逆地漂移出真实类别的支撑集区域。MAGNET 提出了条件自适应的前向加噪定律。首先以标准加噪状态 $\tilde{x}_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$ 评估条件去噪输出与无条件去噪输出的绝对差分向量 $d_t = G_\theta(\tilde{x}_t, t, s, b_t) - G_\theta(\tilde{x}_t, t, \varnothing, \varnothing)$，构建逐坐标归一化的条件相关性重要度掩码 $m_{t,j} = |d_{t,j}| / (\max_k |d_{t,k}| + \varepsilon)$。结合单调递减的时间退火函数 $\rho(t) \in [0, 1]$，构造流形保真高斯先验，进而推导出逐坐标方差受控的前向扰动律： $$\sigma_{t,j}^2 = \operatorname{clip}\left(1 - \kappa \rho(t) m_{t,j},\, \sigma_{\min}^2,\, 1\right), \quad x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t}(\sigma_t \odot \epsilon)$$ 通过在扩散早期对判别性特征明显的流形坐标主动降低扰动方差、在后期逐步回归标准高斯分布，有效减缓了生成轨迹偏离目标流形支撑集的漂移累积。

4. 正交化饱和引导采样：防止流形过特化的残差解耦 CFG 在采样阶段，传统分类器无关引导（CFG）直接沿引导方向施加固定位移外推，不仅容易在有限步数内破坏生成样本的解剖多样性，还可能因为无条件预测与条件预测在平行方向的潜在冲突导致流形偏离。MAGNET 将无条件预测向量 $\hat{v}_\varnothing$ 显式分解为与引导方向 $g = \hat{v}_c - \hat{v}_\varnothing$ 平行和正交的两个分量，剥离平行分量并引入双曲正切饱和控制项： $$\hat{v}_\varnothing^\perp = \hat{v}_\varnothing - \frac{\langle \hat{v}_\varnothing, g \rangle}{\|g\|_2^2 + \varepsilon} g, \quad \hat{v}_\gamma = \hat{v}_\varnothing^\perp + \gamma \tanh(\|g\|_2) g$$ 该策略确保引导方向严格受控，同时利用 $\tanh$ 抑制引导强度过大时的极端值溢出，既保留了类无关的正交拓扑多样性，又实现了稳健的条件导向。

损失函数 / 训练策略¶

MAGNET 的端到端训练联合最小化三个互补目标： $$L = L_{\mathrm{diff}} + \lambda_{\mathrm{bridge}} L_{\mathrm{bridge}} + \lambda_{\mathrm{spec}} L_{\mathrm{spec}}$$ 其中 $L_{\mathrm{diff}}$ 为标准流形扩散匹配损失（采用 min-SNR 加权与纵向受试者逆频重加权以平衡长尾偏置）；$L_{\mathrm{bridge}}$ 为由前向信噪比 $\bar{\alpha}_t$ 加权的平摊桥蒸馏损失（仅作用于有条件分支）；$L_{\mathrm{spec}} = \|\varphi(\widehat{x}_0) - \varphi(x_0)\|_2^2$ 为预测潜变量与真实潜变量之间的谱特征一致性正则项。默认超参数固定为 $\lambda_{\mathrm{bridge}} = \lambda_{\mathrm{spec}} = 0.1$。去噪模型主干采用 4 层 Pre-norm Transformer 编码器（隐藏维度 128，4 个 QK-Normalized 注意力头，GeGLU 维度 512）。采样默认采用 6 步 DDIM 离散求解器。

实验关键数据¶

主实验¶

实验在三个国际公认的神经影像队列上展开：ABIDE（846 例脑扫描，自闭症 vs 对照）、ADNI（2343 例，阿尔茨海默病 vs 对照）以及 OASIS-3（1800 例，纵向认知衰退）。均采用 MSDL 脑图谱划分（39 个 ROI，对应流形图表维度 741），在严格的受试者独立 10 折交叉验证下，以 6 步采样预算对比前沿流形与图生成模型。

数据集	模型	$\alpha$-Precision ↑	$\beta$-Recall ↑	$\alpha,\beta$-F1 ↑	Downstream F1 ↑	采样耗时 (s) ↓
ABIDE	Real Data (Oracle)	0.80 ± 0.04	0.81 ± 0.04	0.80 ± 0.01	0.58 ± 0.06	N/A
	TriangCFM	0.21 ± 0.21	0.24 ± 0.35	0.20 ± 0.29	0.40 ± 0.13	7.16 ± 1.17
	TriangDDPM	0.25 ± 0.04	0.22 ± 0.06	0.23 ± 0.04	0.43 ± 0.08	5.45 ± 0.19
	DiffeoGauss	0.51 ± 0.04	0.23 ± 0.03	0.32 ± 0.03	0.56 ± 0.06	0.06 ± 0.01
	DiffeoCFM	0.81 ± 0.05	0.50 ± 0.04	0.62 ± 0.05	0.52 ± 0.07	9.04 ± 0.86
	GDSS-proj	0.00 ± 0.00	0.28 ± 0.07	0.00 ± 0.00	0.29 ± 0.31	1.05 ± 0.02
	MAGNET (Ours)	0.80 ± 0.04	0.59 ± 0.04	0.68 ± 0.03	0.56 ± 0.06	0.38 ± 0.02
ADNI	Real Data (Oracle)	0.87 ± 0.09	0.87 ± 0.02	0.87 ± 0.05	0.59 ± 0.03	N/A
	TriangCFM	0.23 ± 0.31	0.25 ± 0.35	0.23 ± 0.33	0.37 ± 0.10	21.70 ± 3.35
	TriangDDPM	0.27 ± 0.06	0.30 ± 0.02	0.28 ± 0.03	0.43 ± 0.05	13.92 ± 0.93
	DiffeoGauss	0.02 ± 0.01	0.50 ± 0.08	0.04 ± 0.02	0.28 ± 0.13	1.42 ± 0.40
	DiffeoCFM	0.68 ± 0.12	0.75 ± 0.04	0.71 ± 0.05	0.40 ± 0.09	36.08 ± 5.84
	GDSS-proj	0.57 ± 0.05	0.38 ± 0.07	0.45 ± 0.05	0.53 ± 0.36	1.70 ± 0.01
	MAGNET (Ours)	0.71 ± 0.13	0.84 ± 0.06	0.76 ± 0.06	0.61 ± 0.05	0.86 ± 0.28
OASIS-3	Real Data (Oracle)	0.89 ± 0.04	0.88 ± 0.05	0.88 ± 0.01	0.61 ± 0.06	N/A
	TriangCFM	0.22 ± 0.27	0.26 ± 0.37	0.23 ± 0.32	0.45 ± 0.17	15.29 ± 2.51
	TriangDDPM	0.26 ± 0.08	0.31 ± 0.05	0.28 ± 0.05	0.48 ± 0.03	8.79 ± 1.31
	DiffeoGauss	0.52 ± 0.06	0.31 ± 0.02	0.39 ± 0.03	0.40 ± 0.09	1.88 ± 1.10
	DiffeoCFM	0.63 ± 0.06	0.38 ± 0.04	0.47 ± 0.05	0.53 ± 0.04	23.24 ± 2.55
	GDSS-proj	0.01 ± 0.02	0.04 ± 0.01	0.01 ± 0.01	0.30 ± 0.32	1.44 ± 0.01
	MAGNET (Ours)	0.75 ± 0.07	0.47 ± 0.06	0.57 ± 0.06	0.56 ± 0.04	0.63 ± 0.00

消融实验¶

在 ABIDE 与 ADNI 上针对骨干架构与四大控制模块开展组件消融，采样步数均固定为 6 步。

配置	骨干网络	Subtype	Bridge	Noise	Guidance	$\alpha,\beta$-F1 (ABIDE)	$\alpha,\beta$-F1 (ADNI)	采样时间 (s)
纯向量扁平去噪	Vector	✗	✗	✗	✗	0.26 ± 0.03	0.34 ± 0.04	0.04 ± 0.00
基础图去噪骨干	Graph	✗	✗	✗	✗	0.56 ± 0.05	0.57 ± 0.06	0.36 ± 0.01
+ 仅类别子型	Graph	✓	✗	✗	✗	0.59 ± 0.05	0.67 ± 0.06	0.39 ± 0.01
+ 仅平摊结构桥	Graph	✗	✓	✗	✗	0.58 ± 0.06	0.69 ± 0.11	0.38 ± 0.01
+ 仅相关性保留加噪	Graph	✗	✗	✓	✗	0.59 ± 0.03	0.69 ± 0.07	0.38 ± 0.00
三模块无饱和引导	Graph	✓	✓	✓	✗	0.63 ± 0.04	0.74 ± 0.09	0.38 ± 0.06
完整 MAGNET	Graph	✓	✓	✓	✓	0.68 ± 0.03	0.76 ± 0.06	0.38 ± 0.02

此外，在零样本细粒度临床亚型恢复实验中（模型仅在二元标签上训练，但在测试时通过独立预训练分类器解构为 4 个细分临床阶段），MAGNET 的协方差谱对齐误差 $\lambda$-MSE 在 ADNI 上降至 6.08 ± 0.53，比 DiffeoCFM（14.30）降低超 57%；在 OASIS-3 上降至 1.61 ± 0.97，显著优于所有基线。

关键发现¶

图结构去噪骨干是性能质变的核心：仅将无结构的扁平一维向量去噪替换为脑区 Token 级图 Transformer，$\alpha,\beta$-F1 便从 0.26 骤增至 0.56。证明即便流形图表严格保持正定有效，若缺乏对解剖脑区邻接交互的关系偏置，模型仍无法恢复高维流形支撑集。
三大控制器具备深度协同增强效应：在子型条件、结构桥与相关性噪声中，单一模块均能带来 2–12% 的 F1 提升，但三者集成能使 ADNI 的 $\alpha,\beta$-F1 从 0.57 跃升至 0.74。正交化饱和引导进一步将保真度推至 0.76，且不会引起采样延迟增长。
计算效率颠覆连续流模型：相比需要数十步高阶连续 ODE 积分求解的 DiffeoCFM（ADNI 单折耗时 36.08 秒），MAGNET 借助 6 步离散 DDIM 与平摊结构桥在 0.86 秒内完成生成，获得超 40 倍的采样加速。

亮点与洞察¶

解耦几何有效性与临床真实性：明确指出生成一个数学上合法的正定相关矩阵只是神经科学建模的必要条件，核心瓶颈在于粗粒度监督下如何避免多峰分布被强制均值化，在流形内精确定位临床真实的支撑集。
轻量平摊桥优雅兼顾宏观谱先验与推理速度：不依赖采样时的复杂递归反推（如 Self-Conditioning）或前向模拟，通过独立训练的小型网络平摊预测脑网络高阶谱特征，以极小代价补足了少步采样阶段全局拓扑盲区。
非线性图表下的坐标非对称加噪：洞察到流形坐标参数化与加噪各向同性之间的内在冲突，以条件去噪差分作为敏感度度量，自适应保护关键病理坐标，为非欧几里得几何扩散提供了普适的加噪设计范式。

局限与展望¶

当前图谱分辨率依赖：主实验主要基于 MSDL（39 脑区），虽然在 AAL116 与 CC200（200 脑区）扩展实验中依然保持优势，但当脑区数量攀升至超千个体素团时，全注意力与 Cholesky 坐标参数量呈 $O(d^2)$ 增长，亟需引入稀疏注意力或分层脑图谱分解。
缺乏多模态融合与连续病程建模：目前生成仅依赖结构化功能连接组，未结合结构磁共振（sMRI）、弥散张量成像（DTI）或连续时间维度的纵向衰退轨迹，未来可扩展至多模态脑状态联合生成。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 将归一化 Cholesky 有效流形图表与脑区级图扩散去噪深度结合，并创新性提出解耦粗标签异质性的高斯混合子型、平摊全局结构桥及自适应方差保护机制。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖 ABIDE、ADNI、OASIS-3 三大真实临床队列，设立严格的受试者级 10 折评估，提供多指标支持恢复评估、零样本亚型检验、全模块消融及高分辨率图谱泛化验证。
写作质量: ⭐⭐⭐⭐⭐ 数学定义严密，理论推导扎实（提供两项贝叶斯去噪风险递减命题及其解析证明），实验图表详实且逻辑清晰。
价值: ⭐⭐⭐⭐⭐ 为计算神经科学、精神疾病标记物合成与非欧几里得流形扩散模型提供了极具启发性的高保真生成范式与开源工具。