PaperlayerView exact revision on arXiv

Source-linked AI summary

CARD: Calibration via Agreement in Reverse Diffusion for Out-of-Domain MRI Segmentation

Jiaheng Dai, Weidong Guo, Qingbiao Li, Jie Xu, Yi Guo, Yuanyuan Wang, Zeju Li

arXiv:2608.28681v1cs.CV

TL;DR

Under domain shift, terminal-prediction signals may no longer track segmentation error, motivating reliability evidence from the prediction process itself. CARD uses disagreement between primary and capacity-limited reference diffusion models across reverse trajectories to calibrate confidence without changing labels, lowering mean NLL across four evaluated shift blocks.

  • Problem

    Under domain shift, terminal-prediction signals can lose their relationship with segmentation error, so calibration needs test-time evidence that remains informative across cases and locations.

  • Method

    CARD aggregates disagreement between a primary and capacity-limited reference across shared reverse states, mapping it to a pixel-wise temperature field without updating the segmentor.

  • Results

    CARD lowers mean NLL to 0.405, 0.437, 0.565, and 0.548 across the ACDC-C, M&Ms, prostate, and brain blocks, respectively.

  • Takeaways & Limitations

    Trajectory disagreement remains informative across evaluated artifact, vendor, and site shifts, extending reliability evidence beyond the terminal prediction.

  • Takeaways & Limitations

    CARD assumes primary and reference errors are largely independent; shared errors can produce low disagreement and mask miscalibration.

Abstract

from arXiv · show

Probability calibration aligns model confidence with predictive accuracy, enabling clinicians to identify unreliable segmentation regions. This alignment breaks down under domain shift, where artifacts and unseen protocols produce confident errors. Existing post-hoc methods adapt the correction at test time, conditioning on predictive entropy, the logit pattern, or augmentation response, but each proxy is read from the terminal prediction, the very quantity that shift corrupts. This motivates reliability evidence beyond the terminal prediction, which categorical diffusion provides in two ways. First, a generative shape prior keeps a capacity-limited reference intact when appearance is corrupted, so its disagreement with the primary segmentor highlights primary-model errors. Second, every reverse step yields a class distribution, separating persistent disagreement from transient discrepancy. Aggregated over the trajectory, this disagreement correlates with Dice at 0.788, against 0.521 for a matched discriminative control. We therefore propose CARD (Calibration via Agreement in Reverse Diffusion), which maps the temporal aggregate of this disagreement to a temperature field applied per pixel across all classes, so that confidence changes while the segmentation does not. Across cardiac, prostate and brain MRI shifts, CARD lowers calibration error in 45 of 49 comparisons against the strongest baseline in each setting.

I. INTRODUCTION

CARD addresses overconfident, spatially varying errors under MRI domain shift by using trajectory-based reliability evidence rather than only terminal predictions. It calibrates diffusion segmentation confidence without changing labels and improves calibration across evaluated shifts.

  • MRI domain shifts can produce localized segmentation errors with high-confidence probabilities, making calibration important for reliable use.
  • Existing post-hoc methods use fixed corrections or evidence from terminal outputs, whose relationship to error may change under shift.
  • CARD compares a primary diffusion model with a capacity-limited reference across shared reverse states to obtain test-time reliability evidence.
  • Temporal Aggregation Calibration converts multi-step disagreement into a spatial, input-dependent temperature field without updating the segmentor.
  • 45 of 49 comparisons show lower calibration error for CARD than the strongest baseline across cardiac, prostate, and brain MRI shifts.

II. RELATED WORK

Prior calibration methods range from fixed post-hoc transformations to spatial, output-conditioned, and image-conditioned corrections. Their reliance on source-fitted or terminal/image cues can limit transfer across acquisition and appearance shifts.

  • Training-time calibration changes the optimization objective or data distribution, whereas post-hoc calibration adjusts outputs after training.
  • Temperature scaling applies one learned scalar uniformly, while ensemble temperature scaling increases correction flexibility without changing predicted class ordering.
  • Spatial calibration assigns different corrections to locations, but source-validation fitting may transfer poorly across scanners, protocols, and image appearance.
  • Output-conditioned methods derive temperature corrections from terminal logits or entropy, including spatial and prediction-specific variants.
  • Image-conditioned methods use augmentation statistics, shape residuals, and current-image information to estimate pixel-wise temperature maps.

C. Reliability Estimation for Segmentation

This section situates calibration within reliability estimation and introduces categorical diffusion as a segmentation model whose intermediate reverse states expose prediction evolution. The formulation uses pixel-wise class distributions and deterministic reverse sampling.

  • Uncertainty estimation measures prediction variability or ambiguity, whereas calibration evaluates agreement between class probabilities and empirical accuracy.
  • Reliability methods use image- or pixel-level estimators, including auxiliary references and reconstruction discrepancies, to identify likely segmentation errors.
  • Categorical diffusion produces a pixel-wise class distribution at each reverse denoising step, making prediction evolution observable before the terminal state.
  • The segmentation is represented as a K-class categorical distribution over image pixels, with each pixel-wise probability vector summing to one.
  • Inference starts from a fixed uniform categorical state and uses deterministic DDIM reverse sampling conditioned on the MRI and step index.
  • Temperature scaling divides pixel-wise logits by a positive temperature before softmax, enabling spatial confidence correction while preserving ordering.

B. Agreement-on-the-Line in Categorical Diffusion

CARD uses a capacity-limited diffusion reference and the ordered reverse trajectory to distinguish primary-model failures from stable anatomical structure. Persistent disagreement tracks segmentation quality more strongly than matched dual-U-Net agreement.

  • A capacity-limited reference is intended to preserve dominant anatomical structure while differing enough from the primary to make disagreement informative.
  • The reference is comparatively insensitive to artifact-altered high-frequency appearance because limited capacity retains coarse anatomical modes.
  • Under shift, the primary degrades locally while the reference remains structurally stable, so disagreement localizes primary evidence-driven failures.
  • Trajectory aggregation separates transient disagreement from disagreement that persists as reverse refinement repeatedly reconciles image evidence and anatomical plausibility.
  • The agreement score is bounded in [0, 1], with higher values indicating stronger agreement, while the spatial disagreement map remains available for calibration.
  • 0.788 Pearson correlation with ground-truth Dice exceeds 0.521 for the matched Dual U-Net control in motion corruption.

A. Reference-Guided Dual Diffusion

CARD compares a high-capacity primary denoiser with a capacity-limited reference on shared reverse states, aggregates their disagreement, and converts it into a bounded pixel-wise temperature field for confidence calibration.

  • Reference-Guided Dual Diffusion: The primary and reference denoisers process the same intermediate reverse states, but only the primary advances the sampled trajectory.The reference uses reduced architecture and truncated training to emphasize global anatomical structure over transient local detail.
  • Reference-Guided Dual Diffusion: JS divergence measures symmetric, bounded disagreement between the two denoisers at each sampled reverse step.With natural logarithms, the divergence is bounded by ln 2 across class counts and reverse steps.
  • Temporal Aggregation Calibration: Temporal Aggregation Calibration averages step-wise disagreement across sampled reverse states to reduce transient noise and ambiguity.The aggregate remains bounded by ln 2 because it averages bounded step-wise divergences.
  • Temporal Aggregation Calibration: The aggregate disagreement is mapped to a bounded spatial temperature field using a monotonically increasing logistic function.The mapping preserves disagreement ordering while constraining temperatures between τmin and τmax.
  • Pixel-wise Temperature Scaling: CARD applies one positive temperature to all terminal primary logits at each pixel, changing confidence without changing the predicted segmentation.Constraining τmin ≥1 permits attenuation but prevents confidence increases; mapping parameters are fitted with denoisers frozen.

V. EXPERIMENTAL PROTOCOL

The evaluation spans cardiac, prostate, and brain lesion segmentation, using source cohorts alongside held-out cohorts that represent synthetic artifacts or real scanner, vendor, and acquisition-site shifts.

  • Experimental Protocol: The protocol evaluates three segmentation tasks, each with a source cohort and one or more held-out shifted cohorts.The tasks are cardiac, prostate, and brain lesion segmentation.
  • Cardiac: ACDC-C supplies synthetic cardiac artifact shifts, while M&Ms provides real cross-vendor evaluation for models trained on ACDC.ACDC-C applies bias-field, motion, ghosting, and k-space spike corruptions while preserving anatomy and labels.
  • Prostate: The prostate evaluation uses PROSTATEx as the source cohort and PROMISE12 as a held-out cohort acquired across institutions with different scanners and acquisition parameters.The supplied protocol identifies these cohorts as the source and held-out settings for prostate segmentation.
  • Brain Lesion: The brain-lesion evaluation uses ATLAS, with scanner metadata defining 13 strata in the processed labeled cohort.SiemensTrioTim is identified as the source stratum in the supplied protocol.

B. Baselines

The evaluation compares conventional calibrators with backbone-matched controls and CARD across ROI-based calibration metrics, using source-validation fitting and held-out evaluation. CARD is evaluated as a trajectory-conditioned method against static, output-conditioned, image-conditioned, and terminal-logit controls.

  • Calibration parameters are fitted on source-validation splits, while held-out cohorts are reserved for evaluation.
  • Baselines: Conventional baselines comprise static TS and ETS, output-conditioned LTS, HTS, and PTS, and image-conditioned CalibNet.
  • Backbone-matched controls: Primary+ETS and Primary+HTS calibrate the same primary diffusion segmentor’s terminal logits, isolating the contribution of trajectory disagreement.
  • Calibration Metrics: Metrics include ECE, SCE, ACE, and NLL within a ground-truth-derived ROI, with lower values indicating better calibration.The ROI dilates the foreground mask with a 10-pixel kernel, and ten bins are used for the calibration metrics.
  • Trajectory-conditioned: CARD reports trajectory-conditioned calibration results alongside the baseline groups in Table I.

D. Implementation Details

CARD uses a time-conditioned SwinUNETR diffusion backbone with a fixed Mini reference and a short deterministic reverse trajectory. Training combines diffusion and clean-state objectives, while inference uses five sampled reverse states.

  • The primary segmentor uses a time-conditioned SwinUNETR Base configuration selected by source-validation Dice, while the Mini reference is fixed a priori at 25k steps.
  • Inference uses a five-step deterministic reverse sampler at indices [49, 37, 24, 12, 0] over a 50-step categorical diffusion horizon.
  • Training combines a variational diffusion loss with clean-state cross-entropy, Dice, and boundary terms, while image conditioning is randomly masked with probability 0.2.

VI. RESULTS AND DISCUSSION

Across cardiac, prostate, and brain MRI shifts, CARD improves calibration over conventional and backbone-matched baselines. A cardiac ghosting case further shows that trajectory disagreement localizes confident primary-model errors and guides spatial confidence correction.

  • Cross-task results: CARD lowers mean NLL across all four reported shift blocks: 0.405, 0.437, 0.565, and 0.548.The blocks cover controlled cardiac artifacts, cross-vendor cardiac shift, prostate transfer, and brain scanner shift.
  • Backbone-matched comparison: Against Primary+ETS, CARD reduces mean NLL from 0.559 to 0.405, 0.562 to 0.437, 1.199 to 0.565, and 0.974 to 0.548 across the four blocks.
  • Case study: In a ghosting-corrupted cardiac case, the disagreement map overlaps the primary segmentation error without using labels, and the temperature correction reduces confidence there.
  • Calibration metrics: CARD lowers ECE under spike corruption from 16.22% to 13.81% versus CalibNet and gives lower ECE, SCE, and ACE across prostate and brain groups.
  • Discussion: Calibration robustness under shift depends on whether the deployment-time signal remains aligned with the current error pattern.

B. Reliability Diagnostics Under Artifact Shift

Under artifact shift, trajectory disagreement localizes confident primary-model errors, while calibration is strongest with compact, undertrained references that preserve structure without reproducing the primary’s failures.

  • Reliability diagnostics: The strongest JS responses occur where primary and reference predictions diverge, overlapping the ground-truth error under ghosting corruption.CARD’s resulting temperature correction reduces confidence in that localized region.
  • Reliability diagnostics: Trajectory disagreement tracks segmentation quality both across cases and spatially, wherever local image evidence produces divergent predictions.This links the pixel-level artifact example to aggregate calibration behavior.
  • Reference capacity and training maturity: The lowest NLL generally occurs for Mini, Tiny, and Small references at early or intermediate checkpoints, while later Base and Large checkpoints often worsen calibration.Prostate has a broader low-NLL region, but the Mini reference remains competitive across all tasks.
  • Reference capacity and training maturity: Useful calibration requires a contrast regime: references that are too strong reproduce primary errors, whereas very weak references add disagreement from their own failures.Compact, undertrained references preserve anatomy while keeping predictions sufficiently distinct.
  • Reference capacity and training maturity: 1.74M additional parameters: the Mini reference increases model size modestly, while batch latency rises 15.4%–16.2% across the evaluated tasks.Reported throughputs are 15.88, 15.99, and 22.58 slices/s for ACDC, prostate, and ATLAS.

E. Component Ablation

Component ablations attribute CARD’s calibration gain to genuine primary-reference contrast and full reverse-trajectory aggregation, rather than logit fusion or terminal agreement alone.

  • Component ablation: 13.81% ECE: the full model outperforms Logit Mix at 19.62% and Dual U-Net at 19.22% on ACDC-C spike corruption.These controls test direct logit fusion and paired-model agreement, respectively.
  • Component ablation: 17.78% ECE: removing reference contrast worsens calibration compared with the full model’s 13.81% ECE.The ablation uses an identical denoiser on both paths.
  • Component ablation: The controls attribute the gain specifically to primary-reference model contrast and full reverse-trajectory aggregation, not logit fusion or paired-model agreement alone.This supports trajectory disagreement as the operative calibration signal.
  • Method and implication: CARD uses a post-hoc, test-time temperature field derived from primary-reference disagreement across shared reverse states without updating the segmentor.The broader conclusion frames reverse trajectories as reliability evidence beyond terminal predictions under domain shift.
Loading 2608.28681v1…