Kernel Mean Matching (KMM)

• Kernel Mean Matching (KMM) is a convex quadratic programming approach that aligns source and target distributions in RKHS by matching their empirical kernel mean embeddings.
• It minimizes discrepancies between weighted source and target data under covariate shift, ensuring unbiased estimation of test expectations with finite-sample and asymptotic guarantees.
• Scalable adaptations such as AMKM and residual-based corrections improve computational efficiency and reduce variance in large-scale applications.

Kernel Mean Matching (KMM) is a convex quadratic programming approach for estimating importance weights between probability distributions under scenarios such as covariate shift. KMM operates by aligning empirical kernel mean embeddings in a reproducing kernel Hilbert space (RKHS). It has become a foundational tool in importance-weighted empirical risk minimization, distributional matching for causal inference, and as a building block in scalable, adaptive, and content-addressable generative modeling frameworks.

1. Mathematical Formulation

KMM seeks to minimize the discrepancy between the (possibly weighted) source distribution and the target distribution in RKHS. Given source inputs $\{x_i\}_{i=1}^n \sim P_{\mathrm{tr}}$ and target inputs $\{z_j\}_{j=1}^m \sim P_{\mathrm{te}}$, and letting $\phi: \mathcal{X} \to \mathcal{H}$ be the feature map associated with a kernel $k(\cdot,\cdot)$, the KMM problem is

$$\min_{w \in \mathbb{R}^n} \left\| \frac{1}{n} \sum_{i=1}^n w_i \phi(x_i) - \frac{1}{m} \sum_{j=1}^m \phi(z_j) \right\|_{\mathcal{H}}^2$$

subject to normalization and box constraints:

$\sum_{i=1}^n w_i = n,\quad 0 \leq w_i \leq B$

where $B$ is an upper bound on the true density ratio. This reduces, via the kernel trick, to a quadratic program (QP):

$$\min_{w \in \mathbb{R}^n} \frac{1}{n^2} w^T K w - \frac{2}{n m} \kappa^T w$$

$$\text{s.t. } \mathbf{1}^T w = n, \quad 0 \leq w_i \leq B$$

with $K_{ik} = k(x_i, x_k)$ and $\kappa_i = \sum_{j=1}^m k(x_i, z_j)$ (Yu et al., 2012, Lam et al., 2019). The solution provides empirical estimates $\{w_i\}$ of the Radon–Nikodym derivative $\frac{dP_{\mathrm{te}}}{dP_{\mathrm{tr}}}(x)$.

2. Theoretical Properties

KMM possesses finite-sample and asymptotic guarantees. Under covariate shift ($P_{\mathrm{tr}}(y|x)=P_{\mathrm{te}}(y|x)$), KMM weights permit unbiased estimation of the test expectation:

$$\mathbb{E}_{\mathrm{te}}[Y] = \int m(x) P_{\mathrm{te}}(dx) = \int m(x) \beta(x) P_{\mathrm{tr}}(dx)$$

where $\beta(x)$ is the true importance weight and $m(x) = \mathbb{E}[Y | X=x]$.

Convergence rates depend on the regularity of $m$ with respect to the RKHS and the capacity of the kernel:

• If $m \in \mathcal{H}$, KMM achieves the parametric rate $O(n_{\mathrm{tr}}^{-1/2} + n_{\mathrm{te}}^{-1/2})$ with high probability.
• If $m$ lies in certain ranges of the kernel integral operator, the rate becomes $O(n_{\mathrm{tr}}^{-\theta/(\theta+2)})$ for some $\theta>0$.
• If $m$ is highly irregular, convergence can be logarithmic in sample size (Yu et al., 2012).

Adaptivity is a central property; KMM achieves these rates without prior knowledge of the regularity of $m$ or kernel capacity, and automatically leverages the test-sample distribution.

3. Algorithmic Extensions and Scalable KMM

Standard KMM QP's computational cost is cubic in the source sample size. Several approaches address computational and statistical scalability:

Adaptive Matching of Kernel Means (AMKM)

AMKM operates in two stages:

1. Randomized Subset Optimization: For $T$ repeats, small random subsets of the reference pool are selected. On each, a KMM QP is solved, then further refined to focus on points with highest preliminary weights (high “information potential,” i.e., large $V(S)$).
2. Convex Fusion: The candidate solutions are merged via a low-dimensional convex QP, yielding a nonnegative combination as the final weight vector.

AMKM exhibits error $O(T^{-1/2})+O(n_s^{-1/2})$ and per-iteration cost $O(T n^3 + T n_s^3 + T^3)$, with $n, n_s, T \ll n_r$, dramatically reducing memory and runtime, thus enabling efficient streaming and incremental learning (Cheng et al., 2020).

Empirical results on both small and large-scale datasets (Monks, Ionosphere, Climate, Forest, Letter, CIFAR-100) show that AMKM typically matches or lowers normalized mean squared error (NMSE) compared to full KMM and other advanced variants, with significantly lower computational burden.

Residual-Based Corrections

Practical KMM-based importance weighting can have high variance, especially with limited sample overlap. Combined estimators integrate a control variate