MambaDiff: Mamba-Enhanced Diffusion Model for 3D Medical Image Segmentation

• Published in: IEEE transactions on image processing Vol. 34; p. 1
• Main Authors: Liu, Yu, Feng, Yan, Cheng, Juan, Zhan, Haolin, Zhu, Zhiqin
• Format: Journal Article
• Language: English
• Published: United States IEEE 01.01.2025; The Institute of Electrical and Electronics Engineers, Inc. (IEEE)
• Subjects: 3D Medical Image Segmentation; Algorithms; Brain - diagnostic imaging; Cognition & reasoning; Computational modeling; Computer architecture; Convolutional neural networks; Data models; Databases, Factual; Diffusion model; Diffusion models; Enhanced diffusion; Feature extraction; Hippocampus - diagnostic imaging; Humans; Image segmentation; Imaging, Three-Dimensional - methods; Medical diagnostic imaging; Medical imaging; Perception; Semantics; Source code; State Space Models; Three-dimensional displays; Transformers
• ISSN: 1057-7149; 1941-0042; 1941-0042
• DOI: 10.1109/TIP.2025.3607615

Accurate 3D medical image segmentation is crucial for diagnosis and treatment. Diffusion models demonstrate promising performance in medical image segmentation tasks due to the progressive nature of the generation process and the explicit modeling of data distributions. However, the weak guidance of conditional information and insufficient feature extraction in diffusion models lead to the loss of fine-grained features and structural consistency in the segmentation results, thereby affecting the accuracy of medical image segmentation. To address this challenge, we propose a Mamba-Enhanced Diffusion Model for 3D Medical Image Segmentation. We extract multilevel semantic features from the original images using an encoder and tightly integrate them with the denoising process of the diffusion model through a Semantic Hierarchical Embedding (SHE) mechanism, to capture the intricate relationship between the noisy label and image data. Meanwhile, we design a Global-Slice Perception Mamba (GSPM) layer, which integrates multi-dimensional perception mechanisms to endow the model with comprehensive spatial reasoning and feature extraction capabilities. Experimental results show that our proposed MambaDiff achieves more competitive performance compared to prior arts with substantially fewer parameters on four public medical image segmentation datasets including BraTS 2021, BraTS 2024, LiTS and MSD Hippocampus. The source code of our method is available at https://github.com/yuliu316316/MambaDiff.