# An Uncertainty-Aware Generalization Framework for Cardiovascular Image Segmentation

arXiv: https://arxiv.org/abs/2409.14305 · HTML: https://mlsanity.com/paper/2409.14305/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Placed in Mamba for medical image analysis, one of 10 directions in State-space sequence models (S4, Mamba).
- Builds on 45 of its 68 references that are in the map.
- First posted on arXiv 22 Sep 2024.

## Placed in

- State-space sequence models (S4, Mamba) / [Mamba for medical image analysis](https://mlsanity.com/ssm/mamba-for-medical-image-analysis/)

## Builds on (most cited first, 30 of 45)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [TransUNet: Transformers Make Strong Encoders for Medical Image Segmentation](https://mlsanity.com/paper/2102.04306/) (2021)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation](https://mlsanity.com/paper/2105.05537/) (2021)
- [A Survey on Deep Learning in Medical Image Analysis](https://mlsanity.com/paper/1702.05747/) (2017)
- [Visualizing the Loss Landscape of Neural Nets](https://mlsanity.com/paper/1712.09913/) (2017)
- [A Survey of Uncertainty in Deep Neural Networks](https://mlsanity.com/paper/2107.03342/) (2021)
- [Generalised Dice overlap as a deep learning loss function for highly unbalanced segmentations](https://mlsanity.com/paper/1707.03237/) (2017)
- [Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers](https://mlsanity.com/paper/2110.13985/) (2021)
- [U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation](https://mlsanity.com/paper/2401.04722/) (2024)
- [A survey of loss functions for semantic segmentation](https://mlsanity.com/paper/2006.14822/) (2020)
- [Metrics reloaded: Recommendations for image analysis validation](https://mlsanity.com/paper/2206.01653/) (2022)
- [nnFormer: Interleaved Transformer for Volumetric Segmentation](https://mlsanity.com/paper/2109.03201/) (2021)
- [Deep learning for cardiac image segmentation: A review](https://mlsanity.com/paper/1911.03723/) (2019)
- [Deep Learning in Medical Image Registration: A Review](https://mlsanity.com/paper/1912.12318/) (2019)
- [Improving Generalization in Federated Learning by Seeking Flat Minima](https://mlsanity.com/paper/2203.11834/) (2022)
- [Uncertainty Aware Semi-Supervised Learning on Graph Data](https://mlsanity.com/paper/2010.12783/) (2020)
- [A review: Deep learning for medical image segmentation using multi-modality fusion](https://mlsanity.com/paper/2004.10664/) (2020)
- [Uncertainty-aware multi-view co-training for semi-supervised medical image segmentation and domain adaptation](https://mlsanity.com/paper/2006.16806/) (2020)
- [Preserving Fairness Generalization in Deepfake Detection](https://mlsanity.com/paper/2402.17229/) (2024)
- [How Does Sharpness-Aware Minimization Minimize Sharpness?](https://mlsanity.com/paper/2211.05729/) (2022)
- [Multi-Atlas Segmentation of Biomedical Images: A Survey](https://mlsanity.com/paper/1412.3421/) (2014)
- [ImageCAS: A Large-Scale Dataset and Benchmark for Coronary Artery Segmentation based on Computed Tomography Angiography Images](https://mlsanity.com/paper/2211.01607/) (2022)
- [Friendly Sharpness-Aware Minimization](https://mlsanity.com/paper/2403.12350/) (2024)

## Built on by (most cited first, 3 of 3)

- [LLM-MedQA: Enhancing Medical Question Answering through Case Studies in Large Language Models](https://mlsanity.com/paper/2501.05464/) (2024)
- [MedChat: A Multi-Agent Framework for Multimodal Diagnosis with Large Language Models](https://mlsanity.com/paper/2506.07400/) (2025)
- [Evidence-Based Actor-Verifier Reasoning for Echocardiographic Agents](https://mlsanity.com/paper/2604.06347/) (2026)

## Abstract (arXiv)

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.