Kernel Two-Sample Test (Maximum Mean Discrepancy)

Description

A kernel two-sample test checks whether two sets of data could have come from the same distribution; in assurance it is used to check whether the inputs a deployed model now receives still resemble the data on which it was validated, and to detect when they stop doing so. Its statistic, the maximum mean discrepancy (MMD), is built from a kernel, a function that scores how similar two data points are (commonly a Gaussian function of the distance between them), and in its usual squared form it is the average similarity among points in the reference sample, plus the average similarity among points in the new sample, minus twice the average similarity between points from different samples. With a characteristic kernel such as the Gaussian, the population value is zero only when the two distributions are identical, so the test can respond to any kind of difference, including a change in how features vary together that tests applied to one feature at a time cannot see, and kernels exist for strings and graphs as well as numerical vectors. To obtain a p-value, the two samples are pooled and repeatedly divided at random into groups of the original sizes, the statistic is recomputed for each division, and the p-value is the share of divisions giving a value at least as large as the one observed; high-dimensional inputs such as images are usually first reduced to a short vector, for example the code of an autoencoder or the model's predicted class probabilities. The output, for each window of new data, is the MMD value and its p-value at a stated significance level, which shows whether that window differs detectably from the reference, but not whether the difference affects the model's accuracy.

Example Use Cases

Reliability

A card issuer's model-risk team tests a random sample of each night's transactions, at most one per customer, against a fixed sample from the fraud model's validation period, after standardising the input features. It set the significance level so that, with nothing changed, it expects about one false alarm a quarter, and a significant result triggers a check of the model's accuracy on newly labelled cases before its alert thresholds are relied on the next day.

Before switching on a chest X-ray triage model at a second hospital that uses a different make of scanner, the deploying team compares 500 recent images from the new site with 500 from the validation set. Each image is first reduced to a short vector by an autoencoder trained on validation images. A significant difference sends the model for local validation on labelled cases from the new site instead of relying on its original results.

A customer-service team whose question-answering assistant was evaluated on a fixed test set compares the sentence embeddings of each week's user questions with those of the test set. When a week is flagged, it computes the MMD witness function, which scores each question by how much more it resembles the new week's questions than the test set, reads the top-scoring questions, finds queries about a new product, and adds such questions to the test set.

Limitations

A significant result shows that inputs differ from the reference, not that the model has got worse: Rabanser et al. note that distributions shift constantly and often harmlessly, and with large windows the test can flag differences too small to matter. A drift alarm becomes evidence about performance only when followed by an accuracy check on labelled cases, and an absence of alarms reassures only if the test was shown beforehand to detect the shifts that harm the model.
On raw high-dimensional inputs the test has little power: in Rabanser et al.'s comparison on image data the multivariate test performed poorly without first reducing the dimension. The reduction then decides what can be seen, since a change that leaves the autoencoder code or predicted probabilities unchanged goes undetected. The kernel's width matters too: it is usually set to the median distance between points, which Gretton et al. call a heuristic, the best width being an open question.
The statistic compares every pair of points, so its cost grows with the square of the combined sample size, and the permutation test repeats it many times; Rabanser et al. ran kernel tests only up to 1,000 new samples for this reason. Gretton et al.'s linear-time version is far cheaper but needs much more data for the same power, so high-volume monitoring has to trade window size, detection delay and computing cost against each other.
The permutation p-value assumes that observations within each sample are independent. Consecutive sensor readings, transactions from one customer or successive time steps are correlated, and shuffling them breaks that dependence, so the test signals a difference far more often than its stated level: in Chwialkowski et al.'s experiments on autocorrelated signals the permutation test almost always rejected when nothing had changed. Such data need thinning or a wild bootstrap.
Run on every new window at the conventional 5 per cent level, the test flags about one window in twenty even when nothing has changed, which for daily windows means around 18 false alarms a year. Continuous monitoring needs a stricter level, a correction for repeated testing, or an online version of the test whose threshold is set for an expected run time between false alarms.

Resources

Research Papers

A Kernel Two-Sample Test

Arthur Gretton et al.•Jan 1, 2012

We propose a framework for analyzing and comparing distributions, which we use to construct statistical tests to determine if two samples are drawn from different distributions. Our test statistic is the largest difference in expectations over functions in the unit ball of a reproducing kernel Hilbert space (RKHS), and is called the maximum mean discrepancy (MMD). We present two distribution-free tests based on large deviation bounds for the MMD, and a third test based on the asymptotic distribution of this statistic. The MMD can be computed in quadratic time, although efficient linear time approximations are available. Our statistic is an instance of an integral probability metric, and various classical metrics on distributions are obtained when alternative function classes are used in place of an RKHS. We apply our two-sample tests to a variety of problems, including attribute matching for databases using the Hungarian marriage method, where they perform strongly. Excellent performance is also obtained when comparing distributions over graphs, for which these are the first such tests.

Software Packages

Frouros: an open-source Python library for drift detection in machine learning systems

Apr 1, 2024

Frouros: an open-source Python library for drift detection in machine learning systems.

Tutorials

MNIST dataset (Frouros data drift example: MMD with an autoencoder)

Documentations

Maximum Mean Discrepancy (Alibi Detect documentation)

Oct 16, 2024

Related Techniques

Name	Description	Assurance Goals
Data Validation Checks	Data validation checks apply a written set of rules about what valid data look like to every new batch of data before a model is trained on it or uses it to make predictions, and flag or hold back any batch that breaks them. In Great Expectations (GX Core), a widely used open-source library for this, each rule is an expectation, a declarative assertion such as that a column exists and has a given type, has no missing values or takes values only from an allowed set or within a range (rules on single values can allow a stated share of rows to fail), or keeps a summary statistic (its mean, a quantile, or the divergence of its histogram from one built on reference data) within bounds the team sets. Expectations are grouped into a suite that a checkpoint runs on each batch as it arrives; the result records for every rule whether it passed and how many values broke it, with examples, configured actions can alert the team or update a human-readable report, and the pipeline can be set to stop training or scoring on a batch that failed. Running the same suite on training batches and on the data a deployed model receives catches differences between the two, such as a feature that is always missing in production, before they degrade predictions; unlike out-of-distribution detectors, the checks need no model and judge whole batches against explicit rules rather than scoring single inputs, and unlike runtime monitoring of a deployed model they act before the model sees the data. Unlike a kernel two-sample test, its distribution rules are fixed bounds set by the team, not a statistical test of whether two samples differ.	Reliability
Bootstrap Confidence Intervals for Evaluation Metrics	Bootstrap confidence intervals put a range around a model's measured performance, showing how much the score could differ had another test set of the same size been drawn from the same population. The labelled held-out evaluation set is resampled with replacement thousands of times, each resample holding as many cases as the original so that some cases appear more than once and others not at all, and the metric (an error rate, the area under a curve, a translation score) is recomputed on each resample from the model's stored predictions, with no retraining. The spread of these recomputed values stands in for the sampling distribution of the metric: the percentile interval takes its 2.5th and 97.5th percentiles for a 95 per cent interval, and the bias-corrected and accelerated (BCa) interval shifts those percentiles to correct for bias and skew in that distribution. The statistic can also be the difference between two models scored on the same cases, recomputed on each shared resample, which shows whether an observed improvement is larger than the noise from test-set sampling. The result is the metric with lower and upper bounds at a stated confidence level, so a requirement such as a maximum false-positive rate can be checked against the bound rather than the point estimate.	Reliability
Out-of-Distribution Detector for Neural Networks	ODIN (Out-of-Distribution Detector for Neural Networks) flags inputs that look unlike the data a neural network classifier was trained on, and it works on an already trained network without retraining it. A classifier reports a confidence score, the highest probability it gives to any class, but that score is often high even for unfamiliar inputs. ODIN adjusts the score in two ways so that familiar and unfamiliar inputs separate better. It divides the classifier's raw output scores by a large constant, called the temperature, before turning them into probabilities, and it nudges the input slightly in the direction that raises the score, which lifts familiar inputs more than unfamiliar ones. An input whose adjusted score falls below a chosen threshold is flagged as out-of-distribution, so its prediction can be withheld or referred for review.	Reliability Safety