SegWithU — Uncertainty as Perturbation Energy

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.

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+U_cal), and an anchor → ranking map optimized for error ordering and selective prediction. Only the orange modules are learnable; the backbone is never modified.

Interactive Viewer

Per-class logit volumes produced by the frozen segmentation backbone, rendered as fully interactive Plotly volumes. Click and drag to orbit, scroll to zoom, and hover to read raw logit values.

Class 0 — Backgroundlogit volume

Class 1 — Foregroundlogit volume

Pipeline

A walk through the five stages of SegWithU on a single case — from input volume through tapped features, probe responses, logits, and the resulting entropy map. Use the arrows to flip between two examples.

Input Volume

Raw 3D input fed to the frozen backbone.

Feature Map

Tapped multi-scale feature representation.

Probe Response

Rank-1 posterior probe activations v = ψ(h).

-->

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}
}