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.

  • Problem

    Existing ViM approaches overlook short-range local dependencies and rely on fixed scanning patterns that limit adaptive spatial-context capture in medical image segmentation.

  • Method

    CCViM is a U-shaped architecture whose context clustering layer partitions features into local windows for adaptive clustering and whose CCS6 layer combines clustering with global scanning.

  • Results

    CCViM outperforms existing models across nuclei, skin-lesion, and multi-organ segmentation evaluations on five public medical image segmentation datasets.

  • Takeaways & Limitations

    Combining dynamic local clustering with global scanning provides CCViM with a supported approach for capturing both local spatial context and global information in medical image segmentation.

  • Takeaways & Limitations

    Fixed configurations for scanning directions and CC layers may not adapt to varying lesion sizes, irregular boundaries, and image characteristics.

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.
  • CNNs capture local features but have limited long-range interactions, motivating models that better aggregate global and local representations.
  • VMamba uses cross-scan processing to integrate pixels from different directions and achieve a global receptive field with linear complexity.
  • CCViM combines cross-scan processing with context clustering, dynamically grouping points within local windows to capture local contexts.
  • Experiments across nuclei, skin-lesion, and multi-organ segmentation tasks show CCViM’s superior performance, with five public datasets used in the evaluation.

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.
  • CNN locality helps detect corners, edges, textures, and other small or irregular target features, but local interaction alone is insufficient for complex anatomical structures.
  • CCViM’s CCS6 design combines global scanning directions with context clustering, while its CC layer dynamically captures spatial contextual information.
  • CC views an image as a set of points and uses clustering to capture local topology information adaptively without significant computational complexity.

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.

  • ViM flattens spatial data into one-dimensional tokens and scans them in two directions, but flattening disrupts natural two-dimensional dependencies.
  • VMamba introduces cross-scan processing to bridge one-dimensional scanning and two-dimensional dependencies without compromising the receptive field.
  • CCS6 integrates scanning modules with context clustering to adaptively extract global and local features while capturing spatial context information.
  • Mamba’s selective state-space model makes parameters depend on the input, allowing selective focus or filtering while retaining linear computation.

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.
  • B. Overall Architecture: 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].
  • B. Overall Architecture: 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.
  • B. Overall Architecture: 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.

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.
  • C. Context Clustering Selective State Space Model: 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.
  • C. Context Clustering Selective State Space Model: 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.
  • C. Context Clustering Selective State Space Model: 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.
  • C. Context Clustering Selective State Space Model: 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.

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.
  • E. Post Processing: Sobel operators derive horizontal and vertical gradient maps from the predicted distance maps.These gradients contribute to the subsequent marker-generation process.
  • E. Post Processing: 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.

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.
  • F. Loss Function: 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.
  • F. Loss Function: 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.
  • F. Loss Function: 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.
  • F. Loss Function: Nuclei segmentation is evaluated with DICE, AJI, PQ, DQ, and SQ, while Synapse uses DSC and HD95.PQ combines detection quality and segmentation quality.

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.
  • B. Configuration of Scan Directions and Local Clusters: 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.

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: All experiments use batch size 32, AdamW with initial learning rate 1e-3, cosine annealing, 300 epochs, and ImageNet-pretrained weights on one RTX 3090 GPU.Experiments are implemented with PyTorch.
  • C. Implementation Details: CCViM is compared with prior methods on ISIC17 and ISIC18 skin-lesion segmentation datasets using visual masks and lesion-contour comparisons.The comparisons include ground truth, CCViM predictions, and VM-UNet predictions.
  • C. Implementation Details: CCViM is compared with CNN-, Transformer-, and Mamba-based models on Kumar and CPM17 nuclei segmentation datasets.The evaluation includes visualizations where differently colored nuclear boundaries denote separate instances.
  • C. Implementation Details: On Synapse, CCViM is evaluated against state-of-the-art models for segmenting eight abdominal-organ classes.The reported visual comparisons assess organ segmentation accuracy and edge delineation relative to VM-UNet.

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.
  • CC layer superiority: 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%.
  • 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.
  • 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.
  • Efficiency analysis: The CC operation aggregates local spatial information within small windows, reducing computational cost compared with processing the entire image.

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.
  • Configuration limitations: Small lesions and irregular boundaries may limit the details captured by the CC layer’s local extraction.
  • Configuration limitations: Different clustering centers capture different local details, while scan directions produce different global interactions that affect segmentation performance.
  • Future direction: The authors identify adaptive scan-direction and CC-layer configurations as a future direction for responding to input-specific characteristics.

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.
  • Conclusion: The CC layer partitions features into windows for learnable local clustering and dynamically captures spatial contextual information.
  • Conclusion: The CCS6 layer combines CC with traditional global scanning to capture local and global information.
  • Conclusion: Experiments on nuclei, skin-lesion, and multi-organ segmentation datasets demonstrate promising medical image segmentation performance.
Loading 2501.01618v1…