Anatomically-aware conformal prediction for medical image segmentation with random walks

TL;DR

This work tackles the challenge of uncertainty quantification in medical image segmentation by combining Conformal Prediction with a spatial diffusion mechanism. RW-CP diffuses probability maps through a kNN graph built from foundation-model embeddings, producing anatomically coherent prediction sets that preserve marginal coverage via conformal risk control. The approach improves segmentation quality (higher DSC, lower ASSD/HD95) across ultrasound, MRI, and CT datasets while maintaining statistically valid guarantees. The method offers a model-agnostic enhancement to existing segmentation pipelines, bridging rigorous uncertainty quantification with clinically meaningful anatomically plausible predictions.

Abstract

The reliable deployment of deep learning in medical imaging requires uncertainty quantification that provides rigorous error guarantees while remaining anatomically meaningful. Conformal prediction (CP) is a powerful distribution-free framework for constructing statistically valid prediction intervals. However, standard applications in segmentation often ignore anatomical context, resulting in fragmented, spatially incoherent, and over-segmented prediction sets that limit clinical utility. To bridge this gap, this paper proposes Random-Walk Conformal Prediction (RW-CP), a model-agnostic framework which can be added on top of any segmentation method. RW-CP enforces spatial coherence to generate anatomically valid sets. Our method constructs a k-nearest neighbour graph from pre-trained vision foundation model features and applies a random walk to diffuse uncertainty. The random walk diffusion regularizes the non-conformity scores, making the prediction sets less sensitive to the conformal calibration parameter $λ$, ensuring more stable and continuous anatomical boundaries. RW-CP maintains rigorous marginal coverage while significantly improving segmentation quality. Evaluations on multi-modal public datasets show improvements of up to $35.4\%$ compared to standard CP baselines, given an allowable error rate of $α=0.1$.

Figures (3)

• Figure 1: Examples of predicted masks (red) and ground-truths (blue), and associated probability maps (lighter means higher probability). From left to right, (a) the original predicted mask with no conformal guarantees, (b--d) the conformalized prediction masks with $\alpha=0.1$. Specifically, b) Consema applies a fixed number of dilatations on the original prediction mossina_conformal_2025, c) Standard CRC thresholds the model's raw output probabilities angelopoulos_ConformalRiskControl_2024, and (d), our method RW-CP, thresholds the probabilities after diffusion with a random walk. Columns 4 and 6 show the probability map used, where lighter colour indicates a value close to 1 and darker colour indicates a value close to 0. Our method is able to successfully diffuse the probabilities and uncertainties, making the conformalized prediction masks closer to the ground truth, despite initially incorrect predictions.
• Figure 2: Mean dice score for confidence $(1\,{-}\,\alpha) \in [0.8, 0.95]$, given a calibration set of 20 samples, for the (a) ACDC-RV, (b) ACDC-LV, (c) CAMUS and (d) MSD-Spleen datasets. Our method, RW-CP (red), outperforms other CP methods (blue) across different confidence values.
• Figure 3: Impact of the number of nearest neighbours (1, 5, 10, 20, 50, 80 and 100) used to compute the random walk transition matrix on the model performance and inference time per sample. Considering $k\,{=}\,20$ neighbours for each pixel yields a high dice score while maintaining low inference time.