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

Tianhao Fu1,4,5,6 Austin Wang†,2,5,6 Charles Chen†,1,5,6 Roby Aldave-Garza†,3,5,6 Yucheng Chen5,7
1 University of Toronto 2 McGill University 3 University of Waterloo 4 Vector Institute
5 Project Neura 6 UTMIST 7 Amplimit
† Equal contributions
arXiv 2604.15271 · April 2026
Paper arXiv Code
TL;DR. SegWithU wraps a frozen pretrained segmentation backbone with a lightweight uncertainty head. It models uncertainty as perturbation energy in a compact probe space using rank-1 posterior probes, yielding two voxel-wise maps, a calibration map for probability tempering and a ranking map for error detection, in a single forward pass.

Method

SegWithU treats the segmentation backbone as a fixed predictor and learns a small uncertainty head on tapped intermediate features. Uncertainty is parameterised as perturbation energy: rank-1 probes induce latent perturbations whose effect on output logits is the basis for the epistemic signal. Click any labeled box in the figure below to open the corresponding interactive 3D rendering.

SegWithU architecture diagram. Frozen backbone, multi-scale feature taps, fusion, probe response, delta-logits, residual, epistemic and aleatoric maps, and downstream calibration, anchor, ranking, weight, margin, and entropy maps.
each dashed box is clickable — opens the interactive 3D map

Overview of SegWithU. A frozen segmentation backbone produces logits and a probability map. Multi-scale intermediate features are tapped (purple) and fused into a probe response that yields a probe map. Rank-1 posterior probes induce class-logit perturbations (the residual map), whose variance gives the epistemic map. In parallel, an aleatoric branch reads the fused features. Auxiliary signals — entropy, margin, ambiguity weight — combine with these via lightweight 1×1 convolutions to produce two functionally distinct outputs: a calibration map that tempers logits via z̃ = z / √(1+Ucal), and an anchor → ranking map optimized for error ordering and selective prediction. Only the orange modules are learnable; the backbone is never modified.

BibTeX

@article{fu2026segwithu,
  author    = {Fu, Tianhao and Wang, Austin and Chen, Charles and
              Aldave-Garza, Roby and Chen, Yucheng},
  title     = {SegWithU: Uncertainty as Perturbation Energy for
              Single-Forward-Pass Risk-Aware Medical Image Segmentation},
  journal   = {arXiv preprint arXiv:2604.15271},
  year      = {2026}
}