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

TL;DR

This paper introduces CCViM, a novel method that enhances medical image segmentation by integrating context clustering into vision models to better capture both long-range and local spatial features.

Contribution

The paper proposes a simple, effective context clustering module within vision models to improve the capture of local and global features for medical image segmentation.

Findings

Outperforms state-of-the-art methods on multiple public datasets.

Effectively combines long-range and short-range feature interactions.

Demonstrates superior spatial contextual representation in segmentation tasks.

Abstract

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.