[Paper Review] 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 Techniques4 citations

TL;DR

VMambaMorph proposes a novel 3D multi-modality deformable image registration framework leveraging a hybrid Visual Mamba-CNN architecture with a 3D-optimized cross-scan module and a recursive fine-tuning strategy. It achieves state-of-the-art registration accuracy on a public MR-CT dataset, outperforming VoxelMorph, TransMorph, and MambaMorph with a Dice score of 82.94±2.01 and improved computational efficiency.

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.

Motivation & Objective

• To address the limitations of existing deep learning-based image registration methods in capturing long-range dependencies and handling complex multi-modality deformations.
• To explore the application of the Visual State Space Model (VMamba) with a cross-scan module in 3D volumetric image registration for improved global feature modeling.
• To design a hybrid VMamba-CNN U-Net architecture tailored for 3D multi-modality registration with enhanced feature extraction and deformation field learning.
• To develop a recursive fine-tuning and inference framework that improves registration accuracy while preserving diffeomorphic properties.
• To validate the proposed method on a public MR-CT registration benchmark and compare it with state-of-the-art methods including VoxelMorph, TransMorph, and MambaMorph.

Proposed method

• The framework employs a U-shaped encoder-decoder architecture integrating 3D VMamba blocks with a redesigned 2D cross-scan module adapted for 3D volumetric feature processing.
• A fine-grained feature extractor is introduced to enhance spatial feature representation, improving the model's ability to capture complex anatomical structures across modalities.
• A recursive training and inference strategy is implemented, where the network iteratively refines the deformation field using predicted displacement fields as feedback.
• The loss function combines a normalized cross-correlation dissimilarity metric and a smoothness regularization term to ensure accurate and diffeomorphic transformations.
• The model is trained end-to-end using a public brain MR-CT registration dataset, with optimization via AdamW and learning rate scheduling.
• The deformation field is predicted as a continuous displacement vector field, ensuring diffeomorphic properties through Jacobian determinant regularization.

Figure 1: The Recursive Registration Framework of VMambaMorph.

Experimental results

Research questions

• RQ1 Can the Visual Mamba architecture with a cross-scan module effectively model long-range dependencies in 3D multi-modality medical image registration?
• RQ2 How does the integration of a 3D-optimized VMamba block with 3D CNNs improve registration accuracy compared to standard CNN or Transformer-based approaches?
• RQ3 To what extent does the proposed recursive registration framework enhance the quality of deformation fields and reduce registration errors?
• RQ4 Does the inclusion of a fine-grained feature extractor significantly improve the model’s ability to handle structural and intensity variations across MR and CT images?
• RQ5 How does VMambaMorph compare in performance and efficiency to state-of-the-art methods such as VoxelMorph, TransMorph, and MambaMorph on a standard benchmark?

Key findings

• VMambaMorph achieves a Dice score of 82.94±2.01 on the public MR-CT registration benchmark, outperforming VoxelMorph (62.42±3.29), TransMorph (78.46±2.13), and MambaMorph (82.49±1.99).
• The method reduces the 95% Hausdorff distance to 1.35±0.18 mm, indicating superior spatial alignment accuracy compared to baseline models.
• The network maintains a low percentage of non-positive Jacobian determinants (1.04±0.05%), confirming the preservation of diffeomorphic properties in the deformation field.
• The recursive inference strategy with K=2 achieves the best balance between accuracy and inference time, with a runtime of 0.19 seconds and 3.93 GB GPU memory usage.
• The ablation study confirms that the fine-grained feature extractor improves Dice scores by up to 4.5% compared to the baseline without it.
• Training loss stabilizes faster and reaches lower values in VMambaMorph compared to VoxelMorph, TransMorph, and MambaMorph, indicating better convergence and generalization.

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