Vision Mamba Attention Feature Fusion UNet: An Innovative State Space Model for Accurate Polyp Segmentation

Abstract:

Colorectal cancer (CRC) is a prevalent disease, with polyps serving as its precursors. Accurate polyp segmenta-tion is crucial for early CRC prevention. However, due to the different size of the polyps, the boundaries are not clear. Therefore, accurate segmentation of polyps is a challenging task. This paper proposes vision mamba at-tention feature fusion UNet (VMA-UNet), a U-shaped asymmetric codec structure model grounded in the state space model (SSM). VMA-UNet incorporates attention feature fusion (AFF) in order to enhance the feature representation of small polyps. A new IUD loss function is proposed to address both large polyps and small polyps, and to mitigate the issue of data imbalance. When applied to multiple datasets, VMA-UNet demonstrates robust performance, particularly in small polyp segmentation, showcasing its practical value. The network proposed in this paper overcomes the inherent shortcomings of convolutional neural network (CNN) and Transformers, not only performs well in remote interaction modeling, but also maintains linear computational complexity. Our study introduces a new method for polyp segmentation based on SSM and advances the field.

Key words:Polyp;segmentation. State;Space Model. Colorectal;cancer. VSSA;codec module