SegWithU: Uncertainty as Perturbation Energy for Single-Forward-Pass Risk-Aware Medical Image Segmentation

TL;DR

SegWithU introduces a lightweight, post-hoc uncertainty estimation method for medical image segmentation that enhances reliability without multiple inferences, outperforming existing single-pass methods across multiple datasets.

Contribution

It presents a novel perturbation energy-based uncertainty framework that leverages intermediate features and rank-1 posterior probes for improved single-pass uncertainty estimation.

Findings

Achieves high AUROC/AURC scores on ACDC, BraTS2024, and LiTS datasets.

Produces two voxel-wise uncertainty maps for calibration and error detection.

Maintains segmentation quality while providing reliable uncertainty estimates.

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.