Maximum Mean Discrepancy (MMD)

Maximum mean discrepancy compares two distributions by embedding each one as a single point in a reproducing kernel Hilbert space and measuring the distance between those points.

\[ \mathrm{MMD}^2(P, Q) = \| \mu_P - \mu_Q \|_{\mathcal{H}}^2 = \mathbb{E}[k(x,x’)] - 2\,\mathbb{E}[k(x,y)] + \mathbb{E}[k(y,y’)] \]

With a characteristic kernel, such as the Gaussian, the embedding is injective, so MMD is zero only when the two distributions are identical. Unlike KS shift it is defined in any number of dimensions, which is why it turns up in drift detection on embeddings. Two practical facts about the estimate decide whether the number you compute means anything.

The biased estimator has positive null expectation

Most write-ups give the estimator that replaces every expectation by a full double average, \( \frac{1}{m^2}\sum_i\sum_{i’} k(x_i, x_{i’}) \). Gretton et al. call that one biased, with an upward bias, and then define a second estimator that drops the \( i = i’ \) terms and normalizes by \( m(m-1) \). The difference is not cosmetic. In the code, grams returns the three kernel blocks (within the first sample, within the second, and between them) of a Gaussian kernel on the pooled points, and the two functions turn those blocks into the biased and the unbiased estimate. Both samples are drawn from the same five-dimensional standard normal, so any positive reading is estimator behavior, not drift.

import numpy as np


def grams(X, Y, sigma):
    Z = np.vstack([X, Y])
    sq = np.einsum("ij,ij->i", Z, Z)
    d2 = np.maximum(sq[:, None] + sq[None, :] - 2 * Z @ Z.T, 0.0)
    K = np.exp(-d2 / (2 * sigma ** 2))
    m = len(X)
    return K[:m, :m], K[m:, m:], K[:m, m:]


def mmd2_biased(Kxx, Kyy, Kxy):
    return Kxx.mean() + Kyy.mean() - 2 * Kxy.mean()


def mmd2_unbiased(Kxx, Kyy, Kxy):
    m, n = len(Kxx), len(Kyy)
    return ((Kxx.sum() - np.trace(Kxx)) / (m * (m - 1))
            + (Kyy.sum() - np.trace(Kyy)) / (n * (n - 1))
            - 2 * Kxy.mean())


rng = np.random.default_rng(5)
d, sigma = 5, np.sqrt(5.0)
print("%7s %7s %13s %13s %14s %14s"
      % ("m = n", "reps", "mean biased", "mean unbiased", "biased x m", "share u < 0"))
for m, reps in ((50, 400), (100, 400), (200, 300), (400, 200), (800, 120)):
    bs, us = [], []
    for _ in range(reps):
        g = grams(rng.normal(size=(m, d)), rng.normal(size=(m, d)), sigma)
        bs.append(mmd2_biased(*g))
        us.append(mmd2_unbiased(*g))
    bs, us = np.array(bs), np.array(us)
    print("%7d %7d %13.6f %13.6f %14.4f %13.1f%%"
          % (m, reps, bs.mean(), us.mean(), bs.mean() * m, (us < 0).mean() * 100))

  m = n    reps   mean biased mean unbiased     biased x m    share u < 0
     50     400      0.023101      0.000313         1.1551          55.8%
    100     400      0.011397      0.000023         1.1397          58.0%
    200     300      0.005684     -0.000006         1.1367          56.0%
    400     200      0.002794     -0.000048         1.1174          57.0%
    800     120      0.001425      0.000003         1.1397          60.0%

Both samples come from the same distribution, so the truth is zero. The biased estimator reports 0.023101 at 50 points a side and 0.001425 at 800, and the fourth column shows why: the product of that value and the sample size is flat at about 1.14, consistent with a null expectation of order \(1/m\); individual readings also contain sampling variation. Compare two windows of different sizes with this estimator and the smaller window looks like it drifted.

The unbiased estimator averages zero and, because its null distribution is right-skewed, comes out negative in most windows of this simulation, not necessarily for every kernel and null distribution. A negative squared distance looks like a bug and is not one; it is what an unbiased estimate of zero looks like when the truth sits at the boundary.

The bandwidth decides what you can see

For Gaussian bandwidth \(\sigma\to0\) and distinct pooled observations, off-diagonal kernels vanish: the biased squared estimator tends to \(1/m+1/n\), while the unbiased estimator tends to zero. At \(\sigma\to\infty\), both tend to zero. The biased estimator can also be exactly zero for identical empirical measures; positive null expectation is not a strictly positive result for every sample. Bandwidth selection affects power as well as scale.

The sweep below pairs each bandwidth with a permutation p-value. permutation_p pools the two samples, shuffles which points count as the first sample, recomputes the unbiased estimate with grams and mmd2_unbiased from the previous block, and reports the share of shuffles at least as large as the observed value, counting the observed arrangement itself. The median heuristic sets \( \sigma \) to the square root of the median pairwise squared distance in the pooled data. The block continues the same session.

from scipy.spatial.distance import pdist


def permutation_p(Z, m, sigma, reps=400):
    stat = mmd2_unbiased(*grams(Z[:m], Z[m:], sigma))
    rng, idx, hits = np.random.default_rng(7), np.arange(len(Z)), 0
    for _ in range(reps):
        rng.shuffle(idx)
        hits += mmd2_unbiased(*grams(Z[idx[:m]], Z[idx[m:]], sigma)) >= stat
    return stat, (hits + 1) / (reps + 1)


def median_sigma(Z):
    return float(np.sqrt(np.median(pdist(Z, "sqeuclidean"))))


rng = np.random.default_rng(31)
m = 300
cases = {"mean moved by 0.5": (rng.normal(size=m), rng.normal(0.5, 1.0, m)),
         "same mean, sd 1.0 to 1.6": (rng.normal(size=m), rng.normal(0.0, 1.6, m))}
for name, (x, y) in cases.items():
    Z = np.concatenate([x, y])[:, None]
    med = median_sigma(Z)
    print("%s   median heuristic sigma = %.4f" % (name, med))
    print("%10s %16s %10s" % ("sigma", "MMD^2 unbiased", "perm p"))
    vals = []
    for sigma in (0.001, 0.01, 0.1, med, 10.0, 100.0, 1000.0):
        stat, p = permutation_p(Z, m, sigma)
        vals.append(stat)
        print("%10.4g %16.3e %10.4f%s"
              % (sigma, stat, p, "   <- median" if sigma == med else ""))
    print("largest value / smallest positive value: %.3g"
          % (max(vals) / min(v for v in vals if v > 0)))

mean moved by 0.5   median heuristic sigma = 1.0072
     sigma   MMD^2 unbiased     perm p
     0.001        1.419e-04     0.2643
      0.01        1.806e-03     0.0025
       0.1        1.398e-02     0.0025
     1.007        6.789e-02     0.0025   <- median
        10        2.929e-03     0.0025
       100        2.999e-05     0.0025
      1000        3.000e-07     0.0025
largest value / smallest positive value: 2.26e+05
same mean, sd 1.0 to 1.6   median heuristic sigma = 1.2730
     sigma   MMD^2 unbiased     perm p
     0.001        3.191e-04     0.0524
      0.01        1.446e-03     0.0125
       0.1        7.385e-03     0.0025
     1.273        3.514e-02     0.0025   <- median
        10        8.123e-05     0.1746
       100       -6.134e-07     0.4888
      1000       -6.288e-09     0.4963
largest value / smallest positive value: 433

Two different things move. The value swings across five orders of magnitude for the same data, so a stored threshold on \( \mathrm{MMD}^2 \) is meaningless unless the kernel is stored beside it. And the two alternatives behave differently: a shift in the mean is still detected at \( \sigma = 1000 \), because the permutation null shrinks along with the statistic, while a change in spread is lost once the kernel is wide — at \( \sigma = 100 \) the p-value has walked back to 0.4888. A wide Gaussian kernel is, to first order, a comparison of means.

What it buys you over a per-feature monitor

The reason to pay for a quadratic-time kernel statistic is that it compares the two samples as clouds rather than one axis at a time. The block below tests each feature alone and then both together with the same permutation_p, setting \( \sigma \) by the median heuristic on the pooled points each time.

rng = np.random.default_rng(101)
n = 600
base = rng.normal(size=(n, 2))
a, b = rng.normal(size=n), rng.normal(size=n)
cur = np.column_stack([a, 0.8 * a + 0.6 * b])

print("correlation   baseline %+.4f   current %+.4f"
      % (np.corrcoef(base.T)[0, 1], np.corrcoef(cur.T)[0, 1]))
print("%-24s %9s %14s %9s" % ("tested on", "sigma", "MMD^2", "perm p"))
for label, sl in (("feature 1 alone", [0]), ("feature 2 alone", [1]),
                  ("both features jointly", [0, 1])):
    Z = np.vstack([base[:, sl], cur[:, sl]])
    s = median_sigma(Z)
    stat, p = permutation_p(Z, n, s)
    print("%-24s %9.4f %14.3e %9.4f" % (label, s, stat, p))

correlation   baseline +0.0080   current +0.8245
tested on                    sigma          MMD^2    perm p
feature 1 alone             0.9895      3.688e-04    0.2743
feature 2 alone             0.9862     -1.181e-03    0.8878
both features jointly       1.6417      2.151e-02    0.0025

The generator preserves each feature’s standard-normal marginal distribution. The per-feature tests do not reject their null hypotheses: p of 0.2743 and 0.8878, with the second statistic slightly negative. Tested jointly, the same data gives 0.0025, the smallest value 400 permutations can produce. Nobody had to nominate the interaction; the kernel sees the pair of clouds and the clouds differ.

The price is the \( O(n^2) \) Gram matrix, a bandwidth you now own, and a statistic with no natural units. In practice that means storing the kernel and the window size with the number, using the unbiased estimator, and reading significance from permutations rather than from the raw value.

References

• Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., & Smola, A. (2012). A kernel two-sample test. Journal of Machine Learning Research, 13, 723–773.
• Sejdinovic, D., Sriperumbudur, B., Gretton, A., & Fukumizu, K. (2013). Equivalence of distance-based and RKHS-based statistics in hypothesis testing. The Annals of Statistics, 41(5), 2263–2291.