Source-linked AI summary
Learning Deep Kernels for Non-Parametric Two-Sample Tests
Feng Liu, Wenkai Xu, Jie Lu, Guangquan Zhang, Arthur Gretton, Danica J. Sutherland
TL;DR
The paper addresses the difficulty of testing whether complex, high-dimensional distributions differ when simple kernels are spatially homogeneous. It learns deep neural-network kernels by maximizing test power, proves consistent adaptation, and reports superior performance on benchmark and real-world data, including advantages over cross-entropy-based learning and classifier-based tests.
Problem
Simple spatially homogeneous kernels can map complex, spatially varying distributions to nearby mean embeddings, making their differences difficult to detect.
Method
The method parameterizes kernels with deep neural networks and selects their parameters by maximizing a regularized MMD-to-variance criterion.
Results
Deep kernels outperform simple shallow methods and cross-entropy-based learning across simulated and real-world datasets, with MMD-D generally achieving the best benchmark performance.
Takeaways & Limitations
Learning kernels directly for test power provides a flexible approach for detecting differences in distributions with complex local structure.
Takeaways & Limitations
Classifier-based tests can have higher statistic variance because sign-only decisions discard confidence, and sample reuse during kernel selection violates the stated asymptotic conditions.
Abstract
from arXiv · showhide
We propose a class of kernel-based two-sample tests, which aim to determine whether two sets of samples are drawn from the same distribution. Our tests are constructed from kernels parameterized by deep neural nets, trained to maximize test power. These tests adapt to variations in distribution smoothness and shape over space, and are especially suited to high dimensions and complex data. By contrast, the simpler kernels used in prior kernel testing work are spatially homogeneous, and adaptive only in lengthscale. We explain how this scheme includes popular classifier-based two-sample tests as a special case, but improves on them in general. We provide the first proof of consistency for the proposed adaptation method, which applies both to kernels on deep features and to simpler radial basis kernels or multiple kernel learning. In experiments, we establish the superior performance of our deep kernels in hypothesis testing on benchmark and real-world data. The code of our deep-kernel-based two sample tests is available at https://github.com/fengliu90/DK-for-TST.
1. Introduction
The paper develops non-parametric two-sample tests using deep neural networks to learn flexible kernels that adapt to complex, spatially varying distribution structure. It trains kernel parameters to maximize test power and reports stronger performance than shallow methods and cross-entropy-based learning.
- Motivation: Kernel two-sample tests compare distributions through differences in their kernel mean embeddings, but simple methods struggle with complex distributions.Characteristic kernels identify equal distributions through equal mean embeddings.
- Motivation: A translation-invariant Gaussian kernel can miss differences when multimodal distributions share modes but differ in local substructure.Its uniform behavior within each mode may require many samples to distinguish the distributions.
- Method: The proposed deep kernels use a neural feature map with a simple feature-space kernel and a characteristic input-space kernel.This construction permits flexible behavior that varies across regions of the input space.
- Method: Kernel parameters are selected by maximizing the ratio of MMD to its variance, targeting higher test power at large sample sizes.The paper proves consistent kernel selection for deep kernels and simpler kernel families.
- Contributions: The framework includes classifier-based two-sample tests as a special case while allowing more general model classes and direct test-power optimization.The experiments report that test-power maximization outperforms cross-entropy surrogate learning.
2. MMD Two-Sample Tests
The paper formulates two-sample testing with MMD, estimates it from paired samples, and uses permutation testing to calibrate rejection decisions. It motivates selecting kernels by an estimated power criterion while addressing sample reuse and estimator bias.
- Two-sample testing: Two-sample testing asks whether independent samples come from the same distribution, testing H0: P = Q against H1: P ≠ Q.The procedure computes a statistic and p-value, then rejects H0 when the p-value is below α.
- MMD: MMD measures the distance between kernel mean embeddings of P and Q using point-level similarities defined by a kernel.For characteristic kernels, equal mean embeddings are equivalent to equal distributions.
- MMD: The U-statistic estimator is unbiased for MMD^2 and has nearly minimal variance among unbiased estimators.Its sample expression combines within-sample and cross-sample kernel evaluations.
- Calibration: Permutation testing estimates the null distribution by repeatedly reassigning pooled samples because samples are interchangeable under H0.This supplies the rejection threshold used to calculate test decisions.
- Kernel selection: The power criterion approximately favors kernels that maximize MMD relative to alternative-distribution variance, motivating regularized empirical kernel selection.Selecting and testing on the same samples violates the stated asymptotic conditions and would require retraining under permutations.
- Related approaches: Proxy-task kernel selection depends on whether proxy features capture the differences between P and Q, while simply maximizing MMD can favor uninformative scaling.The paper reports that the latter approach performed notably worse than maximizing the power criterion.
- Related approaches: MMD-GAN objectives differ because a model far from its target may need a weak kernel, whereas a two-sample test seeks a kernel that distinguishes distributions.The paper distinguishes its algorithm, guarantees, and evaluations from MMD-GANs.
3. Limits of Simple Kernels
Simple kernels used in MMD tests are often translation invariant, so one global behavior is suboptimal for distributions whose shape or density varies across space. Deep kernels address this by adapting to local structure.
- Kernel limitations: The criterion ˆJλ can select even a Gaussian kernel lengthscale, but the Blob example illustrates limitations of such simple kernels.The comparison evaluates Gaussian kernels against the deep-kernel family.
- Empirical comparison: The optimal ˆJ and empirical test power are higher for deep kernels than for the compared Gaussian-kernel family.Figure 2c tracks the criterion as the number of samples increases, while Figure 2a reports test power.
- Kernel limitations: Translation-invariant kernels behave identically across space, making a single global choice suboptimal when local directions or densities vary.This limitation applies to Gaussian, Laplace, inverse multiquadric, and automatic relevance determination kernels described here.
- Deep-kernel advantage: Non-translation-invariant deep kernels can adapt their behavior to the different shapes required in different areas.Figure 1c illustrates this spatial adaptation.
4. Relationship to Classifier-Based Tests
The paper recasts classifier-based two-sample tests within the MMD framework and contrasts their objectives with direct test-power optimization. Learned MMD tests can outperform classifier-based approaches in the reported experiments.
- Classifier-based sign tests use classification accuracy and have a simple asymptotically normal null distribution.
- A C2ST-S using classifier f is equivalent to an MMD test with a sign-based kernel.
- Unlike C2ST-S, maximizing the corresponding MMD objective directly maximizes classification accuracy and therefore test power.
- Sign-based statistics ignore classifier confidence, treating highly confident and uncertain predictions alike and potentially increasing variance.
- Learned MMD tests with deep kernels empirically outperform cross-entropy-based learning and achieve higher power than competing methods in Figure 2.
5. Learning Deep Kernels
The method learns flexible kernels by combining neural feature extraction with simple kernels, while retaining characteristicness through an input-space component. Training maximizes a regularized test-power criterion before permutation testing.
- The deep kernel applies a simple kernel to neural-network features and combines it with a characteristic kernel on the input space.
- With positive mixing weight and characteristic input kernel, the resulting deep kernel is characteristic.
- Algorithm 1 separates kernel learning from permutation testing and uses Adam, with minibatches for larger or unequal datasets.
- The experiments report runtime scarcely different from classifier-based two-sample tests, while quadratic testing costs can be reduced using block estimators.
6. Theoretical Analysis
The theoretical analysis establishes consistency for empirical kernel selection under boundedness, Lipschitz, and nonzero-variance conditions. It covers deep kernels, Gaussian bandwidth selection, and multiple kernel learning, while identifying sensitivity to small variances.
- Uniform convergence of empirical MMD and variance estimates implies convergence to the unique best kernel when optimization succeeds.
- Theorem 6 analyzes uniformly bounded, Lipschitz kernel families over bounded parameter sets with variance bounded away from zero.
- The same guarantees cover Gaussian bandwidth selection and multiple kernel learning, including an efficiently available exact empirical maximizer for the latter.
- The consistency framework applies to fully connected and convolutional deep kernels under bounded-domain and network-norm assumptions.
- The ratio criterion is sensitive to errors in very small variances, although regularization and the input-kernel component can reduce problematic selections.
- With perfect optimization, the asymptotic power bound is maximized by training on a specified fraction of the data and testing on the remainder.
7. Experimental Results
The experiments compare MMD-D with optimized Gaussian-kernel, mean-embedding, smooth-characteristic-function, and classifier two-sample tests across synthetic, high-dimensional, physics, and image datasets. MMD-D generally achieves the strongest test power while maintaining reasonable Type I error, and the ablation results favor deep kernels and power-based optimization.
- Benchmark comparisons: MMD-D and C2ST-L are the clear power winners on Blob, with MMD-D stronger in the higher-sample regime and more reliable than C2STs.All methods have expected Type I error rates in the reported Blob experiments.
- Benchmark comparisons: MMD-D generally has the best test power across high-dimensional Gaussian-mixture settings while maintaining reasonable Type I error.Figure 3 varies sample size at d = 10 and dimension at N = 4 000 for null and alternative HDGM settings.
- Benchmark comparisons: On Higgs, MMD-D generally performs best, notably outperforming C2ST while all methods maintain reasonable Type I errors.The comparison uses jet φ-momenta distributions with increasing sample size.
- Benchmark comparisons: MMD-D substantially outperforms competitors on real MNIST versus DCGAN samples while retaining the desired Type I error.The experiments compare true MNIST images with samples from a pretrained DCGAN.
- Benchmark comparisons: MMD-D provides the strongest evidence that CIFAR-10.1 images differ statistically from CIFAR-10 images, while ME also performs well.The study trains on 1 000 images from each dataset and tests on 1 031 images, averaging over ten repetitions.
- Ablation study: The ablation study finds that performance generally improves from S to L to G to D and from cross-entropy optimization to test-power optimization.MMD-D uses a full deep kernel trained for test power, whereas C2ST-S uses a sign kernel trained for cross entropy.
8. Conclusions
The paper concludes that optimizing richer kernels for MMD test power addresses limitations of simple kernels and is theoretically justified, while encompassing classifier-based tests.
- Conclusions: Deep kernels avoid the limitations of translation-invariant kernels for complex-structured distributions and can be optimized even with hundreds of thousands of parameters.The paper reports theoretical guarantees that finite-sample optimization asymptotically selects the most powerful kernel.
- Conclusions: The framework includes classifier-based two-sample tests while allowing more general model classes that can use data more efficiently.Representations are trained directly to maximize test power rather than through a cross-entropy surrogate.
- Conclusions: Practitioners are recommended to use optimized deep kernel methods when checking whether two distributions are the same.This recommendation is framed as an alternative to indirectly training a classifier.
- Estimator caveat: The variance estimator used in the analysis is biased, although the paper states that its bias becomes negligible as sample size increases.The estimator remains attractive because it is always nonnegative and more convenient for implementation and analysis.
- Theoretical conditions: The consistency analysis relies on assumptions including uniformly bounded kernels and a finite-dimensional, Lipschitz kernel parameterization.The paper states that the assumptions are relatively agnostic to kernel form.
Appendix A.4.2 builds this space and its norm for the kernels we use in the paper.
The appendix establishes consistency for selecting kernels by maximizing an empirical test-power criterion under boundedness, finite-dimensionality, and Lipschitz assumptions.
- Scope: The theory applies across different kernel-learning settings once the required assumptions are verified for the relevant parameterization.The appendix explicitly treats Gaussian bandwidth selection and states that assumptions hold for the paper’s kernels.
- Finite-sample guarantee: Theorem 11 provides a finite-sample error bound for the empirical objective under the stated kernel assumptions.The bound is supported by uniform convergence results for the MMD and variance estimators.
- Consistency: Under Assumptions (A) to (C), empirical kernel selection converges in probability to the unique population maximizer of J.The result uses λ = n^-1/3 and requires sufficiently large sample size.
- Test power: With n training points and m testing points, the selected-kernel test retains a high-probability power guarantee relative to the best kernel in the class.The guarantee depends on constants involving ν, L_k, D, R_Ω, and s.
- Sample allocation: When N points are split between training and testing, the asymptotic power upper bound determines an interior choice of n rather than either endpoint.The appendix analyzes the optimizer through stationary-point conditions and asymptotic cases.
A.3. Uniform convergence results
The uniform-convergence appendix controls empirical MMD and variance estimation over kernel parameters, then verifies the regularity conditions for Gaussian, deep, convolutional, and multiple-kernel settings.
- Uniform convergence: Uniform convergence of the empirical squared MMD and variance estimator supplies the concentration results used to prove kernel-selection consistency.The bounds are obtained under boundedness, finite-dimensional parameter space, and Lipschitz parameterization assumptions.
- Proof strategy: The proofs use epsilon-net arguments, bounded differences, McDiarmid’s inequality, and union bounds to control errors across kernel parameters.The covering number depends on the dimension D and radius R_Ω of the parameter space.
- Deep kernels: For deep kernels, the regularity analysis relies on bounded domains, 1-Lipschitz activations, operator-norm control of weights, and assumptions on the kernel form.Convolutional networks are handled by defining parameter norms through the corresponding convolution operators.
- Capacity dependence: The convergence bound includes dependence on network capacity and output-scale terms, linking the guarantee to the number of optimized parameters and feature magnitudes.The appendix describes one term as being weighted roughly equally with the square root of the total number of network parameters.
- Multiple kernel learning: Multiple kernel learning satisfies the same assumptions and permits an exact empirical-objective maximizer through a convex quadratic program.This exact optimization result distinguishes the multiple-kernel setting from the deep setting.
B.1. Details of synthetic datasets
The synthetic-data suite contains four datasets designed to test two-sample methods on rotated blob structures and high-dimensional, heterogeneous distributions.
- Dataset suite: The synthetic benchmarks include Blob datasets, HDGM datasets, and four total dataset specifications summarized in Table 6.HDGM datasets are introduced in this paper, while Blob datasets are established validation benchmarks.
- Blob datasets: The Blob datasets rotate each blob to expose the benefits of nonhomogeneous kernels.This construction preserves the benchmark’s multimodal character while varying local structure across modes.
- HDGM datasets: HDGM-D is a high-dimensional Blob-D analogue with two modes sharing variance but having different covariance.The passage describes this as part of the synthetic-data specifications.
- Distribution parameters: Table 6 specifies the distributions through mode means, including blob means arranged over a two-dimensional grid.The listed blob means range from [0, 0] through [2, 2] in the provided specification.
B.2. Dataset visualization
The experiments compare real and generated MNIST, CIFAR-10 and CIFAR-10.1, using neural-network architectures and specified training procedures for the tested methods.
- Model architectures: C2ST-S and C2ST-L use a deep classifier F trained by minimizing cross entropy, while MMD-O optimizes a kernel lengthscale.The reported implementations use Adam optimization and zero dropout for C2ST-S, C2ST-L and MMD-D.
- Model architectures: MMD-D uses a neural network φω, implemented as a five-layer fully connected network for tabular datasets and a CNN for MNIST and CIFAR.The CNN contains four convolutional layers and one fully connected layer.
- Evaluation protocol: Test-method parameters are searched on a validation set, then test power is computed over 100 test sets of the same size as the training set.The experiments use α = 0.05 and permutation tests for the listed methods.
- Evaluation protocol: The appendix reports average Type I error as sample size increases for Higgs and real-MNIST versus real-MNIST.Higgs averages Type I error across samples drawn from distributions with and without Higgs bosons; MNIST reports average Type I error ± standard errors.
C. Interpretability on CIFAR-10 vs CIFAR-10.1
The paper examines which images reveal the CIFAR-10 versus CIFAR-10.1 difference using mean-embedding test locations and learned kernels. A two-stage deep-kernel procedure improves rejection rates over the basic learned-location setup, while restricting locations to existing images retains similar power and improves interpretability.
- Test locations: Mean Embedding tests compare mean embeddings at selected test locations rather than through their overall norm.The analysis focuses on interpreting the difference between CIFAR-10 and CIFAR-10.1.
- Test locations: Jointly optimizing a deep kernel and test locations failed to find a useful test, motivating a two-stage procedure.The two stages first learn a deep kernel to maximize ˆJλ, then choose test locations with the kernel fixed.
- Rejection rates: 0.415 is the average rejection rate for the best test locations learned by an ME test with L = 1.This result averages 10 experiments on CIFAR-10 versus CIFAR-10.1.
- Rejection rates: 0.637 is the average rejection rate when ME test locations use a deep kernel optimized for an MMD test.The result uses L = 1 and averages 10 experiments.
- Rejection rates: 0.653 is the average rejection rate when the single test location is selected among existing images using the learned deep kernel.Restricting the location to the dataset makes the result interpretable through images with smaller objects of interest than the CIFAR-10 mean.