A minimal, faithful, NumPy-first re-implementation of DUAL — learning diverse kernels for aggregated two-sample testing (MMD: are two samples from the same distribution?) and independence testing (HSIC: are two paired variables dependent?).
References.
- Zhou, Tian, Peng, Lei, Schrab, Sutherland, Liu, "DUAL: Learning Diverse Kernels for Aggregated Two-sample and Independence Testing", NeurIPS 2025. (arXiv:2510.11140; official code: MMD-HSIC-DUAL, full experiments: Aggregated_U_stats)
- Liu, Xu, Lu, Zhang, Gretton, Sutherland, "Learning Deep Kernels for Non-parametric Two-Sample Test", ICML 2020. (arXiv:2002.09116)
Independent educational reproduction, not affiliated with the authors.
Kernel tests live and die by the kernel, and aggregating many kernels only helps if they capture different things. DUAL stacks the per-kernel (paired, second-order) U-statistics into a vector U ∈ R^c, estimates their null covariance Σ̂ on resampled null data, and aggregates with the Mahalanobis statistic T = n² Uᵀ Σ̂⁻¹ U — redundant (correlated) kernels inflate Σ̂, so maximising T over the kernel parameters implicitly learns a diverse pool. Because diverse pools inevitably contain individually-weak kernels, testing adds selection inference: freeze the whitening L = Σ̂^{-1/2} from training and keep only the coordinates of L·U_test whose sign agrees with training (T = n²‖F ∘ L·U_test‖²). A wild bootstrap (Rademacher signs on the h-matrices) calibrates the threshold; sign flips correspond exactly to swapping x_i ↔ y_i (MMD) or the paired x-couples (HSIC), which under H₀ leaves the data distribution invariant — so the test is an exact randomisation test. The same recipe handles HSIC via the pair-up trick, h = ¼·h_X∘h_Y.
dktest/
dual.py # ★ the DUAL method (NumPy): MMDDual + HSICDual
# multivariate U-stats, null covariance, Mahalanobis
# aggregation, analytic-gradient Adam learning of the
# bandwidth pool, selection inference, wild bootstrap
kernels.py # Gaussian/Laplacian kernels, 'agg'/'fuse' bandwidth grids,
# median heuristic (legacy single-kernel modules)
mmd.py # unbiased MMD^2, permutation test, Bonferroni multi-kernel baseline
hsic.py # single-kernel HSIC + permutation independence test
power.py # Liu et al. 2020 power criterion J = MMD^2/sqrt(var_H1+lam)
# with the closed-form H1 variance estimator (Eq. 4-5)
metrics.py # Monte-Carlo Type-I error / power estimation
data.py # synthetic generators incl. the BLOB benchmark
deep_kernel.py # optional PyTorch learned deep kernel (Liu et al. 2020),
# trained with the paper's power criterion
examples/
demo_two_sample.py # 4-test comparison + BLOB benchmark (CPU)
demo_independence.py # HSIC & HSIC-DUAL vs. correlation (CPU)
run_deep_kernel.py # learned deep-kernel test (needs torch)
tests/ # NumPy-only unit tests (no torch, no network)
pip install -r requirements.txt # just numpy
python examples/demo_two_sample.py # a few minutes on CPU
python examples/demo_independence.py
bash scripts/run_tests.sh # 28 unit testsimport numpy as np
from dktest import MMDDual, HSICDual
rng = np.random.default_rng(0)
model = MMDDual().fit(X_train, Y_train, epochs=200, rng=rng) # learn diverse pool
result = model.test(X_test, Y_test, num_boot=200, rng=rng) # exact wild-bootstrap test
print(result["pvalue"], result["reject"])demo_two_sample.py — Gaussian mean/variance shifts (n = 80, d = 5), then the
BLOB benchmark (3×3 grid mixture whose components differ only in local
covariance, n = 150) where the median-heuristic kernel is nearly blind:
Two-sample tests (alpha = 0.05)
Type-I (null) power mean=0.5 power var=0.5
single-kernel MMD 0.025 0.517 1.000
power-selected MMD 0.058 0.258 0.883
Bonferroni-agg MMD 0.017 0.392 1.000
MMD-DUAL (paper) 0.042 0.317 1.000
BLOB benchmark (local covariance difference, rho = 0.7)
Type-I (rho=0) power
single-kernel MMD 0.083 0.100
Bonferroni-agg MMD 0.050 0.050
MMD-DUAL untrained 0.033 0.167
MMD-DUAL trained 0.067 0.367
The BLOB table is the paper's story in miniature: all tests control the Type-I error, but on structured data the single median-heuristic kernel and the Bonferroni aggregation barely see the difference, a fixed diverse pool with covariance-aware aggregation already helps, and learning the pool by maximising T roughly doubles the power again. (MMD-DUAL's Type-I sits near α rather than below it — the wild bootstrap is exact, not conservative like Bonferroni. On the easy unimodal mean-shift a single well-placed kernel remains competitive, as expected.)
demo_independence.py — HSIC catches non-linear dependence that Pearson
correlation cannot see:
Independence tests (alpha = 0.05)
HSIC value independent = 0.0018
HSIC value nonlinear = 0.0286
Pearson r nonlinear = -0.134 (linear correlation misses it)
Type-I (independent) power linear power nonlinear
single-kernel HSIC 0.042 1.000 1.000
HSIC-DUAL (paper) 0.042 1.000 1.000
| Component | Status |
|---|---|
| Paired second-order U-statistics (MMD h-matrix; HSIC pair-up trick h = ¼·h_X∘h_Y) | ✅ matches paper Sec. 3 / official code (unit-tested against naive formulas) |
| Null covariance Σ̂ (paper Eq. 3) on resampled null data (pooled shuffle for MMD) | ✅ matches official code |
| HSIC null resampling: independent permutations that break the x-y pairing | ✅ follows the paper's Eq. 3 (H₀ data). |
| Mahalanobis aggregation T = n²UᵀΣ̂⁻¹U; Adam ascent over Gaussian('agg') + Laplacian('fuse') bandwidth pool | ✅ same objective/optimiser; gradients derived analytically and verified against finite differences |
| Selection inference (sign alignment via frozen Σ̂^{-1/2}) + wild bootstrap with per-replicate masks | ✅ matches official code |
| Liu et al. 2020 power criterion (closed-form H1 variance, Eq. 4-5) for kernel selection and deep-kernel training | ✅ implemented (NumPy + optional torch) |
| Deep/Mahalanobis kernel variants of DUAL, image benchmarks (MNIST/CIFAR/ImageNet) | ❌ out of scope for this CPU reproduction — see the official repos |
Documented deviations from the official torch implementation (see
dktest/dual.py docstring): bandwidths are optimised in log-space
(positivity); the p-value uses the paper's (B+1) threshold convention (exact
finite-sample level); the covariance constant uses the paired sample count n
as written in the paper (scale-only); and the HSIC null resampling breaks the
pairing as the paper requires (see table above).
pip install -r requirements-torch.txt
python examples/run_deep_kernel.pyTrains the Liu et al. (2020) deep kernel k(x,y) = [(1−ε)·k_φ(φ(x),φ(y)) + ε]·k_q(x,y) by maximising the test-power criterion J = MMD²_u / √(σ̂²_H1 + λ) on a training split, then runs a permutation test on a held-out split.
@inproceedings{zhou2025dual,
title = {DUAL: Learning Diverse Kernels for Aggregated Two-sample and Independence Testing},
author = {Zhou, Zhijian and Tian, Xunye and Peng, Liuhua and Lei, Chao and Schrab, Antonin and Sutherland, Danica J. and Liu, Feng},
booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
year = {2025}
}
@inproceedings{liu2020learning,
title = {Learning Deep Kernels for Non-parametric Two-Sample Test},
author = {Liu, Feng and Xu, Wenkai and Lu, Jie and Zhang, Guangquan and Gretton, Arthur and Sutherland, Danica J.},
booktitle = {International Conference on Machine Learning (ICML)},
year = {2020}
}MIT (this reproduction). The original methods and results are the authors'.