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

TL;DR

KMM-CP is a conformal prediction method that uses Kernel Mean Matching to maintain coverage guarantees under covariate shift, with a selective extension for improved stability in low-overlap regions.

Contribution

It introduces a novel KMM-based conformal prediction framework with a selective region extension to handle covariate shift more reliably.

Findings

Reduces coverage gap by over 50% on molecular prediction benchmarks.

Controls bias-variance trade-off in conformal coverage error.

Provides asymptotic coverage guarantees under mild conditions.

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.