Source-linked AI summary
Contrastive uncertainty learning for robust multi-structure segmentation in lower-pelvic MR
Bilal Sidiqi, Yipei Wang, Wen Yan, Shaheer U. Saeed, Daniel C. Alexander, Yipeng Hu
TL;DR
Boundary-ambiguous prostate MRI cases can produce large surface deviations that global overlap metrics do not fully reflect, and direct uncertainty-guided representation learning remains underexplored in supervised multi-structure segmentation. CUL uses uncertainty-guided contrastive learning to structure voxel embeddings; it improves surface distance across structures while maintaining comparable global overlap performance.
Problem
Boundary-ambiguous prostate MRI cases can produce large boundary deviations that global overlap metrics do not fully reflect, while direct uncertainty-guided representation learning in fully supervised multi-structure segmentation remains underexplored.
Source: §1 · p. 2Method
CUL uses predictive uncertainty as a representation-level signal, forming contrastive anchors from low-uncertainty correct predictions and negatives from high-uncertainty or misclassified voxels.
Source: §1 · p. 2, Uncertainty Guided Sampling and Pair… · p. 4Results
Surface distance improves significantly across 6 of 8 anatomical structures (p < 0.001), with a 22× asymmetry between significantly improved and degraded hard cases.
Source: §1 · p. 3Takeaways & Limitations
CUL maintains comparable global overlap while improving boundary-sensitive metrics, with strongest benefits on hard cases characterized by large boundary deviations.
Source: §4 · p. 7
Abstract
from arXiv · showhide
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{\times}$ 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
1 Introduction
Boundary-ambiguous prostate MRI cases can produce large surface deviations that global overlap metrics do not fully reflect, and supervised learning that directly uses uncertainty to structure representations remains underexplored. CUL uses uncertainty-guided contrastive learning to improve surface distance across anatomical structures.
1 Introduction
- Boundary-ambiguous cases can produce large deviations that global overlap metrics do not fully reflect, while standard Dice and cross-entropy losses do not explicitly incorporate confidence.Source: §1 · p. 2
- Prior uncertainty-aware approaches use uncertainty for quality assurance, pseudo-label filtering, loss weighting, or confident-region prototypes; many contrastive segmentation methods do not account for predictive confidence.Source: §1 · p. 2
- Directly using uncertainty to structure embeddings in fully supervised segmentation remains underexplored, particularly for intrinsically ambiguous multi-structure MRI segmentation.Source: §1 · p. 2
- CUL uses predictive uncertainty as a representation-level supervisory signal, aligning low-uncertainty same-class features while preventing ambiguous features from collapsing into confident clusters.The framework forms contrastive pairs from uncertainty-partitioned voxel embeddings, pulling low-uncertainty same-class features together and repelling high-uncertainty or misclassified features.Source: §1 · p. 2, Figure 1 · p. 3
- 6 of 8 anatomical structures show significant surface-distance improvement (p < 0.001), with a 22× asymmetry between significantly improved and degraded hard cases.The reported ablations show 5/8 significant structures without uncertainty guidance and 4/8 when the contrastive term is replaced by boundary loss.Source: §1 · p. 3
2 Method
The method combines a 3D ViT encoder and UperNet decoder with uncertainty-guided contrastive training. It uses voxel-wise predictive uncertainty to select anchors and negatives, then adds a contrastive term to the segmentation objective.
2.1 Model Architecture
- The 3D model uses a masked-autoencoding-pretrained ViT encoder and UperNet decoder, with Dice plus cross-entropy as the baseline segmentation loss.Source: §2.1 · p. 3
- Voxel-wise predictive uncertainty is estimated with five stochastic decoder passes and normalized per sample.The uncertainty map is based on predictive entropy.Source: Predictive Uncertainty Estimation · p. 4
2.2 Contrastive Uncertainty Learning (CUL)
- CUL projects decoder features into a contrastive embedding space and uses a fixed uncertainty threshold and prediction errors to form anchor and negative pools.Anchors are low-uncertainty, correctly predicted voxels; negatives include high-uncertainty or misclassified voxels and are not restricted to other classes.Source: Uncertainty Guided Sampling and Pair… · p. 4, §2.2 · p. 4
- For each anchor, the voxel-wise contrastive term compares a same-class confident positive with sampled negatives, and the full objective adds this term to DiceCE with a learnable weight.Source: Uncertainty-Aware Contrastive Loss · p. 4
3 Experiments and Results
CUL maintains baseline-level Dice while improving boundary-sensitive metrics, with significant surface-distance reductions across most structures. Case-wise analyses show the strongest gains in hard cases, and ablations distinguish contributions from uncertainty guidance and contrastive learning.
Overall Quantitative Performance
- CUL matched baseline Dice (0.8338 vs 0.8335) while lowering mean HD95 (5.93 to 5.41 mm) and mean SD (1.61 to 1.42 mm).Table 1 summarizes performance across the baseline, CUL, and two ablations; the SD reduction is statistically significant.Source: Table 1 · p. 5, §3 · p. 5
- Paired Wilcoxon testing found SD reduction significant (p<0.0001), whereas Dice (p<0.181) and HD95 (p<0.816) median shifts were not.The reduced HD95 standard deviation (8.14 to 4.88 mm) is descriptive; Brown-Forsythe testing did not confirm a variance difference (p=0.538).Source: §3 · p. 5
Case-Wise Performance Distribution and Hard-Case Robustness Analysis
- Significant SD improvements outnumbered degradations 18.8% to 0.9%; hard cases had a mean HD95 reduction of 1.113 mm, versus marginal change in easy cases.Case-wise significance used thresholds of ±0.1 mm for SD and ±0.5 mm for HD95; difficulty strata used the median baseline HD95 threshold.Source: §3 · p. 6, Table 2 · p. 6, Table 3 · p. 6
Structure-wise Analysis
- CUL significantly improved SD in 6/8 structures—bladder, bone, central gland, rectum, transition zone, and neurovascular bundle; obturator internus and seminal vesicle showed no significant change.Dice remained largely unchanged across structures (|∆|<0.004), while HD95 was significant only for rectum.Source: §3 · p. 6, Table 4 · p. 7
Ablation Study
- Removing uncertainty guidance reduced significant SD improvements to 5/8 structures, while replacing contrastive learning with boundary loss reduced them to 4/8.The class-based ablation failed to improve transition zone; the boundary-loss ablation additionally failed to improve bladder.Source: §3 · p. 7
4 Discussion and Conclusion
The paper presents CUL as a supervised framework that uses predictive uncertainty as a representation-level signal for multi-structure MRI segmentation. Experiments report comparable global overlap with improved boundary metrics, especially on hard cases, and connect surface-distance precision to radiotherapy planning.
4 Discussion and Conclusion
- CUL estimates predictive uncertainty as a representation-level supervisory signal to improve geometric robustness in multi-structure MRI segmentation.The framework organizes representations around uncertainty rather than using uncertainty only as a post-hoc reliability estimate.Source: §4 · p. 7
- CUL maintains comparable global overlap while improving boundary-sensitive metrics, with the strongest benefits in hard cases characterized by large boundary deviations.Comparison with two ablations identifies uncertainty guidance as a key driver of transition-zone improvements.Source: §4 · p. 7
- Surface-distance errors can affect radiotherapy planning target volume margins, dose coverage, and organ-at-risk sparing.The discussion also notes precision needs in targeted biopsy and focal therapy workflows.Source: §4 · pp. 7–8
- The qualitative comparison pairs an improved case with a degraded case and highlights correction of extreme boundary outliers in ambiguous regions.Examples include the central gland and both neurovascular bundles.Source: Figure 2 · p. 9