RegMamba: An Improved Mamba for Medical Image Registration
Abstract
:1. Introduction
- A novel registration network, RegMamba, is proposed. By combining the global feature capturing ability of the Mamba module with the local feature extraction advantage of CNN, it not only captures the subtle local changes in the image, but also provide global relationship modeling, thus achieving more efficient and accurate image registration.
- We insert a scaling convolution module into the Mamba module to adaptively scale the local features to alleviate the spatial information that may be lacking in Mamba’s processing of three-bit data.
- The deformable convolutional residual module (DCRM) is able to adaptively adjust the position and shape of the convolution kernel according to the characteristics of the input data to extract clearer and more discriminative local variation features.
2. Related Work
2.1. Prior Works on Medical Image Registration
2.2. Mamba in Vision Tasks
2.3. Diffeomorphic Registration
3. Method
3.1. Network Overview
3.2. Scaled Convolutional Mamba Block
3.3. Deformable Convolutional Residual Module
3.4. Loss Function
4. Experiments
4.1. Data and Pre-Processing
4.2. Evaluation Metrics
4.3. Baseline Methods
4.4. Experimental Settings
4.5. Results
Ablation Study
5. Conclusions
Author Contributions
Funding
Informed Consent Statement
Data Availability Statement
Conflicts of Interest
Appendix A
References
- Avants, B.; Epstein, C.; Grossman, M.; Gee, J. Symmetric diffeomorphic image registration with cross-correlation: Evaluating automated labeling of elderly and neurodegenerative brain. Med. Image Anal. 2008, 12, 26–41. [Google Scholar] [CrossRef] [PubMed]
- Beg, M.F.; Miller, M.I.; Trouvé, A.; Younes, L. Computing Large Deformation Metric Mappings via Geodesic Flows of Diffeomorphisms. Int. J. Comput. Vis. 2005, 61, 139–157. [Google Scholar] [CrossRef]
- Heinrich, M.P.; Maier, O.; Handels, H. Multi-modal Multi-Atlas Segmentation using Discrete Optimisation and Self-Similarities. In Proceedings of the VISCERAL Challenge@ISBI, Brooklyn, NY, USA, 16–19 April 2015. [Google Scholar]
- Modat, M.; Ridgway, G.R.; Taylor, Z.A.; Lehmann, M.; Barnes, J.; Hawkes, D.J.; Fox, N.C.; Ourselin, S. Fast free-form deformation using graphics processing units. Comput. Methods Programs Biomed. 2010, 98, 278–284. [Google Scholar] [CrossRef] [PubMed]
- Balakrishnan, G.; Zhao, A.; Sabuncu, M.R.; Guttag, J.; Dalca, A.V. VoxelMorph: A Learning Framework for Deformable Medical Image Registration. IEEE Trans. Med. Imaging 2019, 38, 1788–1800. [Google Scholar] [CrossRef] [PubMed]
- Chen, J.; Frey, E.C.; He, Y.; Segars, W.P.; Li, Y.; Du, Y. TransMorph: Transformer for unsupervised medical image registration. Med. Image Anal. 2022, 82, 102615. [Google Scholar] [CrossRef] [PubMed]
- Jia, X.; Bartlett, J.; Zhang, T.; Lu, W.; Qiu, Z.; Duan, J. U-Net vs Transformer: Is U-Net Outdated in Medical Image Registration? In Proceedings of the Machine Learning in Medical Imaging, Singapore, 18 September 2022; Lian, C., Cao, X., Rekik, I., Xu, X., Cui, Z., Eds.; Springer Nature: Cham, Switzerland, 2022; pp. 151–160. [Google Scholar]
- Kim, B.; Kim, D.H.; Park, S.H.; Kim, J.; Lee, J.G.; Ye, J.C. CycleMorph: Cycle consistent unsupervised deformable image registration. Med. Image Anal. 2021, 71, 102036. [Google Scholar] [CrossRef]
- Shi, J.; He, Y.; Kong, Y.; Coatrieux, J.L.; Shu, H.; Yang, G.; Li, S. XMorpher: Full Transformer for Deformable Medical Image Registration via Cross Attention. In Proceedings of the Medical Image Computing and Computer Assisted Intervention—MICCAI 2022, Singapore, 18–22 September 2022; Wang, L., Dou, Q., Fletcher, P.T., Speidel, S., Li, S., Eds.; Springer Nature: Cham, Switzerland, 2022; pp. 217–226. [Google Scholar]
- Sokooti, H.; Vos, B.D.; Berendsen, F.; Lelieveldt, B.P.F.; Išgum, I.; Staring, M. Nonrigid Image Registration Using Multi-Scale 3D Convolutional Neural Networks; Springer: Cham, Switzerland, 2017; pp. 232–239. [Google Scholar]
- Xie, Y.; Zhang, J.; Shen, C.; Xia, Y. CoTr: Efficiently Bridging CNN and Transformer for 3D Medical Image Segmentation. In Proceedings of the Medical Image Computing and Computer Assisted Intervention—MICCAI, Strasbourg, France, 27 September–2 October 2021; de Bruijne, M., Cattin, P.C., Cotin, S., Padoy, N., Speidel, S., Zheng, Y., Essert, C., Eds.; Springer Nature: Cham, Switzerland, 2021; pp. 171–180. [Google Scholar]
- Yang, X.; Kwitt, R.; Styner, M.; Niethammer, M. Quicksilver: Fast predictive image registration—A deep learning approach. NeuroImage 2017, 158, 378–396. [Google Scholar] [CrossRef] [PubMed]
- Zhao, S.; Lau, T.; Luo, J.; Chang, E.I.C.; Xu, Y. Unsupervised 3D End-to-End Medical Image Registration with Volume Tweening Network. IEEE J. Biomed. Health Inform. 2020, 24, 1394–1404. [Google Scholar] [CrossRef] [PubMed]
- de Vos, B.D.; Berendsen, F.F.; Viergever, M.A.; Staring, M.; Išgum, I. End-to-End Unsupervised Deformable Image Registration with a Convolutional Neural Network. In Proceedings of the Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support, Quebec City, QC, Canada, 14 September 2017; Cardoso, M.J., Arbel, T., Carneiro, G., Syeda-Mahmood, T., Tavares, J.M.R., Moradi, M., Bradley, A., Greenspan, H., Papa, J.P., Madabhushi, A., et al., Eds.; Springer Nature: Cham, Switzerland, 2017; pp. 204–212. [Google Scholar]
- Li, W.; Zhou, G.; Lin, S.; Tang, Y. PerNet: Progressive and Efficient All-in-One Image-Restoration Lightweight Network. Electronics 2024, 13, 2817. [Google Scholar] [CrossRef]
- Jiao, C.; Yang, T.; Yan, Y.; Yang, A. RFTNet: Region–Attention Fusion Network Combined with Dual-Branch Vision Transformer for Multimodal Brain Tumor Image Segmentation. Electronics 2024, 13, 77. [Google Scholar] [CrossRef]
- Baek, J.H.; Lee, H.K.; Choo, H.G.; Jung, S.h.; Koh, Y.J. Center-Guided Transformer for Panoptic Segmentation. Electronics 2023, 12, 4801. [Google Scholar] [CrossRef]
- Zhang, Y.; Pei, Y.; Zha, H. Learning Dual Transformer Network for Diffeomorphic Registration. In Proceedings of the Medical Image Computing and Computer Assisted Intervention—MICCAI 2021—24th International Conference, Strasbourg, France, 27 September–1 October 2021; de Bruijne, M., Cattin, P.C., Cotin, S., Padoy, N., Speidel, S., Zheng, Y., Essert, C., Eds.; Springer: Berlin/Heidelberg, Germany, 2021; Volume 12904, (Lecture Notes in Computer Science). pp. 129–138. [Google Scholar] [CrossRef]
- Chen, J.; He, Y.; Frey, E.C.; Li, Y.; Du, Y. ViT-V-Net: Vision Transformer for Unsupervised Volumetric Medical Image Registration. arXiv 2021, arXiv:2104.06468. [Google Scholar]
- Gu, A.; Dao, T. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv 2024, arXiv:2312.00752. [Google Scholar]
- Galić, I.; Habijan, M.; Leventić, H.; Romić, K. Machine Learning Empowering Personalized Medicine: A Comprehensive Review of Medical Image Analysis Methods. Electronics 2023, 12, 4411. [Google Scholar] [CrossRef]
- Rueckert, D.; Sonoda, L.; Hayes, C.; Hill, D.; Leach, M.; Hawkes, D. Nonrigid registration using free-form deformations: Application to breast MR images. IEEE Trans. Med. Imaging 1999, 18, 712–721. [Google Scholar] [CrossRef]
- Johnson, H.J.; Christensen, G.E. Landmark and Intensity-Based, Consistent Thin-Plate Spline Image Registration. In Proceedings of the Information Processing in Medical Imaging, Davis, CA, USA, 18–22 June 2001; Insana, M.F., Leahy, R.M., Eds.; Springer: Berlin/Heidelberg, Germany, 2001; pp. 329–343. [Google Scholar]
- Ashburner, J. A fast diffeomorphic image registration algorithm. NeuroImage 2007, 38, 95–113. [Google Scholar] [CrossRef] [PubMed]
- Vercauteren, T.; Pennec, X.; Perchant, A.; Ayache, N. Diffeomorphic demons: Efficient non-parametric image registration. NeuroImage 2009, 45, S61–S72. [Google Scholar] [CrossRef]
- Zhao, S.; Dong, Y.; Chang, E.; Xu, Y. Recursive Cascaded Networks for Unsupervised Medical Image Registration. In Proceedings of the 2019 IEEE/CVF International Conference on Computer Vision (ICCV), Seoul, Republic of Korea, 27 October–2 November 2019. [Google Scholar] [CrossRef]
- Chen, J.; Lu, D.; Zhang, Y.; Wei, D.; Ning, M.; Shi, X.; Xu, Z.; Zheng, Y. Deformer: Towards Displacement Field Learning for Unsupervised Medical Image Registration. In Proceedings of the Medical Image Computing and Computer Assisted Intervention—MICCAI, Singapore, 18–22 September 2022; Wang, L., Dou, Q., Fletcher, P.T., Speidel, S., Li, S., Eds.; Springer Nature: Cham, Switzerland, 2022; pp. 141–151. [Google Scholar]
- Li, Y.X.; Tang, H.; Wang, W.; Zhang, X.F.; Qu, H. Dual attention network for unsupervised medical image registration based on VoxelMorph. Sci. Rep. 2022, 12, 16250. [Google Scholar] [CrossRef]
- Zhu, L.; Liao, B.; Zhang, Q.; Wang, X.; Liu, W.; Wang, X. Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model. arXiv 2024, arXiv:2401.09417. [Google Scholar]
- Vaswani, A.; Shazeer, N.; Parmar, N.; Uszkoreit, J.; Jones, L.; Gomez, A.N.; Kaiser, L.; Polosukhin, I. Attention Is All You Need. arXiv 2023, arXiv:1706.03762. [Google Scholar]
- Yang, C.; Chen, Z.; Espinosa, M.; Ericsson, L.; Wang, Z.; Liu, J.; Crowley, E.J. PlainMamba: Improving Non-Hierarchical Mamba in Visual Recognition. arXiv 2024, arXiv:2403.17695. [Google Scholar]
- Huang, T.; Pei, X.; You, S.; Wang, F.; Qian, C.; Xu, C. LocalMamba: Visual State Space Model with Windowed Selective Scan. arXiv 2024, arXiv:2403.09338. [Google Scholar]
- Pei, X.; Huang, T.; Xu, C. EfficientVMamba: Atrous Selective Scan for Light Weight Visual Mamba. arXiv 2024, arXiv:2403.09977. [Google Scholar]
- Patro, B.N.; Agneeswaran, V.S. SiMBA: Simplified Mamba-Based Architecture for Vision and Multivariate Time series. arXiv 2024, arXiv:2403.15360. [Google Scholar]
- Wang, Z.; Zheng, J.Q.; Zhang, Y.; Cui, G.; Li, L. Mamba-UNet: UNet-Like Pure Visual Mamba for Medical Image Segmentation. arXiv 2024, arXiv:2402.05079. [Google Scholar]
- Ma, J.; Li, F.; Wang, B. U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation. arXiv 2024, arXiv:2401.04722. [Google Scholar]
- Liao, W.; Zhu, Y.; Wang, X.; Pan, C.; Wang, Y.; Ma, L. LightM-UNet: Mamba Assists in Lightweight UNet for Medical Image Segmentation. arXiv 2024, arXiv:2403.05246. [Google Scholar]
- Xing, Z.; Ye, T.; Yang, Y.; Liu, G.; Zhu, L. SegMamba: Long-range Sequential Modeling Mamba For 3D Medical Image Segmentation. arXiv 2024, arXiv:2401.13560. [Google Scholar]
- Ruan, J.; Xiang, S. VM-UNet: Vision Mamba UNet for Medical Image Segmentation. arXiv 2024, arXiv:2402.02491. [Google Scholar]
- Yue, Y.; Li, Z. MedMamba: Vision Mamba for Medical Image Classification. arXiv 2024, arXiv:2403.03849. [Google Scholar]
- Guo, T.; Wang, Y.; Shu, S.; Chen, D.; Tang, Z.; Meng, C.; Bai, X. MambaMorph: A Mamba-based Framework for Medical MR-CT Deformable Registration. arXiv 2024, arXiv:2401.13934. [Google Scholar]
- Arsigny, V.; Commowick, O.; Pennec, X.; Ayache, N. A Log-Euclidean Framework for Statistics on Diffeomorphisms. In Medical Image Computing and Computer-Assisted Intervention: MICCAI, Proceedings of the International Conference on Medical Image Computing and Computer-Assisted Intervention, Copenhagen, Denmark, 1–8 October 2006; Springer: Berlin/Heidelberg, Germany, 2006; pp. 924–931. [Google Scholar]
- Dalca, A.V.; Balakrishnan, G.; Guttag, J.; Sabuncu, M.R. Unsupervised Learning for Fast Probabilistic Diffeomorphic Registration. In Lecture Notes in Computer Science; Springer International Publishing: Berlin/Heidelberg, Germany, 2018; pp. 729–738. [Google Scholar] [CrossRef]
- Jaderberg, M.; Simonyan, K.; Zisserman, A.; Kavukcuoglu, K. Spatial transformer networks. In Proceedings of the 28th International Conference on Neural Information Processing Systems, Cambridge, MA, USA, 7–12 December 2015; Volume 2, pp. 2017–2025. [Google Scholar]
- Gu, A.; Goel, K.; Ré, C. Efficiently Modeling Long Sequences with Structured State Spaces. arXiv 2022, arXiv:2111.00396. [Google Scholar]
- Ramachandran, P.; Zoph, B.; Le, Q.V. Searching for Activation Functions. arXiv 2017, arXiv:1710.05941. [Google Scholar]
- Ba, J.L.; Kiros, J.R.; Hinton, G.E. Layer Normalization. arXiv 2016, arXiv:1607.06450. [Google Scholar]
- Dai, J.; Qi, H.; Xiong, Y.; Li, Y.; Zhang, G.; Hu, H.; Wei, Y. Deformable Convolutional Networks. In Proceedings of the 2017 IEEE International Conference on Computer Vision (ICCV), Venice, Italy, 22–29 October 2017; pp. 764–773. [Google Scholar] [CrossRef]
- Yu, F.; Koltun, V. Multi-Scale Context Aggregation by Dilated Convolutions. arXiv 2016, arXiv:1511.07122. [Google Scholar]
- Shattuck, D.W.; Mirza, M.; Adisetiyo, V.; Hojatkashani, C.; Salamon, G.; Narr, K.L.; Poldrack, R.A.; Bilder, R.M.; Toga, A.W. Construction of a 3D probabilistic atlas of human cortical structures. NeuroImage 2008, 39, 1064–1080. [Google Scholar] [CrossRef] [PubMed]
- Kuang, D.; Schmah, T. FAIM—A ConvNet Method for Unsupervised 3D Medical Image Registration. In Proceedings of the Machine Learning in Medical Imaging, Shenzhen, China, 13 October 2019; Suk, H.I., Liu, M., Yan, P., Lian, C., Eds.; Springer Nature: Cham, Switzerland, 2019; pp. 646–654. [Google Scholar]
- Dice, L.R. Measures of the Amount of Ecologic Association Between Species. Ecology 1945, 26, 297–302. [Google Scholar] [CrossRef]
- Qiu, H.; Qin, C.; Schuh, A.; Hammernik, K.; Rueckert, D. Learning Diffeomorphic and Modality-invariant Registration using B-splines. In Proceedings of the Medical Imaging with Deep Learning, Virtual Event, 19–21 July 2021. [Google Scholar]
- Wang, W.; Xie, E.; Li, X.; Fan, D.P.; Song, K.; Liang, D.; Lu, T.; Luo, P.; Shao, L. PVT v2: Improved baselines with Pyramid Vision Transformer. Comput. Vis. Media 2022, 8, 415–424. [Google Scholar] [CrossRef]
- Zhou, H.; Guo, J.; Zhang, Y.; Yu, L.; Wang, L.; Yu, Y. nnFormer: Interleaved Transformer for Volumetric Segmentation. arXiv 2021, arXiv:2109.03201. [Google Scholar]
- Yu, Z.; Chen, L.; Cheng, Z.; Luo, J. TransMatch: A Transfer-Learning Scheme for Semi-Supervised Few-Shot Learning. arXiv 2020, arXiv:1912.09033. [Google Scholar]
Model | Dice | Time(s) | Parameters | Mult-Adds(G) | |
---|---|---|---|---|---|
Affine | 0.386 ± 0.195 | - | - | - | - |
SyN [1] | 0.639 ± 0.151 | <0.0001 | - | - | - |
NiftyReg [4] | 0.640 ± 0.166 | <0.0001 | - | - | - |
LDDMM [2] | 0.675 ± 0.135 | <0.0001 | - | - | - |
deedsBCV [3] | 0.733 ± 0.126 | 0.147 ± 0.050 | - | - | - |
VoxelMorph-1 [5] | 0.723 ± 0.130 | 1.590 ± 0.339 | 0.114 | 274,387 | 304.05 |
VoxelMorph-2 [5] | 0.726 ± 0.123 | 1.522 ± 0.336 | 0.131 | 301,411 | 398.81 |
VoxelMorph-diff [43] | 0.577 ± 0.165 | <0.0001 | 0.042 | 307,878 | 89.67 |
CycleMorph [8] | 0.730 ± 0.124 | 1.719 ± 0.382 | 0.143 | 361,299 | 49.42 |
MIDIR [53] | 0.736 ± 0.129 | <0.0001 | 0.061 | 266,387 | 47.05 |
ViT-V-Net [19] | 0.728 ± 0.124 | 1.609 ± 0.319 | 0.139 | 9,815,431 | 10.60 |
CoTr [11] | 0.721 ± 0.128 | 1.858 ± 0.314 | 0.478 | 38,684,995 | 1461.61 |
PVT [54] | 0.729 ± 0.135 | 1.292 ± 0.342 | 0.149 | 58,749,007 | 193.61 |
nnFormer [55] | 0.740 ± 0.134 | 1.595 ± 0.358 | 0.916 | 34,415,851 | 686.77 |
TransMorph [6] | 0.746 ± 0.128 | 1.579 ± 0.328 | 0.249 | 46,771,251 | 657.64 |
TransMorph-Bayes [6] | 0.753 ± 0.123 | 1.560 ± 0.333 | 5.942 | 21,205,491 | 657.69 |
TransMorph-diff [6] | 0.594 ± 0.163 | <0.0001 | 0.252 | 46,557,414 | 252.61 |
RegMamba | 0.7665 ± 0.124 | 0.312 ± 0.126 | 0.185 | 2,983,742 | 329.67 |
RegMamba-diff | 0.7627 ± 0.0.123 | <0.0001 | 0.208 | 2,983,742 | 329.67 |
Model | LPBA40 | ||
---|---|---|---|
Dice ↑ | HD95 ↓ | %↓ | |
Initial | 0.686 ± 0.045 | 6.419 ± 1.014 | - |
SyN [1] | 0.821 ± 0.015 | 4.655 ± 0.592 | <0.0001 |
NiftyReg [1] | 0.813 ± 0.012 | 4.589 ± 0.652 | <0.0001 |
VoxelMorph [5] | 0.808 ± 0.017 | 4.804 ± 0.604 | 0.632 ± 0.222 |
CycleMorph [8] | 0.822 ± 0.016 | 4.644 ± 0.623 | 0.030 ± 0.028 |
ViT-V-Net [19] | 0.817 ± 0.018 | 4.851 ± 0.633 | 0.588 ± 0.199 |
TransMorph [6] | 0.825 ± 0.016 | 4.688 ± 0.623 | 0.479 ± 0.143 |
TransMorph-diff [6] | 0.798 ± 0.017 | 5.549 ± 0.582 | <0.0001 |
TransMatch [56] | 0.819 ± 0.018 | 4.576 ± 0.565 | 0.064 ± 0.028 |
RegMamba | 0.839 ± 0.014 | 4.324 ± 0.505 | 0.055 ± 0.022 |
Method | Model | C | DCRM | Mamba | Scale Conv | Dice↑ |
---|---|---|---|---|---|---|
A0 | U-Net | 8 | - | - | - | 76.16 (4.08) |
A1 | RegMamba | 8 | Y | Y | N | 76.92 (3.84) |
A2 | RegMamba | 8 | Y | N | N | 76.99 (3.82) |
A3 | RegMamba | 8 | N | Y | Y | 76.95 (3.87) |
A4 | RegMamba | 8 | Y | Y | Y | 77.06 (3.73) |
B1 | RegMamba | 16 | Y | Y | Y | 77.35 (3.71) |
B2 | RegMamba | 32 | Y | Y | Y | 77.46 (3.70) |
Method | Model | C | DCRM | SCMB | Dice↑ |
---|---|---|---|---|---|
A0 | U-Net | 8 | - | - | 72.6 (1.23) |
A1 | RegMamba | 8 | Y | Y | 76.7 (1.24) |
A2 | RegMamba | 8 | Y | N | 75.3 (1.23) |
A3 | RegMamba | 8 | N | Y | 76.5 (1.24) |
Disclaimer/Publisher’s Note: The statements, opinions and data contained in all publications are solely those of the individual author(s) and contributor(s) and not of MDPI and/or the editor(s). MDPI and/or the editor(s) disclaim responsibility for any injury to people or property resulting from any ideas, methods, instructions or products referred to in the content. |
© 2024 by the authors. Licensee MDPI, Basel, Switzerland. This article is an open access article distributed under the terms and conditions of the Creative Commons Attribution (CC BY) license (https://creativecommons.org/licenses/by/4.0/).
Share and Cite
Hu, X.; Chen, J.; Chen, Y. RegMamba: An Improved Mamba for Medical Image Registration. Electronics 2024, 13, 3305. https://doi.org/10.3390/electronics13163305
Hu X, Chen J, Chen Y. RegMamba: An Improved Mamba for Medical Image Registration. Electronics. 2024; 13(16):3305. https://doi.org/10.3390/electronics13163305
Chicago/Turabian StyleHu, Xin, Jiaqi Chen, and Yilin Chen. 2024. "RegMamba: An Improved Mamba for Medical Image Registration" Electronics 13, no. 16: 3305. https://doi.org/10.3390/electronics13163305