Merging Context Clustering with Visual State Space Models for Medical Image Segmentation

MedISeg is crucial for physicians in diagnosing specific diseases, as it delineates the regions of interest with precision [1, 2, 3]. Automatic MedISeg has been extensively researched to address the limitations of manual segmentation, which is often time-consuming, labor-intensive, and subjective [5, 36, 6]. Consequently, deep learning methods such as convolutional neural networks (CNNs) [17, 18, 19] and vision transformers (ViTs) [37, 30, 24, 20] have been widely applied in MedISeg tasks [38]. For instance, UNet [39] is favored for its simple architecture and robust scalability, which derives many U-shaped models for MedISeg tasks [36]. UNet++ [40] proposes nested encoder-decoder sub-networks with skip connections for MedISeg. TransUnet [20] proposes the first Transformer-based model for MedISeg tasks. Swin-Unet [24], a pure Transformer-based U-shaped model, adopts a hierarchical Swin Transformer [41] block with shifted windows as the encoder and performs effectively in MedISeg. However, CNNs’ performance is limited with the local respective field, and Transformer-based models require quadratic complexity. To overcome these challenges, Mamba [27] has been proposed to solve the computational efficiency problem and modeling long-range dependencies. In this paper, our research is also based on Mamba. Our contribution is to propose a more efficient local interaction mechanism.

Feature interaction matters, and it greatly affects MedISeg’s performance. CNNs have dominated the field of computer vision in recent years, benefiting from some key inductive biases e.g., locality and translation equivalence [42, 14]. The CNN’s locality helps identify distinct and stable points in images, e.g., corners, edges, and textures, which are significant for detecting small and irregularly shaped target features in MedISeg [43, 44]. However, the local feature interaction cannot accurately address the complex and interconnected anatomical structures in MedISeg, which also requires long-range feature interactions to capture the global context and spatial relationships within the medical image [45]. In recent years, ViTs have emerged as effective long-range feature interaction methods in MedISeg tasks, e.g., Swin-Unet [24], UTNet [46], and Missformer [47]. Since ViTs abandon the local bias from CNNs, some researchers combine convolution and attention to acquire both local and global features, e.g., TransUnet [20] and Conformer [22]. Although the above methods inherit the advantages from short-range and long-range features and achieve better performance, the insights and knowledge are still restricted to CNNs and ViTs [48]. Besides, CNNs or ViTs are limited in modeling long-range dependencies or computational efficiency. Therefore, except for Mamba [27], we also attempt to employ a new paradigm for visual representation by using a clustering [48]. A clustering algorithm views an image as a set of points, allowing it to capture local topology information adaptively without introducing significant computational complexity.

Mamba [27], originally designed for the NLP tasks, proposes a novel selective mechanism that focuses on relevant information or filters out irrelevant information. Mamba [27] eliminates the limitations in CNNs and Transformers, and achieves both effectiveness and efficiency. However, due to the intrinsic non-causal characteristic of images, there is a significant obstacle to Mamba’s comprehension of images. Therefore, a series of Mamba-based models have been proposed to address this issue. For instance, the representative ViM [49] flattens spatial data into 1D tokens and scans these tokens in two different directions to address the non-casual problem in the vision domain. However, flattening spatial data into 1D tokens disrupts the natural 2D dependencies. To address this problem, VMamba [25] introduces a cross-scan module to bridge the gap between 1D array scanning and 2D dependencies, effectively resolving the non-causal problem in visual images without compromising the receptive field. Besides, influenced by the success of VMamba [25], VM-UNet [28] introduces a visual Mamba UNet for MedISeg. However, these Mamba-based methods overlook the local features, which are both important in the vision domain. Therefore, LocalMamba [26] introduces a local scan strategy to capture both global and local information. However, such fixed scanning approaches are limited in their ability to dynamically capture spatial relationships. To address this challenge, we propose the CCS6 module, which integrates scanning modules with our CC to adaptively extract both global and local features while capturing spatial context information.

Mamba [27] is a causal visual state space model, which is both effective and efficient, identifying that a key weakness of the structured state space model (SSM) is its inability to perform content-based reasoning. To overcome this, Mamba introduces a selective state-space model (S6). The recurrent formula of SSM can be defined as follows:

where $h(t)\in\mathbb{R}^{N}$ is the intermediate latent state to map $x(t)\in\mathbb{R}$ to $y(t)\in\mathbb{R}$. $\boldsymbol{A}\in\mathbb{R}^{N\times N}$ and $\boldsymbol{B}\in\mathbb{R}^{N\times N}$ are discretized by $\overline{\boldsymbol{A}}=\exp(\Delta\boldsymbol{A})$ and $\overline{\boldsymbol{B}}=(\Delta\boldsymbol{A})^{-1}(\exp(\Delta\boldsymbol{A})-\boldsymbol{E})\cdot\Delta\boldsymbol{B}$. $\Delta$ is the timescale parameter for discretizing the parameters $\boldsymbol{A}$ and $\boldsymbol{B}$. Through discretization, continuous-time system SSM can be integrated into deep-learning models. $\boldsymbol{C}\in\mathbb{R}^{N\times 1}$ is the parameter matrix. $\boldsymbol{E}$ is the identify matrix. On the basis of SSM, S6 lets the SSM parameters be functions of the input, allowing the model to selectively focus on or filter out information. Making the parameters $\Delta,\boldsymbol{B},\boldsymbol{C}$ depend on the input $x$, and the formulas are defined as follows:

where $\mathrm{Linear}_{N},\mathrm{Linear}_{1}$ are fully connected (FC) layers to project the embedding dimension of $x$ to $N$ and $1$. $\mathrm{Broadcasrt}_{D}$ matches the dimension of the output of $x$ to the Parameter. softplus is an activation function. Integrating S6, the Mamba not only achieves linear computation, but also performs excellent in language modeling tasks [27, 50, 51, 52]. However, due to the inherent non-causal characteristics of images, it’s difficult for Mamba-based models [27] to handle image data well. Fortunately, VMamba [25] proposes a cross-scan module to scan the patched and flattened image in four directions, which can address the non-causality of the image without compromising the field of reception and increasing computational complexity.

The overview of CCViM is illustrated in Fig. 2 (a), which is composed of a patch embedding layer, an encoder, a decoder, a final projection layer, and skip connections. CCViM is not a symmetrical structure like UNet [39], but adopts an asymmetric design. First, input the image $\boldsymbol{I}\in\mathbb{R}^{3\times W\times H}$ into the patch embedding layer to divide $\boldsymbol{I}$ into non-overlapping patches of size $4\times 4$, getting a feature map with dimensions of $C\times\frac{W}{4}\times\frac{H}{4}$, where $C$ is the number of feature channels. Then, input the feature map into the encoder network. Each stage in the encoder network is composed of two CCViM blocks and one patch merging layer, while the last stage only has two CCViM blocks without patch merging layer after them. The patch merging layer is for reducing the height and width of the feature maps and increasing the number of channels, and the channels of feature maps in the four stages are $[C,2C,4C,8C]$. Thirdly, the feature maps are translated into the decoder, which also has four stages. Each stage of decoder is composed of one patch expanding layer and two CCViM blocks, while the first stage only has two CCViM blocks without patch expanding layer before them. The patch expanding layer is for increasing the height and width of the feature maps and reducing the number of channels, and the channels of feature maps in the four stages are $[8C,4C,2C,C]$. Furthermore, we simply employ the addition skip connections to capture low-level and high-level features. Besides, we adopt the most foundational Cross-Entropy and Dice loss for our MedISeg tasks. The CCViM block is the core component of CCViM, as shown in Fig. 2 (b). The input is first translated into the normalization layer and then split into two branches. In the first branch, the input passes through a linear layer followed by an activation layer. In the second branch, the input passes through a linear layer, depth-wise convolution, and an activation layer and then translates into the core module of the CCViM block: context clustering selective state space model (CCS6) layer. After normalizing the output of CCS6 layer, multiply it with the output from the first branch to merge the out of the two pathways. Finally, a linear layer is used to project the merged features onto the dimensions of the initial input features to establish residual connections.

Context clustering selective state space model (CCS6) layer is the core component of the CCViM block, which adopts a selective state space model (S6) as its footstone. S6 processes the input data causally, resulting in only capturing vital information within the scanned part of the data, which is difficult for processing non-causal images. Based on which, numerous researches have also proposed various scan strategies to solve this problem well [49, 25, 28, 26]. However, these global scanning methods overlook the local features and the spatial context information in medical images. Furthermore, all these fixed scanning approaches cannot effectively capture spatial relationships adaptively. To overcome these challenges, we propose a CCS6 layer to extract both global and local features while capturing spatial context information in a learnable way. We adopt VMamba’s [25] cross-scan module, which proposes a selective scan mechanism across four different directions. As shown in Fig. 2 (c), we patch and flatten the input image, and scan the flattened image in both horizontal and vertical directions to capture the global information. At the same time, our CC layer performs learnable local context clustering in local windows. Additionally, the CC layer employs two different numbers of clustering centers–4 and 25–to capture varying structural information. Consequently, there are two distinct CC layers and four different scanning directions available for selection. To avoid introducing too much computational complexity, in each CCS6 layer, we select four from the six choices as in [28, 26]. The detailed configuration of these choices will be introduced in section 4.2.

Instead of using convolution or attention to extract information, we use the context cluster (CC) algorithm [48] for local extraction and spatial context information. In our CC layer, we view the image as a set of data points and group all points into clusters. Instead of clustering data points over the entire image, which will bring a significant computational cost. We split the image into distinct windows and limit the clustering operation to a local region. In each cluster, we aggregate the points into a center adaptively. As shown in Fig. 2 (d), after the patch embedding layer, the image $\boldsymbol{I}\in\mathbb{R}^{3\times W\times H}$ is transformed into patched feature maps $\boldsymbol{F}\in\mathbb{R}^{d\times w\times h}$. We convert the patched feature maps into a set of data points $\boldsymbol{P}\in\mathbb{R}^{d\times n}$, where $n=w\times h$ represents the total number of data points, $d$ is the point feature dimension. These points are unordered and disorganized. We then group the data points into several clusters based on their similarities, ensuring that each point is assigned to only one cluster.

For CC layer, We first project $\boldsymbol{P}$ to $\boldsymbol{P}_{S}\in\mathbb{R}^{n\times d^{\prime}}$ to compute the similarity, where $d^{\prime}$ is the new point feature dimension. We evenly propose $t$ centers in each local window and compute the center feature by averaging its $k$ nearest points. Then, we calculate the pair-wise cosine similarity between the resulting $t$ center points and $\boldsymbol{P}_{S}$ to get the similarity matrix $\boldsymbol{S}\in\mathbb{R}^{t\times n}$. Based on the similarity, we allocate each point to the most similar center to get $t$ clusters. Furthermore, each cluster may exhibit a varying number of data points. In exceptional cases, some clusters may contain no data points, indicating that they are redundant. During the clustering, all data points in one cluster are dynamically aggregated into the center point based on the similarities. In a cluster, there is a small set of $m$ data points, represented by $\boldsymbol{P}_{m}\in\mathbb{R}^{m\times d^{\prime}}$. Calculating the similarity between the $m$ data points $\boldsymbol{P}_{m}\in\mathbb{R}^{m\times d^{\prime}}$ and the center, which is represented by $\boldsymbol{s}\in\mathbb{R}^{m}$ and is the subset of similarity matrix $\boldsymbol{S}$. We map the $m$ data points $\boldsymbol{P}_{m}\in\mathbb{R}^{m\times d^{\prime}}$ to a value space to get $\boldsymbol{P}^{\prime}_{m}\in\mathbb{R}^{m\times d^{\prime}}$. We also generate a value center $\boldsymbol{v}_{c}$ of the $\boldsymbol{P}^{\prime}_{m}$ in the value space, which is like the clustering center proposal. Then the aggregated feature $\boldsymbol{g}\in\mathbb{R}^{d^{\prime}}$ is formulated as follow:

where $\alpha$ and $\beta$ are learnable scalars used to scale and shift the similarity, and $\sigma$ is a sigmoid function that re-scales the similarity to the range (0, 1). $\boldsymbol{p}_{i}$ denotes $i$-th point in $\boldsymbol{P}^{\prime}_{m}$. For numerical stability and to emphasize locality, we incorporate the value center $\boldsymbol{v}_{c}$ in Eq. (3). To control the scale, normalizing the aggregated feature by a factor of $T$. Subsequently, adaptively dispatching the aggregated feature $\boldsymbol{g}$ to each data point in a cluster based on the similarity. This approach facilitates the communication of the points in the cluster, enabling them to share features in the cluster. For each data point $\boldsymbol{p}_{i}$ of $\boldsymbol{P}^{\prime}_{m}$ in the cluster, updating it by:

where $\boldsymbol{s}$ denotes similarity, using the same procedures as above. We apply a fully connected (FC) layer to project the feature dimension from the value space dimension $d^{\prime}$ to the original dimension $d$.

In the nuclei segmentation task (i.e., Kumar [31] and CPM17 [32]), we employ the post-processing method watershed algorithm to distinguish between individual nuclei instances. Following research [9], we create horizontal and vertical distance maps by calculating the distances from nuclear pixels to their centers of mass in both the vertical and horizontal directions. Within our model we predict the vertical distance maps $P_{v}$ and horizontal distance maps $P_{h}$. Additionally, we apply the Sobel operator to these distance maps to obtain the horizontal and vertical gradient maps. We then select the maximum value between these gradient maps: $M_{s}=\max(\text{Sobel}(P_{v}),\text{Sobel}(P_{h}))$. This process aids in edge detection by calculating the gradient magnitude, thereby emphasizing areas of significant intensity change. Next, we calculate the markers, $M$, by applying a threshold function to the probability map $P$ and gradient map $M_{s}$, where the markers are defined as $M=\delta(\tau(P,r)-\tau(M_{s},k))$. Here, $\tau(P,r)$ is a threshold function, with $r$ being the threshold value. If the value exceeds $r$, it is set to 1; otherwise, it is set to 0. $\delta$ is used to set negative values to 0. We then obtain the energy landscape $E=(1-\tau(M_{s},k))\times\tau(P,h)$. Finally, $M$ serves as the marker during marker-controlled watershed to determine how to split $\tau(P,h)$, guided by the energy landscape $E$.

For a fair comparison with other state-of-the-art methods, we employ the most fundamental loss functions in medical image segmentation: Cross-Entropy and Dice loss [28, 9, 53]. Cross-entropy ensures pixel-level classification accuracy, while Dice loss addresses the common issue of class imbalance by optimizing the overlap between the predicted and true segmentation [54]. As shown in Eq. (5) and Eq. (LABEL:eq6), we combine Cross-Entropy and Dice loss to balance precise pixel-wise classification with overall segmentation performance.

where, $N$ denotes the total number of samples, and $C$ represents the total number of categories. $y_{i,c}$ is an indicator of ground truth, equals 1 if sample $i$ belongs to category $c$, and 0 if it does not. $\hat{y}_{i,c}$ is the probability that the model predicts sample $i$ as belonging to category $c$. $y_{i}$ and $\hat{y}_{i}$ represent the ground truth and prediction, respectively.

In this paper, we do not apply the redundancy 22 scan strategies [56] in each CCS6 layer, and there are only one, two, or three different scanning directions in each CCS6 layer. To extract local features and capture the spatial context information adaptively, we integrate our CC layer into each CCS6 layer. As shown in Fig. 3, each CCS6 layer contains only four modules, which include one, two, or three different scanning directions and one or two different CC layers. This design avoids adding additional computational overhead compared to previous models. There are a total of four different scanning directions to choose from, including horizontal, horizontal flipped, vertical, and vertical flipped. Given the different sizes and number of image targets, we adopt two different CC layers. One CC layer method has 4 cluster centers in a local region, and the other one has 25 cluster centers in a local region. In nuclei segmentation tasks [31, 32], there are many nuclei in one local region. In skin image segmentation tasks [33, 34], the target is large and may need lower point centers in one local region. We adopt the configuration provided by LocalMamba [26], with our CC lyaer replacing the local scan. Compared with the fixed local scan strategies in LocalMamba [26], our CC layer can dynamically cluster local features and capture spatial information adaptively.

In our all experiments, we set the batch size to 32 and employ AdamW [57] optimizer with an initial learning rate of 1e-3. CosineAnnealingLR [58] is employed as the scheduler. We set the training epochs to 300. All experiments initialize the models with ImageNet [59] pretrained weights. All experiments are conducted on the PyTorch deep learning platform with a single NVIDIA GeForce RTX 3090 GPU.

Although extensive experiments demonstrate the effectiveness of our method in MedISeg tasks, there are two notable limitations. First, while our CC layer may introduce minimal computational costs, this slight addition contributes to promising performance improvements. Second, different medical images may exhibit small lesions and irregular boundaries, which can affect the details captured by our CC layer’s local extraction. Additionally, varying scanning directions and CC layers may influence the performance of MedISeg tasks. However, our current configurations for scanning directions and CC layers are fixed, underscoring the need for more adaptive strategies. As shown in Fig. 8, there are several failure examples from the ISIC18 [34] dataset. We can observe that the fixed configuration of global scan directions and local CC layers may impact performance, particularly in capturing the details of irregular boundaries and complex structures. The model struggles to replicate the intricate features represented in the ground truth, leading to discrepancies in segmentation results. An adaptive configuration could enhance our model’s ability to discern these features more effectively. Given that our CC layers utilize different clustering centers, varying clustering centers can capture different local details. Furthermore, different scan directions can facilitate varying global interactions, which also affect performance in MedISeg tasks. Adopting a more flexible approach could significantly improve segmentation accuracy by allowing the model to better adapt to the unique characteristics of each medical image. Therefore, developing more adaptive configurations for scan directions and CC layers could enhance performance in both global and local extraction. In the future, we will focus on implementing these adaptive strategies, exploring dynamic configuration algorithms that can adjust in real-time based on the input data characteristics.