SegWithU: Uncertainty as Perturbation Energy for Single-Forward-Pass Risk-Aware Medical Image Segmentation
For medical image segmentation practitioners needing reliable uncertainty without repeated inference, SegWithU provides a practical and effective single-forward-pass solution.
SegWithU introduces a post-hoc framework that augments a frozen pretrained segmentation backbone with a lightweight uncertainty head, modeling uncertainty as perturbation energy via rank-1 posterior probes. It achieves state-of-the-art single-forward-pass uncertainty estimation on ACDC, BraTS2024, and LiTS, with AUROC/AURC of 0.9838/2.4885, 0.9946/0.2660, and 0.9925/0.8193 respectively.
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.