UD-Mamba: A pixel-level uncertainty-driven Mamba model for medical image segmentation

Kavli Affiliate: Feng Wang

| First 5 Authors: Weiren Zhao, Feng Wang, Yanran Wang, Yutong Xie, Qi Wu

| Summary:

Recent advancements have highlighted the Mamba framework, a state-space model
known for its efficiency in capturing long-range dependencies with linear
computational complexity. While Mamba has shown competitive performance in
medical image segmentation, it encounters difficulties in modeling local
features due to the sporadic nature of traditional location-based scanning
methods and the complex, ambiguous boundaries often present in medical images.
To overcome these challenges, we propose Uncertainty-Driven Mamba (UD-Mamba),
which redefines the pixel-order scanning process by incorporating channel
uncertainty into the scanning mechanism. UD-Mamba introduces two key scanning
techniques: 1) sequential scanning, which prioritizes regions with high
uncertainty by scanning in a row-by-row fashion, and 2) skip scanning, which
processes columns vertically, moving from high-to-low or low-to-high
uncertainty at fixed intervals. Sequential scanning efficiently clusters
high-uncertainty regions, such as boundaries and foreground objects, to improve
segmentation precision, while skip scanning enhances the interaction between
background and foreground regions, allowing for timely integration of
background information to support more accurate foreground inference.
Recognizing the advantages of scanning from certain to uncertain areas, we
introduce four learnable parameters to balance the importance of features
extracted from different scanning methods. Additionally, a cosine consistency
loss is employed to mitigate the drawbacks of transitioning between uncertain
and certain regions during the scanning process. Our method demonstrates robust
segmentation performance, validated across three distinct medical imaging
datasets involving pathology, dermatological lesions, and cardiac tasks.

| Search Query: ArXiv Query: search_query=au:”Feng Wang”&id_list=&start=0&max_results=3

Read More