AM-UNet: Attention Mamba U-Net for Medical Image Segmentation

Recently, architectures such as Mamba that leverage State Space Models (SSMs) have shown strong potential in rivaling conventional CNN and Transformer architectures. SSM is a deep sequence model known for its power to handle long sequence tasks, efficiently tracking intricate inter-sequence relationships with linear computational overhead. However, previous skip connection methods have not fully bridged the semantic gap across encoder and decoder features, which may lead to insufficient feature fusion and consequently affect fine-detail recovery. To address this problem, we propose the Attention Mamba UNet (AM-UNet), which integrates the traditional U-shaped architecture with Visual State Space (VSS) blocks to exploit richer contextual information. We further enhance the architecture by embedding a novel attention module into the skip connection framework, where dilated convolution broaden the perception range with no extra processing burden, while cross mechanism facilitates optimal feature fusion between the encoder and decoder, mitigating the semantic gap and enabling a more comprehensive understanding of spatial dependencies. Experiments on ISIC17, ISIC18, and ACDC datasets showing that AM-UNet delivers superior results compared to existing methods when applied to medical image segmentation tasks.