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

arXiv: https://arxiv.org/abs/2601.18997 · HTML: https://mlsanity.com/paper/2601.18997/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 21 of its 30 references that are in the map.
- First posted on arXiv 26 Jan 2026.

## Builds on (most cited first, 21 of 21)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [What Uncertainties Do We Need in Bayesian Deep Learning for Computer Vision?](https://mlsanity.com/paper/1703.04977/) (2017)
- [A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges](https://mlsanity.com/paper/2011.06225/) (2020)
- [The Medical Segmentation Decathlon](https://mlsanity.com/paper/2106.05735/) (2021)
- [Medical SAM Adapter: Adapting Segment Anything Model for Medical Image Segmentation](https://mlsanity.com/paper/2304.12620/) (2023)
- [Bayesian Convolutional Neural Networks with Bernoulli Approximate Variational Inference](https://mlsanity.com/paper/1506.02158/) (2015)
- [Bayesian Uncertainty Estimation for Batch Normalized Deep Networks](https://mlsanity.com/paper/1802.06455/) (2018)
- [A review of uncertainty quantification in medical image analysis: probabilistic and non-probabilistic methods](https://mlsanity.com/paper/2310.06873/) (2023)
- [Predictive Inference with Feature Conformal Prediction](https://mlsanity.com/paper/2210.00173/) (2022)
- [Conformal Semantic Image Segmentation: Post-hoc Quantification of Predictive Uncertainty](https://mlsanity.com/paper/2405.05145/) (2024)
- [Deep Q Learning Driven CT Pancreas Segmentation with Geometry-Aware U-Net](https://mlsanity.com/paper/1904.09120/) (2019)
- [Conformal Prediction for Image Segmentation Using Morphological Prediction Sets](https://mlsanity.com/paper/2503.05618/) (2025)
- [ConformalSAM: Unlocking the Potential of Foundational Segmentation Models in Semi-Supervised Semantic Segmentation with Conformal Prediction](https://mlsanity.com/paper/2507.15803/) (2025)
- [Prompt learning with bounding box constraints for medical image segmentation](https://mlsanity.com/paper/2507.02743/) (2025)

## Abstract (arXiv)

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.