MambaMorph: a Mamba-based Framework for Medical MR-CT Deformable Registration

2024-01-25

AI Summary

Existing cross-modal medical image registration methods struggle to balance voxel-level alignment accuracy with clinical deployment efficiency. To address this, we propose MambaReg—a lightweight, end-to-end differentiable registration framework for MR-CT brain image deformable registration—introducing the state space model (Mamba) into medical image registration for the first time. MambaReg decouples long-range dependency modeling from local feature extraction, incorporating a dedicated lightweight encoder and a differentiable deformation field generation module. To support reproducible research, we introduce SR-Reg—the first high-quality, open-source MR-CT brain registration dataset. Extensive experiments demonstrate that MambaReg significantly outperforms state-of-the-art methods on both SR-Reg and the public T1-T2 dataset, achieving substantial improvements in registration accuracy. Moreover, it enables real-time inference (>25 FPS) and reduces GPU memory consumption by over 35%.

Technology Category

• Computer Vision: Multi-modal Vision
• Machine Learning: Multimodal Learning
• Intelligent Robots: Multimodal Perception & Sensor Fusion

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

Abstract

Capturing voxel-wise spatial correspondence across distinct modalities is crucial for medical image analysis. However, current registration approaches are not practical enough in terms of registration accuracy and clinical applicability. In this paper, we introduce MambaMorph, a novel multi-modality deformable registration framework. Specifically, MambaMorph utilizes a Mamba-based registration module and a fine-grained, yet simple, feature extractor for efficient long-range correspondence modeling and high-dimensional feature learning, respectively. Additionally, we develop a well-annotated brain MR-CT registration dataset, SR-Reg, to address the scarcity of data in multi-modality registration. To validate MambaMorph's multi-modality registration capabilities, we conduct quantitative experiments on both our SR-Reg dataset and a public T1-T2 dataset. The experimental results on both datasets demonstrate that MambaMorph significantly outperforms the current state-of-the-art learning-based registration methods in terms of registration accuracy. Further study underscores the efficiency of the Mamba-based registration module and the lightweight feature extractor, which achieve notable registration quality while maintaining reasonable computational costs and speeds. We believe that MambaMorph holds significant potential for practical applications in medical image registration. The code for MambaMorph is available at: https://github.com/Guo-Stone/MambaMorph.

Problem

• Improving accuracy in multi-modality medical image registration
• Addressing data scarcity for MR-CT deformable registration
• Enhancing efficiency in long-range correspondence modeling

Innovation

• Mamba-based module for long-range correspondence
• Lightweight feature extractor for high-dimensional learning
• New annotated MR-CT dataset for multi-modality registration