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

2025-02-04

AI Summary

This work addresses the challenges of ambiguous boundaries and inadequate local feature modeling in medical image segmentation, highlighting the limitations of conventional fixed-order scanning strategies and existing Mamba architectures in capturing high-uncertainty regions—particularly lesion boundaries. To this end, we propose an uncertainty-driven dual-path Mamba architecture: it dynamically schedules sequential and skip-step scanning based on pixel-wise prediction uncertainty and fuses multi-path features via learnable weights. Additionally, a cosine consistency loss is introduced to mitigate feature distortion arising from scanning-mode transitions. Evaluated on three public benchmarks—histopathology, skin lesion, and cardiac imaging—the method achieves state-of-the-art performance, significantly improving boundary localization accuracy and foreground-background interaction modeling. Our approach establishes a novel paradigm for medical image segmentation based on state-space models.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Uncertainty Representations

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• Improves medical image segmentation accuracy.
• Addresses local feature modeling challenges.
• Incorporates channel uncertainty in scanning.

Innovation

• Uncertainty-Driven Mamba model
• Sequential and skip scanning
• Cosine consistency loss function