Mask-Conditional Conformal Prediction: Valid Uncertainty For All Missing Data Mechanisms

Jiarong Fan; Juhyun Park. Thi Phuong Thuy Vo; Nicolas Brunel

stat.MLDec 16, 2025

Mask-Conditional Conformal Prediction: Valid Uncertainty For All Missing Data Mechanisms

Authors: Jiarong Fan, Juhyun Park. Thi Phuong Thuy Vo, Nicolas Brunel

Organizations: LaMME, University of Paris-Saclay, Evry · ENSIIE, Evry · Capgemini Invent, Paris, France

Abstract

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.

Figures & tables

Figure 1: Illustration of the pre-imputation: the calibration data is imputed and then masked by the $M$ of test point such that they share the same pattern. In this way, the split CP can give a prediction set conditional on $M$ . But a distribution shift exists. Unlike CP-MDA, our method leverages all available calibration data without modifying the test point.

	1:	$\displaystyle\forall k\in\widetilde{Id}_{\tt Cal}\text{, set}\bar{M}^{k}=[\max(M^{k}_{i},M^{n+1}_{i}),i=1,...,d]$
	2:	$\displaystyle\widehat{C}_{n,\alpha}^{\text{Nested$ {}^{\star} $}}(\widetilde{X}^{n+1}(M^{n+1})):=\Big\{y:(1-\alpha)(1+\#\widetilde{Id}_{\tt Cal})$
		$\displaystyle>\sum_{k\in\widetilde{Id}_{\tt Cal}}\mathbf{1}\Big[s(\widetilde{X}^{k}(\bar{M}^{k}),Y^{k})<s(\widetilde{X}^{n+1}(\bar{M}^{k}),y)\Big]\Big\}$

Algorithm 1 CP-MDA-Nested ⋆ Zaffran et al. (2024)

Figure 2: Synthetic dataset on MCAR: prediction interval width across 1023 masks (mean and 95% CI).

Figure 3: Synthetic dataset on MCAR: prediction coverage across 1023 masks (mean and 95% CI).

Mechanism	MDA-Exact	MDA-Nested	Weighted CP	ARC CP
MAR	87.8–89.7% / 5.84	87.9–89.8% / 5.67	89.0–91.2% / 5.79	88.3–90.2% / 5.64
MNAR	87.8–89.8% / 5.74	87.8–89.8% / 5.64	89.1–91.0% / 5.79	88.5–90.3% / 5.60

Table 1: Worst-case coverage (95% CI) and average prediction interval width over 31 masks.

Bin of $L_{\text{oracle}}(m)$	ARC-CP cov	ARC-CP wid	WCP cov	WCP wid	Nested cov	Nested wid
$[17.74,20.39]$	0.911	27.92	0.884	27.14	0.926	29.58
$[20.39,21.16]$	0.910	27.99	0.885	27.20	0.927	29.76
$[21.16,21.57]$	0.910	28.04	0.885	27.16	0.928	29.90
$[21.57,22.02]$	0.911	28.08	0.881	27.08	0.927	29.91
$[22.02,22.40]$	0.911	28.20	0.882	27.17	0.928	30.04
$[22.40,22.77]$	0.911	28.14	0.884	27.21	0.927	30.17

Table 2: 100-dimensional synthetic setting: average mask-conditional coverage and interval width, grouped by oracle interval width

L_{\text{oracle}}(m)

. Results of coverage and width are averaged over 1000 masks with 50 replications.

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Synthetic dataset: mean and 99% confidence interval of prediction interval width with respect to 1023 missingness patterns, results by using deterministic imputation.

Figure 5: Synthetic dataset: mean and 99% confidence interval of prediction coverage with respect to 1023 missingness patterns, results by using deterministic imputation.

		$\displaystyle\widehat{C}^{W}_{\alpha}(\widetilde{X}^{n+1}):=\Big\{y:s_{\hat{f}}(\widetilde{X}^{n+1},y)$
		$\displaystyle\leq\mathcal{Q}_{1-\alpha}\bigg(\sum_{i=1}^{n}W^{i}(X^{i}_{\mathrm{obs}(M^{n+1})},y)\delta_{s_{\hat{f}}(\widetilde{X}^{i},Y^{i})}+W^{n+1}(X^{n+1}_{\mathrm{obs}(M^{n+1})},y)\delta_{\infty}\bigg)\Big\}.$		(4)

Appendix

Algorithm 3 Weighted CP with pre-imputation of Calibration

\displaystyle\widehat{C}^{AR}_{\alpha}(\widetilde{X}^{n+1}):=\bigg\{y\in\mathcal{Y}:s_{\hat{f}}(\widetilde{X}^{n+1},y)\leq\mathcal{Q}_{1-\alpha}\left(\sum_{i\in\mathcal{I}_{AR}}\delta_{\hat{S}^{i}_{M^{n+1}}}+\delta_{\infty}\right)\bigg\}.

(10)

Appendix

Algorithm 4 Acceptance-rejection corrected CP with pre-imputation of Calibration

Figure 6: Coverage vs. Correlation of estimated and true density ratios. Each point represents one experiment. High correlation yields near-nominal coverage.

Figure 7: Miscoverage vs estimated bound. Points below

y=x

confirm the bound’s conservativeness.

Method	Synthetic dataset	Concrete dataset
MDA-Exact	17034s	2031s
MDA-Nested	510932s	43225s
Weighted CP	228638s	37983s
ARC CP	21320s	4766s

Appendix

Table 3: Comparison of running time of 4 methods.

Method	Avg. Cov	Avg.Width
MDA-Exact	$92.9\pm 0.7$	$\infty$
MDA-Nested	$91.8\pm 0.5$	$336.3\pm 2.3$
Weighted CP	$91.9\pm 0.6$	$\mathbf{306.4\pm 2.2}$
ARC CP	$93.0\pm 0.6$	$\mathbf{309.9\pm 3.1}$

Appendix

Table 4: Comparison on Bike sharing dataset for different methods. The average coverage and interval width of 15 different missingness patterns are listed.

Method	Coverage	Interval width
MDA-Exact	$94.03-95.01$	$33.0$
MDA-Nested	$94.29-95.15$	$34.1$
ARC CP	$94.04-95.04$	$\mathbf{32.08}$

Appendix

Table 5: Comparison on Bike sharing dataset for different methods. The worst-case coverage and interval width of 32 different missingness patterns are listed.

Figure 8: Synthetic dataset: Comparison of prediction coverage (mean and 99% confidence interval) across 1023 missingness patterns, with and without the proposed correction method.

Figure 9: Synthetic dataset: Comparison of the width of prediction interval (mean and 99% confidence interval) across 1023 missingness patterns, with and without the proposed correction method.

Figure 10: Concrete dataset: mean and 99% confidence interval of prediction coverage with respect to 255 missingness patterns

Mech.	MDA-Exact	MDA-Nested	Weighted CP	ARC CP
MAR	89.1% $\pm$ 1.64% / 29.1	90.8% $\pm$ 1.43% / 31.4	89.1% $\pm$ 1.91% / 29.8	90.6% $\pm$ 1.42% / 28.9

Appendix

Table 6: Worst-case coverage (95% CI) and average prediction interval width over 15 masks, concrete dataset with 4 possible MAR missing covariates

Mech.	MDA-Exact	MDA-Nested	Weighted CP	ARC CP
MNAR	89.0% $\pm$ 1.78% / 30.1	88.7% $\pm$ 1.76% / 30.4	90.0% $\pm$ 1.67% / 30.8	89.0% $\pm$ 1.83% / 30.1

Appendix

Table 7: Worst-case coverage (95% CI) and average prediction interval width over 15 masks, concrete dataset with 4 possible missing covariates

Figure 11: Synthetic dataset with 20% missing values: mean and 99% confidence interval of prediction interval width with respect to 1023 missingness patterns, comparison for whether using correction method.

Figure 12: Synthetic dataset with 20% missing values: mean and 99% confidence interval of prediction coverage with respect to 1023 missingness patterns, comparison for whether using correction method.

Figure 13: Synthetic dataset with 20% missing values and independent features: mean and 99% confidence interval of prediction interval width with respect to 1023 missingness patterns, comparison for whether using correction method.

Figure 14: Synthetic dataset with 20% missing values and independent features: mean and 99% confidence interval of prediction coverage with respect to 1023 missingness patterns, comparison for whether using correction method.

Figure 15: Synthetic dataset: unsmoothed results for mean and 99% confidence interval of prediction width with respect to 1023 missingness patterns.

Figure 16: Synthetic dataset: unsmoothed results for mean and 99% confidence interval of prediction coverage with respect to 1023 missingness patterns.

Conformal Prediction

Conformal prediction wraps predictive models with calibrated sets or intervals that provide statistical coverage guarantees.

See all 364 papers

Beyond Imputation: Mask-Adaptive Conformal Prediction via Tree Embeddings on General Missing Data Mechanisms
Nov 1, 2026Jiarong Fan, Juhyun Park, Thi Phuong Thuy Vo +1Conformal Prediction Learning with Missing Data
LaMME Université Paris-Saclay · LaMME, ENSIIE Université Paris-Saclay · Capgemini Invent
Conformal Prediction Assessment: A Framework for Conditional Coverage Evaluation and Selection
Mar 28, 2026Zheng Zhou, Xiangfei Zhang, Chongguang Tao +1Conformal Prediction
Qiuzhen College, Tsinghua University · Yau Mathematical Sciences Center, Tsinghua University
Isotonic Conformal Inference
Jul 18, 2026Daniel Bensimon, Zhiquan Liu, Sean X. Yu +2Conformal Prediction Uncertainty Calibration
Department of Mathematics and Statistics, McGill University; Mila · Eli Lilly and Company · Corresponding author, Department of Mathematics and Statistics, McGill University; Mila

Nov 1, 2026J/K move · Enter open · S save

Beyond Imputation: Mask-Adaptive Conformal Prediction via Tree Embeddings on General Missing Data Mechanisms

Jiarong Fan, Juhyun Park, Thi Phuong Thuy Vo, Nicolas J-B. Brunel

LaMME Université Paris-Saclay · LaMME, ENSIIE Université Paris-Saclay · Capgemini Invent

Conformal Prediction Learning with Missing Data Regression Prediction Intervals

Open paper AskPDF arXiv

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.

Abstract

Conformal Prediction

Explore similar work

Beyond Imputation: Mask-Adaptive Conformal Prediction via Tree Embeddings on General Missing Data Mechanisms

Conformal Prediction Assessment: A Framework for Conditional Coverage Evaluation and Selection

Isotonic Conformal Inference