SK-VM++: Mamba assists skip-connections for medical image segmentation

Abstract

In medical automatic image segmentation engineering, the U-shaped structure is the primary key framework. And the skip-connection operation in it is an important operation for key fusion of high and low features, which is one of the highlights of the U-shaped architecture. However, the traditional U-shaped architecture usually employs direct concatenation or different variants of convolution-based module composition. The recent emergence of Mamba, based on state-space models (SSMs), has shaken up the traditional convolution and Transformers that have long been the foundational building blocks. In this study, we analyze the impact of Mamba on skip-connection operations for U-shaped architectures and propose a novel skip-connection operation (SK-VM++) combining the UNet++ framework and Mamba. Specifically, Mamba is able to refine the fusion of high and low feature information better than traditional convolution. In addition, SK-VM++ leverages the excellent property of Mamba’s concatenation, making it significantly less sensitive to changes in computational complexity and parameters caused by changes in the number of channels. In particular, the number of channels increases from 64 to 512, and the convolution-based FLOPs and parameters rise by 8.82 and 6.22 times, respectively, compared to our proposed Mamba-based skip-connection operation. In addition, comparing with the most popular nnU-Net and VM-UNet, the DSC of SK-VM++ improves by 2.01% and 1.10% on the ISIC2017 dataset, 1.59% and 9.10% on the CVC-ClinicDB dataset, 1.23% and 18.94% on the Promise12 dataset and 46.25% and 34.01% improvement on the UWF-RHS dataset. The code is available from https://github.com/wurenkai/SK-VMPlusPlus.