MambaClinix: Hierarchical Gated Convolution and Mamba-Based U-Net for Enhanced 3D Medical Image Segmentation

Chenyuan Bian,Nan Xia,Xia Yang,Feifei Wang,Fengjiao Wang,Bin Wei,Qian Dong

2024-09-19

Abstract:Deep learning, particularly convolutional neural networks (CNNs) and Transformers, has significantly advanced 3D medical image segmentation. While CNNs are highly effective at capturing local features, their limited receptive fields may hinder performance in complex clinical scenarios. In contrast, Transformers excel at modeling long-range dependencies but are computationally intensive, making them expensive to train and deploy. Recently, the Mamba architecture, based on the State Space Model (SSM), has been proposed to efficiently model long-range dependencies while maintaining linear computational complexity. However, its application in medical image segmentation reveals shortcomings, particularly in capturing critical local features essential for accurate delineation of clinical regions. In this study, we propose MambaClinix, a novel U-shaped architecture for medical image segmentation that integrates a hierarchical gated convolutional network(HGCN) with Mamba in an adaptive stage-wise framework. This design significantly enhances computational efficiency and high-order spatial interactions, enabling the model to effectively capture both proximal and distal relationships in medical images. Specifically, our HGCN is designed to mimic the attention mechanism of Transformers by a purely convolutional structure, facilitating high-order spatial interactions in feature maps while avoiding the computational complexity typically associated with Transformer-based methods. Additionally, we introduce a region-specific Tversky loss, which emphasizes specific pixel regions to improve auto-segmentation performance, thereby optimizing the model's decision-making process. Experimental results on five benchmark datasets demonstrate that the proposed MambaClinix achieves high segmentation accuracy while maintaining low model complexity.

Image and Video Processing,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper aims to address key challenges in 3D medical image segmentation, particularly in balancing the capture of local features and global dependencies. Specifically:

1. **Balancing Local Features and Global Dependencies**: Convolutional Neural Networks (CNNs) excel at capturing local features but are limited by a finite receptive field in complex clinical scenarios. In contrast, Transformers are effective in modeling long-range dependencies but come with high computational costs, making efficient deployment challenging.

2. **Limitations of Existing Methods**: The Mamba architecture, based on State Space Models (SSM), can efficiently model long-range dependencies but has shown deficiencies in capturing critical local features in medical image segmentation.

To address these issues, the paper proposes MambaClinix, a novel U-shaped architecture that combines Hierarchical Gated Convolutional Networks (HGCN) and Mamba modules, working within an adaptive stage framework. This design improves computational efficiency and enhances high-order spatial interaction capabilities, enabling the model to effectively capture proximal and distal relationships in medical images. Additionally, a region-specific Tversky loss function is introduced to optimize automatic segmentation performance. Experimental results demonstrate that MambaClinix achieves high segmentation accuracy on 5 benchmark datasets while maintaining low model complexity.