An Uncertainty-Aware Generalization Framework for Cardiovascular Image Segmentation

TL;DR

This work addresses the generalization gap in cardiovascular image segmentation by introducing UU-Mamba, an uncertainty-aware generalization framework that extends U-Mamba with a learnable, region/distribution/pixel-level loss and the Sharpness-Aware Minimization (SAM) optimizer. The combination promotes flatter loss landscapes and adaptive emphasis on confident predictions, enabling robust segmentation across cardiac and vascular structures. Extensive experiments on ACDC, ImageCAS, and Aorta datasets demonstrate state-of-the-art performance, with notable improvements in DSC and NSD and strong robustness across regions and modalities. The results underscore the practical potential of uncertainty-aware optimization and SAM for reliable, scalable cardiovascular image analysis in diverse clinical settings.

Abstract

Deep learning models have achieved significant success in segmenting cardiovascular structures, but there is a growing need to improve their generalization and robustness. Current methods often face challenges such as overfitting and limited accuracy, largely due to their reliance on large annotated datasets and limited optimization techniques. This paper introduces the UU-Mamba model, an extension of the U-Mamba architecture, designed to address these challenges in both cardiac and vascular segmentation. By incorporating Sharpness-Aware Minimization (SAM), the model enhances generalization by seeking flatter minima in the loss landscape. Additionally, we propose an uncertainty-aware loss function that integrates region-based, distribution-based, and pixel-based components, improving segmentation accuracy by capturing both local and global features. We expand our evaluations on the ImageCAS (coronary artery) and Aorta (aortic branches and zones) datasets, which present more complex segmentation challenges than the ACDC dataset (left and right ventricles) used in prior work, showcasing the model's adaptability and resilience. Our results confirm UU-Mamba's superior performance compared to leading models such as TransUNet, Swin-Unet, nnUNet, and nnFormer. We also provide a more in-depth assessment of the model's robustness and segmentation accuracy through extensive experiments.

Figures (6)

• Figure 1: Comparison between our method and basic approach. Traditionally, a deep learning model is trained using the Cross-Entropy loss $\mathcal{L}_{CE}$. Our method enhances U-Mamba by utilizing the uncertainty-aware loss $\mathcal{L}_{UA}$, which is optimized via the SAM optimizer over a flattened loss landscape. Evaluation using Dice Similarity Coefficient (DSC), Normalized Surface Dice (NSD) and Mean Squared Error (MSE) shows improvement of our method against basic CNN-based methods.
• Figure 2: Overview of our proposed framework, we encode input images and incorporate a novel uncertainty-aware loss function. Optimization is performed using the Sharpness-Aware Minimization (SAM) optimizer foret2021sharpness, which operates within a flattened loss landscape. Experiments on the ACDC dataset bernard2018deep, ImageCAS dataset zeng2023imagecas, and Aorta dataset imran2024ciskrebs2024volumetric perform 3D heart segmentation on cardiovascular MRI and CT images, delineating each cardiovascular labels.
• Figure 3: Segmentation results for various methods on sample images from the ACDC dataset bernard2018deep. The Mean Squared Error (MSE) between the output segmentation and the ground truth is shown for each method.
• Figure 4: Segmentation results for various methods on sample images from the Aorta dataset imran2024ciskrebs2024volumetric. The Mean Squared Error (MSE) between the output segmentation and the ground truth is shown for each method.
• Figure 5: Segmentation results for various methods on sample images from the ImageCAS dataset zeng2023imagecas. The Mean Squared Error (MSE) between the output segmentation and the ground truth is shown for each method.