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

2026-04-16

AI Summary

Current medical image segmentation methods struggle to simultaneously achieve computational efficiency and reliable uncertainty estimation within a single forward pass, limiting the trustworthiness of clinical decisions. To address this, we propose SegWithU, a lightweight post-processing framework that attaches an uncertainty head to a frozen, pretrained segmentation backbone. By leveraging intermediate features to construct a compact probe space, our method introduces perturbation energy modeling into single-pass inference for the first time, enabling voxel-wise uncertainty estimation without repeated model evaluations via rank-1 posterior probing. This approach jointly optimizes probability calibration and error detection while preserving original segmentation accuracy. SegWithU achieves state-of-the-art single-pass performance on ACDC, BraTS2024, and LiTS, with AUROC/AURC scores of 0.9838/2.4885, 0.9946/0.2660, and 0.9925/0.8193, respectively.

Technology Category

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

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• medical image segmentation
• uncertainty estimation
• single-forward-pass
• risk-aware
• error detection

Innovation

• uncertainty estimation
• perturbation energy
• single-forward-pass
• medical image segmentation
• rank-1 posterior probes