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

TL;DR

UD-Mamba tackles the challenge of capturing local features in medical image segmentation by replacing traditional position-based scanning with pixel-level channel uncertainty-guided selective scanning. It introduces the Uncertainty-Driven Selective Scanning Model (UD-SSM) featuring sequential and skip scans, four learnable weights, and a cosine consistency loss to align forward and backward scans within a linear-time State Space Model framework. The approach demonstrates robust improvements across three medical imaging datasets (DigestPath, ISIC 2018, ACDC) in Dice, IoU, and accuracy while maintaining computational efficiency. This work enhances segmentation reliability in clinical imaging by focusing model attention on uncertain regions and enforcing cross-directional feature consistency, enabling more precise delineation of boundaries and lesions.

Abstract

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.

Figures (6)

• Figure 1: Pixel-level channel uncertainty-based scanning mechanism. (a) Input image; (b) Ground truth; (c) Resulting image obtained from channel-based uncertainty calculations, where pixels with the highest uncertainty are red and the lowest are blue; (d) Feature image sorted by the degree of uncertainty; (e) Previous method using the SS2D liu2024vmamba scanning mechanism; (f) Our UD-SSM scanning mechanism, which includes sequential scanning and skip scanning.
• Figure 2: Performance comparison of ascending and descending methods: the ascending method scans from regions of low uncertainty to high uncertainty, while the descending method scans in the reverse order.
• Figure 3: Illustration of the UD-Mamba architecture, which includes a patch embedding layer, an encoder-decoder with Uncertainty-Driven (UD) Blocks, and a segmentation head. Each UD Block features the Uncertainty-Driven Selective Scanning Model (UD-SSM) for processing input.
• Figure 4: Detailed description of the UD-SSM. I. Describes the uncertainty calculation process based on channel uncertainty in the UD-SSM. II. Explains the scanning expansion operation, which mainly includes two strategies: Sequential scanning ($\text{Scan}^{se}$) and Skip scanning ($\text{Scan}^{sk}$), as well as the subsequent Reweighting operation. III. Presents the S6 block and Recovery processing, while also introduces the calculation process of Consistency Constraints.
• Figure 5: Visual comparisons of segmentation results from UD-Mamba and various other methods are conducted across three different datasets.