Efficient Uncertainty Quantification in Medical Imaging via Mamba State Space Models

Area of Science:

• Medical Imaging
• Computer Vision
• Machine Learning

Background:

• Reliable uncertainty quantification (UQ) is crucial for safety-critical medical imaging.
• Existing UQ methods often compromise computational efficiency or calibration.
• There is a need for efficient and well-calibrated UQ in medical AI.

Purpose of the Study:

• To introduce UQ-Mamba, a novel architecture for native uncertainty quantification within Mamba state space models.
• To achieve efficient and calibrated confidence estimates in medical image classification.
• To enable deployment of automated systems in resource-constrained medical settings.

Main Methods:

• UQ-Mamba embeds UQ into a Mamba state space model using linearized error propagation.
• It provides approximate epistemic and aleatoric uncertainty estimates in a single forward pass.
• Learnable log-variance parameters are propagated through the state transition matrix for calibration.

Main Results:

• UQ-Mamba achieves high accuracy across diverse medical imaging tasks (CT, histopathology, dermoscopy, radiography).
• It demonstrates superior performance compared to baseline models in terms of accuracy and uncertainty calibration (low ECE).
• The architecture is highly parameter-efficient, using significantly fewer parameters than established models like ResNet-50 and EfficientNet-B0.

Conclusions:

• The Mamba state space model propagation mechanism is vital for effective uncertainty decomposition.
• Uncertainty-aware state space models represent a promising direction for parameter-efficient medical image classification.
• UQ-Mamba shows potential for clinical deployment in resource-limited environments, pending further validation.