Factorized Gradients for Scalable Highly-expressive Parametric Diffeomorphisms¶
会议: NeurIPS 2026
arXiv: 2610.34766
论文: https://arxiv.org/abs/2610.34766
领域: 优化/理论
关键词: 微分同胚、连续分段仿射、因子化梯度、离散伴随法、可形变配准
一句话总结¶
本文提出基于离散伴随动量累积的 FG-CPAB 微分同胚变换框架，将空间轨迹积分与参数投影彻底解耦，使反向梯度复杂度从 \(\mathcal{O}(dTN)\) 骤降至 \(\mathcal{O}(TN + dC)\)，并在稀疏节点基与分段检查点加速下实现万倍加速与超 500 倍显存降低，使高分辨率细剖分参数化微分同胚首次在 2D/3D 中具备实用扩展性。
研究背景与动机¶
微分同胚（Diffeomorphism）作为具备光滑可微逆映射的双射变换，能够天然保持拓扑结构并防止空间折叠、撕裂与自相交，因而在医学图像配准、几何形变对齐与神经算子学习中扮演着核心角色。经典构造方式（如稳态速度场 SVF 与大变形微分度量映射 LDDMM）通常在稠密空间网格上离散速度场，导致参数量直接绑定于数据分辨率（二维达数十万、三维达数百万），不仅显存开销沉重，而且易因高维病态损失面引发优化困难。
连续分段仿射（CPAB）变换通过将空间域剖分为若干凸多面体单元（如 2D 三角形或 3D 四面体），在单元内部定义仿射速度场并在边界施加连续性约束，使得速度场完全由紧凑的低维参数矢量 \(\boldsymbol{\theta} \in \mathbb{R}^d\) 控制。这种参数化方式使参数维度 \(d\) 仅取决于剖分精细度而非数据点分辨率，天然具备紧凑性与隐式正则化优势。然而，标准 CPAB 实现（如 libcpab）在计算关于 \(\boldsymbol{\theta}\) 的损失梯度时，采用沿正向轨迹逐参数传递敏感性矩阵的前向敏度传播方式，导致反向传播计算复杂度高达 \(\mathcal{O}(dTN)\)（其中 \(T\) 为时间步数，\(N\) 为空间点数）。
更严重的是，参数维度 \(d\) 随剖分加密呈指数级膨胀，并在 GPU 内核启动配置中直接与线程网格维度耦合，导致现有实现一旦网格精细度稍高（如 2D 网格 \(N_{\text{tess}} > 12\) 或 3D 网格 \(N_{\text{tess}} > 5\)）便因超出硬件启动上限或显存崩溃而无法运行，长期制约了 CPAB 的局部表达能力。核心 idea：利用分段仿射速度场的线性基结构重组离散伴随敏感性，把跨时间步的参数依赖项通过望远级数抵消转化为仅依赖局部位姿的几何动量累积，将全参数依赖后延至单次稀疏基投影中，实现轨迹积分与参数空间维度的彻底解耦。
方法详解¶
整体框架¶
FG-CPAB 整体计算流程由前向轨迹积分、反向离散伴随动量累积以及稀疏基梯度投影三大阶段组成。输入包含变换前的一批空间离散点（或图像坐标）与低维控制参数 \(\boldsymbol{\theta}\)，最终输出变换后的坐标及其端到端损失关于 \(\boldsymbol{\theta}\) 的精确梯度。
首先，模型利用节点稀疏基矩阵将参数矢量 \(\boldsymbol{\theta}\) 解码为每个剖分单元的仿射速度矩阵，并在前向传播中通过预计算矩阵指数或显式积分推进轨迹，同时按固定间隔缓存稀疏检查点以削减显存占用；在反向传播中，算法沿轨迹逆向回放局部段落，通过转置中点雅可比矩阵反传伴随矢量 \(\boldsymbol{\lambda}\)，并将点位置与伴随矢量的外积统计量原子累积至无关于参数维度 \(d\) 的单元动量矩阵 \(\mathbf{M}_c\) 中；最终，所有单元的动量矩阵展平后与预计算的稀疏 CPA 基矩阵 \(\mathbf{B}\) 进行单次矩阵乘法，高效恢复全局参数梯度 \(\nabla_{\boldsymbol{\theta}}\mathcal{L}\)。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入坐标点集 p(0) 与参数 θ"] --> B["阶段 1：前向检查点积分<br/>解算单元矩阵 Ac 并稀疏缓存边界状态"]
B --> C["阶段 2：离散伴随动量累积<br/>回放轨迹段落并累积 d 无关矩阵 Mc"]
C --> D["阶段 3：稀疏基梯度恢复<br/>单次投影 B^T vec({Mc}) 计算 ∇_θ L"]
D --> E["输出变换后坐标与参数梯度 ∇_θ L"]关键设计¶
1. 离散伴随敏感性分解与望远级数消解：消除沿时间步传播的参数耦合
传统 CPAB 对中点积分格式进行微分时，每个时间步的敏感性包含两部分：当前步参数引入的直接偏导项，以及依赖上一时间步位置敏感性的递归传播项 \(\mathbf{J} \frac{\partial \boldsymbol{p}^{(t)}}{\partial \theta_i}\)。由于雅可比矩阵 \(\mathbf{J} = \mathbf{I} + h\mathbf{A}_{c,\text{lin}} + \frac{h^2}{2}\mathbf{A}_{c,\text{lin}}^2\) 沿轨迹不断累积，直接计算必须跟踪完整的 \(N \times d\) 敏感性张量。
本文定义时间步 \(t\) 处的损失伴随梯度 \(\boldsymbol{\lambda}^{(t)} = \frac{\partial \mathcal{L}}{\partial \boldsymbol{p}^{(t)}}\)，其满足转置伴随递归方程 \(\boldsymbol{\lambda}^{(t)} = \mathbf{J}^\top \boldsymbol{\lambda}^{(t+1)}\)。通过将该式左乘敏感性递归式，标量投影 \(s^{(t)} = {\boldsymbol{\lambda}^{(t)}}^\top \frac{\partial \boldsymbol{p}^{(t)}}{\partial \theta_i}\) 恰好满足差分递推关系 \(s^{(t+1)} = d^{(t)} + s^{(t)}\)。沿全轨迹 \(t=0\) 至 \(T-1\) 求和时，左侧与右侧的递归传播项构成严格的望远级数并完全相互抵消。考虑到初始输入点独立于参数（即 \(\frac{\partial \boldsymbol{p}^{(0)}}{\partial \theta_i} = \mathbf{0} \implies s^{(0)}=0\)）且终端伴随为损失梯度（\(s^{(T)} = \frac{\partial \mathcal{L}}{\partial \theta_i}\)），全局损失梯度被完全简化为各步局部贡献之和： $\(\frac{\partial \mathcal{L}}{\partial \theta_i} = \sum_{n=1}^N \sum_{t=0}^{T-1} {\boldsymbol{\lambda}_n^{(t+1)}}^\top \left( h \mathbf{B}_{i,c_n^{(t)}} \widetilde{\boldsymbol{p}}_{n,\text{mid}}^{(t)} + \frac{h^2}{2} \mathbf{A}_{c_n^{(t)},\text{lin}} \mathbf{B}_{i,c_n^{(t)}} \widetilde{\boldsymbol{p}}_n^{(t)} \right)\)$ 这一代数恒等式从根本上消除了参数敏感性在时间维度上的跨步连乘，使梯度计算可以按单元独立聚合。
2. 单元几何动量矩阵因子化：将 d 维投影与轨迹计算分离
上述局部贡献表达式仍显式依赖基分量 \(\mathbf{B}_{i,c}\)。作者利用 Frobenius 内积的循环代数性质 \(\boldsymbol{\lambda}^\top \mathbf{A} \mathbf{B} \widetilde{\boldsymbol{p}} = \langle \mathbf{B}, (\mathbf{A}^\top \boldsymbol{\lambda}) \widetilde{\boldsymbol{p}}^\top \rangle_F\)，将与输入轨迹相关的双线性项完全剥离出来，按剖分单元 \(c \in \{1,\dots,C\}\) 汇聚全部落入该单元的时空轨迹对 \(\mathcal{I}_c = \{(n,t) \mid c_n^{(t)}=c\}\)，构造与参数维度 \(d\) 完全解耦的低维单元动量矩阵 \(\mathbf{M}_c \in \mathbb{R}^{K \times (K+1)}\)： $\(\mathbf{M}_c = \sum_{(n,t) \in \mathcal{I}_c} \left[ h \boldsymbol{\lambda}_n^{(t+1)} {\widetilde{\boldsymbol{p}}_{n,\text{mid}}^{(t)\top}} + \frac{h^2}{2} (\mathbf{A}_{c,\text{lin}}^\top \boldsymbol{\lambda}_n^{(t+1)}) {\widetilde{\boldsymbol{p}}_n^{(t)\top}} \right]\)$ 在整个反向积分阶段，GPU 线程仅需维护微小的 \(K \times (K+1)\) 矩阵（在 2D 中为 \(2 \times 3\)，在 3D 中为 \(3 \times 4\)）并原子累加到全局显存中的 \(\mathbf{M}_c\)。当伴随回传结束后，全参数梯度通过一次紧凑的基投影运算统一恢复： $\(\nabla_{\boldsymbol{\theta}}\mathcal{L} = \mathbf{B}^\top \operatorname{vec}(\{\mathbf{M}_c\}_{c=1}^C)\)$ 从而将反向传播计算复杂度从传统方法的 \(\mathcal{O}(dTN)\) 彻底降至 \(\mathcal{O}(TN + dC)\)，计算量和显存不再受参数维度 \(d\) 的反噬。
3. 节点空间参数化与超稀疏 CPA 基构建：兼顾连续性先验与投影效率
传统 CPAB 直接在堆叠的单元仿射系数空间 \(\mathbb{R}^{CK(K+1)}\) 建立连续性线性方程组并通过密集 SVD 求解零空间，所得基矩阵 \(\mathbf{B}\) 为稠密矩阵，在大剖分尺度下存储和乘法极为沉重。FG-CPAB 转而在网格顶点（Nodal）速度矢量空间 \(\boldsymbol{u} \in \mathbb{R}^{Kn_V}\) 中构建基。
在单纯形剖分中，单元内的仿射场由其 \(K+1\) 个顶点的速度唯一线性插值决定，因此单元仿射系数可表述为局部算子 \(\mathbf{N}_c = \mathbf{X}_c^{-\top} \otimes \mathbf{I}_K\)。全局节点映射矩阵 \(\mathbf{N} \in \mathbb{R}^{D \times Kn_V}\) 自然具备极高稀疏性（每行仅与单元自身的 \(K+1\) 个顶点相关）。由于单纯形节点速度天然保障了跨单元边界的面连续性，原系统的连续性边界约束在节点表示下恒等于零，仅需对域边界的零速度边界条件求解极度简化的零空间矩阵 \(\mathbf{Z}\)。由此导出的基矩阵 \(\mathbf{B} = \mathbf{N}\mathbf{Z}\) 保持结构稀疏性，不仅显著减轻了构建与存储开销（3D 细网格显存从 37.6GB 降至 474MB），更使得 \(\mathbf{B}^\top \operatorname{vec}(\{\mathbf{M}_c\})\) 的投影步骤以稀疏矩阵乘法极速完成。
4. 分段检查点机制：压制单线程私有轨迹显存瓶颈
伴随状态与动量矩阵计算依赖正向状态序列 \(\boldsymbol{p}^{(t)}\)。若每个 GPU 线程显式保存全部 \(T\) 个时间步的轨迹点，单点 \(\mathcal{O}(T)\) 的局部存储会导致严重的寄存器溢出并退化为显存换页，显著拉低 GPU 占用率。
FG-CPAB 引入段长为 \(S\)（默认 \(S=4\)）的分段检查点机制，将长度为 \(T\) 的积分划分为 \(K = \lceil T/S \rceil\) 个段，正向积分仅保留段边界处的 \(K+1\) 个检查点，单线程存储从 \(\mathcal{O}(T)\) 压缩至 \(\mathcal{O}(T/S + S)\)。在反向伴随扫描中，按逆序逐段提取起点边界并重放该段局部轨迹，立即就地计算伴随更新并累加动量，计算完成后即刻释放临时状态，以极微小的正向局部重算代价换取显存占用与寄存器压力的急剧缓解。
实验关键数据¶
主实验¶
论文在合成形变配准基准（Blob 合成图像，固定 1000 步优化预算）以及 IXI 脑部 3D MRI 配准基准（115 个测试受试者）上系统对比了 FG-CPAB、稠密稳态速度场（Dense SVF）、传统非刚体变换（LDDMM、SyN、NiftyReg）及现代深度学习配准方法（VoxelMorph、TransMorph、nnFormer 等）。
| 数据集 / 场景 | 方法 | 几何参数量 | 运行/拟合耗时 | 拓扑违例比例 \(\|J\| \le 0\) (%) | 核心对齐指标 |
|---|---|---|---|---|---|
| 2D 合成配准 (\(512^2\)) | Dense SVF (7 步缩放平方) | 524,288 | \(1.356 \pm 0.007\text{ s}\) | \(0.178\%\) | MSE: \(2.651\times 10^{-3}\) |
| 2D 合成配准 (\(512^2\)) | FG-CPAB (15x15 剖分) | 842 | \(0.525 \pm 0.044\text{ s}\) | \(0.028\%\) | MSE: \(8.853\times 10^{-5}\) |
| 3D 合成配准 (\(128^3\)) | Dense SVF (7 步缩放平方) | 6,291,456 | \(12.906 \pm 0.021\text{ s}\) | \(0.021\%\) | MSE: \(1.676\times 10^{-4}\) |
| 3D 合成配准 (\(128^3\)) | FG-CPAB (15x15x15 剖分) | 8,232 | \(6.140 \pm 0.021\text{ s}\) | \(0.012\%\) | MSE: \(7.255\times 10^{-5}\) |
| IXI 3D MRI (\(160 \times 192 \times 224\)) | LDDMM (经典优化) | - | \(\approx 135\text{ s / pair}\) | \(< 0.0001\%\) | Dice: \(0.680 \pm 0.135\) |
| IXI 3D MRI (\(160 \times 192 \times 224\)) | TransMorph (无监督深度学习) | - | 离线预训练 | \(1.579 \pm 0.328\%\) | Dice: \(0.753 \pm 0.123\) |
| IXI 3D MRI (\(160 \times 192 \times 224\)) | FG-CPAB (\(i=1\) 单变换直接优化) | 20,577 | \(\approx 5\text{ s / pair}\) | \(< 0.0001\%\) | Dice: \(0.760 \pm 0.023\) |
| IXI 3D MRI (\(160 \times 192 \times 224\)) | FG-CPAB (\(i=4\) 级联复合变换) | \(4 \times 20,577\) | \(\approx 20\text{ s / pair}\) | \(< 0.0001\%\) | Dice: \(0.765 \pm 0.024\) |
在 2D/3D 单次变换的前向积分与参数求导延迟对比中（Batch size=32），FG-CPAB 在 2D 下正向 1.69ms、反向 8.64ms、峰值显存 197.97MiB，相较于 VoxelMorph SVF 提速 2.4 倍（显存缩减 6.7 倍），相较于 torch-lddmm 提速 21.4 倍（显存缩减 16 倍）；在 3D 下正向 31.10ms、反向 117.50ms、显存 2351.61MiB，相较 VoxelMorph 提速 2.3 倍且显存降低 6.1 倍，较 torch-lddmm 提速 21.2 倍且显存降低 18.4 倍。
消融实验¶
论文针对基稀疏化、积分步长与检查点分段长度进行了细致的消融评测。
1. 稀疏基 vs. 稠密基在大尺寸剖分下的消融（3D 网格 \(128^3\)）
| 剖分尺寸 \(N_{\text{tess}}\) | 稀疏基前向 (ms) | 稀疏基反向 (ms) | 稀疏基显存 (MiB) | 稠密基前向 (ms) | 稠密基反向 (ms) | 稠密基显存 (MiB) |
|---|---|---|---|---|---|---|
| \(N_{\text{tess}}=2\) | 10.8 | 51.9 | 384 | 10.9 | 51.9 | 384 |
| \(N_{\text{tess}}=10\) | 13.8 | 59.6 | 395 | 14.5 | 60.3 | 501 |
| \(N_{\text{tess}}=15\) | 17.2 | 65.1 | 423 | 26.9 | 74.1 | 6,360 |
| \(N_{\text{tess}}=18\) | 20.5 | 75.4 | 450 | 52.0 | 104.8 | 19,674 |
| \(N_{\text{tess}}=20\) | 23.5 | 83.4 | 474 | 85.0 | 140.3 | 37,678 |
2. 检查点分段大小 \(S\) 对反向时延的影响消融（\(T=50\) 步）
| 检查点分段长度 \(S\) | 分段数 \(K\) | 正向延迟 (ms) | 反向延迟 (ms) | 机制说明 |
|---|---|---|---|---|
| \(S=2\) | 25 | 0.898 | 4.161 | 检查点保存频次过高，元数据维护开销大 |
| \(S=4\) (默认设置) | 13 | 0.882 | 3.860 | 显存缓存与寄存器压力达到最优平衡点 |
| \(S=8\) | 7 | 0.893 | 4.644 | 局部重放段变长，临时存储上升 |
| \(S=16\) | 4 | 0.885 | 5.683 | 线程内部局部存储增大，寄存器溢出增多 |
| \(S=50\) (无检查点全存) | 1 | 0.899 | 7.339 | 占用最大单线程私有内存，反向延迟增加近一倍 |
关键发现¶
- 剖分分辨率与性能正相关：伴随 \(N_{\text{tess}}\) 从 2 增加至 30（对应单元数达到 \(5 \times 30^3 = 135,000\) 个），IXI 数据集上的 Dice 得分持续上升，验证了解锁精细剖分对于提升连续分段仿射模型局部形变表达能力的决定性作用。
- 紧凑参数带来隐式优化平滑性：在相同的优化迭代次数下，具有 8,232 参数的 3D FG-CPAB 能够达到 \(7.255\times 10^{-5}\) 的超低残差，而具备 600 余万自由度的 Dense SVF 在无强正则化下陷入较差局部极小；即使对 Dense SVF 施加平滑正则化，后者也需额外耗费 50,000 至 60,000 次迭代才能追平残差，表明紧凑基约束提供了极其优良的低维优化曲面。
- 复合微分同胚突破稳态限制：CPAB 单变换受限于稳态速度场（无法建模交叉轨迹等非稳态运动），通过将多个 CPAB 级联复合（\(i=1 \to 4\)），在严格保持拓扑（雅可比非正率 \(<0.0001\%\)）的前提下将 Dice 从 0.760 进一步推升至 0.765。
亮点与洞察¶
- 精妙的代数级消去：将离散伴随法与线性基展开相结合，借助望远级数自然抵消了跨时间步的参数累积，把繁重的积分轨迹计算与参数空间投影解耦，从根本上重塑了算法的复杂度天花板。
- 软硬件协同设计：不仅在数学上推导出 \(\mathcal{O}(TN + dC)\) 因子化形式，还针对现代 GPU 特性配套研发了节点空间稀疏基、分段检查点与原子累加融合内核，一举攻克了旧版 libcpab 随剖分加密出现内核启动几何超限崩溃的工程顽疾。
- 即插即用的几何算子表征：在 DIMON 反应扩散方程的神经算子学习任务中，直接以拟合所得的 122 维 CPAB 参数作为几何代码，相比 LDDMM 的高维场 PCA 压缩降维方案，将测试相对 \(L^2\) 误差从 0.0828 压低至 0.0528，证实了其作为低维结构几何编码的泛化潜力。
局限与展望¶
- 网格剖分的维度诅咒：虽然算法对空间维度 \(K\) 具有通用数学表达，但凸多面体剖分的单元总数 \(C\) 随空间维度 \(K\) 呈指数增长，使得当前实现严格局限于 \(K \le 3\) 的物理与视觉任务，若要向更高维状态空间扩展，必须引入自适应网格细化（AMR）或张量积可分离基。
- 正反向数值积分器微弱失配：前向传播采用精确单元矩阵指数以提升计算速度，反向伴随求导基于二阶显式中点法，二者在跨越单元分界面时存在局部 \(\mathcal{O}(h)\) 的理论截断偏差（尽管数值有限差分验证其余弦相似度超过 0.998，未影响一阶优化稳定性）。
- 复合变换阶段数的经验选择：级联多个 CPAB 能够有效提升非稳态表达力，但级联阶段数 \(m\) 目前仍需作为离散超参数手动调优，未来可探索自适应判定非稳态残差并动态增生变换级的机制。
相关工作与启发¶
- vs libcpab / 原始 CPAB: 传统 CPAB 采用前向敏感性跟踪，梯度反传复杂度为 \(\mathcal{O}(dTN)\)，内核几何结构直接绑定参数维度 \(d\)；FG-CPAB 借助伴随动量因子化与稀疏投影将复杂度拉低至 \(\mathcal{O}(TN + dC)\)，计算加速达 \(10^4\) 倍，使得细网格优化成为可能。
- vs Dense SVF (VoxelMorph): 稠密速度场参数量绑定图像分辨率，自由度过高且易产生拓扑违例（雅可比异常率通常在 1.5% 左右）；FG-CPAB 基于剖分网格参数化，自由度缩减两个数量级以上，优化过程更加快速稳健，拓扑违例率低于 0.0001%。
- vs LDDMM: LDDMM 基于非稳态测地线流，形变建模极为灵活但逐对优化极其耗时（单对脑部配准需 135 秒）；FG-CPAB 单对优化仅需约 5 秒，且通过级联复合（\(i=4\)）能够逼近时变速度场的形变能力，在保持高效性的同时兼顾复杂拓扑形变。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [首次在连续分段仿射微分同胚中实现轨迹积分与参数维度的完全解耦，代数推导简洁深刻。]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖 1D/2D/3D、微基准时延、合成优化对比、真实医学图像配准以及 PDE 神经算子学习，消融与分析详实。]
- 写作质量: ⭐⭐⭐⭐⭐ [数学推导严密清晰，附录证明详尽完整，对先验缺陷与软硬件瓶颈定位精准。]
- 价值: ⭐⭐⭐⭐⭐ [攻克了微分同胚领域十余年来制约 CPAB 走向高精细实用的核心算力瓶颈，开源工程价值极高。]