UU-Mamba: Uncertainty-aware U-Mamba for Cardiac Image Segmentation

TL;DR

UU-Mamba is an advanced cardiac MRI segmentation model that combines uncertainty-aware loss and the SAM optimizer to improve accuracy, robustness, and generalization over existing methods.

Contribution

We introduce UU-Mamba, a novel cardiac segmentation model that integrates uncertainty-aware loss and SAM optimizer for enhanced performance and robustness.

Findings

Outperforms state-of-the-art models on ACDC dataset

Achieves higher DSC and lower MSE scores

Demonstrates improved robustness and generalization

Abstract

Biomedical image segmentation is critical for accurate identification and analysis of anatomical structures in medical imaging, particularly in cardiac MRI. Manual segmentation is labor-intensive, time-consuming, and prone to errors, highlighting the need for automated methods. However, current machine learning approaches face challenges like overfitting and data demands. To tackle these issues, we propose a new UU-Mamba model, integrating the U-Mamba model with the Sharpness-Aware Minimization (SAM) optimizer and an uncertainty-aware loss function. SAM enhances generalization by locating flat minima in the loss landscape, thus reducing overfitting. The uncertainty-aware loss combines region-based, distribution-based, and pixel-based loss designs to improve segmentation accuracy and robustness. Evaluation of our method is performed on the ACDC cardiac dataset, outperforming state-of-the-art models including TransUNet, Swin-Unet, nnUNet, and nnFormer. Our approach achieves Dice Similarity Coefficient (DSC) and Mean Squared Error (MSE) scores, demonstrating its effectiveness in cardiac MRI segmentation.