COMPASS: Robust Feature Conformal Prediction for Medical Segmentation Metrics

2025-09-26

AI Summary

In clinical practice, the utility of medical image segmentation models hinges on the accuracy of downstream scalar metrics—such as organ volume—rather than pixel-wise segmentation fidelity, necessitating reliable uncertainty quantification for such metrics. This work proposes COMPASS, a framework that perturbs intermediate deep features along a low-dimensional subspace most sensitive to the target metric, integrating conformal prediction with importance weighting to yield efficient and well-calibrated confidence intervals. Its core innovations are representation-space calibration and a covariate-shift-robust perturbation strategy. Evaluated on four medical segmentation tasks, COMPASS achieves significantly shorter confidence intervals while strictly satisfying marginal coverage guarantees. Moreover, it maintains the desired coverage level under distributional shift, demonstrating robustness to covariate drift.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Computer Vision: Segmentation
• Reasoning under Uncertainty: Uncertainty Representations

Application Category

• User Modeling, Personalization and Recommendation: Metrics for user behavior and evaluating success
• Security and Privacy: Large-scale security measurements
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

In clinical applications, the utility of segmentation models is often based on the accuracy of derived downstream metrics such as organ size, rather than by the pixel-level accuracy of the segmentation masks themselves. Thus, uncertainty quantification for such metrics is crucial for decision-making. Conformal prediction (CP) is a popular framework to derive such principled uncertainty guarantees, but applying CP naively to the final scalar metric is inefficient because it treats the complex, non-linear segmentation-to-metric pipeline as a black box. We introduce COMPASS, a practical framework that generates efficient, metric-based CP intervals for image segmentation models by leveraging the inductive biases of their underlying deep neural networks. COMPASS performs calibration directly in the model's representation space by perturbing intermediate features along low-dimensional subspaces maximally sensitive to the target metric. We prove that COMPASS achieves valid marginal coverage under exchangeability and nestedness assumptions. Empirically, we demonstrate that COMPASS produces significantly tighter intervals than traditional CP baselines on four medical image segmentation tasks for area estimation of skin lesions and anatomical structures. Furthermore, we show that leveraging learned internal features to estimate importance weights allows COMPASS to also recover target coverage under covariate shifts. COMPASS paves the way for practical, metric-based uncertainty quantification for medical image segmentation.

Problem

• Quantifying uncertainty for medical segmentation-derived metrics like organ size
• Improving efficiency of conformal prediction for segmentation-to-metric pipelines
• Generating tighter uncertainty intervals under covariate shifts in medical imaging

Innovation

• Perturbs features along metric-sensitive subspaces
• Calibrates directly in model representation space
• Uses learned features for covariate shift adaptation