KMM-CP: Practical Conformal Prediction under Covariate Shift via Selective Kernel Mean Matching

Siddhartha Laghuvarapu; Rohan Deb; Jimeng Sun

cs.LGMar 27, 2026

KMM-CP: Practical Conformal Prediction under Covariate Shift via Selective Kernel Mean Matching

Authors: Siddhartha Laghuvarapu, Rohan Deb, Jimeng Sun

Organizations: Siebel School of Computing and Data Science, University of Illinois Urbana-Champaign, IL, USA.

Abstract

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.

Figures & tables

Figure 1: Synthetic experiment: (a) Covariate shift with partial support overlap, the target includes regions poorly supported by the source. (b) Classifier-based density-ratio assigns heavy-tailed weights to low-support regions, reducing effective sample size (ESS). (c) Standard KMM enforces bounded weights and matches moments over the full domain, improving stability but still correcting unsupported regions. (d) Selective KMM jointly optimizes source weights and target selection variables, filtering unsupported regions and stabilizing weights. (e) Classifier weighting reduces ESS, KMM lowers moment discrepancy (MMD), and selective KMM achieves both stable ESS and improved moment matching.

\sum_{i=1}^{n}\tilde{w}_{i}\mathbf{1}\{S_{i}\leq t\}\geq 1-\alpha_{\mathrm{conf}}

Algorithm 1 Split Conformal Prediction with (Selective) KMM

Figure 2: Calibration curves under covariate shift (global calibration). The dashed line denotes nominal coverage. Uniform CP undercovers systematically, density-ratio baselines partially correct this, KMM improves alignment, and SKMM most closely tracks the target across datasets.

Dataset	Random	Shifted
AMES	0.0121	0.0844
BBB-Martins	0.0346	0.4579
HIV	0.0044	0.1830
Tox21	0.0249	0.4366
hERG-Karim	0.0254	0.0384

Table 1: Mean absolute deviation (MAD) of uniform conformal prediction under random (no-shift) and shifted splits.

	Uniform		KDE		CoDrug KDE		Logistic		KMM		SKMM
Dataset	G	M	G	M	G	M	G	M	G	M	G	M
AMES	0.166	0.084	0.137	0.046	0.093	0.090	0.102	0.096	0.078	0.079	0.065	0.072
Tox21	0.409	0.437	0.323	0.331	0.103	0.110	0.073	0.089	0.053	0.057	0.015	0.020
BBB-Martins	0.424	0.458	0.364	0.393	0.329	0.381	0.396	0.443	0.319	0.353	0.213	0.270
hERG	0.034	0.038	0.096	0.094	0.031	0.030	0.020	0.022	0.028	0.036	0.009	0.019
HIV	0.185	0.183	0.291	0.283	0.145	0.151	0.297	0.269	0.087	0.098	0.013	0.022

Table 2: Mean absolute deviation (MAD) from nominal coverage under covariate shift for global (G) and Mondrian (M) calibration. Lower is better. Best per column is bold ; second-best is gray bold . SKMM attains the lowest MAD under global calibration and the best overall mean rank across both settings.

Figure 3: Bias–variance proxy (

\mathrm{MMD}+\sqrt{1/\mathrm{ESS}}

) under covariate shift. MAD in brackets.

Figure 4: Selection weights under SKMM.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

\tilde{w}_{i}^{(c)}=\frac{w_{i}}{\sum_{k\in\mathcal{I}_{c}}w_{k}}\hskip 10.00002pt\text{for }i\in\mathcal{I}_{c}

Appendix

Algorithm 2 Mondrian Split Conformal Prediction with KMM / Selective KMM

Dataset	Total $N$	Train	Calib	Test
AMES	7,280	5,096	1,092	1,092
BBB-Martins	2,033	1,423	305	305
HIV	41,132	28,793	6,169	6,170
Tox21	7,267	5,087	1,090	1,090
hERG	13,445	9,411	2,017	2,017

Appendix

Table 3: Dataset statistics and data splits.

Figure 5: Calibration coverage under global conformal prediction across datasets.

Panel of Figure 6 — Figure 6: Figure. (a) illustrates the bias-variance proxy( $\mathrm{MMD}+\sqrt{1/\mathrm{ESS}}$ ) , while (b) displays the emergent sparsity in the weight distribution.

Dataset	Fraction Filtered
AMES	71.2 $\pm$ 5.9
Tox21	58.3 $\pm$ 0.5
BBB-Martins	60.1 $\pm$ 10.1
hERG-Karim	76.7 $\pm$ 5.9
HIV	62.2 $\pm$ 1.4

Appendix

Table 4: Average filtration ratio (%) under two-stage Selective KMM (mean

\pm

std).

Method	0.5	0.55	0.6	0.65	0.7	0.75	0.8	0.85	0.9	0.95	Avg. Rank
Uniform	0.455 $\pm$ 0.000	0.458 $\pm$ 0.004	0.475 $\pm$ 0.006	0.485 $\pm$ 0.004	0.498 $\pm$ 0.005	0.525 $\pm$ 0.014	0.574 $\pm$ 0.034	0.636 $\pm$ 0.040	0.698 $\pm$ 0.034	0.824 $\pm$ 0.048	4.60
KDE	0.774 $\pm$ 0.127	0.788 $\pm$ 0.113	0.799 $\pm$ 0.104	0.810 $\pm$ 0.092	0.833 $\pm$ 0.081	0.848 $\pm$ 0.073	0.866 $\pm$ 0.051	0.898 $\pm$ 0.046	0.915 $\pm$ 0.036	0.933 $\pm$ 0.022	4.80
CoDrugKDE	0.680 $\pm$ 0.041	0.716 $\pm$ 0.040	0.738 $\pm$ 0.042	0.767 $\pm$ 0.035	0.790 $\pm$ 0.034	0.823 $\pm$ 0.042	0.845 $\pm$ 0.033	0.874 $\pm$ 0.027	0.899 $\pm$ 0.033	0.933 $\pm$ 0.027	3.00
Logistic	0.684 $\pm$ 0.058	0.725 $\pm$ 0.068	0.755 $\pm$ 0.051	0.778 $\pm$ 0.037	0.807 $\pm$ 0.030	0.831 $\pm$ 0.036	0.858 $\pm$ 0.044	0.875 $\pm$ 0.028	0.898 $\pm$ 0.027	0.925 $\pm$ 0.030	4.10
KMM	0.664 $\pm$ 0.026	0.694 $\pm$ 0.025	0.726 $\pm$ 0.021	0.750 $\pm$ 0.012	0.775 $\pm$ 0.017	0.799 $\pm$ 0.021	0.822 $\pm$ 0.027	0.861 $\pm$ 0.017	0.885 $\pm$ 0.024	0.915 $\pm$ 0.019	2.60
SKMM	0.630 $\pm$ 0.028	0.669 $\pm$ 0.024	0.703 $\pm$ 0.019	0.734 $\pm$ 0.009	0.761 $\pm$ 0.010	0.791 $\pm$ 0.020	0.819 $\pm$ 0.029	0.862 $\pm$ 0.022	0.883 $\pm$ 0.025	0.927 $\pm$ 0.020	1.90

Appendix

Table 5: AMES — Global calibration (mean

\pm