Challenge Summary U-MedSAM: Uncertainty-aware MedSAM for Medical Image Segmentation

2024-08-03

AI Summary

Quantifying prediction uncertainty remains challenging in medical image segmentation. To address this, we propose U-MedSAM: a framework built upon MedSAM that integrates an uncertainty-aware triple loss—comprising region-level, distribution-level, and pixel-level components—with an automatic loss-weighting mechanism. Furthermore, we adopt Sharpness-Aware Minimization (SharpMin) as the optimizer to guide training toward flat minima, thereby improving generalization and robustness. Evaluated on the CVPR 2024 MedSAM on Laptop Challenge, U-MedSAM achieves state-of-the-art performance, significantly enhancing segmentation accuracy, confidence calibration quality, and cross-domain transferability. Our key contributions are: (i) the first uncertainty-driven, multi-granularity joint loss design for medical segmentation; and (ii) the pioneering successful application of SharpMin in lightweight medical image segmentation.

Technology Category

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

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Medical Image Foundation Models have proven to be powerful tools for mask prediction across various datasets. However, accurately assessing the uncertainty of their predictions remains a significant challenge. To address this, we propose a new model, U-MedSAM, which integrates the MedSAM model with an uncertainty-aware loss function and the Sharpness-Aware Minimization (SharpMin) optimizer. The uncertainty-aware loss function automatically combines region-based, distribution-based, and pixel-based loss designs to enhance segmentation accuracy and robustness. SharpMin improves generalization by finding flat minima in the loss landscape, thereby reducing overfitting. Our method was evaluated in the CVPR24 MedSAM on Laptop challenge, where U-MedSAM demonstrated promising performance.

Problem

• Medical Image Analysis
• Prediction Uncertainty
• Accuracy Improvement

Innovation

• U-MedSAM
• Medical Image Segmentation
• Prediction Uncertainty