[论文解读] VMambaMorph: a Multi-Modality Deformable Image Registration Framework based on Visual State Space Model with Cross-Scan Module

Ziyang Wang, Jian-Qing Zheng+2 more|arXiv (Cornell University)|Apr 7, 2024

Advanced Image and Video Retrieval Techniques被引用 4

一句话总结

VMambaMorph 提出了一种新颖的 3D 多模态可变形图像配准框架，结合了混合视觉 Mamba-CNN 架构、专为 3D 优化的跨扫描模块以及递归微调策略。该方法在公开的 MR-CT 数据集上实现了最先进的配准精度，Dice 分数达到 82.94±2.01，同时计算效率更高。

ABSTRACT

Image registration, a critical process in medical imaging, involves aligning different sets of medical imaging data into a single unified coordinate system. Deep learning networks, such as the Convolutional Neural Network (CNN)-based VoxelMorph, Vision Transformer (ViT)-based TransMorph, and State Space Model (SSM)-based MambaMorph, have demonstrated effective performance in this domain. The recent Visual State Space Model (VMamba), which incorporates a cross-scan module with SSM, has exhibited promising improvements in modeling global-range dependencies with efficient computational cost in computer vision tasks. This paper hereby introduces an exploration of VMamba with image registration, named VMambaMorph. This novel hybrid VMamba-CNN network is designed specifically for 3D image registration. Utilizing a U-shaped network architecture, VMambaMorph computes the deformation field based on target and source volumes. The VMamba-based block with 2D cross-scan module is redesigned for 3D volumetric feature processing. To overcome the complex motion and structure on multi-modality images, we further propose a fine-tune recursive registration framework. We validate VMambaMorph using a public benchmark brain MR-CT registration dataset, comparing its performance against current state-of-the-art methods. The results indicate that VMambaMorph achieves competitive registration quality. The code for VMambaMorph with all baseline methods is available on GitHub.

研究动机与目标

• 解决现有基于深度学习的图像配准方法在捕捉长距离依赖关系以及处理复杂多模态形变方面的局限性。
• 探索视觉状态空间模型（VMamba）结合跨扫描模块在 3D 体积图像配准中的应用，以提升全局特征建模能力。
• 设计一种专为 3D 多模态配准优化的混合 VMamba-CNN U-Net 架构，增强特征提取与形变场学习能力。
• 开发一种递归微调与推理框架，提升配准精度的同时保持微分同胚性质。
• 在公开的 MR-CT 配准基准上验证所提方法，并与 SOTA 方法（包括 VoxelMorph、TransMorph 和 MambaMorph）进行比较。

提出的方法

• 该框架采用 U 形编码器-解码器架构，集成 3D VMamba 模块与重新设计的 2D 跨扫描模块，专为 3D 体积特征处理而适配。
• 引入细粒度特征提取器以增强空间特征表示，提升模型捕捉跨模态复杂解剖结构的能力。
• 实施递归训练与推理策略，网络通过将预测的位移场作为反馈，迭代优化形变场。
• 损失函数结合归一化互相关差异度量与平滑正则化项，确保变换的准确性和微分同胚性。
• 使用公开的脑部 MR-CT 配准数据集进行端到端训练，采用 AdamW 优化并配合学习率调度。
• 形变场被预测为连续的位移向量场，通过雅可比行列式正则化确保其具备微分同胚性质。

Figure 1: The Recursive Registration Framework of VMambaMorph.

实验结果

研究问题

• RQ1 带有跨扫描模块的视觉 Mamba 架构是否能有效建模 3D 多模态医学图像配准中的长距离依赖关系？
• RQ2 将专为 3D 优化的 VMamba 模块与 3D CNN 结合，相较于标准 CNN 或 Transformer 方法，能否显著提升配准精度？
• RQ3 所提出的递归配准框架在多大程度上提升了形变场质量并降低了配准误差？
• RQ4 细粒度特征提取器的引入是否显著增强了模型处理 MR 与 CT 图像间结构与强度差异的能力？
• RQ5 在标准基准上，VMambaMorph 与 VoxelMorph、TransMorph 和 MambaMorph 等 SOTA 方法相比，在性能与效率方面表现如何？

主要发现

• VMambaMorph 在公开的 MR-CT 配准基准上取得 82.94±2.01 的 Dice 分数，优于 VoxelMorph（62.42±3.29）、TransMorph（78.46±2.13）和 MambaMorph（82.49±1.99）。
• 该方法将 95% 的豪斯多夫距离降低至 1.35±0.18 mm，表明其在空间对齐精度方面优于基线模型。
• 网络保持了极低比例的非正雅可比行列式（1.04±0.05%），证实了形变场中微分同胚性质的有效保持。
• 采用 K=2 的递归推理策略在精度与推理时间之间实现了最佳平衡，运行时间为 0.19 秒，GPU 显存占用 3.93 GB。
• 消融实验证实，与无该模块的基线相比，细粒度特征提取器可使 Dice 分数提升最高达 4.5%。
• 与 VoxelMorph、TransMorph 和 MambaMorph 相比，VMambaMorph 的训练损失收敛更快且达到更低值，表明其具有更优的收敛性与泛化能力。

Figure 2: The Architecture of VMambaMorph, and the Details of Visual State Space Block.