SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Felix Biggs*
Centre for Artificial Intelligence
University College London & Inria London
contact@felixbiggs.com
Antonin Schrab*
Gatsby Computational Neuroscience Unit
a.schrab@ucl.ac.uk
Arthur Gretton
arthur.gretton@gmail.com
We propose novel statistics which maximise the power of a two-sample test based on the Maximum Mean Discrepancy (MMD), by adapting over the set of kernels used in defining it. For finite sets, this reduces to combining (normalised) MMD values under each of these kernels via a weighted soft maximum. Exponential concentration bounds are proved for our proposed statistics under the null and alternative. We further show how these kernels can be chosen in a data-dependent but permutation-independent way, in a well-calibrated test, avoiding data splitting. This technique applies more broadly to general permutation-based MMD testing, and includes the use of deep kernels with features learnt using unsupervised models such as auto-encoders. We highlight the applicability of our MMD-FUSE test on both synthetic low-dimensional and real-world high-dimensional data, and compare its performance in terms of power against current state-of-the-art kernel tests.
The fundamental problem of non-parametric two-sample testing consists in detecting the difference between any two distributions having access only to samples from these. Kernel-based tests relying on the Maximum Mean Discrepancy (MMD; Gretton et al., 2012a) as a measure of distance on distributions are well-suited for this framework as they can identify complex non-linear features in the data, and benefit from both strong theoretical guarantees and ease of implementation. This explains their popularity among practitioners and justifies their wide use for real-world applications.
However, the performance of these tests is crucially impacted by the choice of kernel. This is commonly tackled by either: choosing the kernel by some weakly data-dependent heuristic (Gretton et al., 2012a); or splitting off a hold-out set of data for kernel selection, with the other half used for the actual test (Gretton et al., 2012b; Sutherland et al., 2017). This latter method includes training feature extractors such as deep kernels on the first selection half. Both of these methods can incur a significant loss in test power, since heuristics may lead to poor kernel choices, and data splitting reduces the number of data points for the actual test.
Our contribution is to present MMD-based tests which can strongly adapt to the data without data splitting. This comes in two parallel parts: firstly we show how the kernel can be chosen in an
unsupervised fashion using the entire dataset, and secondly we show how multiple such kernels can be adaptively weighted in a single test statistic, optimising test power.
Data Splitting. The data splitting approach selects test parameters on a held-out half of the dataset, and applies the test to the other half. Commonly, this involves optimising a kernel on held-out data in a supervised fashion to distinguish which sample originated from which distribution, either by learning a deep kernel directly (Sutherland et al., 2017; Liu et al., 2020, 2021), or indirectly through the associated witness function (Kübler et al., 2022a,b). Jitkrittum et al. (2016) propose tests which select witness function features (in either spatial or frequency space) on the held-out data, running the analytic representation test of Chwialkowski et al. (2015) on the remaining data.
Our first contribution is to show that it is possible to learn a feature extractor for our test (e.g., a deep kernel) on the entirety of the data in an unsupervised manner while retaining the desired non-asymptotic test level. Specifically, any method can be used that is ignorant of which distribution generated which samples. We can thus leverage many feature extraction methods, from auto-encoders (Hinton and Salakhutdinov, 2006) to recent powerful developments in self-supervised learning (He et al., 2020; Chen et al., 2020a,b,c; Chen and He, 2021; Grill et al., 2020; Caron et al., 2020; Zbontar et al., 2021; Li et al., 2021). Remarkably, our method applies to any permutation-based two-sample test, even non-kernel tests, provided the parameters are chosen in this unsupervised fashion. This includes a very wide array of MMD-based tests, and finally provides a formal justification for the commonly-used median heuristic (Gretton et al., 2012b; Ramdas et al., 2015; Reddi et al., 2015).
Adaptive Kernel Selection. A newer approach originating in Schrab et al. (2023) performs adaptive kernel selection through multiple testing. Aggregating several MMD-based tests with different kernels, each on the whole dataset, results in an overall adaptive test with optimal power guarantees in terms of minimax separation rate over Sobolev balls. Variants of this kernel adaptivity through aggregation have been proposed with linear-time estimators (Schrab et al., 2022b), spectral regularisation (Hagrass et al., 2022), kernel thinning compression (Domingo-Enrich et al., 2023), and in the asymptotic regime (Chatterjee and Bhattacharya, 2023). Another adaptive approach using the entire dataset for both kernel selection and testing is given by Kübler et al. (2020); this leverages the Post Selection Inference framework, but the resulting test suffers from low power in practice.
While our unsupervised feature extraction method is an extremely general technique and potentially powerful, particularly for high-dimensional structured data like images, it is not always sufficient for data where such feature extraction is difficult. This motivates our second contribution, a method for combining whole-dataset MMD estimates under multiple different kernels into single test statistics, for which we prove exponential concentration. Kernel parameters such as bandwidth can then be chosen in a non-heuristic manner to optimise power, even on data with less structure and varying length scales. Using a single statistic also ensures that a single test can be used, instead of the multiple testing approach outlined above, reducing computational expense.
MMD-FUSE. By combining these contributions we construct two closely related MMD-FUSE tests. Each chooses a set of kernels based on the whole dataset in an unsupervised fashion, and then adaptively weights and fuses this (potentially infinite) set of kernels through our new statistics; both parts of this procedure are done using the entire dataset without splitting. On the finite sets of kernels we use in practice, the weighting procedure is done in closed form via a weighted soft maximum. We show these new tests to be well-calibrated and give sufficient power conditions which achieve the minimax optimal rate in terms of MMD. In empirical comparisons, our test compares favourably to the state-of-the-art aggregated tests in terms of power and computational cost.
Outline and Summary of Contributions. In Section 2 we outline our setting, crucial results underlying our work, and alternative existing approaches. Section 3 covers the construction of permutation tests and discusses how we can choose the parameters for any such test in any unsupervised fashion including with deep kernels and by those methods mentioned above. Section 4 introduces and motivates our two proposed tests. Section 5 discusses sufficient conditions for test power (at a minimax optimal rate in MMD), and the exponential concentration of our statistics. Finally, we show that our test compares favourably with a wide variety of competitors in Section 6 and discuss in Section 7.
Two-Sample Testing. The two-sample testing problem is to determine whether two distributions  $p$  and  $q$  are equal or not. In order to test this hypothesis, we are given access to two samples,  $\mathbf{X} := (X_1, \ldots, X_n) \stackrel{\mathrm{iid}}{\sim} p$  and  $\mathbf{Y} := (Y_1, \ldots, Y_m) \stackrel{\mathrm{iid}}{\sim} q$  as tuples of data points with sizes  $n$  and  $m$ . We write the combined (ordered) sample as  $\mathbf{Z} := (Z_1, \ldots, Z_{n+m}) = (\mathbf{X}, \mathbf{Y}) = (X_1, \ldots, X_n, Y_1, \ldots, Y_m)$ .
We define the null hypothesis  $H_0$  as  $p = q$  and the alternative hypothesis  $H_1$  as  $p \neq q$ , usually with a requirement  $\mathbb{D}(p, q) > \epsilon$  for a distance  $\mathbb{D}$  (such as the MMD) and some  $\epsilon > 0$ . A hypothesis test  $\Delta$  is a  $\{0, 1\}$ -valued function of  $Z$ , which rejects the null hypothesis if  $\Delta(Z) = 1$  and fails to reject it otherwise. It will usually be formulated to control the probability of a type I error at some level  $\alpha \in (0, 1)$ , so that  $\mathbb{P}_{p \times p}(\Delta(Z) = 1) \leq \alpha$ , while simultaneously minimising the probability of a type II error,  $\mathbb{P}_{p \times q}(\Delta(Z) = 0)$ . In the above we have used the notation  $\mathbb{P}_{p \times p}$  and  $\mathbb{P}_{p \times q}$  to indicate that the sample  $Z$  is either drawn from the null,  $p = q$ , or the alternative  $p \neq q$ . Similar notation will be used for expectations and variances. When a bound  $\beta \in (0, 1)$  on the probability of a type II error is given (which may depend on the precise formulation of the alternative), we say the test has power  $1 - \beta$ .
Maximum Mean Discrepancy. The Maximum Mean Discrepancy (MMD) is a kernel-based measure of distance between two distributions  $p$  and  $q$ , which is often used for two-sample testing. The MMD quantifies the dissimilarity between these distributions by comparing their mean embeddings in a reproducing kernel Hilbert space (RKHS; Aronszajn, 1950) with kernel  $k(\cdot, \cdot)$ . Formally, if  $\mathcal{H}_k$  is the RKHS associated with kernel function  $k$ , the MMD between distributions  $p$  and  $q$  is the integral probability metric defined by:
$$
\operatorname {M M D} _ {k} (p, q) := \sup  _ {f \in \mathcal {H} _ {k}: \| f \| _ {\mathcal {H} _ {k}} \leq 1} \left(\mathbb {E} _ {X \sim p} [ f (X) ] - \mathbb {E} _ {Y \sim q} [ f (Y) ]\right).
$$
The minimum variance unbiased estimate of  $\mathrm{MMD}_k^2$ , is given by the sum of two U-statistics and a sample average as:
$$
\widehat {\mathrm {M M D}} _ {k} ^ {2} (\mathbf {Z}) := \frac {1}{n (n - 1)} \sum_ {(i, i ^ {\prime}) \in [ n ] _ {2}} k (X _ {i}, X _ {i ^ {\prime}}) + \frac {1}{m (m - 1)} \sum_ {(j, j ^ {\prime}) \in [ m ] _ {2}} k (Y _ {j}, Y _ {j ^ {\prime}}) - \frac {2}{m n} \sum_ {i = 1} ^ {n} \sum_ {j = 1} ^ {m} k (X _ {i}, Y _ {j}),
$$
where we introduced the notation  $[n]_2 = \{(i,i')\in [n]^2:i\neq i'\}$  for the set of all pairs of distinct indices in  $[n] = \{1,\dots ,n\}$ . Tests based on the MMD usually reject the null when  $\widehat{\mathrm{MMD}}_k^2$  exceeds some critical threshold, with the resulting power being greatly affected by the kernel choice. For characteristic kernel functions (Sriperumbudur et al., 2011), it can be shown that  $\mathrm{MMD}_k(p,q) = 0$  if and only if  $p = q$ , leading to consistency results. However, on finite sample sizes, convergence rates of MMD estimates typically have strong dependence on the data dimension, so there are settings in which kernels ignoring redundant or unimportant features (i.e. non-characteristic kernels) will give higher test power in practice than characteristic kernels (which can over-weight redundant features).
Distributions Over Kernels. In our test statistic, we will consider the case where the kernel  $k \in \mathcal{K}$  is drawn from a distribution  $\rho \in \mathcal{M}_{+}^{1}(\mathcal{K})$  (with the latter notation denoting a probability measure; c.f. Benton et al., 2019 in the Gaussian Process literature). This distribution will be adaptively chosen based on the data subject to a regularisation term based on a "prior"  $\pi \in \mathcal{M}_{+}^{1}(\mathcal{K})$  and defined through the Kullback-Liebler divergence:  $\mathrm{KL}(\rho \| \pi) \coloneqq \mathbb{E}_{\rho}[\log (\mathrm{d}\rho /\mathrm{d}\pi)]$  for  $\rho \ll \pi$  and  $\mathrm{KL}(\rho \| \pi) \coloneqq \infty$  otherwise. When constructing these statistics the Donsker-Varadhan equality (Donsker and Varadhan, 1975), holding for any measurable  $g:\mathcal{K}\to \mathbb{R}$ , will be useful:
$$
\sup  _ {\rho \in \mathcal {M} _ {+} ^ {1} (\mathcal {A})} \mathbb {E} _ {\rho} [ g ] - \operatorname {K L} (\rho \| \pi) = \log \mathbb {E} _ {\pi} [ \exp \circ g ]. \tag {1}
$$
This can be further related to the notion of soft maxima. If  $\pi$  is a uniform distribution on finite  $\mathcal{K}$  with  $|\mathcal{K}| = r$ , then  $\mathrm{KL}(\rho \| \pi) \leq \log (r)$ . Setting  $g = tf$  for some  $t > 0$ , Equation (1) relaxes to
$$
\max  _ {k} f (k) - \frac {\log (r)}{t} \leq \frac {1}{t} \log \left(\frac {1}{r} \sum_ {k} e ^ {t f (k)}\right) \leq \max  _ {k} f (k), \tag {2}
$$
$^2$ Kim et al. (2022) note that  $\widehat{\mathrm{MMD}_k^2}$  can equivalently be written as a two-sample U-statistic with kernel  $h_k(x,x';y,y') = k(x,x') + k(y,y') - k(x,y') - k(x',y)$ , which is useful for analysis.
which approximates the maximum with error controlled by  $t$ . Our approach in considering these soft maxima is reminiscent of the PAC-Bayesian (McAllester, 1998; Seeger et al., 2001; Maurer, 2004; Catoni, 2007; see Guedj, 2019 or Alquier, 2021 for a survey) approach to capacity control and generalisation. This framework has raised particular attention recently as it has been used to provide the only non-vacuous generalisation bounds for deep neural networks (Dziugaite and Roy, 2017, 2018; Dziugaite et al., 2021; Zhou et al., 2019; Perez-Ortiz et al., 2021; Biggs and Guedj, 2021, 2022a); it has also been fruitfully applied to other varied problems from ensemble methods (Lacasse et al., 2006, 2010; Masegosa et al., 2020; Wu et al., 2021; Zantedeschi et al., 2021; Biggs and Guedj, 2022b; Biggs et al., 2022) to online learning (Haddouche and Guedj, 2022). Our proofs draw on techniques in that literature for U-Statistics (Lever et al., 2013) and martingales (Seldin et al., 2012; Biggs and Guedj, 2023; Haddouche and Guedj, 2023; Chugg et al., 2023). By considering these soft maxima, we gain a major advantage over the standard approaches, as we can obtain concentration and power results for our statistics without incurring Bonferroni-type multiple testing penalties.
Permutation Tests. The tests we discuss in this paper use permutations of the data  $\mathbf{Z}$  to approximate the null distribution. We begin our discussion in a fairly general form to include other close settings such as independence testing (Albert et al., 2022; Rindt et al., 2021) or wild bootstrap-based two-sample tests (Fromont et al., 2012, 2013; Chwialkowski et al., 2014; Schrab et al., 2023). Let  $\mathcal{G}$  be a group of transformations on  $\mathbf{Z}$ ; in our setting  $\mathcal{G} = \mathfrak{S}_{n + m}$ , denoting the permutation (or symmetric) group of  $[N]$  by  $\mathfrak{S}_N$  and its elements by  $\sigma \in \mathfrak{S}_N, \sigma : [N] \to [N]$ . We write  $g\mathbf{Z}$  for the action of  $g \in \mathcal{G}$  on  $\mathbf{Z}$ ; e.g. defining  $\sigma \mathbf{Z} = (Z_{\sigma(1)}, \ldots, Z_{\sigma(n + m)})$  for the group elements  $\sigma \in \mathfrak{S}_{n + m}$ . We will suppose  $\mathbf{Z}$  is invariant under the action of  $\mathcal{G}$  when the sample is drawn from the null distribution, i.e. if the null is true than  $g\mathbf{Z} =^d\mathbf{Z}$  (by which we denote equality of distribution) for all  $g \in \mathcal{G}$ . This is clearly the case for  $\mathcal{G} = \mathfrak{S}_{n + m}$  in two-sample testing, since under the null  $Z = (Z_1, \ldots, Z_{m + n}) \stackrel{\mathrm{iid}}{\sim} p$  with  $p = q$ , while under the alternative, the permuted sample  $\sigma \mathbf{Z}$  for randomised  $\sigma \sim \mathrm{Uniform}(\mathfrak{S}_{n + m})$  simulates the null distribution.
We can use permutations to construct an approximate cumulative distribution function (CDF) of our test statistic under the null, and choose an appropriate quantile of this CDF as our test threshold, which must be exceeded under the null with probability less than level  $\alpha$ . For this we introduce a quantile operator (analogous to the max and min operators) for a finite set  $\{f(a)\in \mathbb{R}:a\in \mathcal{A}\}$ :
$$
\underset {q, a \in \mathcal {A}} {\text {q u a n t i l e}} f (a) := \inf  \left\{r \in \mathbb {R}: \frac {1}{| \mathcal {A} |} \sum_ {a \in \mathcal {A}} \mathbf {1} \{f (a) \leq r \} \geq q \right\}. \tag {3}
$$
Various different results can be used to choose the threshold giving a correct level; we will here highlight a very general and easy-to-use theorem for constructing permutation tests, and in Section 3 will discuss previously unconsidered implications for using unsupervised feature extraction methods as a part of our tests. Although this result is not new, we believe that its usefulness has been underappreciated in the kernel testing community, and can be more conveniently applied than Romano and Wolf (2005, Lemma 1), which requires exchangeability and is commonly used, e.g. in Albert et al. (2022, Proposition 1) and Schrab et al. (2023, Proposition 1).
Theorem 1 (Hemerik and Goeman, 2018, Theorem 2.). Let  $G$  be a vector of elements from  $\mathcal{G}$ ,  $G = (g_{1}, g_{2}, \ldots, g_{B + 1})$ , with  $g_{B + 1} = \mathrm{id}$  (the identity permutation) for any  $B \geq 1$ . The elements  $g_{1}, \ldots, g_{B}$  are drawn uniformly from  $\mathcal{G}$  either i.i.d. or without replacement (which includes the possibility of  $G = \mathcal{G}$ ). If  $\tau(\mathbf{Z})$  is a statistic of  $\mathbf{Z}$  and  $\mathbf{Z} = ^d g\mathbf{Z}$  for all  $g \in \mathcal{G}$  under the null then
$$
\mathbb{P}_{p\times p,G}\left(\tau (\boldsymbol {Z}) > \operatorname {quantile}  \tau (g\boldsymbol {Z})\right)\leq \alpha .
$$
In other words, if we compare  $\tau(Z)$  with the empirical quantile of  $\tau(gZ)$  as a test threshold, the type I error rate is no more than  $\alpha$ . The potentially complex task of constructing a permutation test reduces to the trivial task of choosing the statistic  $\tau(Z)$ . This result is also true for randomised permutations of any number  $B \geq 1$  without approximation, so an exact and computationally efficient test can be straightforwardly constructed this way.
We finally mention the related approach of the MMD Aggregated test (MMDAgg; Schrab et al., 2023), which combines multiple MMD-based permutation tests with different kernels, and rejects if
3This test is near-exact. Specifically, if the statistic realisations are distinct  $(\tau(gZ) \neq \tau(g'Z)$  for  $g \neq g')$  as is common for continuous data, the level is exactly  $\lfloor (B + 1)\alpha \rfloor / (B + 1)$ .
any of these reject, using distinct quantiles for each kernel. To ensure the overall aggregated test is well-calibrated, these quantiles must be adjusted using a second level of permutations. This incurs additional computational cost, a pitfall avoided by our fused single statistic.
Faster Sub-Optimal Tests. While the main focus of this paper revolves around the kernel selection problem for optimal, quadratic MMD testing, we also highlight the existence of a rich literature on efficient kernel tests which run in linear (or near-linear) time. These speed improvements are achieved using various tools such as: incomplete U-statistics (Gretton et al., 2012b; Yamada et al., 2019; Lim et al., 2020; Kübler et al., 2020; Schrab et al., 2022b), block U-statistics (Zaremba et al., 2013; Deka and Sutherland, 2022), eigenspectrum approximations (Gretton et al., 2009), Nyström approximations (Cherfaoui et al., 2022), random Fourier features (Zhao and Meng, 2015), analytic representations (Chwialkowski et al., 2015; Jitkrittum et al., 2016), deep linear kernels (Kirchler et al., 2020), kernel thinning (Dwivedi and Mackey, 2021; Domingo-Enrich et al., 2023), etc.
The efficiency of these tests usually entail weaker theoretical power guarantees compared to their quadratic-time counterparts, which are minimax optimal $^{4}$  (Kim et al., 2022; Li and Yuan, 2019; Fromont et al., 2013; Schrab et al., 2023; Chatterjee and Bhattacharya, 2023). These optimal quadratic tests are either permutation-based non-asymptotic tests (Kim et al., 2022; Schrab et al., 2023) or studentised asymptotic tests (Kim and Ramdas, 2023; Shekhar et al., 2022; Li and Yuan, 2019; Gao and Shao, 2022; Florian et al., 2023). We emphasise that the parameters of any of these permutation-based two-sample tests can be chosen in the unsupervised way we outline in Section 3. We choose to focus in this work on optimal quadratic-time results, but note that our general approach could be extended to sub-optimal faster tests as well.
3 Learning Statistics for Permutation Tests
Here we discuss how we can improve permutation-based tests by learning parameters for them in an unsupervised manner. The reason that we highlighted Theorem 1 in the previous section is that it holds for any  $\tau$ . For example, we could use the MMD with a kernel chosen based on the data,  $\tau(Z) = \widehat{\mathrm{MMD}}_{k=k(Z)}^{2}(Z)$ ; however, for each permutation  $\sigma$  we would need to re-compute  $\tau(\sigma Z) = \widehat{\mathrm{MMD}}_{k=k(\sigma Z)}^{2}(\sigma Z)$ , so the kernel being used would be different for each permutation. This has two major disadvantages: firstly, it might be computationally expensive to re-compute  $k$  for each permutation, especially for a deep kernel<sup>5</sup>. Secondly, the scale of the resulting MMD could be dramatically different for each permutation, so the empirical quantile might not lead to a powerful test. This second problem is related to the problem of combining multiple different MMD values into a single test which our MMD-FUSE statistics are designed to combat (Section 4; c.f. also MMDAgg, Schrab et al., 2023).
Our Proposal. In two-sample testing we propose to use a statistic  $\tau_{\theta}$  parameterised by some  $\theta$ , where  $\theta$  is fixed for all permutations, but depends on the data in an unsupervised or permutation-invariant way. Specifically, we allow such a parameter to depend on  $\langle Z\rangle \coloneqq \{Z_1,\dots ,Z_{n + m}\}$ , the unordered combined sample. Since our tests will not depend on the internal ordering of  $X$  and  $Y$  (which are assumed i.i.d. under both hypotheses), the only additional information contained in  $Z$  over  $\langle Z\rangle$  is the label assigning  $Z_{i}$  to its initial sample. This is justified since  $\langle Z\rangle = \langle \sigma Z\rangle$  for all  $\sigma \in \mathfrak{S}_{n + m}$ , so setting  $\tau (\boldsymbol {Z}) = \tau_{\theta (\langle \boldsymbol {Z}\rangle)}(\boldsymbol {Z})$  gives a fixed  $\theta$  and statistic for all permutations to use in Theorem 1. The information in  $\langle Z\rangle$  can be used to fine-tune test parameters for any test fitting this setup. This solves both the computation and scaling issues mentioned above.
The above provides a first formal justification for the use of the median heuristic in Gretton et al. (2012a), since it is a permutation-invariant function of the data. However, a far richer set of possibilities are available even when restricting ourselves to these permutation-free functions of  $\mathbf{Z}$ . For example, we can use any unsupervised or self-supervised learning method to learn representations to use as the input to an MMD-based test statistic, while paying no cost in terms of calibration and needing to train such methods only once. Given the wide variety of methods dedicated to feature extraction and dimensionality reduction, this opens up a huge range of possibilities for the design of new and principled two-sample tests. The simplicity and generality of our proposal might lead one to
<sup>4</sup>With the exception of the near-linear test of Domingo-Enrich et al. (2023) which achieves the same MMD separation rate as the quadratic test but under stronger assumptions on the data distributions.
expect that this idea has been used before, but it has not to our best knowledge, underlined by the fact that the median heuristic has been widely used without such justification when one follows from this method. This potentially powerful and widely-applicable possibility represents one of our most practical contributions.
4 MMD-FUSE: Fusing U-Statistics by Exponentiation