Compare Model Input Distributions with Maximum Mean Discrepancy
A model can keep returning valid outputs while the data reaching it gradually changes. An image service may receive photos from new cameras. A text classifier may see a new mix of topics. An embedding pipeline may start processing documents from a different source. Accuracy can degrade even though the model binary, API, and serving code are unchanged.
Maximum mean discrepancy (MMD) gives you a way to compare two samples without reducing each one to a few statistics such as a mean and variance. It uses pairwise similarities to produce a single discrepancy score. That makes it useful as one signal for distribution shift, especially when inputs or representations are vectors.
This article builds MMD from a small example, shows the estimator used in practice, and covers the kernel choices and monitoring mistakes that determine whether the score is informative.
Treat drift as a two-sample comparison
Suppose you saved a reference sample from a period when a model behaved acceptably:
[ X = {x_1, x_2, \ldots, x_n} ]
Later, you collect a current sample:
[ Y = {y_1, y_2, \ldots, y_m} ]
The monitoring question is not necessarily whether every new point resembles an old point. Natural data has variation. The useful question is whether the two samples look as though they came from the same underlying distribution.
For one numeric feature, you could compare histograms or quantiles. For a vector with hundreds of dimensions, checking every coordinate separately can miss changes in relationships between coordinates. Two datasets can have similar per-coordinate means and still have different joint structure.
MMD approaches the problem through a kernel. A kernel is a function (k(a,b)) that assigns a similarity value to two observations. MMD compares three kinds of similarity:
- pairs drawn from the reference sample,
- pairs drawn from the current sample,
- pairs with one point from each sample.
If the two distributions are similar under the chosen kernel, within-sample and cross-sample similarities should be compatible. If cross-sample similarity becomes systematically different, the discrepancy grows.
Start with a Gaussian kernel
A common choice for numeric vectors is the Gaussian radial basis function kernel:
[ k(a,b) = \exp\left(-\frac{|a-b|_2^2}{2\sigma^2}\right) ]
Here, (\sigma) is the kernel bandwidth.
The kernel returns 1 when (a=b). As the Euclidean distance grows, the value moves toward 0. The bandwidth controls the scale at which distance matters.
Consider one-dimensional observations with (\sigma=1):
[ a=0,\qquad b=1 ]
Then:
[ k(0,1)=\exp(-1/2)\approx 0.607 ]
For points farther apart:
[ k(0,3)=\exp(-9/2)\approx 0.011 ]
So the kernel turns geometric distance into a smooth similarity score. MMD then compares averages of these scores rather than comparing raw coordinates directly.
This is already enough for a useful mental model: MMD is small when the reference and current samples have similar patterns of kernel similarity, and larger when those patterns separate.
The population quantity behind MMD
Let (P) be the reference distribution and (Q) the current distribution. Draw independent values (X,X’\sim P) and (Y,Y’\sim Q).
The squared MMD associated with kernel (k) can be written as:
[ \operatorname{MMD}^2(P,Q)
\mathbb{E}[k(X,X’)] + \mathbb{E}[k(Y,Y’)]
2\mathbb{E}[k(X,Y)] ]
Each term has a direct interpretation.
The first measures expected similarity between two reference points. The second does the same for current points. The third measures similarity across the two distributions.
When cross-distribution similarity is lower than the within-distribution terms imply, the result becomes positive.
For suitable characteristic kernels, including the Gaussian kernel on ordinary Euclidean spaces under standard conditions, population MMD is zero exactly when the two distributions are equal. That property is stronger than merely matching a finite set of moments.
The kernel still defines what differences receive practical emphasis. A mathematically discriminating kernel does not make finite samples immune to poor bandwidth choices, noise, or insufficient sample size.
Estimate MMD from finite samples
Production monitoring has samples, not direct access to (P) and (Q). A common unbiased estimator of squared MMD is:
[ \widehat{\operatorname{MMD}}_u^2
\frac{1}{n(n-1)} \sum_{i\ne j} k(x_i,x_j) + \frac{1}{m(m-1)} \sum_{i\ne j} k(y_i,y_j)
\frac{2}{nm} \sum_{i=1}^{n}\sum_{j=1}^{m} k(x_i,y_j) ]
The diagonal terms (k(x_i,x_i)) and (k(y_j,y_j)) are excluded from the first two sums. This removes a finite-sample bias present in the simpler estimator that averages every pair.
One detail can surprise developers: the unbiased estimate of squared MMD can be slightly negative for finite samples. The population quantity is non-negative, but an unbiased estimator can fluctuate below zero. A negative estimate is not evidence for a negative distance.
If your monitoring code needs a non-negative descriptive score, you can use the biased empirical estimator that includes diagonal pairs, or clamp an unbiased estimate for display while retaining the original value for statistical procedures. Do not silently switch estimators and then compare the new values with thresholds calibrated for the old one.
A small numerical example
Take two reference values and two current values:
[ X={0,1},\qquad Y={0,1} ]
With the same kernel and bandwidth for every comparison, the empirical distributions are identical. A biased MMD estimate is exactly zero because all three kernel-average terms cancel.
Now move the current sample:
[ Y={3,4} ]
Within each sample, neighboring points are still one unit apart, so the within-sample similarities remain fairly high. Cross-sample distances are now much larger, making the cross term smaller. MMD therefore increases.
This example exposes a useful distinction from a simple variance check. Both samples have the same spread, but their locations differ. MMD detects the change because the cross-sample geometry changed.
The same mechanism can detect changes more complex than a location shift. With an appropriate kernel and enough data, differences in spread, shape, or combinations of coordinates can affect the score.
Bandwidth controls what the score can see
The Gaussian bandwidth (\sigma) is not a cosmetic parameter.
If (\sigma) is extremely small relative to typical distances, almost every pair of distinct points has kernel value near zero. The similarity matrix then contains little useful structure.
If (\sigma) is extremely large, most pairs have kernel values near one. Distinct distributions can look nearly identical at that scale.
A common starting heuristic sets the bandwidth from a typical pairwise distance, often the median distance among sampled points. It is a heuristic, not a universal optimum. Its result depends on representation scale, dimension, sample composition, and the type of shift you care about.
For monitoring, a more defensible process is to choose bandwidths using representative historical scenarios and then keep the procedure stable. If the bandwidth is recomputed independently on every monitoring window, the measuring scale itself changes along with the data. That can make score trends harder to interpret.
Using several bandwidths can help when relevant shifts occur at different geometric scales. You can compute separate MMD values or combine kernels, but thresholds must be calibrated for the exact statistic you deploy.
Representation choice comes before the kernel
MMD only sees the values you give it.
For raw tabular features, that may be appropriate after sensible scaling. For images, raw pixels may emphasize lighting, cropping, or compression rather than semantic content. For text, token IDs have no useful Euclidean geometry.
A practical AI monitoring system often computes MMD on model representations such as embeddings or intermediate feature vectors. This can make the comparison more aligned with distinctions the model uses.
It also introduces a dependency: if the embedding model changes, the representation space changes. An MMD jump may then reflect the representation update rather than a shift in incoming data. Keep the representation version fixed across a comparison, or establish a fresh baseline after a deliberate representation change.
Feature scaling matters for the same reason. If one coordinate has values around (10^4) and another around (10^{-2}), Euclidean distances may be dominated by the first coordinate. Apply the same fixed preprocessing to reference and current samples.
An MMD score is not a universal drift threshold
A raw value such as MMD² = 0.018 has little meaning by itself. Its scale depends on the kernel, bandwidth, preprocessing, representation, estimator, and data distribution.
Thresholds should therefore come from the monitoring problem rather than from a generic constant.
One practical approach is to create many reference-versus-reference comparisons from periods considered stable. Their score distribution shows the normal variation produced by sampling and routine traffic changes. Candidate thresholds can then be evaluated against known or simulated shifts that matter to the application.
If you need a formal two-sample hypothesis test rather than a descriptive monitor, use a valid test procedure with an appropriate null calibration, such as a permutation-based method. A threshold chosen from an arbitrary MMD value is not automatically a statistical test.
Repeated monitoring creates another issue. Testing every hour gives many opportunities for false alarms. Alerting policy should account for repeated checks, operational tolerance, and whether a single-window spike deserves action.
Quadratic cost can become the bottleneck
The direct estimator compares many pairs. With (n) reference points and (m) current points, the kernel work includes terms on the order of:
[ O(n^2 + m^2 + nm) ]
When (n) and (m) are similar, this is quadratic in sample size.
That is manageable for moderate monitoring batches but can become expensive for large windows or high-dimensional embeddings. The kernel matrix can also consume substantial memory if materialized in full.
Several practical options exist:
- compute kernel sums in blocks instead of storing complete matrices;
- compare fixed-size subsamples when the resulting sensitivity is adequate;
- use linear-time MMD estimators or kernel approximations when scale demands them.
These alternatives trade computation against estimator variance or approximation error. Validate the chosen method on shifts that resemble the failures you need to detect.
Common monitoring mistakes
Treating drift as proof of model failure
MMD measures a distribution discrepancy under a chosen representation and kernel. It does not measure task accuracy.
A large shift can be harmless if the changed features do not affect the decision boundary. A small MMD can coexist with serious degradation if the change is concentrated in a rare but critical subgroup.
Pair drift monitoring with outcome metrics whenever labels or delayed feedback are available. When labels are unavailable, treat MMD as a trigger for investigation rather than a verdict on model quality.
Comparing scores produced by different pipelines
Changing feature normalization, embedding versions, kernels, bandwidths, sample sizes, or estimators can move the score even when the underlying traffic is unchanged.
Version the complete MMD configuration. A monitoring chart is interpretable only when points on the chart represent the same measurement procedure.
Ignoring sample dependence
The standard two-sample formulas are easiest to interpret when observations behave like independent samples from their respective distributions. Real traffic may contain sessions, duplicate events, time-series dependence, or repeated records from the same user or device.
Strong dependence changes the sampling behavior of the statistic and can invalidate a calibration procedure that assumes independent observations. Sampling at the appropriate unit, grouping correlated records, or using a calibration method suited to the data structure can be more important than tuning the kernel.
Using one global score for every operational question
A single MMD value can hide where a shift occurred. If an alert fires, you still need diagnostics.
Useful follow-up views include scores by region or product segment, individual feature summaries, nearest-neighbor examples in representation space, and comparisons across several time windows. MMD can tell you that two samples differ without identifying one simple causal feature.
When MMD fits and when a simpler check is better
MMD is a strong candidate when you need to compare multivariate samples, have a meaningful kernel or vector representation, and want sensitivity beyond a few hand-picked moments.
It is less attractive when a single feature has a clear business interpretation and a simple histogram or quantile comparison answers the operational question. Simpler checks are easier to debug and explain.
MMD also isn’t a substitute for task evaluation. If fresh labels arrive quickly, direct performance metrics such as error rate, precision, recall, or task-specific loss often deserve higher priority. Distribution monitoring is most valuable as an earlier or complementary signal.
Build the monitor around stable comparisons
The useful part of an MMD monitor is not the formula alone. It is the controlled comparison around it.
Choose a representation tied to the model’s inputs or internal features. Fix preprocessing and kernel settings. Measure normal score variation on stable data. Test the monitor against shifts that would matter in operation. Record enough context to investigate an alert.
With those pieces in place, maximum mean discrepancy becomes a practical answer to a specific question: has the geometry of current model inputs moved away from a trusted reference in a way this kernel can detect? That is narrow enough to interpret and broad enough to catch multivariate changes that simpler summaries can miss.