PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yun Zhu, Dong Zhang, Yi Lin, Yifei Feng, Jinhui Tang

arXiv:2501.01618v1cs.CVcs.AI

TL;DR

Medical image segmentation must combine long-range and short-range feature interactions, while existing ViM approaches overlook local dependencies and use fixed scanning patterns. CCViM adds context clustering within a U-shaped vision Mamba architecture to dynamically capture local spatial context alongside global information, and experiments across five public datasets report superior performance over existing models.

Abstract

from arXiv · show

Medical image segmentation demands the aggregation of global and local feature representations, posing a challenge for current methodologies in handling both long-range and short-range feature interactions. Recently, vision mamba (ViM) models have emerged as promising solutions for addressing model complexities by excelling in long-range feature iterations with linear complexity. However, existing ViM approaches overlook the importance of preserving short-range local dependencies by directly flattening spatial tokens and are constrained by fixed scanning patterns that limit the capture of dynamic spatial context information. To address these challenges, we introduce a simple yet effective method named context clustering ViM (CCViM), which incorporates a context clustering module within the existing ViM models to segment image tokens into distinct windows for adaptable local clustering. Our method effectively combines long-range and short-range feature interactions, thereby enhancing spatial contextual representations for medical image segmentation tasks. Extensive experimental evaluations on diverse public datasets, i.e., Kumar, CPM17, ISIC17, ISIC18, and Synapse demonstrate the superior performance of our method compared to current state-of-the-art methods. Our code can be found at https://github.com/zymissy/CCViM.

I. INTRODUCTION

Medical image segmentation requires feature interactions that capture both local detail and broader anatomical context. CCViM addresses this challenge by combining global cross-scan processing with adaptive local context clustering in a U-shaped architecture.

  • Medical image segmentation supports anatomy research, disease diagnosis, treatment planning, and pixel-level delineation of lesions and regions of interest.Source: §I · p. 1
  • CNNs capture local features but have limited long-range interactions, motivating models that better aggregate global and local representations.Source: §I · p. 1
  • VMamba uses cross-scan processing to integrate pixels from different directions and achieve a global receptive field with linear complexity.Source: §I · p. 1, Figure 1 · p. 2
  • CCViM combines cross-scan processing with context clustering, dynamically grouping points within local windows to capture local contexts.Source: §I · p. 2, Figure 1 · p. 2
  • Experiments across nuclei, skin-lesion, and multi-organ segmentation tasks show CCViM’s superior performance, with five public datasets used in the evaluation.Source: §I · p. 2

II. RELATED WORK

Related work shows a progression from local convolutional interactions and global attention toward Mamba-based visual state space models. These approaches motivate context clustering as a way to capture local topology information adaptively.

  • MedISeg methods include CNN- and Transformer-based U-shaped architectures such as UNet, UNet++, TransUnet, and Swin-Unet.Source: A. Medical Image Segmentation (MedISeg) · p. 2
  • CNN locality helps detect corners, edges, textures, and other small or irregular target features, but local interaction alone is insufficient for complex anatomical structures.Source: B. Feature Interaction for MedISeg · p. 2
  • CCViM’s CCS6 design combines global scanning directions with context clustering, while its CC layer dynamically captures spatial contextual information.Source: Figure 2 · p. 3
  • CC views an image as a set of points and uses clustering to capture local topology information adaptively without significant computational complexity.Source: B. Feature Interaction for MedISeg · p. 3

C. Vision Mamba (ViM)

Vision Mamba methods adapt selective state space models to image data, where ordinary causal processing conflicts with images’ non-causal structure. Cross-scan strategies address this issue, while CCS6 further targets adaptive local and global feature extraction.

B. Overall Architecture

CCViM uses an asymmetric U-shaped architecture with hierarchical encoder-decoder stages, patch resizing, skip connections, and a final projection for medical image segmentation.

B. Overall Architecture

  • CCViM comprises patch embedding, an encoder, a decoder, a final projection layer, and skip connections.The architecture is asymmetric rather than a symmetric UNet.Source: B. Overall Architecture · p. 4
  • The encoder has four stages with CCViM blocks and patch merging, while the final stage omits patch merging.Feature channels increase across stages as [C, 2C, 4C, 8C].Source: B. Overall Architecture · p. 4
  • The decoder has four stages with patch expanding layers and CCViM blocks, reducing channels as [8C, 4C, 2C, C].The first decoder stage omits the preceding patch expanding layer.Source: B. Overall Architecture · p. 4
  • Addition skip connections combine low-level and high-level features, while a linear layer projects merged features to the initial input dimension for residual connections.The skip connections are used alongside the decoder pathway.Source: B. Overall Architecture · p. 4

C. Context Clustering Selective State Space Model

The CCS6 layer combines VMamba-style global scanning with context clustering, which locally groups feature points by similarity and adaptively shares aggregated information.

C. Context Clustering Selective State Space Model

  • CCS6 combines four-direction selective scanning with a context-clustering layer to capture global features and local spatial context adaptively.The input is patched and flattened for horizontal and vertical global scans, while CC performs learnable local extraction.Source: C. Context Clustering Selective State… · p. 4
  • The CC layer partitions feature points into local windows rather than clustering across the entire image, reducing the clustering computation.Each point is assigned to one cluster based on similarity.Source: D. Context Clustering · p. 4
  • CC projects points into a new feature dimension, proposes t centers per window, and assigns each point to its most similar center.Center features are computed by averaging k nearest points before cosine similarities are calculated.Source: D. Context Clustering · pp. 4–5
  • Clusters dynamically aggregate their members into center features, allowing variable cluster sizes and potentially empty redundant clusters.The resulting aggregation is represented by g in the value space.Source: D. Context Clustering · p. 5
  • Learnable scaling and shifting transform similarities through a sigmoid, after which aggregated features are normalized and dispatched back to cluster points.The value center is incorporated for numerical stability and locality, and the update uses a fully connected projection.Source: D. Context Clustering · p. 5

E. Post Processing

For nuclei segmentation, the post-processing pipeline predicts distance maps and uses thresholded markers with an energy landscape to guide marker-controlled watershed splitting.

E. Post Processing

  • The nuclei pipeline predicts horizontal and vertical distance maps from nuclear pixels to their centers of mass.It applies this procedure to the Kumar and CPM17 datasets.Source: E. Post Processing · p. 5
  • Sobel operators derive horizontal and vertical gradient maps from the predicted distance maps.These gradients contribute to the subsequent marker-generation process.Source: E. Post Processing · p. 5
  • Thresholding and negative-value suppression produce an energy landscape used by marker-controlled watershed to split nuclei instances.The marker guides the splitting of the thresholded prediction through the energy landscape.Source: E. Post Processing · p. 5

F. Loss Function

CCViM is evaluated across nuclei, skin-lesion, and multi-organ datasets using task-specific segmentation metrics and standard processing configurations.

F. Loss Function

  • The model combines Cross-Entropy and Dice loss to balance pixelwise classification accuracy with overlap optimization under class imbalance.Ground-truth indicators, predicted probabilities, and segmentation predictions define the loss terms.Source: F. Loss Function · p. 5
  • The experiments cover five datasets: Kumar and CPM17 for nuclei, ISIC17 and ISIC18 for skin lesions, and Synapse for multi-organ segmentation.The datasets span microscopy, dermoscopy, and abdominal CT imaging.Source: A. Datasets and Evaluation Metrics · p. 5
  • ISIC17 and ISIC18 contain 2,150 and 2,694 masked images, respectively, and use a 7:3 training-test split with resizing and augmentation.The processing includes 256 × 256 resizing, flips, and rotations.Source: A. Datasets and Evaluation Metrics · pp. 5–6
  • Synapse contains 30 abdominal CT cases and 3,779 axial images covering eight organs, with 18 cases for training and 12 for testing.Images are resized to 224 × 224 with flip and rotation augmentation.Source: A. Datasets and Evaluation Metrics · p. 6
  • Nuclei segmentation is evaluated with DICE, AJI, PQ, DQ, and SQ, while Synapse uses DSC and HD95.PQ combines detection quality and segmentation quality.Source: A. Datasets and Evaluation Metrics · p. 6

B. Configuration of Scan Directions and Local Clusters

Each CCS6 layer uses a limited set of scan directions together with context-clustering layers to extract local features and adaptively capture spatial context.

B. Configuration of Scan Directions and Local Clusters

  • CCViM omits redundant 22 scan strategies, using only one, two, or three scanning directions per CCS6 layer.This reduces the scan configuration within each layer while retaining multiple directional processing options.Source: B. Configuration of Scan Directions… · p. 6
  • Each CCS6 layer contains four modules comprising one to three scan directions and one or two context-clustering layers.The CC layers are integrated to extract local features and capture spatial context adaptively.Source: B. Configuration of Scan Directions… · p. 6

C. Implementation Details

The experiments use standardized training settings and evaluate CCViM across nuclei, skin-lesion, and multi-organ segmentation tasks with visual and quantitative comparisons against established methods.

C. Implementation Details

E. Ablation Analysis

The ablations evaluate context clustering, cluster-center configurations, and computational efficiency. Results support adaptive local clustering alongside global scanning, while maintaining competitive inference speed.

CC layer superiority

  • On Kumar, the CC layer improves PQ, Dice, AJI, DQ, and SQ over LocalVIM by 1.57%, 0.15%, 1.68%, 1.75%, and 0.39%, respectively.Source: E. Ablation Analysis · p. 8
  • On ISIC17, the CC layer improves mIoU and DSC over LocalVIM by 0.29% and 0.18%, and over LocalVMamba by 2.68% and 1.65%.Source: E. Ablation Analysis · p. 8

Cluster-center analysis

  • Combining CC layers with 4 and 25 cluster centers generally improves results on Kumar and ISIC17, although the h-hflip-C4-C25 configuration performs relatively poorly.Source: E. Ablation Analysis · p. 9

Efficiency analysis

  • CCViM achieves 41.16 fps, close to VM-UNet at 41.54 fps and LocalVMamba at 41.01 fps, despite a slight FLOPs increase from the CC layer.Source: E. Ablation Analysis · p. 10
  • The CC operation aggregates local spatial information within small windows, reducing computational cost compared with processing the entire image.Source: E. Ablation Analysis · p. 10

F. Limitation Analysis

The method’s fixed scan directions and clustering configurations limit adaptation to image-specific structures. These constraints particularly affect small lesions, irregular boundaries, and complex details.

Configuration limitations

  • Fixed global scan directions and local CC configurations can impair segmentation of irregular boundaries and complex structures in ISIC18 failure cases.Source: F. Limitation Analysis · p. 10
  • Small lesions and irregular boundaries may limit the details captured by the CC layer’s local extraction.Source: F. Limitation Analysis · p. 10
  • Different clustering centers capture different local details, while scan directions produce different global interactions that affect segmentation performance.Source: F. Limitation Analysis · p. 11

Future direction

  • The authors identify adaptive scan-direction and CC-layer configurations as a future direction for responding to input-specific characteristics.Source: F. Limitation Analysis · p. 11

V. CONCLUSION

The conclusion presents CCViM as an efficient U-shaped architecture for medical image segmentation. Its CC and CCS6 layers combine dynamic local context capture with global scanning across multiple segmentation tasks.

Conclusion

  • CCViM is a U-shaped medical image segmentation architecture that inherits Mamba’s efficiency and effectiveness.Source: §V · p. 11
  • The CC layer partitions features into windows for learnable local clustering and dynamically captures spatial contextual information.Source: §V · p. 11
  • The CCS6 layer combines CC with traditional global scanning to capture local and global information.Source: §V · p. 11
  • Experiments on nuclei, skin-lesion, and multi-organ segmentation datasets demonstrate promising medical image segmentation performance.Source: §V · p. 11

Related papers

Loading 2501.01618v1…