SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Say we have computed several MMD values  $\widehat{\mathrm{MMD}}_k^2$  under different kernels  $k\in \mathcal{K}$ . How might we combine these? One possibility is to perform multiple testing as in the case of MMDAgg. An even simpler approach would simply take the maximum of those values  $\max_k\widehat{\mathrm{MMD}}_k^2$  since Theorem 1 shows that this will not prevent us from controlling the level of our test by  $\alpha$ ; note though that for each permutation we would take this maximum separately. Indeed, Cárcamo et al. (2022) show that for certain kernel choices the supremum of the MMD values with respect to the bandwidth is a valid integral probability metric (Müller, 1997). There are two main problems with this approach: a capacity control issue and a normalisation issue.
Firstly, if the class over which we are maximising is sufficiently rich (for example a complex neural network), then the maximum may be able to effectively memorise the entire sample for each possible permutation, saturating the statistic for every permutation and limiting test power. Any convergence results would need to hold simultaneously for every  $k$  in  $\mathcal{K}$ , and so power results suffer: for finite  $|\mathcal{K}|$  we would incur a sub-optimal Bonferroni correction (see Section 5); while for infinite classes, results would need to involve capacity control quantities like Rademacher complexity. Further, only information from a single maximising "base" kernel can be considered at a time.
Therefore, in both our statistics, we prefer a "soft" maximum, which considers information from every kernel simultaneously and when more than one of the kernels is well-suited (Section 5) it therefore avoids the Bonferroni correction arising from uniform bounds. From Equation (1), our approach is equivalent to using a KL complexity penalty, and is strongly reminiscent of PAC-Bayesian (Section 2) capacity control. We note that other soft maxima could be considered, but our choice makes obtaining exponential concentration inequalities relatively easy (Appendix C), and the dual formulation (Equation (2)) allows us to derive power results in terms of MMD directly.
The second issue is that the MMD estimates might have different scales or variances per kernel which need to be accounted for. In order to be able to meaningfully compare MMD values between each other, these need to be normalised somehow before taking a maximum. We use the common approach of dividing through by a variance-like term (as in "studentised" tests, see Section 2). The specific normaliser is permutation invariant and gives our statistic tight sub-Gaussian null concentration (for well-chosen regularisation parameter  $\lambda$ ; see Appendix C).
Based on the above we introduce the FUSE-1 statistic which uses a log-sum-exp soft maximum, and the FUSE-N which combines this with normalisation. Both statistics use a "prior" distribution on  $\mathcal{K}$ , denoted  $\pi (\langle Z\rangle)$ , which is either fixed independently of the data, or is a function of the data which is invariant under permutation (as discussed in Section 3).
Definition 1. We define the un-normalised (subscript 1 for normaliser of 1) and normalised (subscript  $N$ ) test statistics with parameter  $\lambda >0$ , respectively, as
$$
\widehat {\mathrm {F U S E}} _ {1} (\boldsymbol {Z}) := \frac {1}{\lambda} \log \left(\mathbb {E} _ {k \sim \pi (\langle \boldsymbol {Z} \rangle)} \left[ \exp \left(\lambda \widehat {\mathrm {M M D}} _ {k} ^ {2} (\boldsymbol {X}, \boldsymbol {Y})\right) \right]\right),
$$
$$
\widehat {\mathrm {F U S E}} _ {N} (\boldsymbol {Z}) := \frac {1}{\lambda} \log \left(\mathbb {E} _ {k \sim \pi (\langle \boldsymbol {Z} \rangle)} \left[ \exp \left(\lambda \frac {\widehat {\operatorname {M M D}} _ {k} ^ {2} (\boldsymbol {X} , \boldsymbol {Y})}{\sqrt {\hat {N} _ {k} (\boldsymbol {Z})}}\right) \right]\right),
$$
where  $\widehat{N}_k(\mathbf{Z})\coloneqq \frac{1}{n(n - 1)}\sum_{(i,j)\in [n + m]_2}k(Z_i,Z_j)^2$  is permutation invariant.
Although these statistics appear complex, we note that in the case where  $\pi$  has finite support, their calculation reduces to a log-sum-exp of MMD estimates normalised by  $\widehat{N}_k$ . This is even clearer when  $\pi$  is also uniform on its support, as we consider experimentally; then Equation (2) shows that our statistics reduce to soft maxima, with  $\lambda$  controlling the smoothness or "temperature".
From this, we define the FUSE-1 test (with the FUSE-N test  $\Delta_{\mathrm{FUSE}_N}$  defined analogously) as
$$
\Delta_ {\mathrm {F U S E} _ {1}} (\boldsymbol {Z}) := \mathbf {1} \left\{\widehat {\mathrm {F U S E} _ {1}} (\boldsymbol {Z}) > \underset {1 - \alpha , \sigma \in S} {\text {q u a n t i l e}} \widehat {\mathrm {F U S E} _ {1}} (\sigma \boldsymbol {Z}) \right\} \tag {4}
$$
for sampled set  $S$  of permutations as described above. It compares the test statistic  $\widehat{\mathrm{FUSE}}_1$  with its quantiles under permutation, and rejects if the overall value exceeds a quantile controlled by  $\alpha$ . Note that since  $\widehat{N}_k(Z)$  is permutation invariant, it only needs to be calculated once per kernel in  $\Delta_{\mathrm{FUSE}_N}$  (and not separately for each permutation) as per Section 3. See Appendix A.5 for its time complexity.
Comparison with MMDAgg. MMDAgg is a different way to think about combining multiple kernels and MMD values, but it relies on a framework based on multiple testing. This can be problematic in the case where the number of kernels considered is large, since as this number increases in the multiple testing approach MMDAgg, the level needs to be corrected differently, so its theoretical power behaviour becomes unclear (though empirically MMDAgg retains its power in this setting). By contrast, our proposed FUSE methods avoid multiple testing as a single statistic and quantile are used, avoiding such issues. The FUSE statistics are even defined in the limit of an infinite number of kernels (continuous/uncountable collection of kernels) by considering distributions on the space of kernels, which is not the case for MMDAgg. Moreover, while the MMDAgg approach is only useful for hypothesis testing, having a quantity like FUSE combining multiple kernel-based measures of distance with exponential concentration bounds could be of interest in a wider range of applications.
FUSE-1 and the Mean Kernel. Although the un-normalised test has worse performance in practice than our normalised test, the FUSE-1 statistic is interesting in its own right because of various theoretical properties, as we discuss below. Firstly, we introduce the mean kernel  $K_{\rho}(x,y) = \mathbb{E}_{k\sim \rho}k(x,y)$  under a "posterior"  $\rho \in \mathcal{M}_+^1 (\mathcal{K})$ , which is indeed a reproducing kernel in its own right. In the finite case, this is simply a weighted sum of "base" kernels.
Note that the linearity of  $\widehat{\mathrm{MMD}}_k^2$  and  $\mathrm{MMD}_k^2$  in the kernel  $k$  implies that  $\mathbb{E}_{k\sim \rho}\mathrm{MMD}_k^2 (p,q) = \mathrm{MMD}_{K_\rho}^2 (p,q)$ , and similarly for  $\widehat{\mathrm{MMD}}_k^2$ , with these terms appearing in our power results (Section 5). Combining this linearity with the dual formulation of  $\widehat{\mathrm{FUSE}}_1$  via Equation (1) gives
$$
\widehat {\mathrm {F U S E}} _ {1} (\pmb {Z}) = \sup _ {\rho \in \mathcal {M} _ {+} ^ {1} (\mathcal {K})} \widehat {\mathrm {M M D}} _ {K _ {\rho}} ^ {2} (\pmb {X}, \pmb {Y}) - \frac {\mathrm {K L} (\rho , \pi)}{\lambda}.
$$
This re-states  $\widehat{\mathrm{FUSE}}_1$  in terms of "posterior"  $\rho$ , and makes the interpretation of our statistic as a KL-regularised kernel-learning method clear. In the finite case, our test simply optimises the weightings of the different kernels in a constrained way. We note that for certain infinite kernel sets and choices of prior it is possible to express the mean kernel in closed form. This happens because, e.g. the expectation of a Gaussian kernel with respect to a Gamma prior over the bandwidth is simply a (closed form) rational quadratic kernel. We discuss this point further in Appendix B.
5 Theoretical Power of MMD-FUSE
In this section we outline possible sufficient conditions for our tests to obtain power at least  $1 - \beta$  at given level  $\alpha$ . The conditions will depend on a fixed data-independent "prior"  $\pi \in \mathcal{M}_{+}^{1}(\mathcal{K})$  and thus hold even without unsupervised parameter optimisation. They are stated as requirements for the existence of a "posterior"  $\rho \in \mathcal{M}_{+}^{1}(\mathcal{K})$  with corresponding mean kernel  $K_{\rho}$  as defined in Section 4. In the finite case,  $K_{\rho}$  is simply a weighted sum of kernels, so these requirements are also satisfied under the same conditions for any single kernel, corresponding to the case where the posterior puts all its weight on a single kernel.
Technically, these results require that there is some constant  $\kappa$  upper bounding all of the kernels, and that  $n$  and  $m$  are within a constant multiple (i.e.  $n \leq m \leq cn$  for some  $c \geq 1$ , notated  $n \asymp m$ ). They hold when using randomised permutations provided  $B > c' \alpha^{-2} \log (\beta^{-1})$  for small constant  $c' > 0$ .
Theorem 2 (FUSE-1 Power). Fix prior  $\pi$  independently of the data. For the un-normalised test FUSE-1 with  $\lambda \asymp n / \kappa$  and  $n\asymp m$ , there exists a universal constant  $C > 0$  such that
$$
\exists \rho \in \mathcal {M} _ {+} ^ {1} (\mathcal {K}): \mathrm {M M D} _ {K _ {\rho}} ^ {2} (p, q) > \frac {C \kappa}{n} \left(\frac {1}{\beta} + \log \frac {1}{\alpha} + \mathrm {K L} (\rho , \pi)\right).
$$
is sufficient for FUSE-1 to achieve power at least  $1 - \beta$  at level  $\alpha$ .
A similar result is obtained for FUSE-N under an assumption that the normalising term is well behaved (see assumption in Theorem 3). This requirement will be satisfied for kernels (including most common ones) that tend to zero only in the limit of the data being infinitely far apart.
Theorem 3 (FUSE-N Power). Fix prior  $\pi$  independently of the data such that for all  $k\in \operatorname {supp}(\pi)$  the expectation  $\mathbb{E}_{\boldsymbol{Z}\sim p\times q}[\widehat{N}_k(\boldsymbol {Z})^{-1}] < c / \kappa$  is bounded for some  $c < \infty$ . For the normalised test FUSE-N with  $\lambda \asymp n\asymp m$ , there exists a universal constant  $C > 0$  such that
$$
\exists \rho \in \mathcal {M} _ {+} ^ {1} (\mathcal {K}): \mathrm {M M D} _ {K _ {\rho}} ^ {2} (p, q) > \frac {C \kappa}{n} \left(\frac {1}{\beta^ {2}} + \log \frac {1}{\alpha} + \mathrm {K L} (\rho , \pi)\right).
$$
is sufficient for FUSE-N to achieve power at least  $1 - \beta$  at level  $\alpha$ .
Discussion. The conditions in Theorems 2 and 3 give the optimal  $\mathrm{MMD}^2$  separation rate in  $n$  (Domingo-Enrich et al., 2023, Proposition 2). These results also imply consistency if the prior  $\pi$  assigns non-zero weight to characteristic kernels. Applying these results to uniform priors supported on  $r$  points, the KL penalty can be upper bounded as  $\mathrm{KL}(\rho, \pi) \leq \log(r)$ . Thus in the worst case, where only a single kernel achieves large  $\mathrm{MMD}_k^2(p, q)$ , the price paid for adaptivity is only  $\log(r)$ . In many cases, most of the kernels will give large  $\mathrm{MMD}_k^2(p, q)$ . The posterior will then mirror the prior, and this KL penalty will be even smaller. Thus very large numbers of kernels could be considered, and if all give large MMD values the power would not be greatly affected.
Additional Technical Results. Aside from the presentation of our new statistics and tests, we make a number of technical contributions on the way to proving Theorems 2 and 3, as well as proving some additional results. In particular, we give exponential concentration bounds for our statistics under permutations and the null, which do not require bounded kernels. This refined analysis requires the construction of a coupled Rademacher chaos and concentration thereof. We obtain intermediate results using variances from the proofs of Theorems 2 and 3 that could be used in future work to obtain power guarantees under alternative assumptions such as Sobolev spaces (Schrab et al., 2023). Finally, we prove exponential concentration for  $\widehat{\mathrm{FUSE}}_1$  under the alternative and bounded kernels, requiring the proof of a "PAC-Bayesian" bounded differences-type concentration inequality. See the appendix for a more detailed overview.
We compare the test power of MMD-FUSE-N ( $\lambda = \sqrt{n(n - 1)}$ ) against various MMD-based kernel selective tests (see Section 1 for details) using: the median heuristic (MMD-Median; Gretton et al., 2012a), data splitting (MMD-Split; Sutherland et al., 2017), analytic Mean Embeddings and Smooth Characteristic Functions (ME & SCF; Jitkrittum et al., 2016), the MMD Deep kernel (MMD-D; Liu et al., 2020), Automated Machine Learning (AutoML; Kübler et al., 2022b), kernel thinning to (Aggregate) Compress Then Test (CTT & ACTT; Domingo-Enrich et al., 2023), and MMD Aggregated (Incomplete) tests (MMDAgg & MMDAggInc; Schrab et al., 2023, 2022b). Additional details and code link for experimental reproducibility are provided in Appendix A.
Distribution on Kernels. We choose our kernel prior distribution  $\pi$  as uniform over a collection of Gaussian,  $k_{\gamma}^{g}(x,y) = \exp \left(-\| x - y\|_{2}^{2} / 2\gamma^{2}\right)$ , and Laplace,  $k_{\gamma}^{\ell}(x,y) = \exp \left(-\sqrt{2}\| x - y\|_{1} / \gamma\right)$ , kernels with various bandwidths  $\gamma >0$ . These bandwidths are chosen as the uniform discretisation of the interval between half the  $5\%$  and twice the  $95\%$  quantiles (for robustness) of  $\{\| z - z'\| _r:z,z'\in Z\}$ , with  $r\in \{1,2\}$ , respectively. This choice is similar to that of Schrab et al. (2023, Section 5.2), who empirically show that ten points for the discretisation is sufficient (Schrab et al., 2023, Figure 6), which we verify in Appendix A.4. This set of distances is permutation-invariant, so Theorem 1 guarantees a well-calibrated test even though the kernels are data-dependent.
Mixture of Gaussians. Our first experiments (Figure 1) consider multimodal distributions  $p$  and  $q$ , each a 2-dimensional mixture of four Gaussians with means  $(\pm \mu, \pm \mu)$  with  $\mu = 20$  and diagonal covariances. For  $p$ , the four components all have unit variances, while for  $q$  we vary the standard deviation  $\sigma$  of one of the Gaussians,  $\sigma = 1$  corresponds to the null hypothesis  $p = q$ . Intuitively, an appropriate kernel bandwidth to distinguish  $p$  from  $q$  would correspond to that separating Gaussians
Under this condition  $\mathrm{KL}(\nu, \pi) < \infty$  with  $\nu$  the restriction of  $\pi$  to characteristic kernels.  $K_{\nu}$  is characteristic so  $\mathrm{MMD}_{K_{\nu}}(p,q) > 0 \iff p \neq q$ , and our condition lower bound tends to zero with  $n \to \infty$ .
with standard deviations 1 and  $\sigma$ . This is significantly smaller than the median bandwidth which scales with the distance  $\mu$  between modes.
Perturbed Uniform. In Figure 1, we report test power for detecting perturbations on uniform distributions in one and two dimensions. We vary the amplitude  $a$  of two perturbations from  $a = 0$  (null) to  $a = 1$  (maximum value for the density to remain non-negative). A similar benchmark was first proposed by Schrab et al. (2023, Section 5.5) and considered in several other works (Schrab et al., 2022b; Hagrass et al., 2022; Chatterjee and Bhattacharya, 2023). Different bandwidths are required to detect different amplitudes of the perturbations.
Galaxy MNIST. We examine performance on real-world data in Figure 1, through galaxy images (Walmsley et al., 2022) in dimension  $d = 3 \times 64 \times 64 = 12288$  captured by a ground-based telescope. These consist of four classes: 'smooth and cigar-shaped', 'edge-on-disk', 'unbarred spiral', and 'smooth and round'. One distribution uniformly samples images from the first three categories, while the other does the same with probability  $1 - c$  and uniformly samples a 'smooth and round' galaxy image with probability of corruption  $c \in [0,1]$ . The null hypothesis corresponds to the case  $c = 0$ .
CIFAR 10 vs 10.1. The aim of this experiment is to detect the difference between images from the CIFAR-10 (Krizhevsky, 2009) and CIFAR-10.1 (Recht et al., 2019) test sets. This is a challenging problem as CIFAR-10.1 was specifically created to consist of new samples from the CIFAR-10 distribution so that it can be used as an alternative test set for models trained on CIFAR-10. Samples from the two distributions are presented in Figure 6 in Appendix A.3 (Liu et al., 2020, Figure 5). This
benchmark was proposed by Liu et al. (2020, Table 3) who introduced the deep MMD test MMD-D and the MMD-Split test (here referred to as MMD-O to point out that their implementation has been used rather than ours). They also compare to ME and SCF, as well as to C2ST-L and C2ST-S (Lopez-Paz and Oquab, 2017) which correspond to Classifier Two-Sample Tests based on Sign or
Linear kernels. For the tests slitting the data, 1000 images from both datasets are used for parameter selection and/or model training, and 1021 other images from each distributions are used for testing. Consequently, tests avoiding data splitting are given the full 2021 images from CIFAR-10.1 and 2021 images sampled from CIFAR-10.
Experimental Results of Figure 1. We observe similar trends in all eight experiments in Figure 1: MMD-FUSE matches the power of state-of-the-art MMDAgg while being computationally faster (both theoretically and practically). These two tests consistently obtain the highest power in every experiment, except when increasing the number of Galaxy MNIST images where MMD-D obtains higher power. However, we observe in Table 1 that MMD-FUSE outperforms MMD-D on another image data problem, also with large sample size. On synthetic data, the deep kernel test MMD-D surprisingly only obtains power similar to MMD-Split in most experiments (even lower in the 2-dimensional perturbed uniform experiment). The two near-linear aggregated variants ACTT and MMDAggInc trade-off a small portion of test power for computational efficiency, with the former outperforming the latter for large sample sizes. The importance of kernel selection is emphasised by the fact that the two tests using the median bandwidth (MMD-Median and CTT) achieve very low power. While the linear-time SCF test, based in the frequency domain, attains high power in the Mixture of Gaussians experiments, it has low power in the three other experiments. Its spatial domain variant ME performs better on Perturbed Uniform  $d \in \{1,2\}$  experiments but in general still has reduced power compared to both linear and quadratic time alternatives. Finally, the AutoML test performs well for fixed sample size  $m = n = 500$  (first row of Figure 1), but its power compared to other tests considerably deteriorates as the sample size increases (second row of Figure 1). Overall, MMD-FUSE achieves state-of-the-art performance across all experiments on both low-dimensional synthetic data and high-dimensional real-world data.
Experimental Results of Table 1. We report in Table 1 the power achieved by each test on the CIFAR 10 vs 10.1 experiment, which is averaged over 1000 repetitions. We observe that MMD-FUSE performs the best and obtains power 0.937, which means that out of 1000 repetitions, it was 937 times able to distinguish between samples from CIFAR-10 and from CIFAR-10.1. This demonstrates that the images in CIFAR-10.1 do not come from the same distribution as those in CIFAR-10.
Experimental Results of Figure 7. We observe in Figure 7 of Appendix A.4 that MMD-FUSE can achieve higher power than MMDAgg in some additional perturbed uniform experiment. We also note that using a relatively small number of kernels (e.g., 10) is enough to capture all the required information, and that the test power is retained when further increasing the number of kernels.
In this work, we propose MMD-FUSE, an MMD-based test which fuses kernels through a soft maximum and a method for learning general two-sample testing parameters in an unsupervised fashion. We demonstrate the empirical performance of MMD-FUSE and show that it achieves the optimal MMD separation rate guaranteeing high test power. This optimality holds with respect to the sample size and likely also for the logarithmic dependence in  $\alpha$ , but we believe the dependence on  $\beta$  could be improved in future work; a general question is whether lower bounds in terms of  $\alpha$  and  $\beta$  can be proved. Obtaining separation rates in terms of the  $L^2$ -norm between the densities (Schrab et al., 2023) may also be possible but challenging since this distance is independent of the kernel.
An open question is in explaining the significant empirical power advantage of the normalised test over its un-normalised variant, which is currently not reflected in the derived rates. The importance of this normalisation is clear when considering kernels with different bandwidths, leading to vastly different scaling in the un-normalised permutation distributions. Work here could begin with finite-sample concentration guarantees for our normalised statistic or other "studentised" variants, some of which might obtain better performance.
Future work could also examine computationally efficient variants of MMD-FUSE, by either relying on incomplete  $U$ -statistics (Schrab et al., 2022b) and leading to suboptimal rates, or by relying on recent ideas of kernel thinning (Domingo-Enrich et al., 2023) which can lead to the same optimal rate under stronger assumptions on the data distributions, or by considering the permutation-free approach of Shekhar et al. (2022). Finally, our two-sample MMD fusing approach could be extended to the HSIC independence framework (Gretton et al., 2005, 2008; Albert et al., 2022) and to the KSD goodness-of-fit setting (Chwialkowski et al., 2016; Liu et al., 2016; Schrab et al., 2022a).
We would like to thank the anonymous reviewers and area chair for their thorough reading of our work, for their constructive feedback, and for engaging in the discussions, all of which have helped to improve the paper. Felix Biggs and Antonin Schrab both gratefully acknowledge the support from the U.K. Research and Innovation under the EPSRC grant EP/S021566/1. Arthur Gretton acknowledges support from the Gatsby Charitable Foundation.
Overview of Appendices
Appendix A We give further experimental details, discuss time complexity of our results while graphing run-times.
Appendix B We discuss how  $\widehat{\mathrm{FUSE}}_1$  can be expressed in a simple form for certain uncountable priors.
Appendix C We prove exponential concentration bounds for both our statistics under the null and permutations.
Appendix D We prove exponential concentration bounds for the  $\widehat{\mathrm{FUSE}}_1$  statistic under the alternative hypothesis.
Appendix E We prove the power results stated in Section 5 and other interesting intermediate results.
A Additional Experimental Details
A.1 Code and licenses
Our implementation of MMD-FUSE in Jax (Bradbury et al., 2018), as well as the code for the reproducibility of the experiments, are made publicly available at:
https://github.com/antoninschrab/mmdfuse-paper
Our code is under the MIT License, and we also implement ourselves the MMD-Median (Gretton et al., 2012a) and MMD-Split (Sutherland et al., 2017) tests. For ME & SCF (Jitkrittum et al., 2016), MMD-D (Liu et al., 2020), AutoML (Kübler et al., 2022b), CTT & ACTT (Domingo-Enrich et al., 2023), MMDAgg & MMDAggInc (Schrab et al., 2023, 2022b), we use the implementations of the respective authors, which are all under the MIT license.
The experiments were run on an AMD Ryzen Threadripper 3960X 24 Cores 128Gb RAM CPU at  $3.8\mathrm{GHz}$  and on an NVIDIA RTX A5000 24Gb Graphics Card, with a compute time of a couple of hours.
A.2 Test parameters
In general, we use the default test parameters recommended by the authors of the tests (listed above). For the ME and SCF tests, ten test locations are chosen on half of the data. AutoML is run with the recommended training time limit of one minute.