SegWithU: Uncertainty as Perturbation Energy for Single-Forward-Pass Risk-Aware Medical Image Segmentation
Tianhao Fu, Austin Wang, Charles Chen, Roby Aldave-Garza, Yucheng Chen
TLDR
SegWithU provides efficient, single-pass uncertainty estimation for medical image segmentation by modeling uncertainty as perturbation energy, improving reliability.
Key contributions
- Introduces SegWithU, a post-hoc framework for single-forward-pass uncertainty estimation in medical segmentation.
- Models uncertainty as perturbation energy using rank-1 posterior probes in a compact probe space.
- Generates two distinct voxel-wise uncertainty maps: one for calibration and one for error ranking.
- Achieves state-of-the-art uncertainty performance on ACDC, BraTS2024, and LiTS datasets.
Why it matters
This paper addresses the critical need for reliable uncertainty estimation in medical image segmentation without sacrificing efficiency. SegWithU offers a practical, high-performing solution that enhances clinical decision support by providing robust error detection and calibration. Its single-pass nature makes it suitable for real-world applications.
Original Abstract
Reliable uncertainty estimation is critical for medical image segmentation, where automated contours feed downstream quantification and clinical decision support. Many strong uncertainty methods require repeated inference, while efficient single-forward-pass alternatives often provide weaker failure ranking or rely on restrictive feature-space assumptions. We present $\textbf{SegWithU}$, a post-hoc framework that augments a frozen pretrained segmentation backbone with a lightweight uncertainty head. SegWithU taps intermediate backbone features and models uncertainty as perturbation energy in a compact probe space using rank-1 posterior probes. It produces two voxel-wise uncertainty maps: a calibration-oriented map for probability tempering and a ranking-oriented map for error detection and selective prediction. Across ACDC, BraTS2024, and LiTS, SegWithU is the strongest and most consistent single-forward-pass baseline, achieving AUROC/AURC of $0.9838/2.4885$, $0.9946/0.2660$, and $0.9925/0.8193$, respectively, while preserving segmentation quality. These results suggest that perturbation-based uncertainty modeling is an effective and practical route to reliability-aware medical segmentation. Source code is available at https://github.com/ProjectNeura/SegWithU.
📬 Weekly AI Paper Digest
Get the top 10 AI/ML arXiv papers from the week — summarized, scored, and delivered to your inbox every Monday.