Paper deep dive
KMM-CP: Practical Conformal Prediction under Covariate Shift via Selective Kernel Mean Matching
Siddhartha Laghuvarapu, Rohan Deb, Jimeng Sun
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 98%
Last extracted: 3/31/2026, 1:55:23 AM
Summary
KMM-CP is a conformal prediction framework designed to handle covariate shift by using Kernel Mean Matching (KMM) to align calibration and test distributions in an RKHS. By minimizing MMD under explicit weight constraints, it avoids unstable density-ratio estimation and improves coverage stability. The paper further introduces a selective extension that jointly optimizes source weights and target selection variables to focus on regions of reliable support overlap, significantly reducing coverage gaps in high-dimensional settings like molecular property prediction.
Entities (5)
Relation Signals (3)
KMM-CP → addresses → Covariate Shift
confidence 100% · KMM-CP, a conformal prediction framework based on Kernel Mean Matching (KMM) for covariate-shift correction.
KMM-CP → uses → Kernel Mean Matching
confidence 100% · We propose KMM-CP, a conformal prediction framework based on Kernel Mean Matching (KMM)
Kernel Mean Matching → minimizes → Maximum Mean Discrepancy
confidence 95% · KMM aligns the weighted calibration distribution... by minimizing the Maximum Mean Discrepancy (MMD)
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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 this https URL.
Tags
Links
- Source: https://arxiv.org/abs/2603.26415v1
- Canonical: https://arxiv.org/abs/2603.26415v1
Trouble viewing inline? Open PDF directly →
Full Text
89,634 characters extracted from source content.
Expand or collapse full text
KMM-CP: Practical Conformal Prediction under Covariate Shift via Selective Kernel Mean Matching Siddhartha Laghuvarapu Siebel School of Computing and Data Science, University of Illinois Urbana-Champaign, IL, USA. Rohan Deb Siebel School of Computing and Data Science, University of Illinois Urbana-Champaign, IL, USA. Jimeng Sun 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 1 Introduction Reliable uncertainty quantification is essential for deploying machine learning models in high-stakes domains such as scientific discovery and healthcare. Conformal Prediction (CP) vovk2005 provides a principled framework for constructing prediction sets with finite-sample marginal coverage guarantees under minimal distributional assumptions. In standard split conformal prediction, a trained model is evaluated on a held-out calibration set, whose prediction residuals are used to determine a quantile threshold for forming prediction sets. These guarantees rely on exchangeability between calibration and test data—an assumption that is frequently violated in real-world settings due to distribution shift tibshirani2019, barber2023conformal. Extensions of conformal prediction to non-exchangeable or distribution-shift settings have been studied in recent work. A common and practically relevant form of distribution shift is covariate shift, where the marginal distribution of inputs changes while the conditional distribution Y∣XY X remains stable sugiyama2007covariate. This setting is especially relevant in scientific discovery, where the data-generating process is consistent under comparable experimental conditions laghuvarapu2023codrug, fannjiang2022conformal. Under covariate shift, weighted conformal prediction restores validity by reweighting calibration samples using the density ratio between test and calibration covariate distributions tibshirani2019. When the ratio is known or accurately estimated, coverage guarantees are recovered. In practice, density ratio estimation can be unstable, particularly in moderate or high-dimensional settings and under limited support overlap. Ratios may be estimated via probabilistic classification bickel2009discriminative, direct density-ratio fitting methods such as KLIEP or LSIF sugiyama2008kliep, kanamori2009least, or kernel density estimation silverman1986density. Classifier-based methods depend on well-calibrated probability estimates and can yield extreme weights when the test distribution assigns mass to regions poorly represented in calibration data, while kernel density approaches suffer from the curse of dimensionality. In both cases, heavy-tailed or highly concentrated importance weights reduce the effective sample size (ESS), leading to unstable quantile estimation and degraded coverage. In this work, we propose KMM-CP, a framework for conformal prediction under covariate shift based on Kernel Mean Matching (KMM) huang2007correcting. Rather than explicitly estimating density ratios, KMM aligns the weighted calibration distribution with the test covariate distribution in a reproducing kernel Hilbert space (RKHS) by minimizing the Maximum Mean Discrepancy (MMD) gretton2012kernel. The resulting weights are obtained via a constrained optimization that enforces boundedness, thereby controlling variance and improving stability in high-dimensional settings. Although KMM-based reweighting has appeared in recent work almeida2025high, its connection to conformal coverage has not been systematically analyzed. In particular, prior work does not explicitly relate moment-matching quality and effective sample size to coverage behavior, nor compare KMM with alternative density-ratio estimators from a stability perspective. We further introduce a selective extension of KMM for low-overlap or disjoint-support settings. When test regions lack calibration support, enforcing global moment matching can induce extreme weights and collapse the ESS. Our formulation jointly optimizes calibration weights and target selection variables, restricting correction to regions of shared support. This improves the bias–variance tradeoff and limits uncertainty quantification to regimes where shift correction is reliable. Empirically, we evaluate KMM-CP on molecular property prediction tasks with substantial covariate shift. In high-dimensional settings, our method consistently reduces coverage gap and improves efficiency relative to classifier- and kernel density–based baselines, demonstrating robustness in low-overlap regimes. Our main contributions are as follows. Main Contributions 1. We propose KMM-CP, a practical framework for conformal prediction under covariate shift that uses Kernel Mean Matching (KMM) to align calibration and test distributions via kernel mean embeddings, improving stability without explicit density estimation. 2. We introduce a selective extension for low-overlap regimes that restricts correction to regions of shared support, improving effective sample size and calibration reliability. 3. We provide analysis connecting moment-matching quality and effective sample size to conformal coverage under distributional shift, explaining the improved stability of KMM-based reweighting. 4. We evaluate KMM-CP on molecular property prediction benchmarks with substantial covariate shift, achieving over 50% reduction in coverage gap compared to existing baselines. 2 Preliminaries 2.1 Conformal Prediction Framework Conformal Prediction (CP) is a framework for constructing prediction sets with finite-sample coverage guarantees. We consider a classification setup where each data point is Z=(X,Y)Z=(X,Y) with X∈ℝdX ^d and Y∈[K]=1,…,KY∈[K]=\1,…,K\. Beyond a point estimate from a base classifier f, we want a confidence level encoded as a prediction set C^(X)⊆[K] C(X) [K]. The main goal is valid coverage: for a target level α∈(0,1)α∈(0,1) (e.g. 0.1), the set should contain the true label with at least 1−α1-α probability. Formally, for a new test point (XN+1,YN+1)(X_N+1,Y_N+1), C C is (1−α)(1-α)-valid if ℙYN+1∈C^(XN+1)≥1−α.P \Y_N+1∈ C(X_N+1) \≥ 1-α. (1) Split conformal prediction. In practice, the coverage guarantee in (1) is typically attained using split (inductive) conformal prediction, which partitions the data into a training set and a calibration set DcalD_cal. A base predictor is fit on the training set, and conformity scores on the calibration set are used to choose a threshold so that the resulting set-valued predictor attains the target coverage. Under exchangeability of the calibration and test points, the procedure is distribution-free and achieves the validity condition (1); see vovk2005. 2.2 Conformal Prediction under Covariate Shift Standard conformal prediction assumes calibration and test data are i.i.d. (or exchangeable), an assumption that is rarely realistic in practice. In many scientific settings, however, while the marginal distribution of covariates X may shift between calibration and test time, the conditional distribution Y∣XY X remains stable. For example, molecular properties such as target activity or physicochemical characteristics are governed by underlying physical laws and typically remain consistent under comparable experimental conditions. This motivates the covariate shift setting, where the marginal law of X changes but the conditional law of Y∣XY X is invariant. Formally, let Dcal=(Xi,Yi)i=1ND_cal=\(X_i,Y_i)\_i=1^N denote the calibration sample and (XN+1,YN+1)(X_N+1,Y_N+1) the test point. Under covariate shift, PY∣Xcal=PY∣XtestP^cal_Y X=P^test_Y X, so the joint laws decompose as Pcal(dx,dy) P^cal(dx,dy) =PXcal(dx)PY∣Xcal(dy∣x), =P_X^cal(dx)\,P_Y X^cal(dy x), (2) Ptest(dx,dy) P^test(dx,dy) =PXtest(dx)PY∣Xcal(dy∣x). =P_X^test(dx)\,P_Y X^cal(dy x). (3) The calibration data satisfy (Xi,Yi)∼i.i.d.Pcal(X_i,Y_i) .i.d. P^cal, while (XN+1,YN+1)∼Ptest(X_N+1,Y_N+1) P^test. Importance weighting under covariate shift. Under covariate shift, one typically assumes PXtest≪PXcalP_X^test P_X^cal, so the likelihood ratio w(x):=dPXtestdPXcal(x)w(x):= dP_X^testdP_X^cal(x) is well-defined. Tibshirani et al. [tibshirani2019] view this as a special case of weighted exchangeability: independent (but not identically distributed) samples are weighted exchangeable with weights given by appropriate Radon–Nikodym derivatives [tibshirani2019, Lemma 2]. In particular, if Z1,…,ZN∼PcalZ_1,…,Z_N P^cal and ZN+1∼PtestZ_N+1 P^test, then one may take wi≡1w_i≡ 1 for i≤Ni≤ N and wN+1(x,y)=w(x)w_N+1(x,y)=w(x). Weighted conformal prediction then forms a cutoff as a weighted quantile of the calibration scores, where the weight w(x)w(x) of the candidate test point enters through the weighted conformal rank/quantile construction [tibshirani2019, Lemma 3 and Theorem 2]. Using this weighted cutoff (instead of the usual unweighted cutoff) yields prediction sets with the desired marginal coverage under PtestP^test. Theorem 1 (Coverage under covariate shift [tibshirani2019, Theorem 2]). Assume the covariate shift model (2)–(3) and PXtest≪PXcalP_X^test P_X^cal. Let w(x)=dPXtestdPXcal(x)w(x)= dP_X^testdP_X^cal(x), and construct the weighted conformal prediction set C C using calibration nonconformity scores and the weighted-quantile rule of tibshirani2019 with weights wi≡1w_i≡ 1 for i≤Ni≤ N and wN+1(x,y)=w(x)w_N+1(x,y)=w(x). Then ℙtestYN+1∈C^(XN+1)≥1−α,P_test\Y_N+1∈ C(X_N+1)\≥ 1-α, where the probability is over (XN+1,YN+1)∼Ptest(X_N+1,Y_N+1) P^test (and the randomness in the calibration sample). Remark 1. The guarantee above is exact when w is known. If w is replaced by an estimate w w (e.g., learned from unlabeled test covariates), finite-sample distribution-free coverage is no longer guaranteed and instead depends on the accuracy of w w [tibshirani2019]. 2.3 Density Ratio Estimation In practice, estimating the density ratio under covariate shift is non-trivial. Let psp_s and ptp_t denote the source (calibration) and target (test) marginal densities of X, respectively, and define the density ratio w(x)=pt(x)ps(x).w(x)= p_t(x)p_s(x). A common strategy is classifier-based density-ratio estimation, where a probabilistic classifier distinguishes source (D=0D=0) from target (D=1D=1) samples. Let η(x)=ℙ(D=1∣X=x)η(x)=P(D=1 X=x) and πs,πt _s, _t denote class priors. Then w(x)=πsπtη(x)1−η(x).w(x)= _s _t η(x)1-η(x). Another approach is kernel density estimation (KDE), which separately estimates psp_s and ptp_t via p^s(x)=1ns∑i=1nsKh(x−xi(s)),p^t(x)=1nt∑i=1ntKh(x−xi(t)), p_s(x)= 1n_s _i=1^n_sK_h(x-x_i^(s)), p_t(x)= 1n_t _i=1^n_tK_h(x-x_i^(t)), and forms w^(x)=p^t(x)/p^s(x) w(x)= p_t(x)/ p_s(x). Classifier-based methods can be sensitive to probability calibration and may produce high-variance importance weights under limited support overlap [shimodaira2000, sugiyama2012], while KDE-based approaches suffer from the curse of dimensionality in moderate or high dimensions [wand1995, silverman1986]. We instead employ Kernel Mean Matching (KMM), which estimates importance weights by directly matching kernel mean embeddings of the weighted source and target distributions, avoiding explicit density estimation [huang2007]. 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. 3 Kernel Mean Matching 3.1 Kernel Mean Matching under Covariate Shift Let xii=1n∼PS\x_i\_i=1^n P_S and zjj=1m∼PT\z_j\_j=1^m P_T denote source and target samples, related by a covariate shift. Kernel Mean Matching (KMM) [NIPS2006_a2186a7] seeks to correct this shift by constructing nonnegative weights w1,…,wnw_1,…,w_n such that the reweighted empirical source measure PSw:=1n∑i=1nwiδxiP_S^w:= 1n _i=1^nw_i\, _x_i (4) approximates (after normalization) PTP_T. Let ℋH be a Reproducing Kernel Hilbert Space (RKHS) with a characteristic kernel k and feature map φ . For probability measures Q and R on ℝdR^d, define the associated MMD by MMD(Q,R):=sup‖h‖ℋ≤1|X∼Q[h(X)]−X∼R[h(X)]|. (Q,R):= _\|h\|_H≤ 1 |E_X Q[h(X)]-E_X R[h(X)] |. KMM solves the empirical Maximum Mean Discrepancy (MMD) minimization problem: minw‖1n∑i=1nwiφ(xi)−1m∑j=1mφ(zj)‖ℋ2 _w \| 1n _i=1^nw_i (x_i)- 1m _j=1^m (z_j) \|_H^2 (5) subject to the following for some ϵ<1ε<1 0≤wi≤B,|1n∑i=1nwi−1|≤ϵ.0≤ w_i≤ B, | 1n _i=1^nw_i-1 |≤ε. (6) 3.2 Coverage Guarantees under KMM Reweighting Let s:ℝd×[K]→ℝs:R^d×[K] be a conformity score. For t∈ℝt define the conditional score CDF gt(x):=ℙ(s(x,Y)≤t|X=x), g_t(x):=P (s(x,Y)≤ t\, |\,X=x ), where (X,Y)∼PS(X,Y) P_S (equivalently (X,Y)∼PT(X,Y) P_T) since PY∣XS=PY∣XTP^S_Y X=P^T_Y X. The target (population) CDF of the score is FT(t) F_T(t) :=ℙ(X,Y)∼PT(s(X,Y)≤t)=X∼PXT[gt(X)]. :=P_(X,Y) P_T (s(X,Y)≤ t )=E_X P^T_X [g_t(X) ]. Given weights wiw_i from KMM, define w~i:=wi∑j=1nwj,∑i=1nw~i=1, w_i:= w_i _j=1^nw_j, _i=1^n w_i=1, and the corresponding weighted empirical covariate measure P^S,Xw:=∑i=1nw~iδxi. P^\,w_S,X:= _i=1^n w_i\, _x_i. This weighted empirical covariate measure induces a reweighted (population) score CDF Fw(t):=X∼P^S,Xw[gt(X)]=∑i=1nw~igt(xi). F_w(t):=E_X P^\,w_S,X [g_t(X) ]= _i=1^n w_i\,g_t(x_i). The first step toward coverage analysis is to control the population discrepancy supt|Fw(t)−FT(t)| _t|F_w(t)-F_T(t)| by the MMD discrepancy between the covariate measures. Theorem 2 (CDF stability under MMD control). Assume that for each t∈ℝt there exists g~t∈ℋ g_t such that supt∈ℝ‖g~t‖ℋ≤Bℋ,supt∈ℝ‖gt−g~t‖∞≤ε∞. _t \| g_t\|_H≤ B_H, _t \|g_t- g_t\|_∞≤ _∞. Then, for any probability measure Q on ℝdR^d, supt∈ℝ|X∼Q[gt(X)] _t |E_X Q[g_t(X)] −X∼PXT[gt(X)]| -E_X P^T_X[g_t(X)] | ≤BℋMMD(Q,PXT)+2ε∞. ≤ B_H\,MMD(Q,P^T_X)+2\, _∞. In particular, with Q=P^S,XwQ= P^\,w_S,X we obtain supt∈ℝ|Fw(t)−FT(t)|≤BℋMMD(P^S,Xw,PXT)+2ε∞. _t |F_w(t)-F_T(t)|≤ B_H\,MMD( P^\,w_S,X,P^T_X)+2\, _∞. Using this we can show the following. Theorem 3 (Coverage error for KMM-weighted split conformal). Let (X1,Y1),…,(Xn,Yn)∼i.i.d.PS(X_1,Y_1),…,(X_n,Y_n) .i.d. P_S be a calibration sample, let Z1,…,Zm∼i.i.d.PXTZ_1,…,Z_m .i.d. P_X^T denote unlabeled target covariates used to compute the KMM weights, and let (Xn+1,Yn+1)∼PT(X_n+1,Y_n+1) P_T be an independent test point. Define F^w(t):=∑i=1nw~i 1s(Xi,Yi)≤t, F_w(t):= _i=1^n w_i\,1\s(X_i,Y_i)≤ t\, and ESS:=1/∑i=1nw~i2ESS:=1/ _i=1^n w_i^2. Let the weighted (1−α)(1-α)-quantile be q^1−α:=inft∈ℝ:F^w(t)≥1−α, q_1-α:= \t : F_w(t)≥ 1-α \, and define the set-valued predictor C^(x):=y∈[K]:s(x,y)≤q^1−α. C(x):= \y∈[K]:s(x,y)≤ q_1-α \. Assume that with probability 11 the calibration scores s(X1,Y1),…,s(Xn,Yn)s(X_1,Y_1),…,s(X_n,Y_n) are all distinct and that the conditions of Theorem 2 hold with constants BℋB_H and ε∞ _∞. Then for any δ∈(0,1)δ∈(0,1), conditional on X1,…,Xn,Z1,…,ZmX_1,…,X_n,Z_1,…,Z_m, with probability at least 1−δ1-δ over Y1,…,YnY_1,…,Y_n we have |ℙ(X,Y)∼PT(Y∈C^(X))−(1−α)| |P_(X,Y) P_T (Y∈ C(X) )-(1-α) | ≤BℋMMDℋ(P^S,Xw,PXT)+2ε∞+(5+12log1δ)1ESS. \!\!≤ B_H\,MMD_H( P^\,w_S,X,P_X^T)+2\, _∞\!+\! (5+ 12 1δ ) 1ESS. Corollary 1 (Asymptotic validity of KMM-weighted split conformal). Assume the conditions of Theorem 3. Allow the approximation error in Theorem 2 to depend on (n,m)(n,m), and write it as ε∞,n,m _∞,n,m. If, as n,m→∞n,m→∞, MMDℋ(P^S,Xw,PXT)→0,ESS→∞,ε∞,n,m→0, _H ( P^\,w_S,X,P_X^T )→ 0,ESS→∞, _∞,n,m→ 0, in probability, then the following holds in probability ℙ(X,Y)∼PT(Y∈C^(X))→1−α _(X,Y) P_T (Y∈ C(X) )→ 1-α Theorem 3 reveals a clear bias–variance decomposition: the MMD term controls systematic distributional mismatch between the reweighted calibration distribution and the target distribution, while the 1/ESS1/ ESS term captures variance induced by weight concentration. Accurate covariate-shift correction therefore requires both small moment discrepancy and stable (non-degenerate) weights. Kernel Mean Matching is well-suited to this objective. Unlike classifier- or KDE-based density-ratio estimation, which can yield heavy-tailed or highly concentrated weights under limited support overlap, KMM minimizes an RKHS discrepancy under explicit box and mass constraints. These constraints control weight variance, preserve effective sample size, and stabilize quantile estimation. As a result, KMM directly targets the two terms governing coverage in our bound—moment alignment and ESS—without explicit density estimation. We further analyze these stability considerations in relation to other density estimation procedures in Appendix C and empirically validate them. 3.3 Selective Kernel Mean Matching via Joint Target Selection When the target distribution PTP_T is not well-supported by the source distribution PSP_S, the standard density ratio diverges. Forcing the KMM optimizer to match moments over the entire target domain requires extreme source weights, which collapses the Effective Sample Size (ESS) and inflates the variance of our coverage bounds (Theorem 3). To address this, we incorporate a selective mechanism that restricts conformal correction to regions of shared support. Specifically, we introduce test-instance selection weights αj∈[0,1] _j∈[0,1] applied to the target samples zjz_j, alongside the standard source weights wi∈[0,B]w_i∈[0,B] applied to xix_i. We jointly optimize the source weights and the target selection variables to match the selected target moments. This formulation is related to double-weighting strategies for covariate shift adaptation segovia2023double, though our objective and constraints differ and are tailored to conformal calibration rather than risk minimization. minw,α‖1n∑i=1nwiφ(xi)−1m∑j=1mαjφ(zj)‖ℋ2 _w,α \| 1n _i=1^nw_i (x_i)- 1m _j=1^m _j (z_j) \|_H^2 (7) subject to the constraints: 0≤wi≤B, 0≤ w_i≤ B, |1n∑i=1nwi−1m∑j=1mαj|≤ϵ, | 1n _i=1^nw_i- 1m _j=1^m _j |≤ε, (8) 0≤αj≤1, 0≤ _j≤ 1, 1m∑j=1mαj≥τ. 1m _j=1^m _j≥τ. Here, τ∈(0,1)τ∈(0,1) enforces a acceptance yield (the fraction of target points we wish to make predictions for). The KKT conditions suggest that when gradients clearly favor inclusion or exclusion, many selection variables αj _j may approach the boundaries 0 or 11 (see Appendix B). Empirically, we observe such sparsity: the learned αj _j often concentrate near the extremes, particularly for target points poorly supported by the source distribution. This induced sparsity directly controls the two terms in the coverage error bound: 1. Bias Reduction (MMD Control): By restricting inference to the selected subset, moment matching is concentrated on regions of shared support. In these regions, bounded source weights suffice to approximate the selected target moments, reducing the effective MMD between PSwP_S^w and induced target distribution P~T P_T. 2. Variance Reduction (ESS↑ESS ): Excluding unsupported target points prevents excessive weight concentration on a few samples, yielding a more uniform weight distribution and larger Effective Sample Size (ESS), thereby stabilizing the C21/ESSC_2 1/ESS term. 4 Kernel Mean Matching Based Conformal Prediction We describe the practical implementation of split conformal prediction under covariate shift using KMM and its selective extension, focusing on classification. We also discuss class-conditional calibration; details are provided in Appendix D. KMM-Based Split Conformal Prediction. Given labeled data D and unlabeled test covariates zjj=1mz_j_j=1^m, we split D into training and calibration sets, train a predictor on the training split, and compute conformity scores on the calibration data. KMM is then applied between calibration and test covariates to obtain importance weights, which are used to compute a weighted empirical quantile of the calibration scores. Selective KMM. The double-weighted formulation in 3.3 induces sparsity in the selection variables αj _j via complementary slackness, concentrating mass near the boundaries. In finite samples, the solutions cluster near 0 and 11 but are not exactly discrete. We therefore adopt a two-stage procedure: first obtain αj _j from the double-weighted objective, then threshold the learned α values to retain sufficiently large-α test points, and finally rerun standard KMM on the retained subset. Coverage applies conditionally to the selected test instances under the assumptions of Section 3.1. The unified procedure is summarized in Algorithm 1. Algorithm 1 Split Conformal Prediction with (Selective) KMM 0: Labeled data =(xi,yi)i=1ND=\(x_i,y_i)\_i=1^N, test covariates zjj=1m\z_j\_j=1^m, miscoverage level αconf _conf, kernel k, weight bound B, threshold θ (optional) 0: Prediction sets C^(zj) C(z_j) 1: Split D into training and calibration sets cal=(xi,yi)i=1nD_cal=\(x_i,y_i)\_i=1^n 2: Train predictor f and define conformity score S(x,y)S(x,y) 3: Compute calibration scores Si=S(xi,yi)S_i=S(x_i,y_i) 4: if Selective KMM then 5: Solve double-weighted KMM to obtain αj _j 6: sel←zj:αj≥θT_sel←\z_j: _j≥θ\ (or top τmτ m) 7: else 8: sel←zjj=1mT_sel←\z_j\_j=1^m 9: end if 10: Solve KMM between calD_cal and selT_sel to obtain wiw_i 11: Normalize w~i=wi/∑k=1nwk w_i=w_i/ _k=1^nw_k 12: Compute weighted quantile q^1−αconf q_1- _conf satisfying ∑i=1nw~iSi≤t≥1−αconf _i=1^n w_i1\S_i≤ t\≥ 1- _conf 13: for zj∈selz_j _sel do 14: C^(zj)=y:S(zj,y)≤q^1−αconf C(z_j)=\y:S(z_j,y)≤ q_1- _conf\ 15: end for Connection to Selective Conformal Inference. Our formulation is conceptually related to selective conformal inference angelopoulos2021learn, which guarantees coverage on a data-dependent subset of test instances. Classical selective methods typically assume exchangeability and use predictive confidence to define this subset. In contrast, our selection mechanism is driven by support mismatch under covariate shift: guarantees are restricted to regions with stable reweighting and reliable density-ratio estimation. Selective guarantees are particularly relevant in high-stakes domains. In molecular discovery, conformal prediction is used to prioritize active compounds for experimental validation, where missing truly active molecules is costly; abstaining in low-support regions is therefore practical. Similarly, in healthcare, models may defer out-of-distribution cases to clinicians while maintaining calibrated guarantees on well-supported patient groups. 5 Experimental Details In this section, we evaluate whether the proposed methods correct covariate shift and restore conformal coverage on small-molecule property prediction tasks, where conformal prediction is widely used to prioritize compounds for wet-lab validation astigarraga2025conformal. These benchmarks derive from real-world drug discovery pipelines huang2021therapeutics and are standard testbeds for molecular uncertainty quantification laghuvarapu2023codrug. 5.1 Datasets We conduct experiments on widely evaluate molecular property prediction benchmarks fooladi2025evaluating from Therapeutics Data Commons (TDC) huang2021therapeutics. We evaluate on five classification tasks: Tox21 (multi-target toxicity from high-throughput assays), AMES (mutagenicity prediction), hERG (cardiac toxicity via potassium channel inhibition), B-Martins (blood–brain barrier penetration), and HIV (bioactivity via inhibition of HIV replication). These tasks span toxicity and bioactivity prediction and exhibit substantial structural diversity across molecular scaffolds. In such regimes, reliable uncertainty quantification is critical: overconfident predictions on novel compounds can lead to unsafe or inefficient decisions, while overly conservative predictions hinder prioritization. Conformal prediction offers a principled mechanism for controlling predictive uncertainty, making these benchmarks well-suited for evaluating robustness under covariate shift. 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. 5.2 Baselines and Methods We compare against the following approaches: 1. Vanilla Conformal Prediction. Standard split conformal prediction assuming exchangeability. 2. KDE Density Ratio Estimation (64D). Importance weights via kernel density estimation in 64-dimensional learned feature space. 3. KDE Density Ratio Estimation (8D). KDE using an 8-dimensional linear projection to assess sensitivity to dimensionality, as suggest by CoDruglaghuvarapu2023codrug. We refer to this as CoDrugKDE. 4. Logistic Classifier-Based Density Ratio Estimation. Logistic regression domain classifier distinguishing calibration and test samples. 5. KMM: Kernel moment matching with bounded importance weights but no selective filtering. 6. SKMM (KMM + α Filtering) Joint optimization of calibration weights (β) and target selection variables (α), corresponding to our selective KMM framework. 5.3 Implementation Details Covariate Shift Construction. To simulate realistic structural shift, we consider Fingerprint split wu2018moleculenet, following prior work laghuvarapu2023codrug. In this split, molecules are represented using binary embeddings (fingerprints) and partitioned based on distance from the dataset centroid in fingerprint space. The 15% farthest molecules form the test set, while the remainder constitutes the calibration pool. This induces reduced scaffold overlap and support mismatch, reflecting practical discovery settings where new chemical series differ from historical data. Model Training. All predictive models use the AttentiveFP graph neural network architecture attentivefp. Models are trained with five independent random seeds per dataset. Conformal calibration and evaluation are performed separately for each run, and results are averaged. Weight Construction. For weighting-based methods, covariate representations are extracted from the final hidden layer of AttentiveFP (64-dimensional embeddings). These embeddings serve as inputs for density ratio estimation and kernel moment matching, reflecting realistic high-dimensional representation learning regimes. For the CoDrugKDE baseline, we evaluate an 8-dimensional linear projection to assess sensitivity to dimensionality. Conformal Setup. Equation 1 describes the standard marginal coverage guarantee of conformal prediction, where coverage holds over the joint distribution of (X,Y)(X,Y). In addition to global calibration, we also evaluate class-conditional (Mondrian) conformal prediction, which computes quantile thresholds separately within predefined groups—in our case, class labels. Mondrian conformal prediction guarantees coverage conditional on class membership, i.e., ℙ(YN+1∈C^(XN+1)∣YN+1=c)≥1−αP\! (Y_N+1∈ C(X_N+1) Y_N+1=c )≥ 1-α for each class c. For molecular activity tasks, we calibrate on the active class (C1), while for toxicity tasks we calibrate on the non-toxic class (C0), prioritizing reliable coverage for outcomes most critical to downstream drug discovery decisions. Additional details of this are provided in Appendix D. KMM Hyperparameters. For kernel mean matching (KMM), the kernel bandwidth is selected using the maximum test power criterion following prior work gretton2012kernel. The weight upper bound is set to B=30B=30, and the target selection fraction is fixed at α=0.5α=0.5, corresponding to selecting 50% of the test molecules. These hyperparameters were not tuned per dataset. Additional implementation details and hyperparameters are provided in Appendix E. All reported results are averaged over five independent random runs and detailed error bars are included in Appendix F. 6 Results and Discussion 6.1 Coverage Under Covariate Shift We first evaluate vanilla (uniform) conformal prediction under both no-shift and fingerprint-based covariate shift. Results are summarized in Table 1. To quantify calibration error, we report the mean absolute deviation (MAD) from nominal coverage. Specifically, for target coverage levels 1−α∈0.5,0.55,…,0.91-α∈\0.5,0.55,…,0.9\, MAD=1K∑k=1K|cov^(1−αk)−(1−αk)|,MAD= 1K _k=1^K | cov(1- _k)-(1- _k) |, where cov^(1−αk) cov(1- _k) denotes empirical coverage at level 1−αk1- _k, and K is the number of evaluated levels. This metric aggregates deviation across coverage levels and provides a measure of miscalibration. Under random (no-shift) splits, uniform conformal prediction achieves low MAD (e.g., 0.0044 on HIV and 0.0121 on AMES), confirming nominal coverage under exchangeability. Under fingerprint-based covariate shift, calibration deteriorates sharply: MAD exceeds 0.43 on Tox21 and 0.45 on B-Martins, and increases by more than an order of magnitude even on milder shifts such as HIV. These results demonstrate that exchangeability violations substantially degrade conformal validity in realistic molecular settings, motivating explicit covariate-shift correction. Table 1: Mean absolute deviation (MAD) of uniform conformal prediction under random (no-shift) and shifted splits. Dataset Random Shifted AMES 0.0121 0.0844 B-Martins 0.0346 0.4579 HIV 0.0044 0.1830 Tox21 0.0249 0.4366 hERG-Karim 0.0254 0.0384 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. 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 B-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 Mean Rank 5.4 5.0 5.0 4.4 3.4 3.6 4.0 4.2 2.2 2.6 1.0 1.2 6.2 Global and Mondrian Calibration Table 2 reports mean absolute deviation (MAD) from nominal coverage under both global and Mondrian conformal calibration. Uniform conformal prediction performs poorly under covariate shift, exhibiting the highest deviation across nearly all datasets, with mean ranks of 5.4 (global) and 5.0 (Mondrian). In particular, uniform MAD exceeds 0.40 on Tox21 and B-Martins, confirming that exchangeability violations severely degrade calibration. Density-ratio baselines partially mitigate this effect but remain inconsistent. Logistic weighting substantially improves Tox21 (0.073 vs. 0.409 under uniform), yet remains unstable on B-Martins (0.396) and HIV (0.297). KDE-based approaches are particularly sensitive to the dimensionality of the representation, with MAD exceeding 0.30 on multiple datasets under global calibration. Even the CoDrug KDE variant remains above 0.30 in B-Martins and fails to consistently match nominal coverage. In contrast, KMM (no filtering) consistently reduces deviation across datasets. On Tox21 (global), MAD decreases from 0.073 (logistic) to 0.053 under KMM, with similar improvements in B-Martins (0.319) and HIV (0.087), indicating more stable moment alignment under covariate shift. The selective extension (SKMM) further improves calibration, achieving the lowest MAD on all five benchmarks under global calibration—reducing MAD to 0.015 in Tox21, 0.009 in hERG, and 0.013 in HIV. Under Mondrian calibration, SKMM remains competitive and attains the best mean rank (1.2), showing that its gains extend to class-conditional settings. Furthermore, figure 2 corroborates these observations: SKMM consistently tracks the diagonal target line across coverage levels, while uniform and KDE baselines exhibit systematic undercoverage under shift. KMM narrows this gap, and SKMM further stabilizes calibration across both toxicity and activity prediction tasks. We observe similar trends global calibration, depicted in Appendix E. Overall, the results demonstrate that moment-matching-based reweighting improves calibration stability under covariate shift, and that selective filtering further enhances robustness in low-overlap regimes. On average across datasets under global calibration, SKMM achieves a 55% reduction in MAD relative to CoDrugKDE, highlighting substantial improvements in coverage stability. 6.3 Bias–Variance and Selective Filtering Figure 3: Bias–variance proxy (MMD+1/ESSMMD+ 1/ESS) under covariate shift. MAD in brackets. Figure 4: Selection weights under SKMM. To understand the source of improvement under covariate shift, we examine the dominant terms in our coverage bound, namely MMD+1/ESSMMD+ 1/ESS, which capture bias (distributional mismatch) and variance (weight concentration), respectively. Figure 3 reports this bias–variance proxy across methods. Uniform conformal prediction exhibits large combined error due to both high MMD and lack of shift correction. The ordering of MMD+1/ESSMMD+ 1/ESS across methods closely aligns with empirical calibration performance, with KMM and SKMM achieving the smallest values Selective KMM (SKMM) further reduces this quantity by restricting correction to regions of shared support, achieving the smallest combined error across datasets. To better understand this behavior, Figure 4 visualizes the learned selection weights α. The distribution of α concentrates near the boundaries, with approximately half of the test instances retained (as controlled by τ=0.5τ=0.5). The learned sparsity is consistent across datasets and remains stable under shift. and exhibit the same qualitative trends. 7 Conclusion We presented KMM-CP, a framework for conformal prediction under covariate shift that leverages Kernel Mean Matching with a selective extension for limited-support regimes. By combining bounded moment matching with target-aware filtering, the proposed method improves coverage stability relative to standard density-ratio approaches, particularly in high-dimensional molecular representation spaces. Our analysis clarifies how moment alignment and effective sample size govern coverage under shift, and experiments across multiple molecular property benchmarks demonstrate more reliable coverage under structural distribution shift. We acknowledge several limitations. The selective filtering step reduces the number of evaluated samples, and performance depends on the quality of learned representations and kernel hyperparameters. Although our synthetic shift construction reflects realistic structural mismatch, real-world deployment may involve more complex distribution shifts. Finally, while we demonstrate effectiveness on small-molecule property prediction tasks, generalization to other domains and data modalities remains to be validated. References Appendix Appendix A Proofs of Theoretical Results A.1 Proof of Theorem 2 Fix a probability measure Q on ℝdR^d and define FQ(t):=X∼Q[gt(X)]. F_Q(t):=E_X Q[g_t(X)]. Then, for any fixed t∈ℝt , FQ(t)−FT(t)=X∼Q[gt(X)]−X∼PXT[gt(X)]. F_Q(t)-F_T(t)=E_X Q[g_t(X)]-E_X P_X^T[g_t(X)]. Add and subtract g~t g_t to obtain |FQ(t)−FT(t)| |F_Q(t)-F_T(t) | ≤|X∼Q[g~t(X)]−X∼PXT[g~t(X)]| ≤ |E_X Q[ g_t(X)]-E_X P_X^T[ g_t(X)] | +|X∼Q[gt(X)−g~t(X)]|+|X∼PXT[gt(X)−g~t(X)]|. + |E_X Q[g_t(X)- g_t(X)] |+ |E_X P_X^T[g_t(X)- g_t(X)] |. Since |gt(x)−g~t(x)|≤‖gt−g~t‖∞|g_t(x)- g_t(x)|≤\|g_t- g_t\|_∞ for all x, and both Q and PXTP_X^T are probability measures, we have |X∼Q[gt(X)−g~t(X)]|≤‖gt−g~t‖∞,|X∼PXT[gt(X)−g~t(X)]|≤‖gt−g~t‖∞. |E_X Q[g_t(X)- g_t(X)] |≤\|g_t- g_t\|_∞, |E_X P_X^T[g_t(X)- g_t(X)] |≤\|g_t- g_t\|_∞. For the first term, if ‖g~t‖ℋ=0\| g_t\|_H=0, then g~t=0 g_t=0 in ℋH and hence |X∼Q[g~t(X)]−X∼PXT[g~t(X)]|=0. |E_X Q[ g_t(X)]-E_X P_X^T[ g_t(X)] |=0. Otherwise, define ht:=g~t‖g~t‖ℋ, h_t:= g_t\| g_t\|_H, so that ‖ht‖ℋ=1\|h_t\|_H=1. By the definition of MMDℋMMD_H, |X∼Q[g~t(X)]−X∼PXT[g~t(X)]| |E_X Q[ g_t(X)]-E_X P_X^T[ g_t(X)] | =‖g~t‖ℋ|X∼Q[ht(X)]−X∼PXT[ht(X)]| =\| g_t\|_H |E_X Q[h_t(X)]-E_X P_X^T[h_t(X)] | ≤‖g~t‖ℋMMDℋ(Q,PXT). ≤\| g_t\|_H\,MMD_H(Q,P_X^T). Combining the bounds gives |FQ(t)−FT(t)|≤‖g~t‖ℋMMDℋ(Q,PXT)+2‖gt−g~t‖∞. |F_Q(t)-F_T(t)|≤\| g_t\|_H\,MMD_H(Q,P_X^T)+2\|g_t- g_t\|_∞. Taking the supremum over t∈ℝt and using the uniform bounds supt∈ℝ‖g~t‖ℋ≤Bℋ,supt∈ℝ‖gt−g~t‖∞≤ε∞, _t \| g_t\|_H≤ B_H, _t \|g_t- g_t\|_∞≤ _∞, we obtain supt∈ℝ|FQ(t)−FT(t)|≤BℋMMDℋ(Q,PXT)+2ε∞. _t |F_Q(t)-F_T(t)|≤ B_H\,MMD_H(Q,P_X^T)+2 _∞. Finally, applying this bound with Q=P^S,XwQ= P^\,w_S,X and recalling that Fw(t)=X∼P^S,Xw[gt(X)], F_w(t)=E_X P^\,w_S,X[g_t(X)], yields supt∈ℝ|Fw(t)−FT(t)|≤BℋMMDℋ(P^S,Xw,PXT)+2ε∞. _t |F_w(t)-F_T(t)|≤ B_H\,MMD_H( P^\,w_S,X,P_X^T)+2 _∞. This proves the theorem. A.2 Proof of Theorem 3 Proof. Throughout the proof, condition on X1,…,Xn,Z1,…,ZmX_1,…,X_n,Z_1,…,Z_m so the weights w~1,…,w~n w_1,…, w_n are deterministic and the labels Y1,…,YnY_1,…,Y_n remain independent. 1. Reduce coverage error to a CDF error and an overshoot term. By definition of C C, ℙ(X,Y)∼PT(Y∈C^(X))=ℙ(X,Y)∼PT(s(X,Y)≤q^1−α)=FT(q^1−α). _(X,Y) P_T (Y∈ C(X) )=P_(X,Y) P_T (s(X,Y)≤ q_1-α )=F_T( q_1-α). Since F^w(q^1−α)≥1−α F_w( q_1-α)≥ 1-α by definition of q^1−α q_1-α, FT(q^1−α)−(1−α) F_T( q_1-α)-(1-α) =(FT(q^1−α)−F^w(q^1−α))+(F^w(q^1−α)−(1−α)), = (F_T( q_1-α)- F_w( q_1-α) )+ ( F_w( q_1-α)-(1-α) ), and therefore |FT(q^1−α)−(1−α)|≤supt∈ℝ|FT(t)−F^w(t)|+(F^w(q^1−α)−(1−α)). |F_T( q_1-α)-(1-α) |≤ _t |F_T(t)- F_w(t)|+ ( F_w( q_1-α)-(1-α) ). Because F^w F_w is right-continuous and has jumps of size exactly w~i w_i at the distinct points s(Xi,Yi)s(X_i,Y_i), the distinct-scores assumption implies that every jump has size at most maxiw~i _i w_i. Moreover, by minimality of q^1−α q_1-α we have F^w(t)<1−α F_w(t)<1-α for all t<q^1−αt< q_1-α, hence the overshoot satisfies 0≤F^w(q^1−α)−(1−α)≤max1≤i≤nw~i. 0≤ F_w( q_1-α)-(1-α)≤ _1≤ i≤ n w_i. Therefore |ℙ(X,Y)∼PT(Y∈C^(X))−(1−α)|≤supt∈ℝ|FT(t)−F^w(t)|+max1≤i≤nw~i. |P_(X,Y) P_T (Y∈ C(X) )-(1-α) |≤ _t |F_T(t)- F_w(t)|+ _1≤ i≤ n w_i. 2. Decompose supt|FT(t)−F^w(t)| _t|F_T(t)- F_w(t)| into bias and variance. By the triangle inequality, supt∈ℝ|FT(t)−F^w(t)|≤supt∈ℝ|FT(t)−Fw(t)|+supt∈ℝ|Fw(t)−F^w(t)|. _t |F_T(t)- F_w(t)|≤ _t |F_T(t)-F_w(t)|+ _t |F_w(t)- F_w(t)|. The first term is controlled by Theorem 2: supt∈ℝ|FT(t)−Fw(t)|≤BℋMMDℋ(P^S,Xw,PXT)+2ε∞. _t |F_T(t)-F_w(t)|≤ B_H\,MMD_H( P^\,w_S,X,P_X^T)+2\, _∞. It remains to control supt|Fw(t)−F^w(t)| _t|F_w(t)- F_w(t)|. 3. An expectation bound for the weighted empirical process. Define Δ:=supt∈ℝ|F^w(t)−Fw(t)|=supt∈ℝ|∑i=1nw~i(s(Xi,Yi)≤t−gt(Xi))|. := _t | F_w(t)-F_w(t) |= _t | _i=1^n w_i (1\s(X_i,Y_i)≤ t\-g_t(X_i) ) |. Let Y1′,…,Yn′Y_1 ,…,Y_n be an independent copy of Y1,…,YnY_1,…,Y_n conditional on X1,…,XnX_1,…,X_n (and thus also on the weights). Define F^w′(t):=∑i=1nw~i 1s(Xi,Yi′)≤t. F_w (t):= _i=1^n w_i\,1\s(X_i,Y_i )≤ t\. By Jensen’s inequality and Fw(t)=[F^w(t)∣X1,…,Xn,Z1,…,Zm]F_w(t)=E[ F_w(t) X_1,…,X_n,Z_1,…,Z_m], [Δ|X1,…,Xn,Z1,…,Zm] [ \, |\,X_1,…,X_n,Z_1,…,Z_m ] ≤[supt|F^w(t)−F^w′(t)||X1,…,Xn,Z1,…,Zm]. [ _t | F_w(t)- F_w (t) |\, |\,X_1,…,X_n,Z_1,…,Z_m ]. Introduce i.i.d. Rademacher signs ε1,…,εn _1,…, _n independent of everything else. By standard symmetrization, [Δ|X1,…,Xn,Z1,…,Zm]≤2[supt∈ℝ|∑i=1nw~iεi 1s(Xi,Yi)≤t||X1,…,Xn,Z1,…,Zm]. [ \, |\,X_1,…,X_n,Z_1,…,Z_m ]≤ 2\,E [ _t | _i=1^n w_i\, _i\,1\s(X_i,Y_i)≤ t\ |\, |\,X_1,…,X_n,Z_1,…,Z_m ]. Condition further on X1,…,Xn,Z1,…,Zm,Y1,…,YnX_1,…,X_n,Z_1,…,Z_m,Y_1,…,Y_n and write si:=s(Xi,Yi)s_i:=s(X_i,Y_i). Let π be a permutation such that sπ(1)≤⋯≤sπ(n)s_π(1)≤·s≤ s_π(n). Then supt∈ℝ|∑i=1nw~iεi 1si≤t|=max0≤k≤n|∑j=1kw~π(j)επ(j)|. _t | _i=1^n w_i\, _i\,1\s_i≤ t\ |= _0≤ k≤ n | _j=1^k w_π(j)\, _π(j) |. Define Mk:=∑j=1kw~π(j)επ(j)M_k:= _j=1^k w_π(j)\, _π(j) with M0:=0M_0:=0. By Doob’s L2L^2 maximal inequality, [max0≤k≤nMk2|X1,…,Xn,Z1,…,Zm,Y]≤4[Mn2|X1,…,Xn,Z1,…,Zm,Y]=4∑i=1nw~i2. [ _0≤ k≤ nM_k^2\, |\,X_1,…,X_n,Z_1,…,Z_m,Y ]≤ 4\,E [M_n^2\, |\,X_1,…,X_n,Z_1,…,Z_m,Y ]=4 _i=1^n w_i^2. Taking square roots and using Jensen’s inequality yields [supt∈ℝ|∑i=1nw~iεi 1si≤t||X1,…,Xn,Z1,…,Zm,Y]≤2∑i=1nw~i2. [ _t | _i=1^n w_i\, _i\,1\s_i≤ t\ |\, |\,X_1,…,X_n,Z_1,…,Z_m,Y ]≤ 2 _i=1^n w_i^2. Combining the last two displays and unconditioning on Y gives [Δ|X1,…,Xn,Z1,…,Zm]≤4∑i=1nw~i2=41ESS. [ \, |\,X_1,…,X_n,Z_1,…,Z_m ]≤ 4 _i=1^n w_i^2=4 1ESS. 4. A high-probability bound via bounded differences. View Δ as a function of the independent variables Y1,…,YnY_1,…,Y_n conditional on X1,…,Xn,Z1,…,ZmX_1,…,X_n,Z_1,…,Z_m. If one replaces YiY_i by another value Yi♯Y_i (leaving all other YjY_j fixed), then for every t the quantity F^w(t) F_w(t) changes by at most w~i w_i, and Fw(t)F_w(t) is unchanged. Hence Δ changes by at most w~i w_i. McDiarmid’s inequality implies that for any u>0u>0, ℙ(Δ−[Δ∣X,Z]≥u|X,Z)≤exp(−2u2∑i=1nw~i2)=exp(−2u2ESS), ( -E[ X,Z]≥ u\, |\,X,Z )≤ (- 2u^2 _i=1^n w_i^2 )= (-2u^2\,ESS ), where X denotes (X1,…,Xn)(X_1,…,X_n) and Z denotes (Z1,…,Zm)(Z_1,…,Z_m). Taking u:=12ESSlog1δu:= 12ESS 1δ yields that with probability at least 1−δ1-δ, Δ≤41ESS+12ESSlog1δ. ≤ 4 1ESS+ 12ESS 1δ. On the event from item 4, we combine items 1 and 2 to obtain |ℙ(X,Y)∼PT(Y∈C^(X))−(1−α)| |P_(X,Y) P_T (Y∈ C(X) )-(1-α) | ≤supt|FT(t)−Fw(t)|+supt|Fw(t)−F^w(t)|+maxiw~i ≤ _t|F_T(t)-F_w(t)|+ _t|F_w(t)- F_w(t)|+ _i w_i ≤BℋMMDℋ(P^S,Xw,PXT)+2ε∞+Δ+maxiw~i. ≤ B_H\,MMD_H( P^\,w_S,X,P_X^T)+2\, _∞+ + _i w_i. Finally, maxiw~i≤∑i=1nw~i2=1/ESS _i w_i≤ _i=1^n w_i^2= 1/ESS, so item 4 gives Δ+maxiw~i≤51ESS+12ESSlog1δ. + _i w_i≤ 5 1ESS+ 12ESS 1δ. Substituting this completes the proof. ∎ A.3 Proof of Corollary 1 Proof. Fix η>0η>0 and δ∈(0,1)δ∈(0,1). By Theorem 3, conditional on X1,…,Xn,Z1,…,ZmX_1,…,X_n,Z_1,…,Z_m, with probability at least 1−δ1-δ over Y1,…,YnY_1,…,Y_n, |ℙ(X,Y)∼PT(Y∈C^(X))−(1−α)| |P_(X,Y) P_T (Y∈ C(X) )-(1-α) | ≤BℋMMDℋ(P^S,Xw,PXT)+2ε∞,n,m+(5+12log1δ)1ESS. ≤ B_H\,MMD_H ( P^\,w_S,X,P_X^T )+2\, _∞,n,m+ (5+ 12 1δ ) 1ESS. Define the random bound Rn,m(δ):=BℋMMDℋ(P^S,Xw,PXT)+2ε∞,n,m+(5+12log1δ)1ESS. R_n,m(δ):=B_H\,MMD_H ( P^\,w_S,X,P_X^T )+2\, _∞,n,m+ (5+ 12 1δ ) 1ESS. Then, ℙ(|ℙ(X,Y)∼PT(Y∈C^(X)) ( |P_(X,Y) P_T (Y∈ C(X) ) −(1−α)|>η) -(1-α) |>η ) =[ℙ(|ℙ(X,Y)∼PT(Y∈C^(X))−(1−α)|>η|X1,…,Xn,Z1,…,Zm)] =E [P ( |P_(X,Y) P_T (Y∈ C(X) )-(1-α) |>η\, |\,X_1,…,X_n,Z_1,…,Z_m ) ] ≤[Rn,m(δ)>η+ℙ(the bound above fails|X1,…,Xn,Z1,…,Zm)] [1\R_n,m(δ)>η\+P (the bound above fails\, |\,X_1,…,X_n,Z_1,…,Z_m ) ] ≤ℙ(Rn,m(δ)>η)+δ, (R_n,m(δ)>η )+δ, where the probability is over X1,…,Xn,Z1,…,Zm,Y1,…,YnX_1,…,X_n,Z_1,…,Z_m,Y_1,…,Y_n. By assumption, Rn,m(δ)→0R_n,m(δ)→ 0 in probability for each fixed δ, hence ℙ(Rn,m(δ)>η)→0P(R_n,m(δ)>η)→ 0. Taking lim sup and then letting δ↓0δ 0 yields limn,m→∞ℙ(|ℙ(X,Y)∼PT(Y∈C^(X))−(1−α)|>η)=0. _n,m→∞P ( |P_(X,Y) P_T (Y∈ C(X) )-(1-α) |>η )=0. ∎ Appendix B Justification for Selective KMM Theorem 4 (Conditional boundary rule under a yield constraint). Fix w and consider the optimization over α∈ℝmα ^m minα _α J(α) J(α) s.t. 0≤αj≤1for all j∈[m], 0≤ _j≤ 1 all j∈[m], 1m∑j=1mαj≥τ, 1m _j=1^m _j≥τ, where J is convex and differentiable. Assume τ<1τ<1, so Slater’s condition holds. Let α⋆α be an optimal solution, and let (λ⋆,μ⋆,ν⋆)(λ ,μ ,ν ) be KKT multipliers, where λj⋆≥0 _j ≥ 0 and μj⋆≥0 _j ≥ 0 correspond to the constraints −αj≤0- _j≤ 0 and αj−1≤0 _j-1≤ 0, and ν⋆≥0ν ≥ 0 corresponds to the yield constraint τ−1m∑j=1mαj≤0τ- 1m _j=1^m _j≤ 0. Then for each j∈[m]j∈[m] the following statements hold: 1. If ∂J∂αj(α⋆)−ν⋆m>0 ∂ J∂ _j(α )- ν m>0, then αj⋆=0 _j =0. 2. If ∂J∂αj(α⋆)−ν⋆m<0 ∂ J∂ _j(α )- ν m<0, then αj⋆=1 _j =1. 3. If ∂J∂αj(α⋆)−ν⋆m=0 ∂ J∂ _j(α )- ν m=0, then αj⋆ _j may be interior or boundary; in particular, any interior coordinate αj⋆∈(0,1) _j ∈(0,1) must satisfy the equality. Proof. Since τ<1τ<1, there exists a strictly feasible point, for example αj0=(τ+1)/2α^0_j=(τ+1)/2 for all j, which satisfies 0<αj0<10<α^0_j<1 and 1m∑jαj0=(τ+1)/2>τ 1m _jα^0_j=(τ+1)/2>τ. Hence Slater’s condition holds, and KKT conditions are necessary and sufficient for optimality. The KKT conditions assert the existence of multipliers (λ⋆,μ⋆,ν⋆)(λ ,μ ,ν ) satisfying: 1. (Primal feasibility) 0≤αj⋆≤10≤ _j ≤ 1 for all j and 1m∑j=1mαj⋆≥τ 1m _j=1^m _j ≥τ. 2. (Dual feasibility) λj⋆≥0 _j ≥ 0, μj⋆≥0 _j ≥ 0 for all j, and ν⋆≥0ν ≥ 0. 3. (Complementary slackness) λj⋆αj⋆=0,μj⋆(αj⋆−1)=0,ν⋆(τ−1m∑j=1mαj⋆)=0. _j _j =0, _j ( _j -1)=0, ν (τ- 1m _j=1^m _j )=0. 4. (Stationarity) For each j∈[m]j∈[m], ∂J∂αj(α⋆)−ν⋆m−λj⋆+μj⋆=0. ∂ J∂ _j(α )- ν m- _j + _j =0. Fix an index j∈[m]j∈[m] and define the shifted gradient gj⋆:=∂J∂αj(α⋆)−ν⋆m. g_j := ∂ J∂ _j(α )- ν m. By stationarity, gj⋆=λj⋆−μj⋆. g_j = _j - _j . We now analyze cases using complementary slackness: 1. If αj⋆=0 _j =0, then μj⋆(αj⋆−1)=μj⋆(−1)=0 _j ( _j -1)= _j (-1)=0 implies μj⋆=0 _j =0, hence gj⋆=λj⋆≥0g_j = _j ≥ 0. 2. If αj⋆=1 _j =1, then λj⋆αj⋆=λj⋆=0 _j _j = _j =0, hence gj⋆=−μj⋆≤0g_j =- _j ≤ 0. 3. If αj⋆∈(0,1) _j ∈(0,1), then λj⋆αj⋆=0 _j _j =0 and αj⋆>0 _j >0 imply λj⋆=0 _j =0, and μj⋆(αj⋆−1)=0 _j ( _j -1)=0 and αj⋆−1≠0 _j -1≠ 0 imply μj⋆=0 _j =0. Therefore gj⋆=0g_j =0. These implications yield the desired conditional boundary rule: 1. If gj⋆>0g_j >0, then αj⋆ _j cannot equal 11 (which would force gj⋆≤0g_j ≤ 0) and cannot be interior (which would force gj⋆=0g_j =0). Hence αj⋆=0 _j =0. 2. If gj⋆<0g_j <0, then αj⋆ _j cannot equal 0 (which would force gj⋆≥0g_j ≥ 0) and cannot be interior (which would force gj⋆=0g_j =0). Hence αj⋆=1 _j =1. 3. If gj⋆=0g_j =0, then the KKT conditions allow αj⋆ _j to be interior (in which case necessarily λj⋆=μj⋆=0 _j = _j =0) or boundary (for example αj⋆=0 _j =0 with λj⋆=0 _j =0 or αj⋆=1 _j =1 with μj⋆=0 _j =0). This proves the final claim. ∎ Corollary 2 (Thresholding structure for SKMM target selection). Fix w and consider the SKMM subproblem in α∈ℝmα ^m with yield constraint minα _α J(α):=‖1n∑i=1nwiφ(xi)−1m∑j=1mαjφ(zj)‖ℋ2 J(α):= \| 1n _i=1^nw_i (x_i)- 1m _j=1^m _j (z_j) \|_H^2 s.t. 0≤αj≤1for all j∈[m], 0≤ _j≤ 1 all j∈[m], 1m∑j=1mαj≥τ, 1m _j=1^m _j≥τ, where τ<1τ<1. Let α⋆α be an optimal solution, and let ν⋆≥0ν ≥ 0 be a KKT multiplier for the yield constraint. Define the shifted partial derivatives gj⋆:=∂J∂αj(α⋆)−ν⋆m,j∈[m]. g_j := ∂ J∂ _j(α )- ν m, j∈[m]. Then: 1. If gj⋆>0g_j >0, then αj⋆=0 _j =0. 2. If gj⋆<0g_j <0, then αj⋆=1 _j =1. 3. If αj⋆∈(0,1) _j ∈(0,1), then gj⋆=0g_j =0. Equivalently, there exists a scalar threshold t⋆:=ν⋆/mt :=ν /m such that for every j, αj⋆=0⇒∂J∂αj(α⋆)≥t⋆,αj⋆=1⇒∂J∂αj(α⋆)≤t⋆,0<αj⋆<1⇒∂J∂αj(α⋆)=t⋆. _j =0\ \ ∂ J∂ _j(α )≥ t , _j =1\ \ ∂ J∂ _j(α )≤ t , 0< _j <1\ \ ∂ J∂ _j(α )=t . Proof. First note that for the SKMM objective, J(α)J(α) is a convex quadratic function of α because J(α) J(α) =‖μw−μα‖ℋ2,μw:=1n∑i=1nwiφ(xi),μα:=1m∑j=1mαjφ(zj), = \| _w- _α \|_H^2, _w:= 1n _i=1^nw_i (x_i), _α:= 1m _j=1^m _j (z_j), and expanding yields J(α)=‖μw‖ℋ2−2m∑j=1mαj⟨μw,φ(zj)⟩ℋ+1m2∑j=1m∑l=1mαjαlk(zj,zl), J(α)=\| _w\|_H^2- 2m _j=1^m _j _w, (z_j) _H+ 1m^2 _j=1^m _l=1^m _j _l\,k(z_j,z_l), whose Hessian with respect to α is (2/m2)KTT(2/m^2)\,K_T, where KTT=(k(zj,zl))j,lK_T=(k(z_j,z_l))_j,l is positive semidefinite. Hence J is convex and differentiable. Since τ<1τ<1, the point αj0=(τ+1)/2α^0_j=(τ+1)/2 is strictly feasible, so Slater’s condition holds and KKT conditions are necessary and sufficient for optimality. Therefore Theorem 4 applies directly to this choice of J, yielding the three implications stated in the corollary. The equivalent threshold formulation follows by defining t⋆:=ν⋆/mt :=ν /m. ∎ Appendix C Instability of Classifier-Based Density Ratios A common alternative to KMM for estimating the density ratio w(x)=PT(x)/PS(x)w(x)=P_T(x)/P_S(x) is probabilistic classification. By training a classifier p^(x) p(x) to predict the probability that a sample x belongs to the target domain, the estimated weight is given by w^(x)=p^(x)/(1−p^(x)) w(x)= p(x)/(1- p(x)). The following proposition demonstrates why this approach can fail under covariate shift with limited support. Proposition 1 (Unbounded MMD under Bounded Classifier Error). Let the true density ratio be w∗(x)=PT(x)/PS(x)w^*(x)=P_T(x)/P_S(x) and let w^(x) w(x) be the density ratio estimated via a classifier p^(x) p(x). For any ϵ>0ε>0 and M>0M>0, there exists a sufficiently small δ and a sample containing such points such that MMD(PSw^,PT)>MMMD(P_S w,P_T)>M. Proof. Assuming equal prior class probabilities for the source and target domains, the true conditional probability of a sample belonging to the target domain is p∗(x)=PT(x)/(PT(x)+PS(x))p^*(x)=P_T(x)/(P_T(x)+P_S(x)). The true density ratio is exactly w∗(x)=p∗(x)/(1−p∗(x))w^*(x)=p^*(x)/(1-p^*(x)). Consider a region in the covariate space where the source distribution has very low mass but the target distribution is well-supported. In this region, PS(x)→0P_S(x)→ 0, which implies p∗(x)→1p^*(x)→ 1. Let p∗(x)=1−δp^*(x)=1-δ for some sufficiently small δ>0δ>0. The true weight is: w∗(x)=1−δ≈1δ.w^*(x)= 1-δ≈ 1δ. Assume our classifier slightly overestimates the probability of the target class such that p^(x)=1−δ2 p(x)=1-δ^2. The absolute error of our classifier is: |p^(x)−p∗(x)|=|(1−δ2)−(1−δ)|=δ−δ2<δ.| p(x)-p^*(x)|=|(1-δ^2)-(1-δ)|=δ-δ^2<δ. By choosing δ≤ϵδ≤ε, the classifier satisfies the bounded error condition. However, the estimated weight in this region becomes: w^(x)=1−δ2δ2≈1δ2. w(x)= 1-δ^2δ^2≈ 1δ^2. The deviation of the estimated weight from the true weight is: |w^(x)−w∗(x)|=|1−δ2δ2−1−δ|=1−δ2.| w(x)-w^*(x)|= | 1-δ^2δ^2- 1-δ |= 1-δ^2. As δ→0δ→ 0, the difference between the estimated weight and the true weight grows as (1/δ2)O(1/δ^2). Recall the empirical MMD objective in the RKHS ℋH with feature map φ : MMD(PSw^,PT)=‖1n∑i=1nw^(xi)φ(xi)−1m∑j=1mφ(zj)‖ℋ.MMD(P_S w,P_T)= \| 1n _i=1^n w(x_i) (x_i)- 1m _j=1^m (z_j) \|_H. Because w^(xi) w(x_i) grows unboundedly as δ→0δ→ 0 for points in this low-density region, the term w^(xi)φ(xi) w(x_i) (x_i) will dominate the empirical source mean. Thus, for any arbitrary bound M, there exists a sufficiently small δ and a sample containing such points such that MMD(PSw^,PT)>MMMD(P_S w,P_T)>M. This demonstrates that minimizing a classifier’s probability estimation error does not safely bound the MMD or the variance of the weights, leading to degraded conformal prediction coverage. KMM explicitly resolves this by enforcing the bound wi≤Bw_i≤ B and directly matching the moments. ∎ Appendix D Mondrian Conformal Prediction D.1 Mondrian Conformal Prediction Mondrian Conformal Prediction (MCP) extends standard conformal prediction by conditioning calibration on a predefined partition of the data vovk2012conditional. Instead of enforcing marginal coverage over the entire population, MCP guarantees coverage within each subgroup defined by a class label or other stratification variable. Let g(X)∈g(X) denote a grouping function. In classification tasks, a common choice is g(X,Y)=Yg(X,Y)=Y, leading to class-conditional calibration. For each group c∈c , define the group-specific conformity scores Si:g(Xi,Yi)=c.\S_i:g(X_i,Y_i)=c\. Let ncn_c denote the number of calibration points in group c. The group-specific empirical CDF is Fc(t)=1nc∑i:g(Xi,Yi)=cSi≤t.F_c(t)= 1n_c _i:g(X_i,Y_i)=c1\S_i≤ t\. The Mondrian conformal threshold for group c is then q^1−α(c)=inft:Fc(t)≥1−α. q_1-α^(c)= \t:F_c(t)≥ 1-α \. The resulting prediction set for a new input (XN+1,YN+1)(X_N+1,Y_N+1) with group c=g(XN+1)c=g(X_N+1) satisfies ℙ(YN+1∈C^(XN+1)∣g(XN+1)=c)≥1−α,P (Y_N+1∈ C(X_N+1) g(X_N+1)=c )≥ 1-α, under exchangeability within each group. In imbalanced molecular datasets, class-conditional calibration prevents dominant classes from overwhelming rare but important categories (e.g., active or toxic compounds). D.2 Weighted Mondrian Conformal Prediction Under covariate shift, we extend MCP by incorporating importance weights within each group. Let wiw_i denote the calibration weights (e.g., obtained via KMM or SKMM). Define the normalized group-specific weights w~i(c)=wi∑j:g(Xj,Yj)=cwj,for g(Xi,Yi)=c. w_i^(c)= w_i _j:g(X_j,Y_j)=cw_j, g(X_i,Y_i)=c. The weighted group CDF becomes Fcw(t)=∑i:g(Xi,Yi)=cw~i(c)Si≤t.F_c^w(t)= _i:g(X_i,Y_i)=c w_i^(c)1\S_i≤ t\. The Mondrian weighted quantile is q^1−α(c),w=inft:Fcw(t)≥1−α. q_1-α^(c),w= \t:F_c^w(t)≥ 1-α \. The prediction set is constructed exactly as in the global weighted case, except that quantiles are computed separately within each group. Algorithm 2 Mondrian Split Conformal Prediction with KMM / Selective KMM 0: Labeled dataset =(xi,yi)i=1ND=\(x_i,y_i)\_i=1^N, unlabeled test covariates zjj=1m\z_j\_j=1^m, miscoverage level αconf _conf, kernel k(⋅,⋅)k(·,·), weight bound B, selection fraction τ (optional) 0: Prediction sets C^(zj) C(z_j) 1: Split D into training set trainD_train and calibration set cal=(xi,yi)i=1nD_cal=\(x_i,y_i)\_i=1^n 2: Train predictor f on trainD_train 3: Define conformity score S(x,y)S(x,y) based on f 4: if Selective KMM is enabled then 5: Solve double-weighted KMM to obtain selection weights αj _j 6: Define selected subset selT_sel 7: else 8: Set sel=zjj=1mT_sel=\z_j\_j=1^m 9: end if 10: Solve standard KMM between calD_cal and selT_sel to obtain source weights wiw_i 11: for each class c do 12: Let ℐc=i:yi=cI_c=\i:y_i=c\ 13: Normalize weights within class: w~i(c)=wi∑k∈ℐcwkfor i∈ℐc w_i^(c)= w_i _k _cw_k i _c 14: Compute class-conditional quantile q^1−α(c) q_1-α^(c): ∑i∈ℐcw~i(c)Si≤t≥1−αconf _i _c w_i^(c)1\S_i≤ t\≥ 1- _conf 15: end for 16: for each zj∈selz_j _sel do 17: Let c^=predicted class of zj c=predicted class of z_j 18: Output C^(zj)=y:S(zj,y)≤q^1−α(c^) C(z_j)=\y:S(z_j,y)≤ q_1-α^( c)\ 19: end for Practical Choice of Groups. In our experiments, we use class-conditional Mondrian calibration: g(X,Y)=Yg(X,Y)=Y. For toxicity datasets, we calibrate within the non-toxic class to avoid undercoverage of rare toxic compounds; for activity datasets, we calibrate within the active class. This choice reflects domain-specific asymmetry in decision importance. In Algorithm 2, we summarize the application of KMM-CP for class conditional guarentees under covariate shift. Appendix E Additional Experimental Details E.1 Code Availability All code used to reproduce the experiments is publicly available at: https://github.com/siddharthal/KMM-CP The repository contains complete training scripts, data preprocessing pipelines, and exact random seeds for all five runs per experiment. All results reported in the paper are fully reproducible. E.2 Dataset Statistics Table 3 summarizes dataset sizes and splits. Table 3: Dataset statistics and data splits. Dataset Total N Train Calib Test AMES 7,280 5,096 1,092 1,092 B-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 E.3 Training Details and Hyperparameters All models use the AttentiveFP attentivefp graph neural network architecture implemented via the dgllife library dgllife, with training performed in PyTorch. Optimization. • Optimizer: Adam • Learning rate: 1×10−31× 10^-3 • Learning rate scheduler: StepLR (step_size=1, γ=0.8γ=0.8) Training Loop. • Maximum epochs: 50 • Batch size: 128 • Early stopping: patience = 5 (based on training loss) Model Architecture. • Graph embedding dimension: 512 • Hidden dimension (classifier): 64 • GNN layers (AttentiveFP steps): 3 • Output dimension: 2 (binary classification) • Total parameters: 10M Training time Model training takes approximately 5–25 minutes per dataset on a single NVIDIA A100 GPU, depending on dataset size. Loss Function. • Cross-entropy loss • Class weights enabled (inverse class frequency) E.4 Conformal Prediction Parameters Non-conformity Score. Non-conformity scores are defined using the predicted class probabilities from the trained AttentiveFP model. Feature Representation. Penultimate-layer embeddings from the trained model are used as feature representations for density-ratio estimation and KMM. E.5 Implementation of KMM Kernel Mean Matching (KMM) optimization is solved using cvxopt library. Weight Constraints. • Weight upper bound: B=30B=30 • Target selection fraction (SKMM): τ=0.5τ=0.5 • Second-stage filtering threshold on α: 0.2 These hyperparameters are fixed across all datasets and are not tuned per dataset. Kernel Bandwidth Selection. We use an RBF kernel with bandwidth selected via a median-distance heuristic gretton2012kernel. Specifically, we compute the median pairwise distance between calibration and test embeddings and evaluate bandwidths from the set: σ∈0.01, 0.1, 0.5, 1, 2×median_distance.σ∈\0.01,\,0.1,\,0.5,\,1,\,2\×median\_distance. The bandwidth maximizing standardized MMD (z-scored across permutations) is selected. This strategy balances sensitivity to local discrepancies and stability under high-dimensional embeddings, and is commonly used in kernel two-sample testing. Computation Time. KMM optimization is performed on CPUs and requires approximately: • 6 minutes for the smallest dataset, • up to 50 minutes for the largest dataset. E.6 Compute Resources All experiments were conducted on a cluster with: • 4 NVIDIA A100 GPUs, • 64 CPU cores. E.7 Implementation of Density-Ratio Baselines KDE Baseline. For the KDE-based density-ratio baseline, we estimate source and target marginal densities using a Gaussian kernel. The bandwidth is selected via a held-out validation split by maximizing the log-likelihood over candidate values 0.01,0.1,1.0,10\0.01,0.1,1.0,10\. The same bandwidth is used for both source and target density estimation. The density ratio is then formed as w^(x)=p^t(x)/p^s(x) w(x)= p_t(x)/ p_s(x). Logistic Regression Baseline. For the classifier-based density-ratio baseline, we train a logistic regression classifier to distinguish calibration and test samples using 64-dimensional embeddings extracted from the final layer of the trained predictor. The classifier is implemented using the scikit-learn library with default regularization settings. Density ratios are obtained from the predicted probabilities via the standard odds transformation. Appendix F Additional results F.1 Calibration coverage plots for global CP Figure 5 presents calibration curves under covariate shift for global conformal prediction across all datasets. Each curve plots realized coverage against target coverage levels from 0.50.5 to 0.950.95, averaged over five random runs. The dashed line indicates perfect calibration. Under covariate shift, uniform conformal prediction systematically undercovers, with deviations increasing at higher target coverage levels. Density-ratio baselines partially correct this behavior but exhibit residual miscalibration and variability across datasets. In particular, classifier- and KDE-based weighting methods often fail to fully align with the nominal coverage line at moderate-to-high coverage levels. In contrast, KMM produces curves that more closely track the nominal diagonal, indicating improved distributional alignment. The proposed selective KMM (SKMM) further tightens this alignment, yielding the smallest deviations from nominal coverage across datasets. The improvement is consistent across both toxicity and activity prediction tasks, and persists across the full range of coverage levels. These plots visually corroborate the mondran CP trends and ranking results reported in the main text. Figure 5: Calibration coverage under global conformal prediction across datasets. F.2 Additional Bias-Variance and Selective Filtering plots In this section, we provide supplemental plots to support the analysis detailed in Section 6.3 for the remaining datasets: AMES, hERG, and HIV.Figure 6(a) presents the bias-variance decomposition across these benchmarks, while Figure 6(b) illustrates the corresponding α distribution plots. Overall, the bias-variance trends observed here closely mimic the results shown in Figure 3; specifically, the KMM-based methods continue to exhibit higher stability across varying shift intensities. Similarly, the α distributions demonstrate emergent sparsity across nearly all experimental settings. However, the hERG_Karim dataset exhibits a denser α distribution; this is potentially attributable to a higher intrinsic overlap in the covariate support. (a) Bias-Variance (AMES) (b) α Distribution (AMES) Figure 6: Figure. (a) illustrates the bias-variance proxy(MMD+1/ESSMMD+ 1/ESS) , while (b) displays the emergent sparsity in the weight distribution. F.3 Filtered Fraction Table 4 reports the average filtration ratio under global and Mondrian calibration, defined as the percentage of test instances retained after the selective KMM filtering stage (mean ± standard deviation across five runs). Recall that the target selection fraction is set to τ=0.5τ=0.5, corresponding to retaining approximately 50% of the test instances. Across datasets, the observed retention rates range from roughly 58% to 77%, indicating that the learned selection variables αj _j concentrate mass on a substantial but restricted subset of the test distribution. Also note that SKMM does not collapse to extreme filtering: a majority of test instances are retained across all benchmarks. This confirms that the method focuses correction on regions where density-ratio estimation is stable, rather than aggressively discarding data. Table 4: Average filtration ratio (%) under two-stage Selective KMM (mean ± std). Dataset Fraction Filtered AMES 71.2 ± 5.9 Tox21 58.3 ± 0.5 B-Martins 60.1 ± 10.1 hERG-Karim 76.7 ± 5.9 HIV 62.2 ± 1.4 F.4 Full Calibration Results Across Datasets The following tables report full calibration curves (mean ± std over five runs) for all datasets under both global and Mondrian conformal prediction across target coverages from 0.50.5 to 0.950.95. Best results per column are shown in bold, and second best in gray bold. Overall trends. Across all datasets and coverage levels, SKMM consistently achieves the lowest or near-lowest deviation from nominal coverage. This pattern holds for both global and class-conditional (Mondrian) calibration. KMM (Beta-only) typically ranks second, while classifier- and KDE-based baselines exhibit larger systematic deviations, particularly at higher target coverage levels. Global calibration. Under global conformal prediction, uniform weighting frequently exhibits pronounced undercoverage as the target coverage increases. Density-ratio baselines partially mitigate this effect but remain unstable across datasets. In contrast, KMM substantially improves calibration alignment, and SKMM further reduces deviation across nearly all coverage levels. The average rank columns confirm this consistency: SKMM achieves the best mean rank on every dataset under global calibration. Mondrian calibration results show analogous results for Mondrian conformal prediction. The same qualitative behavior persists: SKMM maintains the strongest alignment with nominal coverage across coverage levels, followed by KMM. Importantly, improvements are not restricted to a specific coverage region; gains remain stable from moderate (0.60.6–0.70.7) to high (0.90.9–0.950.95) target coverage. Table 5: AMES — Global calibration (mean ± std). Best is bold; second best is gray bold. 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± 0.000 0.458± 0.004 0.475± 0.006 0.485± 0.004 0.498± 0.005 0.525± 0.014 0.574± 0.034 0.636± 0.040 0.698± 0.034 0.824± 0.048 4.60 KDE 0.774± 0.127 0.788± 0.113 0.799± 0.104 0.810± 0.092 0.833± 0.081 0.848± 0.073 0.866± 0.051 0.898± 0.046 0.915± 0.036 0.933± 0.022 4.80 CoDrugKDE 0.680± 0.041 0.716± 0.040 0.738± 0.042 0.767± 0.035 0.790± 0.034 0.823± 0.042 0.845± 0.033 0.874± 0.027 0.899± 0.033 0.933± 0.027 3.00 Logistic 0.684± 0.058 0.725± 0.068 0.755± 0.051 0.778± 0.037 0.807± 0.030 0.831± 0.036 0.858± 0.044 0.875± 0.028 0.898± 0.027 0.925± 0.030 4.10 KMM 0.664± 0.026 0.694± 0.025 0.726± 0.021 0.750± 0.012 0.775± 0.017 0.799± 0.021 0.822± 0.027 0.861± 0.017 0.885± 0.024 0.915± 0.019 2.60 SKMM 0.630± 0.028 0.669± 0.024 0.703± 0.019 0.734± 0.009 0.761± 0.010 0.791± 0.020 0.819± 0.029 0.862± 0.022 0.883± 0.025 0.927± 0.020 1.90 Table 6: AMES — Mondrian calibration (mean ± std). Best is bold; second best is gray bold. 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.470± 0.006 0.499± 0.014 0.530± 0.025 0.569± 0.035 0.606± 0.045 0.642± 0.033 0.683± 0.035 0.732± 0.038 0.809± 0.048 0.916± 0.055 4.00 KDE 0.585± 0.080 0.605± 0.083 0.613± 0.094 0.642± 0.096 0.658± 0.112 0.704± 0.090 0.758± 0.085 0.808± 0.053 0.817± 0.064 0.847± 0.086 2.80 CoDrugKDE 0.685± 0.031 0.710± 0.034 0.735± 0.039 0.764± 0.038 0.792± 0.032 0.820± 0.028 0.838± 0.033 0.865± 0.031 0.901± 0.026 0.926± 0.030 4.10 Logistic 0.691± 0.052 0.722± 0.033 0.746± 0.033 0.777± 0.038 0.795± 0.041 0.819± 0.043 0.847± 0.029 0.865± 0.031 0.894± 0.020 0.926± 0.030 4.70 KMM 0.669± 0.020 0.698± 0.024 0.723± 0.029 0.744± 0.030 0.767± 0.030 0.803± 0.026 0.834± 0.029 0.855± 0.027 0.881± 0.024 0.926± 0.028 3.00 SKMM 0.655± 0.013 0.680± 0.020 0.706± 0.022 0.740± 0.026 0.765± 0.032 0.801± 0.029 0.837± 0.032 0.857± 0.028 0.889± 0.028 0.936± 0.026 2.40 Table 7: B_Martins — Global calibration (mean ± std). Best is bold; second best is gray bold. 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.148± 0.053 0.179± 0.072 0.226± 0.069 0.243± 0.078 0.266± 0.081 0.288± 0.083 0.322± 0.077 0.374± 0.064 0.439± 0.087 0.659± 0.110 6.00 KDE 0.205± 0.094 0.226± 0.112 0.254± 0.116 0.294± 0.112 0.317± 0.114 0.362± 0.125 0.396± 0.142 0.442± 0.173 0.532± 0.195 0.698± 0.230 3.80 CoDrugKDE 0.207± 0.098 0.244± 0.100 0.275± 0.100 0.308± 0.102 0.347± 0.089 0.390± 0.084 0.446± 0.093 0.495± 0.108 0.629± 0.080 0.770± 0.055 2.40 Logistic 0.180± 0.068 0.216± 0.066 0.243± 0.077 0.268± 0.078 0.289± 0.079 0.317± 0.073 0.354± 0.067 0.400± 0.059 0.467± 0.072 0.693± 0.084 5.00 KMM 0.187± 0.061 0.242± 0.108 0.266± 0.119 0.291± 0.111 0.331± 0.096 0.372± 0.084 0.462± 0.125 0.559± 0.110 0.719± 0.112 0.838± 0.038 2.60 SKMM 0.296± 0.121 0.375± 0.171 0.401± 0.166 0.424± 0.157 0.464± 0.142 0.530± 0.135 0.559± 0.138 0.652± 0.094 0.685± 0.091 0.812± 0.086 1.20 Table 8: B_Martins — Mondrian calibration (mean ± std). Best is bold; second best is gray bold. 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.189± 0.005 0.202± 0.002 0.208± 0.003 0.216± 0.004 0.222± 0.004 0.236± 0.010 0.252± 0.010 0.287± 0.022 0.367± 0.069 0.683± 0.088 5.70 KDE 0.214± 0.015 0.226± 0.022 0.241± 0.020 0.269± 0.020 0.282± 0.025 0.308± 0.040 0.343± 0.072 0.408± 0.086 0.473± 0.121 0.632± 0.173 3.90 CoDrugKDE 0.218± 0.014 0.235± 0.020 0.248± 0.026 0.265± 0.027 0.291± 0.041 0.315± 0.044 0.357± 0.067 0.432± 0.037 0.511± 0.062 0.687± 0.075 2.50 Logistic 0.202± 0.003 0.209± 0.006 0.218± 0.005 0.226± 0.011 0.237± 0.013 0.252± 0.019 0.272± 0.018 0.311± 0.047 0.386± 0.065 0.658± 0.090 5.00 KMM 0.219± 0.013 0.227± 0.018 0.242± 0.015 0.264± 0.023 0.279± 0.013 0.353± 0.078 0.437± 0.027 0.496± 0.070 0.604± 0.089 0.686± 0.085 2.50 SKMM 0.272± 0.037 0.312± 0.031 0.329± 0.028 0.383± 0.091 0.433± 0.094 0.453± 0.112 0.510± 0.112 0.575± 0.082 0.604± 0.092 0.672± 0.095 1.40 Table 9: HIV — Global calibration (mean ± std). Best is bold; second best is gray bold. 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.376± 0.063 0.406± 0.063 0.439± 0.064 0.472± 0.064 0.506± 0.064 0.544± 0.063 0.581± 0.065 0.626± 0.065 0.684± 0.061 0.779± 0.049 4.40 KDE 0.329± 0.124 0.346± 0.126 0.361± 0.132 0.379± 0.137 0.401± 0.144 0.418± 0.160 0.450± 0.168 0.482± 0.188 0.517± 0.203 0.577± 0.255 5.60 CoDrugKDE 0.405± 0.101 0.443± 0.105 0.478± 0.110 0.510± 0.108 0.545± 0.112 0.585± 0.121 0.625± 0.116 0.674± 0.110 0.727± 0.089 0.803± 0.091 3.30 Logistic 0.997± 0.002 0.997± 0.002 0.997± 0.002 0.997± 0.002 0.997± 0.002 0.997± 0.002 0.997± 0.002 0.997± 0.002 0.997± 0.002 0.997± 0.002 4.50 KMM 0.455± 0.058 0.493± 0.056 0.537± 0.048 0.571± 0.042 0.610± 0.042 0.648± 0.033 0.683± 0.031 0.729± 0.035 0.787± 0.031 0.863± 0.034 2.20 SKMM 0.510± 0.053 0.556± 0.056 0.596± 0.050 0.641± 0.052 0.680± 0.051 0.735± 0.043 0.787± 0.040 0.829± 0.034 0.880± 0.031 0.945± 0.028 1.00 Table 10: HIV — Mondrian calibration (mean ± std). Best is bold; second best is gray bold. 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.385± 0.063 0.415± 0.064 0.448± 0.063 0.480± 0.066 0.512± 0.067 0.546± 0.068 0.581± 0.068 0.619± 0.067 0.667± 0.064 0.732± 0.060 4.40 KDE 0.341± 0.124 0.358± 0.123 0.373± 0.130 0.392± 0.132 0.411± 0.139 0.427± 0.150 0.457± 0.159 0.479± 0.171 0.515± 0.177 0.552± 0.201 5.60 CoDrugKDE 0.409± 0.096 0.445± 0.103 0.481± 0.106 0.509± 0.109 0.544± 0.106 0.574± 0.107 0.615± 0.119 0.659± 0.115 0.705± 0.112 0.769± 0.097 3.40 Logistic 0.969± 0.001 0.969± 0.001 0.969± 0.001 0.969± 0.001 0.969± 0.001 0.969± 0.001 0.969± 0.001 0.969± 0.001 0.969± 0.001 0.972± 0.009 4.20 KMM 0.452± 0.059 0.488± 0.060 0.527± 0.057 0.565± 0.048 0.600± 0.044 0.640± 0.039 0.674± 0.035 0.714± 0.038 0.759± 0.032 0.826± 0.030 2.30 SKMM 0.505± 0.059 0.548± 0.054 0.593± 0.053 0.631± 0.048 0.676± 0.053 0.713± 0.048 0.770± 0.040 0.816± 0.043 0.861± 0.034 0.910± 0.031 1.10 Table 11: Tox21 — Global calibration (mean ± std). Best is bold; second best is gray bold. 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.142± 0.014 0.159± 0.012 0.187± 0.018 0.209± 0.018 0.236± 0.019 0.291± 0.020 0.352± 0.009 0.460± 0.015 0.586± 0.010 0.760± 0.029 5.80 KDE 0.249± 0.021 0.270± 0.023 0.308± 0.023 0.331± 0.023 0.358± 0.017 0.397± 0.023 0.434± 0.032 0.480± 0.037 0.564± 0.028 0.709± 0.070 5.20 CoDrugKDE 0.408± 0.117 0.444± 0.116 0.480± 0.125 0.537± 0.129 0.582± 0.145 0.628± 0.132 0.682± 0.137 0.769± 0.066 0.846± 0.038 0.906± 0.027 3.90 Logistic 0.570± 0.063 0.613± 0.076 0.671± 0.094 0.731± 0.117 0.780± 0.101 0.839± 0.071 0.902± 0.036 0.918± 0.034 0.931± 0.021 0.940± 0.009 2.80 KMM 0.454± 0.027 0.509± 0.021 0.551± 0.005 0.591± 0.008 0.626± 0.015 0.687± 0.006 0.748± 0.011 0.807± 0.013 0.847± 0.015 0.905± 0.013 2.30 SKMM 0.471± 0.024 0.519± 0.014 0.574± 0.033 0.643± 0.069 0.695± 0.056 0.761± 0.028 0.798± 0.018 0.834± 0.024 0.892± 0.014 0.956± 0.004 1.00 Table 12: Tox21 — Mondrian calibration (mean ± std). Best is bold; second best is gray bold. 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.137± 0.017 0.154± 0.013 0.175± 0.010 0.202± 0.017 0.225± 0.019 0.257± 0.019 0.317± 0.013 0.386± 0.009 0.518± 0.014 0.623± 0.011 6.00 KDE 0.247± 0.019 0.268± 0.021 0.305± 0.023 0.326± 0.029 0.353± 0.016 0.386± 0.018 0.429± 0.032 0.469± 0.030 0.542± 0.039 0.669± 0.078 5.00 CoDrugKDE 0.399± 0.121 0.438± 0.115 0.473± 0.126 0.523± 0.130 0.575± 0.149 0.627± 0.145 0.678± 0.139 0.751± 0.104 0.845± 0.038 0.906± 0.027 4.00 Logistic 0.590± 0.062 0.636± 0.086 0.692± 0.100 0.760± 0.108 0.807± 0.087 0.873± 0.041 0.903± 0.034 0.916± 0.033 0.925± 0.020 0.934± 0.008 2.80 KMM 0.448± 0.022 0.504± 0.033 0.546± 0.005 0.585± 0.017 0.616± 0.010 0.675± 0.025 0.743± 0.012 0.809± 0.012 0.860± 0.016 0.907± 0.014 2.20 SKMM 0.466± 0.025 0.521± 0.013 0.570± 0.029 0.627± 0.051 0.688± 0.048 0.767± 0.017 0.796± 0.018 0.830± 0.016 0.891± 0.015 0.951± 0.003 1.00 Table 13: hERG_Karim — Global calibration (mean ± std). Best is bold; second best is gray bold. 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.564± 0.026 0.600± 0.024 0.643± 0.034 0.692± 0.032 0.734± 0.037 0.776± 0.045 0.820± 0.044 0.862± 0.037 0.912± 0.024 0.958± 0.014 4.30 KDE 0.700± 0.180 0.700± 0.180 0.700± 0.180 0.700± 0.180 0.700± 0.180 0.709± 0.170 0.733± 0.151 0.737± 0.149 0.755± 0.157 0.807± 0.177 5.50 CoDrugKDE 0.551± 0.037 0.592± 0.036 0.636± 0.045 0.690± 0.042 0.737± 0.057 0.779± 0.060 0.822± 0.053 0.863± 0.046 0.907± 0.032 0.949± 0.021 3.90 Logistic 0.535± 0.043 0.578± 0.039 0.622± 0.043 0.674± 0.037 0.719± 0.028 0.765± 0.030 0.813± 0.035 0.861± 0.027 0.910± 0.027 0.957± 0.015 2.50 KMM 0.545± 0.027 0.587± 0.038 0.645± 0.045 0.686± 0.057 0.728± 0.055 0.776± 0.055 0.812± 0.060 0.844± 0.074 0.883± 0.066 0.930± 0.065 3.60 SKMM 0.523± 0.046 0.570± 0.053 0.606± 0.067 0.645± 0.067 0.687± 0.058 0.747± 0.053 0.795± 0.046 0.846± 0.028 0.899± 0.018 0.944± 0.016 1.20 Table 14: hERG_Karim — Mondrian calibration (mean ± std). Best is bold; second best is gray bold. 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.594± 0.044 0.629± 0.040 0.664± 0.037 0.696± 0.036 0.729± 0.037 0.768± 0.038 0.812± 0.032 0.853± 0.019 0.899± 0.015 0.941± 0.009 4.30 KDE 0.545± 0.129 0.545± 0.129 0.591± 0.115 0.612± 0.115 0.620± 0.120 0.642± 0.125 0.642± 0.125 0.662± 0.140 0.684± 0.137 0.716± 0.071 4.30 CoDrugKDE 0.571± 0.043 0.612± 0.045 0.649± 0.042 0.681± 0.044 0.720± 0.045 0.759± 0.041 0.801± 0.035 0.840± 0.026 0.887± 0.014 0.929± 0.010 3.20 Logistic 0.546± 0.046 0.590± 0.043 0.631± 0.045 0.675± 0.036 0.716± 0.033 0.758± 0.030 0.811± 0.024 0.861± 0.018 0.909± 0.018 0.950± 0.011 2.50 KMM 0.572± 0.058 0.609± 0.057 0.656± 0.062 0.694± 0.072 0.726± 0.075 0.764± 0.085 0.802± 0.090 0.831± 0.093 0.868± 0.089 0.914± 0.069 4.40 SKMM 0.548± 0.078 0.589± 0.074 0.625± 0.069 0.666± 0.067 0.706± 0.063 0.753± 0.062 0.797± 0.055 0.839± 0.040 0.884± 0.029 0.931± 0.023 2.30