MSVM-UNet: Multi-Scale Vision Mamba UNet for Medical Image Segmentation

2D-Selective-Scan block performs selective scanning along four scanning paths on the input feature map to capture global contextual information and long-range dependencies. Specifically, the 2D input feature map first undergoes a linear layer, a depth-wise convolution operation, and an activation function. Then, further feature extraction is carried out through the 2D-Selective-Scan (SS2D) operation. Finally, the output is obtained after another layer normalization and a linear projection. As shown in Fig. 2(a), SS2D first flattens the 2D input along four different scanning paths to obtain four 1D sequences. These sequences are then fed into the S6 blocks [11] for selective scanning to model long-range dependencies. Finally, the four 1D sequences are restored to their original 2D form and summed to produce the output. The definition of the SS2DBlock is given by Equation (3):

where $z_{i}$ and $\hat{z}_{i}$ represent the input and output feature maps of the SS2DBlock in the $i^{th}$ stage, respectively. $C_{1}(\cdot)$ represents a linear projection used to double the channel dimension. $DWConv(\cdot)$ denotes a depth-wise convolution with a kernel size of $3\times 3$. $\sigma_{1}$ represents a $SiLU(x)=x\cdot sigmoid(x)$ activation function. $SS2D(\cdot)$ denotes 2D-Selective-Scan operation. $C_{2}(\cdot)$ represents another linear projection used to reduce the channel dimension by half.

As depicted in Fig. 2(b), we introduce convolution operations in the feed-forward network to aggregate information from these four diagonal directions. Additionally, to effectively capture detailed information and multi-scale feature representations in hierarchical features, we employ a set of convolution operations with different kernel sizes. To avoid introducing excessive parameters and computational overhead, we use depth-wise convolutions, which are parameter and computation efficient, to implement the MS-FFN. As shown in Fig. 1(d), the MS-FFN consists of two linear layers and a set of parallel depth-wise convolution layers. The definitions of MS-FFN are given by Equations (4) and (5):

where $m_{i}$ and $\hat{m}_{i}$ represent the input and output tensors of the MS-FFN in the $i^{th}$ stage, respectively, and $m_{i}^{{}^{\prime}}$ denote the output of the first linear transformation. $C_{3}(\cdot)$ represents a linear projection used to expand the channel dimension by four times. $DWConv_{ks}(\cdot)$ denotes a depth-wise convolution with a kernel size of $ks$. $KS$ defines a group of parallel convolution kernels with values of $\{1\times 1,3\times 3,5\times 5\}$. $\sigma_{2}(\cdot)$ represents a $GELU(x)=x\cdot\Phi(x)$ activation function. $C_{4}(\cdot)$ represents another linear projection that reduces the channel dimension back to the input dimension.

The dataset consists of 30 abdominal CT scans with 3779 axial contrast-enhanced abdominal CT images. Each CT volume contains 85 to 198 slices of $512\times 512$ pixels, with a voxel spatial resolution of ($[0.54\sim 0.54]\times[0.98\sim 0.98]\times[2.5\sim 5.0]$) $mm^{3}$. Similar to TransUNet [2], we randomly divide the dataset into 18 cases for training and 12 cases for testing. We segment only 8 types of abdominal organs: aorta, gallbladder, left kidney, right kidney, liver, pancreas, spleen, and stomach.

The dataset contains 100 cardiac MRI scan images, each of which consists of three sub-organs: right ventricle (RV), myocardium (Myo), and left ventricle (LV). Following the TransUNet [2], we split the data into 70 cases for training, 10 for validation, and 20 for testing.

In our experiments, all models were implemented based on the Pytorch 2.0.0 framework, and all training was conducted on an NVIDIA GeForce RTX 3090 GPU. We initialized the backbone network using pretrained weights from the ImageNet-1k dataset. To reduce overfitting and enhance the model’s generalization capability, we employed extensive data augmentation techniques, including resizing the input images to $224\times 224$, horizontal and vertical flips, random rotations, Gaussian noise, Gaussian blur, and contrast enhancement. We set the batch size to 32 and used the AdamW optimizer to train the network for a maximum of 300 epochs. The initial learning rate was set to 5e-4 and was decayed during training using a cosine annealing schedule. Due to the varying difficulty levels of different datasets and to reduce overfitting, we set different weight decay for different datasets: 1e-3 for the Synapse dataset and 1e-4 for the ACDC dataset. Additionally, the reported FLOPs and parameter counts of the models were calculated using calflops with an input size of $224\times 224\times 3$. We used a combination of Dice and cross-entropy loss functions to train the network, defined as follows:

where $\alpha=0.6$ and $1-\alpha=0.4$ are the weights for the Dice loss $L_{DICE}$ and the cross-entropy loss $L_{CE}$, respectively.

Following the commonly adopted metrics to measure model performance, we used the Dice SCore (DSC) and the 95% Hausdorff Distance (HD95) to assess our model’s performance on the Synapse and ACDC datasets. The DSC and HD are calculated according to Equations (9) and (10):

where the $X$ and $Y$ denote the ground truth and segmented maps, respectively, and $d(x,y)$ represents the distance between points $x$ and $y$. HD95 is the $95^{th}$ percentile of the distances of the boundaries of $X$ and $Y$.

As shown in Table I, compared to various types of methods, our proposed MSVM-UNet achieved the best average DSC of 85.00% and HD95 of 14.75mm. Specifically, compared to CNN-based methods (such as 2D D-LKA Net), our method improved the DSC and HD95 by 0.73% and 5.29mm, respectively; by 1.37% and 0.93mm compared to transformer-based methods (such as PVT-EMCAD-B2); and by 2.62% and 1.47mm compared to mamba-based methods (such as VM-UNet). Additionally, for the small organs, there were 0.09% and 1.80% DSC improvements in the gallbladder and pancreas, respectively, compared to the second-best methods, and a 1.57% improvement in the stomach for the large organ. This is because MSVM-UNet can effectively capture both long-range dependencies between pixels and local contextual relationships simultaneously. Due to the introduction of multi-scale convolution operations, MSVM-UNet not only handles organs of varying shapes and sizes effectively but also better locates organ boundaries.

As shown in Table II, we present the performance of our method compared to the aforementioned types of methods on MRI medical images, where our proposed MSVM-UNet achieved the best average DSC of 92.58%. Additionally, for the three categories in the ACDC dataset (RV, Myo, and LV), our method achieved the best DSC results of 91.00%, 90.35%, and 96.39%, respectively. This demonstrates the generalizability of our method, as it performs well on different modalities of medical image data (MRI and CT).

We conducted a series of experiments on the Synapse multi-organ dataset to understand the impact of different components in the MSVM-UNet decoder. We assessed the effects of replacing specific modules in the decoder with our proposed modules. As shown in Table III, significant performance improvements were observed after substituting the original modules with our proposed ones, with only a minimal increase in computational overhead and the number of parameters. Specifically, compared to the decoder with VSS blocks and patch expanding layers, using our proposed modules improved the DSC and HD95 by 2.2% and 12.54mm, respectively, with only an additional 0.42G FLOPs and 0.25M parameters. This demonstrates both the effectiveness and efficiency of our proposed MSVM-UNet.

To explore the effectiveness of our proposed LKPE, we conducted experiments on the Synapse multi-organ dataset by using the original decoder with transposed convolution, the UpSample block [3], the patch expanding layer, and the LKPE layer as upsampling layers, respectively, to evaluate their individual performance. As shown in Table IV, we reported the performances and overheads corresponding to different upsampling methods. To provide a clearer comparison of the computational overhead and the number of parameters of different upsampling operations, we only report the FLOPs and parameters of the decoder. Compared to the original patch expanding layer, LKPE improved the DSC and HD95 by 1.5% and 13.08mm, respectively, while only introducing an additional 0.12G FLOPs and 0.06M parameters. Moreover, it can be observed that methods aggregating channel and spatial information generally achieve better performance, further demonstrating the effectiveness of our proposed LKPE.

We also conducted additional experiments on the Synapse multi-organ dataset to explore the impact of different multi-scale convolution kernels in the depth-wise convolution of the MSVSS block. As shown in Table V, the performances for various multi-scale convolution kernels are reported. Similarly, the FLOPs and the number of parameters are provided only for the decoder. To avoid excessive computational overhead, we designed three sets of kernels: $1\times 1$ and $7\times 7$ in the first set, $3\times 3$ and $5\times 5$ in the second, and a union of both in the third. As the number and size of convolution kernels increase, we found that the performance decreases. Based on these observations, we selected $[1,3,5]$ as the default multi-scale convolution kernels in the MSVSS block.

To investigate the impact of different encoder depths on model performance, we conducted two sets of experiments on the Synapse multi-organ dataset to study the effects of varying encoder scales. As shown in Table VI, there is a slight decrease in performance as the encoder becomes larger and deeper. Since both sets of experiments used the same setup, we hypothesize that this minor performance drop may be due to slight overfitting caused by the increased complexity of the model. Based on these observations and considering the computational overhead and the number of parameters, we chose the tiny version of the encoder as the default scale.

The features of the corresponding layers of the original decoder and our MSVM-UNet decoder are visualized in Fig. 5. We compute the average of all channels in the feature map and then produce the heatmap using Matplotlib. It is evident from Fig. 5 that our method helps in handling organs of varying sizes and shapes and in obtaining more discriminative feature representations.