Towards Real-Time Medical Image Segmentation via Axial Channel Attention and Multi-scale Contextual Mamba
摘要
The rapid advancement of handheld medical devices has increased the demand for accurate yet lightweight medical image segmentation models. While current CNN and Transformer architectures are effective, they often struggle to balance global context modeling with computational efficiency, a critical factor for deployment on resource-constrained devices. Recently, the Mamba architecture, based on the State Space Model (SSM), has emerged as a promising approach for capturing long-range dependencies with linear computational complexity. Motivated by this potential, we present LiteHAM-Net, an ultra-lightweight U-Net variant featuring only 0.21M parameters and approximately 1 GFLOP of computational cost. The model incorporates two primary novel Mamba-based modules: the Mamba Axial Channel Attention (MACA) Encoder and the Multi-scale Contextual Mamba (MSCM) Block. These modules collaboratively extract rich semantic features by effectively leveraging global and multi-scale information. Furthermore, we introduce a Hybrid Convolution Attention Fusion (HCAF) block at the bottleneck to effectively refine features while keeping the model lightweight. Results from extensive experiments on the PH2 and DSB datasets indicate that LiteHAM-Net not only substantially reduces computational resource requirements but also surpasses existing lightweight models in segmentation accuracy. These findings underscore the potential of hybrid architectures combining Mamba, convolution, and attention mechanisms for real-time medical image segmentation applications. The source code is available at: https://github.com/HUY-BK/LiteHAM-Net .