MambaMIR: An Arbitrary-Masked Mamba for Joint Medical Image Reconstruction and Uncertainty Estimation

Jiahao Huang,Liutao Yang,Fanwen Wang,Yang Nan,Angelica I. Aviles-Rivero,Carola-Bibiane Schönlieb,Daoqiang Zhang,Guang Yang

2024-06-26

Abstract:The recent Mamba model has shown remarkable adaptability for visual representation learning, including in medical imaging tasks. This study introduces MambaMIR, a Mamba-based model for medical image reconstruction, as well as its Generative Adversarial Network-based variant, MambaMIR-GAN. Our proposed MambaMIR inherits several advantages, such as linear complexity, global receptive fields, and dynamic weights, from the original Mamba model. The innovated arbitrary-mask mechanism effectively adapt Mamba to our image reconstruction task, providing randomness for subsequent Monte Carlo-based uncertainty estimation. Experiments conducted on various medical image reconstruction tasks, including fast MRI and SVCT, which cover anatomical regions such as the knee, chest, and abdomen, have demonstrated that MambaMIR and MambaMIR-GAN achieve comparable or superior reconstruction results relative to state-of-the-art methods. Additionally, the estimated uncertainty maps offer further insights into the reliability of the reconstruction quality. The code is publicly available at <a class="link-external link-https" href="https://github.com/ayanglab/MambaMIR" rel="external noopener nofollow">this https URL</a>.

Image and Video Processing,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper aims to address the issues of medical image reconstruction and uncertainty estimation. Specifically, the authors propose a new method based on the Mamba model—MambaMIR, and its Generative Adversarial Network (GAN) variant MambaMIR-GAN, for joint medical image reconstruction and uncertainty estimation. MambaMIR inherits the advantages of the Mamba model, such as linear complexity, global receptive field, and dynamic weights, and introduces a new arbitrary mask mechanism to adapt to image reconstruction tasks, thereby providing randomness for subsequent Monte Carlo-based uncertainty estimation.

Experimental results show that MambaMIR and its variants achieve comparable or better reconstruction results than existing state-of-the-art methods in various medical image reconstruction tasks, such as fast MRI and sparse-view CT (SVCT). Additionally, the estimated uncertainty maps further provide insights into the reliability of the reconstructed image quality. MambaMIR-GAN, due to its adversarial training strategy, outperforms MambaMIR in terms of perceptual quality. Overall, these models represent significant advancements in the field of medical image reconstruction and demonstrate good scalability and computational efficiency potential.