Conformal prediction (CP) provides a distribution-free framework for uncertainty quantification with finite-sample marginal coverage guarantees. Yet, missing values introduce distributional heterogeneity, where marginal coverage can hide systematic undercoverage for certain missingness patterns. At the same time, existing methods based on distances or smoothness require a continuous representation of partially observed covariates. Although imputation restores continuity, it obscures the missingness patterns and introduces task-dependent overhead. We develop a new CP framework that is adaptive to missingness patterns, without imputing covariates. We rely on data augmentation by feature embedding that accounts for missingness. Our framework represents each input with missing values by a tree embedding together with its raw missingness mask, and then learns a regularized quantile score function. Varying the choices for the penalty function and the augmentation enables us to characterize three practically relevant validity levels: exact coordinate-wise validity, exact validity on observed masks, and approximate missingness-conditional guarantees for general mechanisms (MCAR, MAR, MNAR). Experiments on synthetic and real datasets support these guarantees and show that the proposed method reduces miscoverage risk while maintaining informative intervals, especially in non-MCAR settings. These results highlight the practical value of uncertainty quantification that adapts to missingness.
Figures & tables
Mechanism
M1=1
M2=1
M3=1
M4=1
M5=1
M6=1
MCAR
92.3±2.1
92.2±1.6
91.7±2.0
92.4±1.1
91.5±2.2
92.1±2.2
MAR
92.8±1.1
92.7±1.9
92.6±1.8
N/A
N/A
N/A
PMIX
91.7±1.9
92.1±2.3
90.8±2.0
93.9±1.1
92.4±2.2
92.9±0.9
Table 1: Coordinate-wise coverage estimates with d=6 and nominal level 90% . Entries are mean percentage coverage ± 95% CI half-width over 10 seeds.
Figure 1: Empirical shortfall risk and average interval length across missingness mechanisms. Our method is competitive with the baselines on MCAR/MAR and gives a stronger coverage–length trade-off on the mask-heterogeneous PMIX setting. Here ntree=256 .
Table 3: Data splits used in the main experiments. The reported mask-wise metrics are computed on the pattern-stratified test set.
Setting
Coverage
Shortfall
Interval length
MCAR
0.9013→0.9162
0.0062→0.0031
7.505→8.120
MAR
0.9020→0.9120
0.0013→0.0002
5.776→6.008
PMIX
0.9068→0.9145
0.0234→0.0131
9.728→10.917
MIMIC-IV
0.9278→0.9225
0.0090→0.0152
107.882→114.814
Appendix
Table 4: Sensitivity to the tuning protocol (10-seed means). Each entry is calibration-based tuning → training-only tuning. Synthetic settings use d=6 , p=0.2 , and nominal coverage 0.9 .
Figure 6: Choice of the best λ over different numbers of trees and replications
Method
d=6
d=100
ARC-CP
0.0054
0.0079
MDA-Exact
0.0057
0.0073
Ours (the exact dual optimization)
1.0594
2.2980
Weighted CP
0.2316
0.4034
MDA-Nested
0.8275
Did not finish
Appendix
Table 5: Runtime per test point in seconds.
Calibration size
100
200
500
1000
2000
Time (s)
0.2454
0.3585
0.7130
1.0594
2.0533
Appendix
Table 6: Runtime scaling for our exact method on the d=6 synthetic setting.
Figure 7: Risk and average interval length for all m∈Mcal . For the mask observed in calibration data, we have a near-zero risk for miscoverage, and the average interval length is the smallest.
Figure 8: Risk and average interval length for all m∈/Mcal . For the mask not observed in calibration data, the risk is slightly bigger, and the interval length is also largely augmented, although it is still the shortest interval, showing that our method can be adapted to "easy" masks and "difficult" masks.
Figure 9: Empirical coverage over different embedding tree models.
Figure 10: Empirical interval length over different embedding tree models.
Figure 11: Empirical shortfall risk and average interval length for the high-dimensional synthetic experiment ( d=100 , p=0.2 ) under MCAR and PMIX mechanisms. Our method keeps shortfall risk near zero for both mechanisms while maintaining informative intervals.
Conformal prediction (CP) offers a principled framework for uncertainty quantification, but it fails to guarantee mask-conditional coverage when faced with missing covariates. In addressing the heterogeneity induced by various missing patterns, Mask-Conditional Valid (MCV) Coverage has emerged as a more desirable property than Marginal Coverage. In this work, we adapt split CP to handle missing values by proposing a preimpute-mask-then-correct framework that can offer valid coverage. We show that our method provides guaranteed Marginal Coverage and Mask-Conditional Validity for general missing data mechanisms. A key component of our approach is a reweighted conformal prediction procedure that corrects the prediction sets after distributional imputation (multiple imputation) of the calibration dataset, making our method compatible with standard imputation pipelines. We derive two algorithms and prove that they achieve both marginal validity and MCV. We evaluate them on synthetic and real-world datasets. It reduces significantly the width of prediction intervals w.r.t standard MCV methods, while maintaining the target guarantees.
Jiarong Fan, Juhyun Park. Thi Phuong Thuy Vo, Nicolas Brunel
LaMME, University of Paris-Saclay, Evry · ENSIIE, Evry · Capgemini Invent, Paris, France
In regression, conformal prediction often suffers from inefficiency under heteroscedasticity and skewness due to fixed, non-adaptive interval centering. We propose CoCP, a new method that parameterizes prediction intervals by a center m(x) and a radius h(x), and alternates between learning h(x) via smooth quantile regression on folded residuals and refining m(x) using a smooth interval loss. This corrects mis-centering and drives the interval toward high-density regions. Beyond finite-sample marginal validity via split-conformal calibration, we prove that CoCP also attains asymptotic conditional coverage and optimal interval length if the base estimators are consistent. Experiments demonstrate that CoCP yields tighter intervals and achieves improved empirical conditional-coverage reliability. Code is available at https://github.com/jjoeysu/cocp.
Xiaoyi Su, Zhixin Zhou, Rui Luo
Department of Systems Engineering, City University of Hong Kong, Hong Kong SAR, China · Alpha Benito Research, Los Angeles, USA
Uncertainty quantification is essential for deploying machine learning models in high-stakes domains such as scientific discovery and healthcare. Conformal Prediction (CP) provides finite-sample coverage guarantees under exchangeability, an assumption often violated in practice due to distribution shift. Under covariate shift, restoring validity requires importance weighting, yet accurate density-ratio estimation becomes unstable when training and test distributions exhibit limited support overlap. We propose KMM-CP, a conformal prediction framework based on Kernel Mean Matching (KMM) for covariate-shift correction. We show that KMM directly controls the bias-variance components governing conformal coverage error by minimizing RKHS moment discrepancy under explicit weight constraints, and establish asymptotic coverage guarantees under mild conditions. We then introduce a selective extension that identifies regions of reliable support overlap and restricts conformal correction to this subset, further improving stability in low-overlap regimes. Experiments on molecular property prediction benchmarks with realistic distribution shifts show that KMM-CP reduces coverage gap by over 50% compared to existing approaches. The code is available at https://github.com/siddharthal/KMM-CP.
Siddhartha Laghuvarapu, Rohan Deb, Jimeng Sun
Siebel School of Computing and Data Science, University of Illinois Urbana-Champaign, IL, USA.