TMU: Transmission-Enhanced Mamba-UNet for Medical Image Segmentation

Xiongfeng Yang,Ziyang Luo,Yanlin Wu,Xueshuo Xie,Li Nan,Tao Li
DOI: https://doi.org/10.1007/978-981-97-5609-4_33
2024-01-01
Abstract:In the field of medical image segmentation, the Mamba-UNet is seen as a diamond in the rough due to its robust capability in capturing long-range interactions within images while maintaining linear computational complexity. However, the existing Mamba-based U-shaped networks utilize direct skip connections, which limit the exploration of features at different scales. To address this issue, we propose the Transmission-Enhanced Mamba-UNet (TMU) by incorporating a DCA (Double Cross Attention) block between the encoder and decoder, aiming to enhance skip connections in mamba-based networks. This design elevates segmentation performance by introducing an attention mechanism into the skip connection, effectively fusing features from different layers and improving the model's capacity to capture intricate details and contextual information. We also explored the performance of DCA blocks with different inputs and outputs to find the best combination. Experimental results on the publicly available ACDC dataset and ISIC dataset demonstrate that TMU outperforms the Mamba-UNet model across all evaluation metrics when utilizing pretrained models. Especially in IoU, the TMU model with pretrained weights saw improvements of 1.56% on the ACDC dataset and 0.68% on the ISIC dataset. Identically, without pretrained weights, it exhibited improvements of 4.73% and 0.67%.
What problem does this paper attempt to address?