Contrastive uncertainty learning for robust multi-structure segmentation in lower-pelvic MR
Authors: Bilal Sidiqi, Yipei Wang, Wen Yan, Shaheer U. Saeed, Daniel C. Alexander, Yipeng Hu
Organizations: UCL Hawkes Institute, University College London, London, UK · Department of Medical Physics and Biomedical Engineering, University College London, London, UK · Centre for Bioengineering, School of Engineering and Materials Science, Queen Mary University of London, London, United Kingdom · Digital Environment Research Institute, Queen Mary University of London, London, United Kingdom · Department of Computer Science, University College London, London, UK
Boundary delineation of anatomical structures in prostate MRI is often compromised by regional ambiguity. Automated segmentation methods relying on global overlap loss functions tend to be insensitive to surface deviations, leading to significant errors in ambiguous regions. Most existing uncertainty-aware methods treat uncertainty as a post-hoc quality indicator or prediction filter. We propose Contrastive Uncertainty Learning (CUL), a supervised framework that uses predictive uncertainty to directly guide representation learning. Using Monte Carlo Dropout, CUL identifies voxel-wise uncertainty and structures the embedding geometry accordingly: low-uncertainty, correctly predicted features form stable voxels while high-uncertainty or misclassified features are actively repelled, preventing feature collapse at difficult boundaries. We evaluate CUL on a 589-scan 3D prostate MRI multi-structure segmentation task (8 anatomical structures) using a ViT with UperNet decoder architecture. Average Dice (83.38%) remains comparable to the baseline (83.35%), while surface distance decreases significantly (p<0.001) across 6 of 8 anatomical structures. A hard-case analysis reveals a 22× asymmetry between significantly improved and degraded cases (18.8% vs 0.9%). Ablation studies show that removing uncertainty-guided sample selection eliminates improvements for the transition zone, while replacing the contrastive objective with a boundary loss additionally removes bladder improvements, demonstrating that both uncertainty-guided sampling and the contrastive objective are essential for the observed gains. These findings demonstrate that the proposed uncertainty-guided representation learning is a useful mechanism for reducing the risk of extreme errors in safety-critical medical segmentation
Figures & tables
Figure 1: Overview of the proposed Contrastive Uncertainty Learning (CUL) framework. Prostate MRI is processed by a segmentation network to produce voxel wise predictions and an associated uncertainty map. Voxel embeddings are partitioned into low and high uncertainty regions and used to construct contrastive pairs. Low-uncertainty same-class features are pulled together, while high-uncertainty or misclassified features are repelled, improving boundary robustness.
Metric
Baseline
CUL
Class-Based
Classic BL
Dice ( ↑ )
0.8335±0.0647
0.8338±0.0651
0.8331±0.0637
0.8338±0.0634
HD95(mm) ( ↓ )
5.9343±8.1414
5.4130±4.8829
5.3736±4.6558
5.6253±5.2306
SD(mm) ( ↓ )
1.6076±2.5101
1.4190±1.8135
1.4385±1.8873
1.4641±1.9188
Table 1: Overall segmentation performance for Baseline, CUL, and the two ablations. SD(mm) is significantly improved which is consistent with CUL’s design: predictive uncertainty concentrates at local boundary regions, so improvement manifests in surface-sensitive metrics rather than in global overlap (Dice) or gross deviation extremes (HD95(mm)).
Metric
Improved / Degraded (%)
Sig+ / Sig- (%)
Ratio
Dice ( ↑ )
64 / 53 (54.7 / 45.3)
–
–
HD95(mm) ( ↓ )
66 / 51 (56.4 / 43.6)
12 / 8 (10.3 / 6.8)
1.5 ×
SD(mm) ( ↓ )
79 / 38 (67.5 / 32.5)
22 / 1 (18.8 / 0.9)
22 ×
Table 2: Case-wise performance distribution (n=117). Significant improvement/degradation thresholds: SD ± 0.1 mm, HD95 ± 0.5 mm.
Subset
n
Δ HD95 (mm)
Easy ( < 4.0032)
58
≈0
Hard ( ≥ 4.0032)
59
−1.1131
Table 3: HD95(mm) change stratified by baseline difficulty (median=4.0032).
Figure 2: Qualitative comparison between an improved case (top) and a degraded case (bottom). The top row highlights CUL’s role as geometric “insurance”, where it effectively corrects extreme boundary outliers in cases characterised by high regional ambiguity (e.g. central gland and both neurovascular bundles). CUL organises representations to stabilise boundaries by anchoring confident regions and isolating ambiguity. This visualisation demonstrates the framework’s capacity to correct significant spatial deviations while maintaining the integrity of stable predictions.
Medical image segmentation delineates organs, lesions, and other structures in scans to support measurement, diagnosis, treatment planning, and research.
Uncertainty quantification complements model predictions by characterizing their reliability, which is essential for high-stakes decision making such as medical image segmentation. However, most existing methods reduce uncertainty to a scalar confidence estimate, leaving its spatial distribution semantically underconstrained. In this work, we focus on uncertainty interpretability, namely, whether estimated uncertainty behaves in a human-understandable manner with respect to sources of ambiguity. We identify three perception-aligned principles requiring the spatial distribution of uncertainty to reflect: (1) image contrast between structures, (2) severity of image corruption, and (3) geometric complexity in anatomical structures. Accordingly, we develop a principle-guided uncertainty supervision framework (PriUS) based on evidential learning, in which the corresponding supervision objectives are explicitly enforced during training. We further introduce quantitative metrics to measure the consistency between predicted uncertainty and image attributes that induce ambiguity. Experiments on ACDC, ISIC, and WHS datasets showed that, compared with state-of-the-art methods, PriUS produced more consistent uncertainty estimates while maintaining competitive segmentation performance.
An Sui, Yuzhu Li, Gunter Schumann +2
School of Data Science, Fudan University, Shanghai, China · Institute of Science and Technology for Brain-Inspired Intelligence, Fudan University, Shanghai, China · National Heart and Lung Institute, Imperial College London, London, United Kingdom
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 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.
Tianhao Fu, Austin Wang, Charles Chen +2
University of Toronto, Toronto, ON, Canada · Vector Institute, Toronto, ON, Canada · Project Neura, Toronto, ON, Canada +4
Medical image segmentation models often report high benchmark accuracy under ideal imaging conditions, yet their failures under clinical degradation can be quiet: sensor noise, patient motion, low- resolution acquisition, and contrast variability may all alter model behavior without producing an obvious warning. We present a reproducible framework for evaluating uncertainty-aware segmentation under con- trolled clinical degradation. Our experiments use a synthetic multimodal brain tumor MRI cohort generated with a biophysical phantom simulator that follows the BraTS protocol. We train U-Net and Attention U-Net baselines for multi-class tumor sub-region segmentation and augment both models with Monte Carlo dropout to estimate per-voxel uncertainty. Across eight clinically motivated corruption types at five severity levels, we measure segmentation accuracy, calibration, failure detection, and selective prediction coverage. On clean data, Attention U-Net achieves a whole-tumor Dice of 0.990; under severe Gaussian noise, its performance falls to 0.089. Predictive uncertainty rises with degradation and tracks segmentation error (Pearson r = 0.53 under severity-3 Gaussian noise), allowing us to flag failures with an AUROC of 0.843. These results argue for uncertainty-aware inference as a practical safety layer in physician-in-the-loop radiology workflows. We release the code, trained models, and evaluation protocol to support direct reproduction.
Pranav Kaliaperumal, Manisha Kaliaperumal
Department of Computer Science University of Colorado Denver Aurora, CO, USA · Department of Biology (Pre-Medical) Creighton University Omaha, NE, USA