CONSIGN: Conformal Segmentation Informed by Spatial Groupings via Decomposition

TL;DR

CONSIGN introduces a conformal prediction method that leverages spatial groupings to produce more accurate and interpretable uncertainty estimates in image segmentation, especially for medical imaging.

Contribution

It presents a novel CP-based approach that incorporates spatial correlations, improving uncertainty quantification over existing methods.

Findings

Consistent improvement in uncertainty estimation metrics across datasets.

Enhanced interpretability of prediction sets with spatial structure consideration.

Applicable to any pre-trained segmentation model with multiple outputs.

Abstract

Most machine learning-based image segmentation models produce pixel-wise confidence scores that represent the model's predicted probability for each class label at every pixel. While this information can be particularly valuable in high-stakes domains such as medical imaging, these scores are heuristic in nature and do not constitute rigorous quantitative uncertainty estimates. Conformal prediction (CP) provides a principled framework for transforming heuristic confidence scores into statistically valid uncertainty estimates. However, applying CP directly to image segmentation ignores the spatial correlations between pixels, a fundamental characteristic of image data. This can result in overly conservative and less interpretable uncertainty estimates. To address this, we propose CONSIGN (Conformal Segmentation Informed by Spatial Groupings via Decomposition), a CP-based method that incorporates spatial correlations to improve uncertainty quantification in image segmentation. Our method generates meaningful prediction sets that come with user-specified, high-probability error guarantees. It is compatible with any pre-trained segmentation model capable of generating multiple sample outputs. We evaluate CONSIGN against two CP baselines across three medical imaging datasets and two COCO dataset subsets, using three different pre-trained segmentation models. Results demonstrate that accounting for spatial structure significantly improves performance across multiple metrics and enhances the quality of uncertainty estimates.