Reliable Risk Estimation with Kernel Mean Matching under Dependent Sampling

Alexey Zaytsev; Diana Koldasbayeva; Egor Serov

cs.LGOct 13, 2026

Reliable Risk Estimation with Kernel Mean Matching under Dependent Sampling

Authors: Alexey Zaytsev, Diana Koldasbayeva, Egor Serov

Organizations: Applied AI Institute

Abstract

Kernel Mean Matching (KMM) estimates risk under covariate shift by reweighting labelled source samples so that their kernel mean embedding matches that of unlabeled target samples. Existing finite-sample guarantees are usually stated for independent samples, although many applications involve temporal, spatial, or clustered dependence. We give a finite-sample KMM analysis under dependent sampling by replacing nominal sample sizes with effective sample sizes in the concentration steps used by the classical proof. The approach uses Bernstein inequalities for geometrically strongly mixing processes and derives a model-based effective sample size. The resulting bound separates overlap, dependence, and smoothness effects with specific constants derived under a Gaussian process assumption. This analysis yields an additional tradeoff: smoother losses preserve more input dependence and therefore reduce effective sample size, while also reducing function complexity and empirical-to-population risk gaps. Applying the result to prediction yields confidence intervals that achieve near-nominal coverage in synthetic experiments and substantially reduce undercoverage on real monitoring data.

Figures & tables

Reference	Setting	Relation to Ours
( Huang et al., 2006 ; Gretton et al., 2009 )	Covariate shift KMM	Introduces KMM objective
( Yu and Szepesvári, 2012 )	Finite-sample bound	i.i.d. proof template
( Roberts et al., 2017 ; Koldasbayeva and Zaytsev, 2025 )	Spatio-temporal CV	Dependence motivation
( Merlevède et al., 2009 )	$\beta$ -mixing processes	Scalar concentration tools
This paper	Dependent KMM	Explicit bounds via $n_{\mathrm{eff}}$

Table 1: Positioning relative to related theory.

Figure 1: Validation of the GP-induced effective sample size. The variance-based empirical

n_{\mathrm{eff}}

aligns with our closed-form prediction (diagonal line) for both Target

Y

(point estimation) and Risk (loss evaluation) across matched GP realizations.

	Target Mean ( $Y$ )		Target Risk ( $L$ )
Method	Coverage	Radius	Coverage	Radius
Unweighted	0.017 ∗	0.0362	0.080 ∗	0.0019
Nominal $n$	0.917 ∗	0.0626	0.853 ∗	0.0055
Block-Bootstrap	1.000	0.4538	0.967	0.0133
Input $n_{\mathrm{eff}}$	1.000	0.1696	0.987	0.0150
GP $\bm{n_{\mathrm{eff}}}$ [Ours]	1.000	0.1380	0.970	0.0120

Table 2: Confidence Interval Calibration under Dependent Sampling. Coverage and interval radius for target-mean and target-risk estimation. Asterisks ( ∗ ) indicate significant empirical undercoverage below the nominal 0.95 level (

<0.925

). Bold denotes the best valid practical estimator.

$\bm{n}$	Method	Coverage	Radius	WS $\downarrow$
48	Nominal $n$	0.773 ∗	0.0048	0.0275
	Input $n_{\mathrm{eff}}$	0.910 ∗	0.0128	0.0290
	GP $\bm{n_{\mathrm{eff}}}$ [Ours]	0.943	0.0101	0.0260
64	Nominal $n$	0.773 ∗	0.0046	0.0218
	Input $n_{\mathrm{eff}}$	0.980	0.0124	0.0258
	GP $\bm{n_{\mathrm{eff}}}$ [Ours]	0.947	0.0098	0.0217

Table 3: Sample-Size Scaling of Target Risk Intervals. Nominal intervals systematically under-cover across all sample sizes. Asterisks ( ∗ ) indicate significant undercoverage below nominal 0.95 (

<0.925

). Bold denotes the best valid practical estimator.

Dataset: PM 2.5 (Temporal Shift)
	Target Mean ( $Y$ )		Target Risk ( $L$ )
Method	Coverage	Radius	Coverage	Radius
Unweighted	0.604 ∗	0.0189	0.852 ∗	0.0016
Nominal $n$	0.848 ∗	0.0672	0.900 ∗	0.0056
Block-Bootstrap	0.884 ∗	0.0832	0.916 ∗	0.0068
Input $n_{\mathrm{eff}}$	0.928	0.1359	0.952	0.0104

Table 4: Confidence Interval Calibration under Dependent Sampling on Real Datasets. Empirical coverage and interval radius for target-mean and target-risk estimation. Asterisks ( ∗ ) indicate significant empirical undercoverage below the relaxed nominal

0.92

level. Bold denotes the best valid practical estimator (tightest radius among methods with valid coverage).

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Controls	Source of Value
$n_{\mathrm{eff}}(t)$	Scalar average at deviation $t$	Mixing Bernstein bound
$n_{\mathrm{eff}}^{\mathrm{tr}}$	Source scalar effective size (moderate regime)	Long-run variance proxy
$n_{\mathrm{eff,RKHS}}^{\mathrm{tr}},n_{\mathrm{eff,RKHS}}^{\mathrm{te}}$	RKHS empirical feature-matching (MMD)	Geometric block size
$n_{\mathrm{eff}}^{\mathrm{GP}}(h)$	GP-modeled conditional loss average	Closed-form GP covariance

Appendix

Table 5: Taxonomy of effective sample sizes used in the analysis.

	Target Mean ( $Y$ )				Target Risk ( $L$ )
Method	MAE	Coverage	Radius	WS $\downarrow$	MAE	Coverage	Radius	WS $\downarrow$
Unweighted	0.1138	0.017 ∗	0.0362	3.182	0.0048	0.080 ∗	0.0019	0.119
Nominal $n$	0.0196	0.917 ∗	0.0626	0.134	0.0025	0.853 ∗	0.0055	0.024
Block-Bootstrap	0.0196	1.000	0.4538	0.908	0.0025	0.967	0.0133	0.026
Input $n_{\mathrm{eff}}$	0.0196	1.000	0.1696	0.339	0.0025	0.987	0.0150	0.031
GP $n_{\mathrm{eff}}$ (Ours)	0.0196	1.000	0.1380	0.276	0.0025	0.970	0.0120	0.026

Appendix

Table 6: Full Benchmark Results including Point Estimation and Winkler Score. Point estimates (MAE) are identical for Nominal

n

, Block-Bootstrap, Input

n_{\mathrm{eff}}

, GP

n_{\mathrm{eff}}

, and the Strict Bound, as they rely on the exact same KMM weights. An asterisk ( ∗ ) indicates empirical undercoverage significantly below the nominal 0.95 level (

<0.925

). Winkler Score (WS, lower is better) penalizes both interval width and undercoverage. Bold highlights the best valid practical estimator (tightest radius and lowest WS while maintaining valid coverage). Strict Thm 1 and Oracle methods are excluded from best-method highlighting.

Figure 2: Validation under Deterministic Misspecification. Empirical variance-based

n_{\mathrm{eff}}

vs. our predicted

n_{\mathrm{eff}}^{\mathrm{GP}}

for both Target

Y

and Risk when evaluated on a deterministic, non-GP harmonic surface. The surrogate formula remains highly predictive.

Figure 3: GP Realizations (Full Grid). Ratio of the empirical to predicted $n_{\mathrm{eff}}$ across diverse GP length-scales ( $h$ ) and input dependencies ( $\ell_{x}$ ). Ratios near $1.0$ indicate close calibration. The approximation tracks the empirical effective sample size across much of the grid, with larger deviations in some strong-dependence and smooth-function regimes.

Figure 4: Deterministic Misspecification (Full Grid). Ratio of the empirical to predicted $n_{\mathrm{eff}}$ when the true function is a deterministic harmonic surface with varying frequency $\omega$ . The GP surrogate remains informative beyond the matched stochastic-GP setting, although its calibration varies across dependence and frequency regimes.

$\bm{n}$	Method	Coverage	Radius	Winkler Score $\downarrow$
32	Unweighted	0.717 ∗	0.0026	0.0201
	Nominal $n$	0.647 ∗	0.0047	0.0394
	Input $n_{\mathrm{eff}}$	0.893 ∗	0.0123	0.0325
	GP $n_{\mathrm{eff}}$ (Ours)	0.850 ∗	0.0098	0.0313
48	Unweighted	0.727 ∗	0.0022	0.0143
	Nominal $n$	0.773 ∗	0.0048	0.0275

Appendix

Table 7: Sample-size scaling of Target Risk intervals (

n\in\{32,\dots,256\}

). Bold highlights the best valid practical estimator (tightest radius and lowest Winkler Score among methods maintaining valid coverage). Asterisks ( ∗ ) indicate significant empirical undercoverage below the nominal 0.95 level (

<0.925

). Strict Thm 1 is excluded from best-method highlighting due to worst-case conservatism.

$\bm{\ell_{x}}$	Method	Coverage	Radius	Winkler Score $\downarrow$
1.5	Unweighted	0.843 ∗	0.0016	0.0061
	Nominal $n$	0.840 ∗	0.0041	0.0161
	Input $n_{\mathrm{eff}}$	0.953	0.0078	0.0176
	GP $n_{\mathrm{eff}}$ (Ours)	0.923 ∗	0.0062	0.0159
2.0	Unweighted	0.843 ∗	0.0016	0.0068
	Nominal $n$	0.870 ∗	0.0043	0.0143

Appendix

Table 8: Hyperparameter Ablation: Input Dependence (

\ell_{x}\in[1.5,6.0]

) on Target Risk. Asterisks ( ∗ ) indicate significant empirical undercoverage below the nominal 0.95 level (

<0.925

). Bold highlights the best valid practical estimator (tightest radius and lowest Winkler Score among methods maintaining valid coverage). Strict Thm 1 is excluded from best-method highlighting due to worst-case conservatism.

$\bm{h}$	Method	Coverage	Radius	Winkler Score $\downarrow$
1.00	Unweighted	0.173 ∗	0.0023	0.0749
	Nominal $n$	0.910 ∗	0.0062	0.0194
	Input $n_{\mathrm{eff}}$	0.993	0.0167	0.0338
	GP $n_{\mathrm{eff}}$ (Ours)	0.983	0.0129	0.0271
1.25	Unweighted	0.047 ∗	0.0018	0.1597
	Nominal $n$	0.917 ∗	0.0072	0.0183

Appendix

Table 9: Hyperparameter Ablation: Functional Smoothness (

h\in[1.0,3.0]

) on Target Risk. Asterisks ( ∗ ) indicate significant empirical undercoverage below the nominal 0.95 level (

<0.925

). Bold highlights the best valid practical estimator (tightest radius and lowest Winkler Score among methods maintaining valid coverage).

$\bm{B}$	Method	Coverage	Radius	Winkler Score $\downarrow$
5.0	Unweighted	0.767 ∗	0.0016	0.0100
	Nominal $n$	0.780 ∗	0.0030	0.0154
	Input $n_{\mathrm{eff}}$	0.987	0.0082	0.0168
	GP $n_{\mathrm{eff}}$ (Ours)	0.963	0.0065	0.0141
7.5	Unweighted	0.767 ∗	0.0016	0.0100
	Nominal $n$	0.803 ∗	0.0035	0.0147

Appendix

Table 10: Hyperparameter Ablation: Clipping Bound (

B\in[5.0,30.0]

) on Target Risk. Asterisks ( ∗ ) indicate significant empirical undercoverage below the nominal 0.95 level (

<0.925

). Bold highlights the best valid practical estimator (tightest radius and lowest Winkler Score among methods maintaining valid coverage).

Data Kernel	Method	Coverage	Radius	WS $\downarrow$
RBF (Smooth)	Unweighted	0.110 ∗	0.0022	0.113
	Nominal $n$	0.903 ∗	0.0076	0.023
	Input $n_{\mathrm{eff}}$	0.993	0.0206	0.042
	GP $\bm{n_{\mathrm{eff}}}$ (Ours)	0.987	0.0147	0.031
Matérn-5/2	Unweighted	0.623 ∗	0.0019	0.019
	Nominal $n$	0.850 ∗	0.0059	0.029

Appendix

Table 11: Kernel Type Ablation (Roughness Stress Test) on Target Risk. Asterisks ( ∗ ) denote significant empirical undercoverage below the nominal 0.95 level (

<0.925

). GP

n_{\mathrm{eff}}

maintains valid interval coverage even when the true generating process is non-differentiable, consistently outperforming Input

n_{\mathrm{eff}}

in Winkler Score (WS). Note that the slight undercoverage on Matérn-5/2 is matched by the Oracle i.i.d. benchmark, indicating inherent finite-sample variance rather than a failure of the GP correction.

Dependence Profile	Method	Coverage	Radius	WS $\downarrow$
Stationary	Unweighted	0.700 ∗	0.0020	0.013
	Nominal $n$	0.900 ∗	0.0048	0.014
	Input $n_{\mathrm{eff}}$	0.993	0.0120	0.024
	GP $\bm{n_{\mathrm{eff}}}$ (Ours)	0.990	0.0099	0.020
Abrupt Switch	Unweighted	0.717 ∗	0.0020	0.014
	Nominal $n$	0.880 ∗	0.0046	0.014

Appendix

Table 12: Robustness to Non-Stationary Covariates on Target Risk. Asterisks ( ∗ ) indicate significant empirical undercoverage below the nominal 0.95 level (

<0.925

). Covariates are generated via a non-stationary Gibbs kernel with mean dependence

\bar{\ell}_{x}=2.5

. Plugging the trajectory-average correlation length into our stationary formula maintains valid coverage (

\geq 0.983

) and yields the lowest Winkler Score (WS) among valid estimators across all non-stationary regimes.

Figure 5: Controlled Overlap Stress Test on California Housing Covariates. Increasing

\gamma

progressively separates source and target distributions. Left: Oracle and KMM weight ESS decrease as overlap deteriorates. Right: KMM controls target-mean error under moderate shift, but its bounded weights cannot fully reproduce the density ratio under severe overlap loss. Weight ESS measures weight concentration and is distinct from the dependence-induced

n_{\mathrm{eff}}

. Error bars show one standard error across 100 repetitions.

$\bm{\gamma}$	MMD 2	Oracle ESS	KMM ESS	Oracle Max $\bm{w}$	Unweighted MAE	Oracle MAE	KMM MAE
0.0	0.0007	256.0	218.3	1.0	0.0380	0.0380	0.0381
0.5	0.0048	198.3	154.8	4.3	0.0461	0.0359	0.0355
1.0	0.0176	87.0	71.5	15.5	0.0789	0.0371	0.0354
1.5	0.0393	38.7	31.5	36.6	0.1203	0.0465	0.0365
2.0	0.0704	22.1	26.9	58.9	0.1634	0.0653	0.0460
3.0	0.1549	9.8	26.3	92.8	0.2403	0.1013	0.1260

Appendix

Table 13: Semi-Synthetic Overlap Stress Test on California Housing Covariates. Increasing

\gamma

separates source and target distributions. ESS denotes weight effective sample size (

\operatorname{ESS}_{w}

), not the dependence-induced

n_{\mathrm{eff}}

. MAE is computed over 100 repetitions.

Figure 6: Effect of the KMM clipping bound

B

under deteriorating overlap. Increasing

B

allows KMM to achieve better source–target balance and reduces residual estimation bias, but at the cost of increasingly concentrated weights and lower weight ESS. When the population density-ratio bound

B^{*}

is moderate (

\gamma=1

), the results stabilize once

B

reaches

B^{*}

.Under severe overlap loss (

\gamma=2,3

B^{*}

is much larger than the considered clipping levels, exposing a pronounced balance–stability trade-off. Weight ESS here measures weight concentration and is distinct from the dependence-induced

n_{\mathrm{eff}}

used elsewhere in the paper.

Data	Estimand	Unweighted	Log-IW	KMM
NOAA	Target Mean ( $\mu_{T}$ )	0.0503	0.0389	0.0379
PM 2.5	Target Mean ( $\mu_{T}$ )	0.0208	0.0270	0.0155
EPA Ozone	Target Mean ( $\mu_{T}$ )	0.0290	0.0245	0.0239
NOAA	Target Risk ( $R_{T}(f)$ )	0.00209	0.00205	0.00213
PM 2.5	Target Risk ( $R_{T}(f)$ )	0.00214	0.00205	0.00244
EPA Ozone	Target Risk ( $R_{T}(f)$ )	0.00137	0.00109	0.00120

Appendix

Table 14: Point Estimation on the Three-Domain Real-Data Benchmark. Mean Absolute Error (MAE) is reported for target response mean

\mu_{T}=\mathbb{E}_{T}[Y]

and squared-error risk

R_{T}(f)=\mathbb{E}_{T}[(Y-f(X))^{2}]

. Bold indicates the lowest MAE within each dataset–estimand pair.

Kernel Methods

Kernel methods use similarity functions and reproducing kernel spaces to model nonlinear relationships, quantify uncertainty, and compare data.

See all 227 papers

KMM-CP: Practical Conformal Prediction under Covariate Shift via Selective Kernel Mean Matching
Mar 27, 2026Siddhartha Laghuvarapu, Rohan Deb, Jimeng SunWeighted Conformal Prediction Covariate Shift
Siebel School of Computing and Data Science, University of Illinois Urbana-Champaign, IL, USA.
A Kernel Nonconformity Score for Multivariate Conformal Prediction
Apr 23, 2026Louis Meyer, Wenkai XuConformal Prediction Kernel Methods
University of Warwick
On Finite-sample Concentration of Median of Incomplete U-Statistics
May 30, 2026Nong Minh Hieu, Antoine LedentRobust Estimation U-Statistics
Singapore Management University, School of Computing and Information Systems

Mar 27, 2026J/K move · Enter open · S save

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

Siddhartha Laghuvarapu, Rohan Deb, Jimeng Sun

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

Weighted Conformal Prediction Covariate Shift Distribution Shift Robustness Kernel Methods

Open paper AskPDF arXiv

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.

Abstract

Kernel Methods

Explore similar work

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

A Kernel Nonconformity Score for Multivariate Conformal Prediction

On Finite-sample Concentration of Median of Incomplete U-Statistics