A Kernel Two-Sample Test

Authors: Arthur Gretton, Karsten M. Borgwardt, Malte J. Rasch, Bernhard Schölkopf, Alexander Smola, Alex J. Smola
Published: 2012
Journal: Journal of Machine Learning Research
Link: PDF

Summary

This paper provides a comprehensive treatment of the Maximum Mean Discrepancy (MMD) as a statistic for two-sample testing. It establishes theoretical properties, efficient estimators, and practical guidelines for kernel selection.

Key Contributions

1. Maximum Mean Discrepancy (MMD)

The MMD between distributions $P$ and $Q$ is:

$MMD[\mathcal{F}, P, Q] = \sup_{f \in \mathcal{F}} \left( \mathbb{E}_P[f(x)] - \mathbb{E}_Q[f(y)] \right)$

For unit ball in RKHS $\mathcal{H}$: $MMD^2[\mathcal{H}, P, Q] = \|\mu_P - \mu_Q\|_{\mathcal{H}}^2$

2. Characteristic Kernels

Definition: Kernel $K$ is characteristic if $MMD[\mathcal{H}, P, Q] = 0 \iff P = Q$.

Examples of characteristic kernels:

• Gaussian RBF: $K(x,y) = \exp(-\|x-y\|^2/2\sigma^2)$
• Laplacian: $K(x,y) = \exp(-\|x-y\|_1/\sigma)$
• Matérn family

3. Unbiased MMD Estimator

Given $X = \{x_1, \ldots, x_n\} \sim P$ and $Y = \{y_1, \ldots, y_m\} \sim Q$:

$\widehat{MMD}_u^2 = \frac{1}{n(n-1)}\sum_{i \neq j}h(x_i, x_j, y_i, y_j)$

where: $h(x_i, x_j, y_i, y_j) = K(x_i, x_j) + K(y_i, y_j) - K(x_i, y_j) - K(x_j, y_i)$

4. Asymptotic Distribution

Under $H_0: P = Q$: $n \cdot \widehat{MMD}_u^2 \xrightarrow{d} \sum_{l=1}^{\infty} \lambda_l (z_l^2 - 2)$

where $z_l \sim \mathcal{N}(0, 2)$ and $\lambda_l$ are eigenvalues of the covariance operator.

Statistical Test Framework

Hypothesis Test

$H_0: P = Q \quad \text{vs} \quad H_1: P \neq Q$

Test Statistic Variants

1. Quadratic-time statistic (U-statistic): $\widehat{MMD}_u^2 = \frac{1}{n(n-1)}\sum_{i \neq j} K(x_i, x_j) + \frac{1}{m(m-1)}\sum_{i \neq j} K(y_i, y_j) - \frac{2}{nm}\sum_{i,j} K(x_i, y_j)$

2. Linear-time statistic: $\widehat{MMD}_l^2 = \frac{2}{n}\sum_{i=1}^{n/2} h(x_{2i-1}, x_{2i}, y_{2i-1}, y_{2i})$

Threshold Selection

Permutation test (recommended for finite samples):

1. Compute observed $\widehat{MMD}^2$
2. Pool samples, permute, compute $\widehat{MMD}^2$
3. Repeat $B$ times
4. Reject if observed exceeds $(1-\alpha)$ quantile

Asymptotic test (large samples):

1. Estimate eigenvalues $\hat{\lambda}_l$
2. Simulate null distribution
3. Use appropriate quantile

Implementation

import numpy as np
from scipy.spatial.distance import cdist

def rbf_kernel(X, Y, sigma=1.0):
    """Compute RBF kernel matrix."""
    sq_dist = cdist(X, Y, 'sqeuclidean')
    return np.exp(-sq_dist / (2 * sigma**2))

def mmd_squared_u_statistic(X, Y, kernel_func):
    """
    Compute unbiased MMD^2 U-statistic.
    """
    n, m = len(X), len(Y)
    
    K_XX = kernel_func(X, X)
    K_YY = kernel_func(Y, Y)
    K_XY = kernel_func(X, Y)
    
    # Remove diagonal for unbiased estimate
    term1 = (K_XX.sum() - np.trace(K_XX)) / (n * (n - 1))
    term2 = (K_YY.sum() - np.trace(K_YY)) / (m * (m - 1))
    term3 = 2 * K_XY.mean()
    
    return term1 + term2 - term3

def mmd_linear_time(X, Y, kernel_func):
    """
    Compute linear-time MMD estimator.
    """
    n = min(len(X), len(Y))
    n = n - (n % 2)  # Make even
    
    X, Y = X[:n], Y[:n]
    
    # Pair up samples
    h_values = []
    for i in range(0, n, 2):
        x1, x2 = X[i:i+1], X[i+1:i+2]
        y1, y2 = Y[i:i+1], Y[i+1:i+2]
        
        h = (kernel_func(x1, x2)[0,0] + kernel_func(y1, y2)[0,0]
             - kernel_func(x1, y2)[0,0] - kernel_func(x2, y1)[0,0])
        h_values.append(h)
    
    return np.mean(h_values)

def mmd_permutation_test(X, Y, kernel_func, n_permutations=1000, alpha=0.05):
    """
    Perform MMD two-sample test with permutation.
    """
    # Observed statistic
    observed = mmd_squared_u_statistic(X, Y, kernel_func)
    
    # Pool and permute
    combined = np.vstack([X, Y])
    n = len(X)
    
    null_statistics = []
    for _ in range(n_permutations):
        perm = np.random.permutation(len(combined))
        X_perm = combined[perm[:n]]
        Y_perm = combined[perm[n:]]
        null_statistics.append(mmd_squared_u_statistic(X_perm, Y_perm, kernel_func))
    
    # Compute p-value
    p_value = np.mean(np.array(null_statistics) >= observed)
    
    return {
        'statistic': observed,
        'p_value': p_value,
        'reject': p_value < alpha,
        'threshold': np.quantile(null_statistics, 1 - alpha)
    }

Kernel Selection

Median Heuristic

Set bandwidth to median of pairwise distances: $\sigma = \text{median}\{\|x_i - x_j\|: i < j\}$

Multiple Kernel MMD

Use multiple bandwidths for robustness: $K_{sum}(x, y) = \sum_{j=1}^J K_{\sigma_j}(x, y)$

Optimized Kernel Selection

Maximize test power by optimizing kernel parameters on held-out data.

Applications

1. Model Validation: Test if generator matches real data distribution
2. Feature Selection: Test which features differ between groups
3. Change Detection: Detect distribution shifts in streaming data
4. Domain Adaptation: Measure domain discrepancy

Power Analysis

Test power depends on:

• Sample size: Power increases with $n$
• Kernel bandwidth: Optimal $\sigma$ depends on distribution separation
• Dimensionality: Power decreases in high dimensions

For well-chosen kernel: $\text{Power} \to 1 \text{ as } n \to \infty$

Citation

@article{gretton2012kernel,
  title={A kernel two-sample test},
  author={Gretton, Arthur and Borgwardt, Karsten M and Rasch, Malte J 
          and Sch{\"o}lkopf, Bernhard and Smola, Alexander},
  journal={Journal of Machine Learning Research},
  volume={13},
  pages={723--773},
  year={2012}
}

Further Reading

• Borgwardt, K., et al. (2006). Integrating structured biological data by kernel MMD
• Sutherland, D., et al. (2017). Generative models and model criticism via optimized MMD
• Liu, F., et al. (2020). Learning deep kernels for non-parametric two-sample tests