跳转至

A Steerable Deep Network for Model-Free Diffusion MRI Registration

会议: NeurIPS 2026
arXiv: 2501.04794
论文: arXiv
领域: 医学图像
关键词: 扩散磁共振成像、可操纵卷积、特殊欧氏群 SE(3)、非刚性配准、最大均值差异

一句话总结

提出首个基于 SE(3)-等变可操纵 UNet 的免模型(model-free)扩散磁共振成像(dMRI)非刚性配准框架,将配准建模为位置-方向联合空间上的等变微分同胚,并结合傅里叶空间最大均值差异损失直接对齐原始信号与集成平均传播子。

研究背景与动机

扩散磁共振成像(dMRI)通过测量水分子在不同扩散梯度方向与强度下的弥散行为来探查脑部神经微观结构。由于水分子扩散具有强烈的空间方向依赖性,每个体素内的信号不仅反映解剖位置信息,更包含方向特征。在常规结构像(如 T1/T2 sMRI 或 CT)中,配准目标是估计三维标量图像之间的空间形变;而在 dMRI 中,一旦空间网格发生非刚性变形,各体素内的方向信号或局部微观结构几何朝向必须依据形变雅可比矩阵(Jacobian)进行物理一致的局部重定向(reorientation),以保持白质纤维束等微观方向与解剖形变同步。

现有 dMRI 配准方法几乎全部依赖于预先估计的“派生表征”(derived representations),例如扩散张量(DT)、纤维方向分布函数(fODF)、张量混合模型(GMM)或分数阶各向异性(FA)图与纤维朝向图(TOM)。这种将方向信息转嫁到值域(codomain)的做法面临两个核心痛点:其一,任何派生模型都引入了先验假设(如高斯扩散假设),在拟合过程中不可避免地丢失了原始扩散信号中丰富的非高斯或多纤维交叉细节;其二,由于方向信息被置于值域,配准形变场依然仅由三维空间位置驱动,未能直接利用原始信号的高维方向耦合关系来联合引导非刚性配准过程。

为了打破对派生表征的依赖并彻底避免复杂脆弱的显式雅可比重定向,本文提出了全新的切入视角:将扩散磁共振成像的配准域从单纯的三维欧氏空间拉回到包含位置与方向的联合空间。核心 idea:将原始 dMRI 配准建模为位置与方向联合空间 \(\mathbb{R}^3 \oplus \mathbb{R}^3\) 上的 SE(3)-等变微分同胚,通过隐式满足等变性天然消除显式雅可比重定向,并借助高斯核诱导的傅里叶空间 MMD 损失在无需参数化假设的前提下隐式对齐水分子集成平均传播子(EAP)。

方法详解

整体框架

本文方法直接接收经仿射预对齐与 \(b=0\) 归一化的原始扩散加权衰减信号对(浮动图像 \(E_m\) 与固定图像 \(E_f\))。在物理上,原始 dMRI 的定义域是空间位置 \(p\) 与扩散梯度波矢量 \(q\) 强耦合的流形空间 \(\Omega = \mathbb{R}^3 \times \mathbb{S}^2\)。模型首先利用一个深度为 3 的 SE(3)-等变 UNet 预测联合流形上的等变速度场 \(v(p, g)\);接着,借助流形上的比例与平方(scaling-and-squaring)李代数积分层将速度场积分推进为等变微分同胚形变场 \(\Phi\);随后,专用的球面空间变换器(Spatial Transformer)利用三线性插值与球面高斯径向基函数(RBF)插值完成浮动图像在位置与方向上的联合形变变换;最后,通过傅里叶空间的最大均值差异(MMD)损失与速度场能量正则项指导网络端到端无监督优化。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入对齐原始 dMRI<br/>浮动图像 Em 与固定图像 Ef"] --> B["SE(3)-等变 UNet<br/>SO(5) 可操纵隐式核生成速度场"]
    B --> C["李代数比例与平方积分层<br/>将 6D 速度场推进为等变微分同胚 Φ"]
    C --> D["双域空间变换器<br/>p-空间三线性 + q-空间球面 RBF 插值"]
    D --> E["傅里叶空间 MMD 损失<br/>加权 L2 隐式匹配集成平均传播子 EAP"]

关键设计

1. diff-pq 配准范式与 SE(3)-等变微分同胚:消除显式雅可比重定向

传统标量配准(anat-p)与派生表征配准(diff-p)均将形变映射定义为三维空间域 \(\mathbb{R}^3 \to \mathbb{R}^3\) 上的位移场,这导致 diff-p 范式在形变后必须使用局部雅可比变换 \(d\Phi_p\) 对微观方向做显式重定向。本文提出 diff-pq 配准范式,将原始信号定义在位置空间(p-space)与扩散梯度空间(q-space)通过刚体变换物理耦合的空间 \(\mathbb{R}^3 \oplus \mathbb{R}^3\) 上。对于单壳(single-shell)信号,定义域限制在流形 \(\Omega = \mathbb{R}^3 \times \mathbb{S}^2\)。在群动作 \((t, r) \in \mathrm{SE}(3)\) 下,流形上的变换满足: $$ (t, r) \triangleright (p, q) = (rp + t, rq) $$ 若形变场 \(\Phi: \Omega \to \Omega\) 满足与 \(\mathrm{SE}(3)\) 群动作的可交换性(即等变微分同胚),则空间旋转与位移自然诱导了方向空间的严格同步转动,从数学原理上彻底免除了传统方法中对雅可比矩阵求导与离散有限应变(finite strain)重定向的繁琐过程。依据李群微分拓扑理论,只要网络输出的速度场 \(X\) 为等变向量场,由其生成的积分流即为严格的等变微分同胚。

2. 基于等变 MLP 的 SO(5) 可操纵卷积:跨位置与方向的几何表征学习

由于常规卷积网络无法在五维黎曼流形 \(\Omega\) 上直接保持 \(\mathrm{SE}(3)\) 等变性且标准正则群表示无法输出切向量场,本文将规范等变可操纵卷积(Gauge Equivariant Steerable Convolutions)拓展至 dMRI 领域。流形上的特征图被形式化为张量场,卷积核 \(K: \mathbb{R}^5 \to \mathbb{R}^{d_{\mathrm{out}} \times d_{\mathrm{in}}}\) 需满足 \(\mathrm{SO}(5)\)-可操纵约束: $$ K(t^{-1}v) = \rho_{\mathrm{out}}(t^{-1}) K(v) \rho_{\mathrm{in}}(t), \quad \forall t \in \mathrm{SO}(5) $$ 为避免解析求解高维偏微分约束基函数的复杂性,本文采用隐式神经网络核参数化方案,构建以标准表示 \(\rho_{\mathrm{st}}\) 为输入、张量积表示 \(\rho_{\mathrm{in}} \otimes \rho_{\mathrm{out}}\) 为输出的 \(\mathrm{SO}(5)\)-等变多层感知机(MLP)。通过在初始化阶段利用李代数生成元将群等变约束分解为一组线性系统求解,使网络在单次初始化开销下获得完全连续的等变核,搭建出包含 3 层下采样与上采样的等变 UNet 骨干。

3. 球面李代数比例与平方积分及对偶空间变换器:流形上的光滑形变推进

网络输出的切向量场 \(v(p, g) \in T_{(p,g)}\Omega\) 采用 6 维切向量表示,前 3 维为空间平移速度 \(v_{\mathbb{R}^3}\),后 3 维为嵌于 \(\mathbb{R}^3\) 中的切向量 \(v_{\mathbb{S}^2}\)。流形上的局部指数映射解析形式为: $$ \exp(v(p, g)) = \left( p + v_{\mathbb{R}^3},\; \cos(|v_{\mathbb{S}^2}|) g + \sin(|v_{\mathbb{S}^2}|) \frac{v_{\mathbb{S}^2}}{|v_{\mathbb{S}^2}|} \right) $$ 通过设置迭代步数 \(N=4\) 的 Scaling-and-Squaring 算法,递推计算 \(\Phi^{(j+1)} = \Phi^{(j)} \circ \Phi^{(j)}\),从而高精度逼近全局微分同胚 \(\Phi = \exp(v)\)。在得到形变场后,空间变换器通过三线性插值采样空间位置 \(p\),并在扩散梯度方向 \(g\) 上利用对映对称的球面高斯径向基函数(RBF,核宽度 \(\sigma=0.1\))进行加权角插值: $$ E(p, g) = \frac{\sum_j w_j E(p, g_j)}{\sum_j w_j}, \quad w_j = \exp\left( -\frac{(\arccos(|g^\top g_j|))^2}{2\sigma^2} \right) $$ 保证了离散球面采样下连续形变场的反向传播与对映对称物理守恒。

损失函数 / 训练策略

dMRI 配准的终极临床与解剖目标在于对齐微观水分子集成平均传播子(EAP)概率密度 \(P_p(r)\)。由于归一化衰减信号 \(E(p, q) = S(p, q) / S(p, 0)\) 恰好是 EAP 的傅里叶变换(即特征函数),本文根据再生核希尔伯特空间(RKHS)理论与 Bochner 定理,选择高斯测度 \(\Lambda \sim \mathcal{N}(0, \sigma^2 I_3)\) 作为对偶权重。根据积分性质,EAP 之间的最大均值差异(MMD)在傅里叶对偶下等价于加权 \(L_2\) 范数。最终目标函数定义为:

\[ \mathcal{L}(E_f, E_m \circ \Phi) = C \sum_{p, q} |E_f(p, q) - [E_m \circ \Phi](p, q)|^2 \cdot e^{-\frac{\|q\|^2}{2\sigma^2}} \|q\|^2 + \lambda \sum_{p, q} \|v(p, q)\|^2 \]

其中 \(C = (2\pi\sigma^2)^{-3/2}\) 为归一化常数,\(\lambda\) 为速度场光滑正则化系数。在训练策略上,使用带动量(Nesterov)的随机梯度下降(SGD),初始学习率为 0.001,每 100 个 epoch 衰减减半,总共训练 500 个 epoch,batch size 设为 1。

实验关键数据

主实验

实验基于人类连接组计划(Human Connectome Project, HCP)年轻成年人数据集,选取 400 例高质量去噪配准的 dMRI 数据(349 训练、25 验证、25 测试、1 固定目标)。测试阶段使用 TractSeg 分割出的 72 条全脑白质纤维束真实解剖标注,计算形变后图像与目标图像之间的平均 Dice 相似系数。

方法 (Method) 输入模态 (Modality) 表征类型 (Representation) 纤维束平均 Dice
Ours raw dMRI 免模型原始信号 (Model-Free) 0.7468
mrregister fODF 经典 SyN 派生表征优化 0.7601
VoxelMorph FA 标量特征图 0.7126
DDMReg FA, TOM 标量 + 向量特征图 0.7417
MVCReg fODF 一阶流形卷积网络 0.7317
MVVSReg fODF 二阶流形卷积网络 0.7493

跨配准范式与计算机制消融分析

下表系统对比了主流深度学习配准方法在输入数据依赖、是否依赖派生模型、重定向处理机制以及计算复杂度维度的差异:

配准方法 输入依赖与离散度 显式重定向需求 模型先验依赖 (Model-Free) 72 束白质纤维 Dice
VoxelMorph FA 标量图 否 (忽略方向) 否 (依赖张量拟合) 0.7126
MVCReg 1 阶流形 fODF 需李代数局部变换 否 (依赖约束球面反褶积) 0.7317
DDMReg FA + 72束 TOM 需向量重定向 否 (依赖多阶段提取) 0.7417
Ours 单壳原始 dMRI 否 (等变流场隐式保证) 是 (无模型预设) 0.7468
MVVSReg 2 阶流形 fODF 需高阶张量重定向 否 (计算/显存开销极大) 0.7493

关键发现

  • 在所有无需二阶高阶流形卷积的深度学习框架中,本文直接基于原始 dMRI 数据的配准性能达到最优(Dice 0.7468),超过了利用工程派生特征的 DDMReg(0.7417)以及一阶流形卷积网络 MVCReg(0.7317)。
  • 尽管 MVVSReg 的 Dice 指标略高 0.0025(0.7493 vs 0.7468),但其代价是必须在预处理中运行耗时的多组织约束超分辨球面反褶积(CSD)生成高阶 fODF 流形特征,且网络推理与显存消耗随流形张量阶数急剧爆炸;本文完全省去了所有离线特征提取开销。
  • 棋盘格(Checkerboard)可视化分析证实,SE(3)-等变性能够在剧烈空间扭曲过程中完美维持神经白质束的解剖连续性,消除了基线方法中常见的纤维断裂伪影。

亮点与洞察

  • 范式升维化解重定向难题:创造性地将配准域升维至位置-方向耦合空间 \(\mathbb{R}^3 \oplus \mathbb{R}^3\),利用几何微分拓扑中“等变切向量场积分流必为等变微分同胚”的定理,完全绕过了二十年来困扰 dMRI 配准的显式局部雅可比重定向计算。
  • Bochner 定理与特征函数 MMD 桥接物理传播子:证明了归一化信号衰减与 EAP 之间的傅里叶变换对应关系,通过高斯测度下的对偶 \(L_2\) 范数优雅地隐式最小化 EAP 概率分布的 MMD 距离,理论严密且实现极度轻量。
  • 即插即用的通用免模型优势:无需研究人员在数据管线早期预先选定特定弥散模型(如 DT 或 fODF),输出的形变场可无缝迁移至任意后续下游微观分析。

局限与展望

  • 实验受限于显存限制,将空间分辨率由原始 \(145 \times 174 \times 145\) 下采样至 \(72 \times 88 \times 72\),全分辨率配准能力尚待扩展。
  • 当前网络实现主要针对单壳(single-shell)扩散数据;对于多 b-值(multi-shell)或弥散谱成像(DSI),需要进一步引入沿径向尺度的完整 \(\mathbb{R}^3 \oplus \mathbb{R}^3\) 离散化架构。
  • 未来有望应用于阿尔茨海默病、帕金森等伴随微观白质退行的神经退行性疾病队列分析中。

相关工作与启发

  • vs DDMReg (IEEE TMI 2021):DDMReg 依赖 FA 标量图与 TractSeg 预先抽取的 72 组向量朝向图(TOM),使用标准 CNN 并在后处理中显式旋转向量;本文直接输入原始信号,无任何解剖先验特征工程,端到端实现隐式几何重定向。
  • vs MVCReg / MVVSReg (IPMI 2021/2023):MVCReg 系列必须依赖 fODF 平方根密度流形数据,并采用复杂的流形卷积;本文将方向性从值域抽离至定义域,利用可操纵群等变卷积在更纯粹的原始数据层面上对齐分布。
  • vs mrregister (NeuroImage 2011):基于对称归一化(SyN)的经典迭代优化方法在白质束重合度上依然具有微弱的数值优势(0.7601),但其单例优化耗时长达数十分钟至数小时,而本文深度网络在单次前向传播内数秒即可完成全脑配准。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次提出 model-free 的 diff-pq 等变配准框架与傅里叶 MMD 损失,构思极其精巧。
  • 实验充分度: ⭐⭐⭐⭐ 覆盖 HCP 400 例被试与 72 条白质纤维束评测,但由于显存原因做了空间下采样且缺少更多消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰优美,从微分流形到群表示论逻辑严密自洽。
  • 价值: ⭐⭐⭐⭐⭐ 为免模型 dMRI 配准与几何深度学习在神经影像中的应用确立了全新基准。