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

Siddhartha Laghuvarapu, Rohan Deb, Jimeng Sun
Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, PMLR 337:3237-3261, 2026.

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v337-laghuvarapu26a, title = {KMM-{CP}: Practical Conformal Prediction under Covariate Shift via Selective Kernel Mean Matching}, author = {Laghuvarapu, Siddhartha and Deb, Rohan and Sun, Jimeng}, booktitle = {Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence}, pages = {3237--3261}, year = {2026}, editor = {Perković, Emilija and Malinsky, Daniel}, volume = {337}, series = {Proceedings of Machine Learning Research}, month = {17--21 Aug}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v337/main/assets/laghuvarapu26a/laghuvarapu26a.pdf}, url = {https://proceedings.mlr.press/v337/laghuvarapu26a.html}, 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.} }
Endnote
%0 Conference Paper %T KMM-CP: Practical Conformal Prediction under Covariate Shift via Selective Kernel Mean Matching %A Siddhartha Laghuvarapu %A Rohan Deb %A Jimeng Sun %B Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence %C Proceedings of Machine Learning Research %D 2026 %E Emilija Perković %E Daniel Malinsky %F pmlr-v337-laghuvarapu26a %I PMLR %P 3237--3261 %U https://proceedings.mlr.press/v337/laghuvarapu26a.html %V 337 %X 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.
APA
Laghuvarapu, S., Deb, R. & Sun, J.. (2026). KMM-CP: Practical Conformal Prediction under Covariate Shift via Selective Kernel Mean Matching. Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, in Proceedings of Machine Learning Research 337:3237-3261 Available from https://proceedings.mlr.press/v337/laghuvarapu26a.html.

Related Material