A Witness Two-Sample Test

Jonas M. Kübler Wittawat Jitkrittum Bernhard Schölkopf Krikamol Muandet

Max Planck Institute for Intelligent Systems Tübingen, Germany Google Research
Max Planck Institute for Intelligent Systems Tübingen, Germany Max Planck
Institute for Intelligent Systems Tübingen, Germany

###### Abstract

The Maximum Mean Discrepancy (MMD) has been the state-of-the-art nonparametric
test for tackling the two-sample problem. Its statistic is given by the
difference in expectations of the witness function, a real-valued function
defined as the mean of kernel evaluations on a set of basis points. Typically
the kernel is optimized on a training set, and hypothesis testing is performed
on a separate test set to avoid overfitting (i.e., control type-I error). That
is, the test set is used to simultaneously estimate the expectations and
define the basis points, while the training set only serves to select the
kernel and is discarded. In this work, we propose to use the training set to
also define the weights and the basis points for better data efficiency. We
show that 1) the new test is consistent and has a well-controlled type-I
error; 2) the optimal witness function is given by a precision-weighted mean
in the reproducing kernel Hilbert space associated with the kernel; and 3) the
test power of the proposed test is comparable or exceeds that of the MMD and
other modern tests, as verified empirically on challenging synthetic and real
problems (e.g., Higgs data).

##  1 INTRODUCTION

We tackle the classic _two-sample problem_ : given two samples, do they differ
significantly enough that we can conclude they originate from two different
distributions? This is a common task in many life sciences such as
bioinformatics and cancer diagnosis (Borgwardt et al., 2006). To decide this,
one can perform a _two-sample test_ , whose goal is to reject the _null
hypothesis_ "the probability distributions are the same" in favor of the
_alternative hypothesis_ "the probability distributions are not the same"
based on data (Lehmann and Romano, 2005). To quantitatively assess this, one
defines a _test statistic_ and estimates its value on the observed samples. If
we know (or are able to simulate) the distribution of this test statistic
under the null, we can reject the null if the observed value is significantly
larger than what we would expect if the null was true. Traditional hypothesis
tests have test statistics that are defined a priori. A simple example are
$t$\- or $z$-tests, which only test whether the empirical means of both
samples differ significantly (Lehmann and Romano, 2005). However, such a
simple approach is not sufficient to detect differences of distributions with
the same mean but, for example, different variance, skewness, or kurtosis.

To detect any differences between two distributions we focus on two categories
of tests closely tied to machine learning, but note that various other methods
exist (Friedman and Rafsky, 1979; Chen and Friedman, 2017). The former first
transforms data into a high-dimensional feature space based on a pre-defined
feature map, e.g., kernel function. The test statistics can then be defined in
terms of the embeddings of the two distributions in the feature space
(Harchaoui et al., 2008b; Gretton et al., 2012a). The second approach instead
learns to distinguish the two distributions by training a classifier, e.g.,
via a deep neural network. Based on the learned model, the test statistics is
then computed on an independent set of samples, e.g., through data splitting
(Friedman, 2003; Kim et al., 2021; Lopez-Paz and Oquab, 2017; Cheng and
Cloninger, 2019).

The popular kernel two-sample test based on the _Maximum Mean Discrepancy_
(MMD) in principle does not require data splitting and is completely
determined a priori by a positive definite kernel function (Gretton et al.,
2012a). However, recent research has shown that optimizing the kernel function
on a held-out dataset increases the power of the MMD-based tests (Gretton et
al., 2012b; Sutherland et al., 2017; Liu et al., 2020; Kirchler et al., 2020).
Thus most modern MMD-based tests are used as two-stage procedures with data
splitting, although it is in principle possible to use the entire dataset for
kernel selection and testing Fromont et al. (2012, 2013); Kübler et al.
(2020).111Schrab et al. (2021) recently proposed an aggregated MMD two-sample
test working without data splitting.

To obtain maximally significant results in the testing phase, we advocate that
in a “two-stage” two-sample test, it is more appropriate to learn a test
statistic that is as problem-specific as possible. For the MMD tests, this
means that we advocate to learn a one-dimensional witness function and not a
kernel. To formalize this, we propose a general two-stage witness two-sample
test (WiTS test). The introduced WiTS test has the following properties:

  * •

The test statistic is the difference in means of a one-dimensional function
called the _witness function_ and is thus asymptotically normal under _both_
the null and alternative hypotheses. This allows for a simple theoretical
treatment (cf. Theorem 1 and Proposition 1).

  * •

Compared to Sutherland et al. (2017) and Liu et al. (2020), the WiTS test has
a simpler test power criterion as a training objective and test thresholds can
be simulated more efficiently (cf. Section 3 & Eq. (7)).

  * •

The WiTS tests empirically outperform the benchmark tests of Liu et al. (2020)
and classification-based tests on challenging synthetic and real problems,
e.g., Higgs data (cf. Figure 2).

The rest of the paper is organized as follows. Section 2 reviews MMD based
two-sample tests with a focus on the witness function and discusses our
motivation. We then present the general WiTS test framework in Section 3,
followed by a specific example in Section 4. Next, we discuss related work in
detail in Section 5. Finally, Section 6 provides the empirical results
comparing the proposed WiTS tests to existing ones on several benchmark
datasets. The code to reproduce the experiments is published under
https://github.com/jmkuebler/wits-test.

##  2 BACKGROUND AND MOTIVATION

#### Notation and definitions.

Let $X,Y$ be random variables with probability distributions $P$ and $Q$ on
$\mathcal{X}\subseteq\mathbb{R}^{d}$, respectively. In this work, we aim to
test the null hypothesis $H_{0}:P=Q$ against the alternative $H_{1}:P\neq Q$
based on samples $\mathbb{X}=\\{x_{1},\dots,x_{n}\\}$ and
$\mathbb{Y}=\\{y_{1},\dots,y_{m}\\}$ drawn i.i.d. from $P$ and $Q$,
respectively. Rejecting $H_{0}$ although it is true creates a type-I error,
whereas not rejecting the null when it is false creates a type-II error.
Desirable testing procedures should minimize the type-II error rate, while
controlling the type-I error rate at a significance level $\alpha$ (or below).
When we consider data splitting, we use
$\mathbb{X}_{\text{tr}},\mathbb{X}_{\text{te}}$ and
$\mathbb{Y}_{\text{tr}},\mathbb{Y}_{\text{te}}$ to denote the disjoint
training and test sets with $n=n_{\text{tr}}+n_{\text{te}}$,
$m=m_{\text{tr}}+m_{\text{te}}$. We define the shorthands
$[n]:=\\{1,\ldots,n\\}$, $\mathbb{Z}=\\{\mathbb{X},\mathbb{Y}\\}$,
$\mathbb{Z}_{\text{tr}}=\\{\mathbb{X}_{\text{tr}},\mathbb{Y}_{\text{tr}}\\}$
and
$\mathbb{Z}_{\text{te}}=\\{\mathbb{X}_{\text{te}},\mathbb{Y}_{\text{te}}\\}$.

Although most of our analysis applies to more general function spaces, we will
consider a reproducing kernel Hilbert space (RKHS) $\mathcal{H}$ with positive
definite kernel $k:\mathcal{X}\times\mathcal{X}\to\mathbb{R}$ (Schölkopf and
Smola, 2002). By the Riesz representation theorem, we have that
$f(x)=\braket{f,k(x,\cdot)}$ for all $x\in\mathcal{X}$ and $f\in\mathcal{H}$.
We assume that

$\textbf{(A1):
}{\mathbb{E}\left[{k(X,X)}\right]}<\infty,\,{\mathbb{E}\left[{k(Y,Y)}\right]}<\infty$

holds. (A1) ensures the kernel mean embeddings of $P$ and $Q$ exist, i.e.,
$\mu_{P}={\mathbb{E}\left[{k(X,\cdot)}\right]},\mu_{Q}={\mathbb{E}\left[{k(Y,\cdot)}\right]}$,
and that we can write ${\mathbb{E}\left[{f(X)}\right]}=\braket{f,\mu_{P}}$ for
all $f\in\mathcal{H}$ (Muandet et al., 2017). For a sample $\mathbb{X}$, we
define the empirical mean embedding as
$\mu_{\mathbb{X}}=\frac{1}{|\mathbb{X}|}\sum_{x\in\mathbb{X}}k(x,\cdot)$.

#### MMD and witness function.

A popular class of two-sample tests are based on the _Maximum Mean
Discrepancy_ (MMD) (Gretton et al., 2012a). The MMD of two distributions with
respect to the unit ball of $\mathcal{H}$ is defined as (Gretton et al.,
2012a, Eq. (1)): $\text{MMD}=\sup_{f\in\mathcal{H},\|f\|\leq
1}\left\\{{\mathbb{E}\left[{f(X)}\right]}-{\mathbb{E}\left[{f(Y)}\right]}\right\\}.$
The function that _witnesses_ the MMD is
$\text{argmax}_{f\in\mathcal{H},\|f\|\leq
1}\\{{\mathbb{E}\left[{f(X)}\right]}-{\mathbb{E}\left[{f(Y)}\right]}\\}=(\mu_{P}-\mu_{Q})/\|\mu_{P}-\mu_{Q}\|$
(Gretton et al., 2012a, Sec. 2.3). We define its unnormalized version as
$h_{k}^{P,Q}=\mu_{P}-\mu_{Q}$ and obtain

$\displaystyle\begin{aligned}
\text{MMD}^{2}&=\braket{\mu_{P}-\mu_{Q},\mu_{P}-\mu_{Q}}=\braket{\mu_{P}-\mu_{Q},h_{k}^{P,Q}}\\\
&={\mathbb{E}\left[{h_{k}^{P,Q}(X)}\right]}-{\mathbb{E}\left[{h_{k}^{P,Q}(Y)}\right]}.\end{aligned}$
(1)

With a _characteristic_ kernel (Sriperumbudur et al., 2010), $\mu_{P}=\mu_{Q}$
if and only if $P=Q$. Hence, the squared MMD (1) can be used to test the
hypothesis $H_{0}:P=Q$ against $H_{1}:P\neq Q$.

#### MMD-BOOT test statistics.

We can estimate the squared MMD (1) by replacing the witness $h_{k}^{P,Q}$ and
the expectations in (1) with their empirical counterparts
$h_{k}^{\mathbb{Z}}=\mu_{\mathbb{X}}-\mu_{\mathbb{Y}}$ and

$\displaystyle\begin{aligned}
&\widehat{\text{MMD}}_{{\textsc{boot}}}^{2}(\mathbb{Z}|k)=\frac{1}{n}\sum_{x\in\mathbb{X}}h_{k}^{\mathbb{Z}}(x)-\frac{1}{m}\sum_{y\in\mathbb{Y}}h_{k}^{\mathbb{Z}}(y)\\\
&=\Braket{\frac{1}{n}\sum_{x\in\mathbb{X}}k(x,\cdot)-\frac{1}{m}\sum_{y\in\mathbb{Y}}k(y,\cdot),h_{k}^{\mathbb{Z}}(\cdot)}\\\
&=\frac{1}{n^{2}}\sum_{x,x^{\prime}\in\mathbb{X}}k(x,x^{\prime})+\frac{1}{m^{2}}\sum_{y,y^{\prime}\in\mathbb{Y}}k(y,y^{\prime})\\\
&\quad-\frac{2}{nm}\sum_{x\in\mathbb{X},y\in\mathbb{Y}}k(x,y).\end{aligned}$
(2)

The latter expression is a sum of $V$-statistics and up to the biased terms
where $x=x^{\prime}$ or $y=y^{\prime}$ equals the unbiased $U$-statistic which
is the standard MMD estimate (Gretton et al., 2012a). The witness itself
depends on the same data $\mathbb{Z}$ used to evaluate the test statistic (2)
and compute the test threshold, the null distribution has to be simulated via
permutation of the samples (or bootstrapping) (Gretton et al., 2012a). Thus,
we refer to this approach as mmd-boot.222Our naming convention should
emphasizes that the asymptotic distribution cannot be evaluated in closed-form
and hence we necessarily need to simulate it. Note, however, that in practice
often permutations are used (Sutherland et al., 2017) and it is not necessary
to completely simulate the distribution from scratch.

#### OPT-MMD-BOOT test statistics.

A drawback of mmd-boot is that the kernel $k$ has to be chosen a priori before
observing the data. Kernel choice, however, critically affects the performance
of MMD based two-sample tests (Gretton et al., 2012b; Sutherland et al., 2017;
Liu et al., 2020; Kübler et al., 2020; Jitkrittum et al., 2016). It is thus
common to split the data into
$\mathbb{Z}=(\mathbb{Z}_{\text{tr}},\mathbb{Z}_{\text{te}})$ and optimize the
kernel only on the held-out set $\mathbb{Z}_{\text{tr}}$. For the moment,
without specifying how the kernel is optimized, we denote the resulting
optimized kernel as $k_{\text{tr}}$ with a subscript $\mathrm{tr}$ to indicate
that it depends on the training data. After optimizing the kernel, the
standard mmd-boot test is conducted on $\mathbb{Z}_{\text{te}}$ with the
optimized kernel $k_{\text{tr}}$ (Sutherland et al., 2017; Liu et al., 2020).
Hence, the empirical expectations and witness function in (2) are still
dependent on the same data $\mathbb{Z}_{\text{te}}$, and the null distribution
still has to be bootstrapped, for the same reason as in the case of mmd-boot.
We will refer to this approach as opt-mmd-boot with the test statistic

$\displaystyle\begin{aligned} &\widehat{\text{MMD}}_{{\textsc{opt-
boot}}}^{2}(\mathbb{Z}_{\text{te}}|k_{\text{tr}})\\\
&=\frac{1}{n_{\text{te}}}\sum_{x\in\mathbb{X}_{\text{te}}}h_{k_{\text{tr}}}^{\mathbb{Z}_{\text{te}}}(x)-\frac{1}{m_{\text{te}}}\sum_{y\in\mathbb{Y}_{\text{te}}}h_{k_{\text{tr}}}^{\mathbb{Z}_{\text{te}}}(y).\end{aligned}$
(3)

#### Our Motivation.

This is the starting point of our investigations: Although the kernel is
optimized, it is still a multidimensional representation of the data. While
this makes the test statistic applicable to other problems (Liu et al., 2020;
Kirchler et al., 2020), features that contain little information about the
differences of $P$ and $Q$ will mainly add noise to the test statistic.
Generally, the noisier the test statistic, the harder it is to obtain
significant test results. Motivated by this drawback, we propose to formulate
a test statistic that is more specific to the observed difference in
$\mathbb{Z}_{\text{tr}}$. Being more specific to the training data (that is
all we know about $P$ and $Q$), comes at the risk of overfitting, which we
mitigate via regularization and model selection (cf. Section 3). Specifically
for MMD, after the kernel is optimized, _we define the witness directly on the
training data_ by replacing $h_{k_{\text{tr}}}^{\mathbb{Z}_{\text{te}}}$ with
$h_{k_{\text{tr}}}^{\mathbb{Z}_{\text{tr}}}=\frac{1}{n_{\text{tr}}}\sum_{x\in\mathbb{X}_{\text{tr}}}k_{\text{tr}}(x,\cdot)-\frac{1}{m_{\text{tr}}}\sum_{y\in\mathbb{Y}_{\text{tr}}}k_{\text{tr}}(y,\cdot)$.
We call this opt-mmd-witness:

$\displaystyle\begin{aligned} &\widehat{\text{MMD}}_{{\textsc{opt-
witness}}}^{2}\left(\mathbb{Z}_{\text{te}}|h_{k_{\text{tr}}}^{\mathbb{Z}_{\text{tr}}}\right)\\\
&=\frac{1}{n_{\text{te}}}\sum_{x\in\mathbb{X}_{\text{te}}}h_{k_{\text{tr}}}^{\mathbb{Z}_{\text{tr}}}(x)-\frac{1}{m_{\text{te}}}\sum_{y\in\mathbb{Y}_{\text{te}}}h_{k_{\text{tr}}}^{\mathbb{Z}_{\text{tr}}}(y).\end{aligned}$
(4)

This test statistic comes with numerous advantages. Firstly, the expectations
(defined via $\mathbb{Z}_{\text{te}}$) are now independent of the witness
function (defined via $\mathbb{Z}_{\text{tr}}$). Thus, the test statistic is
asymptotically normal. Secondly, as we will see in the following sections, (4)
allows us to compute asymptotic test thresholds in closed form and allows for
a simpler derivation of a test power criterion than in the case of opt-mmd-
boot (Sutherland et al., 2017; Liu et al., 2020). Lastly, our empirical
results suggest that opt-mmd-witness outperforms opt-mmd-boot on datasets
considered in Liu et al. (2020).

##  3 WITNESS TWO-SAMPLE TEST (WiTS TEST)

Similar to (4), the WiTS tests we propose are by design two-stage procedures:
In _Stage I_ , we learn the witness function $h$ with the training data
$\mathbb{Z}_{\text{tr}}$. This ensures that $h$ is independent of the test
data $\mathbb{Z}_{\text{te}}$, used in _Stage II_ to define a test statistic

$\displaystyle\hat{\tau}(\mathbb{Z}_{\text{te}}|h)\propto\frac{1}{n_{\text{te}}}\sum_{x\in\mathbb{X}_{\text{te}}}h(x)-\frac{1}{m_{\text{te}}}\sum_{y\in\mathbb{Y}_{\text{te}}}h(y).$
(5)

We reject the null hypothesis $H_{0}:P=Q$ if the observed value is larger than
a test threshold. We start presenting Stage II and analyze the test’s
asymptotic power for a given function $h$. Then, we will use this test power
criterion as the objective when optimizing the witness function in Stage I.

###  3.1 Stage II - Testing with the Witness Function

We start with a basic result on asymptotic normality of empirical means
((Serfling, 1980), Proof in App. A.1).

######  Theorem 1 (Asymptotic normality of WiTS test).

For a witness function $h:\mathcal{X}\to\mathbb{R}$, let
$\sigma_{P}^{2}:=\text{Var}[h(X)]$ and $\sigma_{Q}^{2}:=\text{Var}[h(Y)]$ such
that $0<\sigma^{2}_{P},\sigma^{2}_{Q}<\infty$. Let
$\\{X_{i}\\}_{i\in[n]}\overset{\text{i.i.d.}}{\sim}P$,
$\\{Y_{j}\\}_{j\in[m]}\overset{\text{i.i.d.}}{\sim}Q$, and
$c:=\frac{n}{n+m}\in(0,1)$ as $n+m\to\infty$. Denote by
$\bar{h}_{P}:={\mathbb{E}\left[{h(X)}\right]}$ and
$\bar{h}_{Q}:={\mathbb{E}\left[{h(Y)}\right]}$. We define the empirical means
$\hat{h}^{n}_{P}:=\frac{1}{n}\sum_{i\in[n]}h(X_{i})$,
$\hat{h}^{m}_{Q}:=\frac{1}{m}\sum_{i\in[m]}h(Y_{i})$ and denote the sample
variance as
$\hat{\sigma}^{2}_{c}(h):=\hat{\sigma}^{2}_{P}/c+\hat{\sigma}^{2}_{Q}/(1-c)$.
Then

$\displaystyle\frac{\sqrt{n+m}}{\hat{\sigma}_{c}(h)}\left[\left(\hat{h}^{n}_{P}-\bar{h}_{P}\right)-\left(\hat{h}^{m}_{Q}-\bar{h}_{Q}\right)\right]\overset{d}{\to}\mathcal{N}\left(0,1\right).$

For any fixed $h$ and for sufficiently large sample sizes, we can thus work
with the asymptotic distribution of test statistics of the form
$\tau(\cdot|h)$ in Eq. (5) to compute test thresholds and derive an asymptotic
test-power objective for choosing $h$ based on the training data
$\mathbb{Z}_{\text{tr}}$ in Stage I. Data splitting ensures that $h$ is
independent of $\mathbb{Z}_{\text{te}}$, which is necessary for Theorem 1 to
hold. In the following, to make the comparison between different choices of
$h$ easier, we consider the standardized test statistic on the test samples
$\mathbb{Z}_{\text{te}}$

$\displaystyle\begin{aligned}
&\tau(\mathbb{Z}_{\text{te}}|h)=\sqrt{n_{\text{te}}+m_{\text{te}}}\frac{\frac{1}{n_{\text{te}}}\sum\limits_{x\in\mathbb{X}_{\text{te}}}h(x)-\frac{1}{m_{\text{te}}}\sum\limits_{y\in\mathbb{Y}_{\text{te}}}h(y)}{\hat{\sigma}_{c}(h)},\end{aligned}$

where $c=\frac{n_{\text{te}}}{n_{\text{te}}+m_{\text{te}}}$ and
$\hat{\sigma}_{c}(h)$ is the empirical estimate of the pooled variance as in
Theorem 1 based on $\mathbb{Z}_{\text{te}}$. To control the type-I error at a
significance level $\alpha$, we need to find a _test threshold_ $t_{\alpha}$
such that $P(\tau(\mathbb{Z}_{\text{te}}|h)>t_{\alpha}|H_{0})\leq\alpha$. By
Theorem 1, we can define the threshold to be the $(1-\alpha)$ quantile of the
asymptotic null distribution. Under the null hypothesis we have
$\bar{h}_{P}=\bar{h}_{Q}$ and obtain $t_{\alpha}=\Phi^{-1}(1-\alpha)$ where
$\Phi^{-1}$ denotes the inverse CDF of the standard normal.

Note that we only consider a "one-sided" test, since we choose $h$ in stage I
with the appropriate sign, i.e., such that it has larger expectation under
$\mathbb{X}_{\text{tr}}$ than under $\mathbb{Y}_{\text{tr}}$. A "two-sided"
test ignores this and may lead to a reduction in test power.

We reject the null hypothesis $H_{0}:P=Q$ if
$\tau(\mathbb{Z}_{\text{te}}|h)>t_{\alpha}$. As an advantage of the asymptotic
normality under the alternative and the closed form of the threshold of our
test, we can write the asymptotic type-II error rate in closed form, similar
as in Gretton et al. (2012b, Sec. 3):

$\displaystyle\begin{aligned} &P(\tau(\mathbb{Z}_{\text{te}}|h)<t_{\alpha})\\\
&\approx\Phi\left(\Phi^{-1}(1-\alpha)-\sqrt{n_{\text{te}}+m_{\text{te}}}\,\frac{\bar{h}_{P}-\bar{h}_{Q}}{\sigma_{c}(h)}\right).\end{aligned}$
(6)

An important consideration in designing a hypothesis test is test consistency.
A hypothesis test is called consistent, if for a fixed alternative hypothesis,
its test power converges to one as sample size goes to infinity. With (6), we
can characterize for which functions $h$ the statistic $\tau_{h}$ leads to a
consistent test.

######  Proposition 1 (Consistency of WiTS test).

Assume $0<\sigma_{c}(h)<\infty$, where $\sigma_{c}(h)$ is defined in Theorem
1. A WiTS test based on $h$ is consistent against a fixed alternative
hypothesis $P\neq Q$ if and only if $\bar{h}_{P}>\bar{h}_{Q}$.

Proposition 1 ensures that, for a given alternative hypothesis, our proposed
test will eventually (in the limit of the sample size) reject the null
hypothesis $H_{0}$ when it is false. Associated with this notion is the _test
power_ , the probability that the test rejects $H_{0}$ when it is false; this
quantity is equivalent to $1-$ type-II error. Defining the _signal-to-noise_
ratio $\text{SNR}(h)=\frac{\bar{h}_{P}-\bar{h}_{Q}}{\sigma_{c}(h)}$, it
follows from (6) that the asymptotic test power of our test is

$\displaystyle\beta_{h}\approx
1-\Phi\left(\Phi^{-1}(1-\alpha)-\sqrt{n_{\text{te}}+m_{\text{te}}}\,\text{SNR}(h)\right).$
(7)

Since $\Phi$ increases monotonically, the test power grows monotonically with
the _signal-to-noise_ ratio (SNR).

###  3.2 Stage I - Finding an Optimal Witness

We now propose an objective to find an optimal witness function. Based on our
test power consideration, we argue that in the first stage one should find a
witness by maximizing a, possibly regularized, empirical estimate of the SNR
in (7). Let $\mathcal{F}$ be a function class containing candidates for the
witness. We propose using the witness $\hat{h}_{\lambda}$ defined as

$\displaystyle\begin{aligned}
&\hat{h}_{\lambda}=\mathop{\mathrm{argmax}}\limits_{f\in\mathcal{F}}\frac{\bar{f}_{\mathbb{X}_{\text{tr}}}-\bar{f}_{\mathbb{Y}_{\text{tr}}}}{\sigma^{\mathbb{Z}_{\text{tr}}}_{c,\lambda}(f)},\\\
&\text{with
}\bar{f}_{\mathbb{X}_{\text{tr}}}=\frac{1}{n_{\text{tr}}}\sum_{x\in\mathbb{X}_{\text{tr}}}f(x),\;\bar{f}_{\mathbb{Y}_{\text{tr}}}=\frac{1}{m_{\text{tr}}}\sum_{y\in\mathbb{Y}_{\text{tr}}}f(y),\end{aligned}$
(8)

and
$\sigma_{c,\lambda}^{\mathbb{Z}_{\text{tr}}}(f)=((\sigma_{c}^{\mathbb{Z}_{\text{tr}}}(f)^{2}+\lambda\Omega(f))^{\frac{1}{2}}$,
where $\sigma_{c}^{\mathbb{Z}_{\text{tr}}}(f)$ corresponds to
$\hat{\sigma}_{c}(h)$ defined in Theorem 1 and $\Omega$ is a regularizer. We
remark that the optimal witness is generally not uniquely defined since the
SNR is invariant to rescaling the function. Correctly rejecting $H_{0}$ when
it is false is at the core of hypothesis testing. Our choice of maximizing the
SNR in (7) is in line with this principle: it leads to a test that maximizes
the asymptotic test power. By contrast, while other objectives such as
classification loss(Kim et al., 2021; Lopez-Paz and Oquab, 2017), softmax loss
(Cheng and Cloninger, 2019), or the MMD statistic itself (Gretton et al.,
2012a), can be used to learn the witness function, their relationship to the
test power may be indirect.

#### OPT-MMD-Witness.

A closely related objective to our SNR in (7) was used in previous work
(Sutherland et al., 2017; Liu et al., 2020) to find a good kernel for a opt-
mmd-boot test, see (3). For a given kernel $k$, Liu et al. (2020, Eq.(3))
derive the training objective as
$J(P,Q|k)=\text{MMD}^{2}(P,Q|k)/\sigma_{H_{1}}(P,Q|k)$ where
$\sigma^{2}_{H_{1}}(P,Q|k)$ is the asymptotic variance of the MMD estimate
under the alternative hypothesis. In Appendix A.5, we examine this quantity in
more detail, and show that $J(P,Q|k)=1/\sqrt{2}\,\text{SNR}(h_{k}^{P,Q})$. For
a given class of kernels and corresponding (empirical) MMD witnesses, this
implies that selecting the optimal witness according to our SNR criterion
leads to the same function as first optimizing the kernel with the $J$
criterion and defining the witness afterwards.

Algorithm 1 WiTS test with kfda-witness

1:Input: $\mathbb{X},\mathbb{Y},$ $\alpha$, paramGrid, $r$

2:$\mathbb{X}_{\text{tr}},\mathbb{X}_{\text{te}},\mathbb{Y}_{\text{tr}},\mathbb{Y}_{\text{te}}\leftarrow$
RandomSplit($\mathbb{X},\mathbb{Y},r$)

3:$\\#$ Optionally perform model selection

4:$k,\lambda\leftarrow$ GridSearchCV(paramGrid, $\mathbb{Z}_{\text{tr}}$)

5:$\\#$ Stage I - Optimize Witness

6:$h\leftarrow$ kfdaWitness($\mathbb{Z}_{\text{tr}},k,\lambda$)
$\triangleright$ App. Alg.2

7:$\\#$ Stage II - Test

8:return: witnessTest($\mathbb{Z}_{\text{te}},h,\alpha$)

9:function witnessTest($\mathbb{Z}_{\text{te}},h(\cdot),\alpha$, $B=200$)

10: $h_{\mathbb{Z}_{\text{te}}}\leftarrow[h(z)\text{ for $z$ in
}\mathbb{Z}_{\text{te}}]$

11:
$\tau\leftarrow\textsc{mean}(h_{\mathbb{Z}_{\text{te}}}[:n_{\text{te}}])-\textsc{mean}(h_{\mathbb{Z}_{\text{te}}}[n_{\text{te}}:])$

12: $p\leftarrow 0$ $\triangleright$ simulate $p$-value via permutations

13: for $i$ in $[B]$ do

14:
$h_{\mathbb{Z}_{\text{te}}}\leftarrow\textsc{Permute}(h_{\mathbb{Z}_{\text{te}}})$

15: if
$\textsc{mean}(h_{\mathbb{Z}_{\text{te}}}[:n_{\text{te}}])-\textsc{mean}(h_{\mathbb{Z}_{\text{te}}}[n_{\text{te}}:])\geq\tau$
then

16: $p\leftarrow p+1/B$

17: if $p\leq\alpha$ then return: 1 else return: 0

#### Model Selection and Optimization.

The choice of function class $\mathcal{F}$ and regularization parameter
$\lambda$ affects the learned witness in (8). We therefore recommend that
practitioners use standard tools for model selection such as cross-validation
(CV) for finding suitable “hyperparameters” and to validate that the learned
witness actually has a high SNR. CV ensures that the witness actually learns
the differences between $P$ and $Q$ and does not solely overfit the training
data. Model-selection on $\mathbb{Z}_{\text{tr}}$ is legit since in Stage II
we only use $\mathbb{Z}_{\text{te}}$, which are independent of
$\mathbb{Z}_{\text{tr}}.$ While this is also possible in classifier two-sample
tests (Lopez-Paz and Oquab, 2017), in the standard mmd-boot this is not done.

Our objective (7) can be used with a variety of function classes
$\mathcal{F}$. For instance, $\mathcal{F}$ can be defined based on an RKHS, or
parameterized by a deep neural network. Note that optimization methods to
maximize (8) are generally function class specific, and may require an
iterative procedure. However, when $\mathcal{F}$ is an RKHS, we can derive the
closed-form solution to (8), as shall be explained in Section 4. Algorithm 1
shows the general procedure for the two-stage WiTS test.

#### Permutation-based Thresholds.

For our theoretical analysis we used the asymptotic threshold. However, the
witness is also chosen in a data-dependent manner. Thus, we generally
recommend to simulate the threshold via permutations in order to ensure type-I
error control at finite sample size. In this case, for simplicity and ease of
implementation, we compute the test statistic without normalization and simply
take the difference in means. We first compute the value of the witness
function on all points in $\mathbb{Z}_{\text{te}}$ and store it in an array.
Then we compute the simplified test statistic by taking the difference in
means of $\mathbb{X}_{\text{te}}$ and $\mathbb{Y}_{\text{te}}$ (as computed
from the array that stores all the witness evaluations). We then iterate over
$B\in\mathbb{N}$ permutation runs to estimate the $p$-value of the computed
test statistic. For each run, we permute the array storing the witness
evaluations, and then compute the difference in means of the first
$n_{\text{te}}$ and the last $m_{\text{te}}$ entries. If this is larger or
equal than the test statistic on the original partition, this contributes
$1/B$ to the p-value. After all permutations, if the p-value is smaller than
$\alpha$, we reject (see Alg. 1). This correctly controls type-I errors, as
under the null, the initial partition can be thought of as being itself a
random permutation of the data. Since for this procedure we only need to
compute the witness once on each data point the overall cost is
$\mathcal{O}((n_{\text{te}}+m_{\text{te}})B)$. Note that simulating the null
for mmd-boot instead has cost
$\mathcal{O}((n_{\text{te}}+m_{\text{te}})^{2}B)$ (Liu et al., 2020, Sec. 5).

##  4 KFDA-WITNESS

In this section, we consider the function class in (8) to be an RKHS, and show
that this choice leads to a closed form solution for the optimal witness. To
start, let $\mathcal{H}$ be an RKHS associated with a positive definite kernel
$k$ (see Section 2). Additionally to the mean embeddings $\mu_{P},\mu_{Q}$, we
define the (centered) covariance operator
$\Sigma_{P}={\mathbb{E}\left[{k(X,\cdot)\otimes
k(X,\cdot)}\right]}-\mu_{P}\otimes\mu_{P}$ (analogously for $Q$) whose
existence is ensured by Assumption (A1) (Muandet et al., 2017, Sec. 3). For
any function in the RKHS we then have
${\mathbb{E}\left[{f(X)}\right]}=\braket{\mu_{P},f}$ and
$\text{Var}[f(X)]=\braket{f,\Sigma_{P}f}$, and analogously for $Q$. We define
the pooled covariance operator
$\Sigma=\frac{\Sigma_{P}}{c}+\frac{\Sigma_{Q}}{1-c}$. Then for all
$f\in\mathcal{H}$ with non-zero variance we have

$\displaystyle\text{SNR}(f)=\frac{\braket{\mu_{P}-\mu_{Q},f}}{\braket{f,\Sigma
f}^{\frac{1}{2}}},$ (9)

where SNR is defined in (7). This objective corresponds to Kernel Fisher
discriminant analysis (KFDA)’s learning objective (Mika et al., 1999). For
singular covariance operator the SNR can diverge, and for infinite-dimensional
RKHS, the empirical estimation of the covariance operator is ill-posed. In the
following, we therefore consider a regularized ($\lambda>0$) version of (9)
and call its solution _(regularized) KFDA witness_ :

$\displaystyle
h_{\lambda}=\mathop{\mathrm{argmax}}\limits_{\begin{subarray}{c}f\in\mathcal{H}\end{subarray}}\frac{\braket{\mu_{P}-\mu_{Q},f}}{\braket{f,(\Sigma+\lambda\mathop{I})f}^{\frac{1}{2}}}.$
(10)

The solution of (10) is given by the solution to the generalized eigenvalue
problem $(\Sigma+\lambda\mathop{I})h_{\lambda}=\gamma(\mu_{P}-\mu_{Q})$ (Mika,
2003, Sec.3.2), thus

$\displaystyle
h_{\lambda}=\gamma(\Sigma+\lambda\mathop{I})^{-1}(\mu_{P}-\mu_{Q}),$ (11)

where $\gamma>0$ is an arbitrary positive constant we fix to $1$, unless
stated otherwise. We will refer to the test with the witness function
$h_{\lambda}$ as the kfda-witness test.

Next, we show how we can estimate the KFDA-witness with the training data.

#### Estimation of the KFDA Witness.

Let
$\mathbb{Z}_{\text{tr}}=\\{x_{1},\dots,x_{n_{\text{tr}}},y_{1},\dots,y_{m_{\text{tr}}}\\}$
denote the pooled training sample and $K$ denote the kernel matrix such that
$K_{ij}=k(z_{i},z_{j})$ for $i,j\in[n_{\text{tr}}+m_{\text{tr}}]$. Further, we
define
$\delta=(\frac{1}{n_{\text{tr}}},\dots,\frac{1}{n_{\text{tr}}},-\frac{1}{m_{\text{tr}}},\dots,-\frac{1}{m_{\text{tr}}})^{\top}\in\mathbb{R}^{n_{\text{tr}}+m_{\text{tr}}}$.
For $l\in\\{n_{\text{tr}},m_{\text{tr}}\\}$, we define the idempotent
centering matrix $P_{l}=\mathop{I}_{l}-l^{-1}\bm{1}_{l}\bm{1}_{l}^{\top}$,
where $\mathop{I}_{l}$ denotes the identity operator and $\bm{1}_{l}$ the $l$
dimensional vector with all ones. With this we define the
$(n_{\text{tr}}+m_{\text{tr}})\times(n_{\text{tr}}+m_{\text{tr}})$ matrix
$N_{c}=\begin{pmatrix}\frac{1}{c}P_{n_{\text{tr}}}&0\\\
0&\frac{1}{1-c}P_{m_{\text{tr}}}\end{pmatrix}.$ Using the representer theorem
(Schölkopf et al., 2001), we can empirically estimate the KFDA witness (more
detail in App. A.3) as

$\displaystyle\hat{h}_{\lambda}(\cdot)=\sum_{i=1}^{n_{\text{tr}}+m_{\text{tr}}}\hat{\alpha}_{i}k(z_{i},\cdot),$
(12)
$\displaystyle\hat{\alpha}=\left(\frac{KN_{c}K}{n_{\text{tr}}+m_{\text{tr}}}+\lambda
K\right)^{-1}K\delta.$ (13)

$\hat{h}_{\lambda}(\cdot)$ can be viewed as a precision-weighted (inverse
covariance) mean of the embeddings of the basis points
$\mathbb{Z}_{\text{tr}}$ in the RKHS. Since
$\mu_{\mathbb{X}_{\text{tr}}},\mu_{\mathbb{Y}_{\text{tr}}},$ and
$\hat{\Sigma}$ are consistent estimates of $\mu_{P},\mu_{Q}$, and $\Sigma$,
for fixed regularization, we have $\hat{h}_{\lambda}\to
h_{\lambda}=(\Sigma+\lambda\mathop{I})^{-1}(\mu_{P}-\mu_{Q})$ (see Appendix
A.4). For the asymptotic witness $h_{\lambda}$ we can compute the difference
in expectation under $P$ and $Q$ in closed form:
$\bar{h}_{\lambda,P}-\bar{h}_{\lambda,Q}=\braket{\mu_{P}-\mu_{Q},h_{\lambda}}=\braket{\mu_{P}-\mu_{Q},(\Sigma+\lambda\mathop{I})^{-1}(\mu_{P}-\mu_{Q})}$.
This difference is positive, and hence by Proposition 1 we obtain a consistent
WiTS test, if and only if $\mu_{P}\neq\mu_{Q}$. We can ensure this for
arbitrary $P\neq Q$ by using a _characteristic_ kernel (Sriperumbudur et al.,
2010), the same condition as for MMD based tests.

Despite asymptotic consistency, the test power at finite sample size depends
on the _splitting ratio_ $r\in(0,1)$, i.e., $n_{\text{tr}}=\lceil rn\rceil$
and $n_{\text{te}}=n-n_{\text{tr}}$ and accordingly for the sample from $Q$.
Based on our experimental results, we observe that, for a fixed kernel $k$,
fixed regularization $\lambda>0$ and sufficiently large sample size, the
splitting ratio $r=1/2$ appears to give the highest test power in many cases,
compared to other values of $r$. Generally, identifying the optimal splitting
ratio remains an open problem. We observe (middle panel of Fig. 1) that if we
include model selection in stage I, it is favorable to use more than half of
the data for the first stage, i.e., $r>1/2$. However, since we cannot quantify
how much "_more_ " data we should use, we generally recommend using a 50/50
split.

The cost of computing the exact solution $\hat{\alpha}$ in (12) is
$\mathcal{O}((n_{\text{tr}}+m_{\text{tr}})^{2})$ in space (storing the kernel
matrix) and $\mathcal{O}((n_{\text{tr}}+m_{\text{tr}})^{3})$ time (matrix
inversion). In Appendix C, we adopt recent advances in large-scale kernel
machines (Rudi et al., 2017; Meanti et al., 2020) to obtain approximate
solutions with lower time and space complexity and thus scale to large
datasets. Using the Nyström approximation (Williams and Seeger, 2000) to
approximate the solution and approximately solving it with conjugate gradient,
we obtain a complexity of $\mathcal{O}((n_{\text{tr}}+m_{\text{tr}})Mt+M^{3})$
in time and $\mathcal{O}(M^{2})$ in space, where $M$ denotes the number of
Nyström centers and $t$ the number of conjugate gradient iterations. For stage
II we then only need $(n_{\text{te}}+m_{\text{te}})M$ kernel evaluations to
compute the test statistic. This makes our approach scalable to large-scale
dataset.333After acceptance of this work, Chatalic et al. (2022) proposed a
Nyström approximation of the kernel mean embedding to speed up the MMD
estimation.

#### Connection of opt-mmd-witness and kfda-witness.

To emphasize the relationship between optimizing the MMD and using KFDA,
consider a fixed kernel $k$ and denote by $\mathcal{A}$ the set of bounded
positive operators on $\mathcal{H}_{k}$. We consider the nonparametric class
of kernels
$\mathcal{K}=\\{{k}_{A}|k_{A}(x,y)=\braket{Ak(x,\cdot),Ak(y,\cdot)},A\in\mathcal{A}\\}$.
For this class of kernels, we show in App. A.6 that using opt-mmd-witness
leads to the same witness function as using kfda-witness.

####  kfda-boot.

It turns out that KFDA-like test statistics were considered before (Harchaoui
et al., 2008b), but in settings without data splitting. Indeed, for a fixed
$k$ and $\lambda>0$, we can use the whole data, i.e., $\mathbb{X},\mathbb{Y}$
for learning the witness
$(\hat{\Sigma}+\lambda)^{-1}(\mu_{\mathbb{X}}-\mu_{\mathbb{Y}})$ and computing
the test statistic (empirical mean difference). The test statistic thus is
$\tau_{\textsc{kfda-
boot}}=\braket{\mu_{\mathbb{X}}-\mu_{\mathbb{Y}},(\hat{\Sigma}+\lambda)^{-1}(\mu_{\mathbb{X}}-\mu_{\mathbb{Y}})},$
and we call its population version $\text{KFDA}^{2}(P,Q|k,\lambda)$. This, is
the test statistic as studied by Harchaoui et al. (2008b). As for mmd-boot,
the same data is used for estimating the witness and computing the mean
difference, hence Theorem 1 does not hold anymore. We thus need to bootstrap
the null distribution via permutations of the samples; thus, we refer to it as
kfda-boot. kfda-boot has similar drawbacks as mmd-boot: 1\. simulating the
null distribution via permutations has cost $\mathcal{O}((n+m)^{3}B)$ for
$B\in\mathbb{N}$ draws from the null distribution; and 2\. we have to fix $k$
and $\lambda$ a priori, and their choices strongly affect the test power.
Harchaoui et al. (2008b) do not provide guidance for how to choose $k$ and
$\lambda$.

Table 1: Overview of kernel-based two-sample tests. a priori means that the
kernel/regularization is chosen independently of the data. The present work
proposes the "witness" methods.

Method | kernel choice |  reg. $\lambda$ | witness obj. | witness estim. | test data | threshold
---|---|---|---|---|---|---
kfda-witness(proposed) | CV | CV | SNR | $\mathbb{Z}_{\text{tr}}$ | $\mathbb{Z}_{\text{te}}$ | analytic
kfda-boot(Harchaoui et al., 2008b) | a priori | a priori | SNR |  $\mathbb{Z}$ (implicit) | $\mathbb{Z}$ | bootstrap
mmd-boot(Gretton et al., 2012a) | a priori | - | MMD |  $\mathbb{Z}$ (implicit) | $\mathbb{Z}$ | bootstrap
opt-mmd-witness(proposed) |  $J$ with $\mathbb{Z}_{\text{tr}}$ | - | MMD | $\mathbb{Z}_{\text{tr}}$ | $\mathbb{Z}_{\text{te}}$ | analytic
opt-mmd-boot(Sutherland et al., 2017) |  $J$ with $\mathbb{Z}_{\text{tr}}$ | - | MMD |  $\mathbb{Z}_{\text{te}}$ (implicit) | $\mathbb{Z}_{\text{te}}$ | bootstrap

Figure 1: Instructive experiments on "Blobs" dataset. Left: Fixed kernel and
fixed regularization for sample size $n=m=100$. Middle: For multiple candidate
kernels ($\mathcal{K}_{10}$) kernel optimization becomes more important and
the difference of kfda-witness and opt-mmd-witness becomes smaller. Further,
opt-mmd-witness already outperforms opt-mmd-boot. Right: Same kernels as in
the middle figure and $r=1/2$. All the tests are consistent, i.e., converge to
power equal 1.

##  5 RELATED WORK

Besides the kernel-based tests we discussed so far, Chwialkowski et al. (2015)
proposed tests based on _smooth characteristic functions_ (SCF), and projected
_mean embeddings_ (ME) of the distributions where the mean embeddings are
projected to $J$-dimensional Euclidean vectors for $J\in\mathbb{N}$. In fact,
the normalized ME statistic in (Chwialkowski et al., 2015, Eq. 13) can be seen
as a variant of the KFDA where the function classes is restricted by the $J$
projection directions. Note that for a finite-dimensional RKHS and without
regularization, kfda-boot corresponds to the Hotelling’s $T^{2}$ statistic
(Hotelling, 1931). Jitkrittum et al. (2016) improve this approach by
optimizing the features in the first stage. However, they also discard the
training data after learning the $J$ projection directions. Kirchler et al.
(2020) propose to learn a deep finite-dimensional representation of the data
and to use this for a subsequent MMD or KFDA test. However, their training
objective does not directly maximize the test power (Kirchler et al., 2020,
Sec. 3.1.1). Liu et al. (2020) propose a deep version of opt-mmd-boot. They
learn a deep-kernel (mmd-d) of the form
$k_{\omega}(x,x^{\prime})=\left[(1-\epsilon)\kappa(\phi_{\omega}(x),\phi_{\omega}(x^{\prime})))+\epsilon\right]q(x,x^{\prime})$,
where $\epsilon\in(0,1)$, $\kappa$ and $q$ are Gaussian kernels and
$\phi_{\omega}$ is a deep representation optimized via the criterion $J$, see
App. A.5. They also consider a version called mmd-o which is
$k_{\omega}(x,x^{\prime})=\kappa(\phi_{\omega}(x),\phi_{\omega}(x^{\prime}))$
and conclude that learning a full kernel (they advocate mmd-d) is better than
learning a one-dimensional representation.

Most of the aforementioned works focus on developing a practical testing
procedure for a specific dataset at hand. However, there also exist more
theoretical work on the statistical optimality of different kernel-based
approaches. Balasubramanian et al. (2021) show that a _moderated_ MMD approach
(which is related to KFDA) leads to optimal rates when testing against local
alternatives. A similar discussion can be found in the long version of
Harchaoui et al. (2008a, Sec.5.1). This resonates our findings, that a witness
based on KFDA is more powerful than simply using the MMD witness. Furthermore,
Li and Yuan (2019) show how the choice of scaling parameter in Gaussian
kernels affects the statistical optimality. However, such theoretically
optimal tests oftentimes are unpractical to use. Balasubramanian et al.
(2021), for examples requires, the eigendecomposition of the kernel function,
which generally is hard to obtain. Furthermore, without data splitting also
these works cannot find a good kernel function.

Since our proposed witness function is one-dimensional, it is closely related
to classification based two-sample tests (Friedman, 2003; Kim et al., 2021;
Lopez-Paz and Oquab, 2017; Cheng and Cloninger, 2019; Cai et al., 2020).
Lopez-Paz and Oquab (2017) proposed learning a deep classifier and using its
classification accuracy as test statistic. We refer to this as c2st-s, where s
stands for sign. The method has two drawbacks. First, classification loss does
optimize the 0-1 loss, whereas we directly maximize test power (Lopez-Paz and
Oquab, 2017, Remark 2). Second, it only uses the sign of the classification
function and thus neglects information by weighting all points equally. Cheng
and Cloninger (2019) address the second issue by considering the network’s
output before thresholding the function into a classifier. They train with a
softmax loss, which also does not directly address test power. The connections
of these methods to kernel-based tests were also thoroughly discussed by Liu
et al. (2020) and, in accordance, we refer to the approach of Cheng and
Cloninger (2019) as c2st-l.

##  6 EXPERIMENTS

Figure 2: Benchmark experiments adapted from Liu et al. (2020) Left: Blobs,
Right: HIGGS. Computing the MMD witness after kernel optimization and
performing a witness test (mmd-d-witness) improves the test power over mmd-d.
Directly learning the kfda-witness also leads to high power.

We empirically assess the test power of the proposed WiTS tests in two
settings. First, we perform instructive experiments to highlight the
differences of the methods summarized in Table 1. Second, we perform benchmark
experiments on two challenging datasets and compare the performance of the
introduced WiTS tests (kfda-witness and opt-mmd-witness) to the benchmarks
(mmd-d, mmd-o, me, scf, c2st-s, c2st-l) introduced in Section 5. For the
benchmarks, we reuse the implementation provided by Liu et al. (2020) without
changing any hyperparameters. Throughout our experiments we set the level
$\alpha=0.05$. App. B contains experiments for correct type-I error control.
The shaded regions contain $\pm$ one standard deviation of the
estimates.444Note that in Fig. 2 we used different approaches to estimate the
rejection rates, see Appendix B. This explains that at the same rejection rate
we can have differently large errors.

#### Instructive experiments.

In Figure 1, we consider a Blobs dataset (Gretton et al., 2012b) where $P$ and
$Q$ are mixtures of nine anisotropic 2-d Gaussians with $Q$ having the
covariance matrix rotated by an angle $\theta=\pi/4$, see Figure 5 in the
appendix. For the left panel of Fig. 1, we consider a single Gaussian kernel
$k_{\sigma}(x,x^{\prime})=\exp{(-\lVert x-x^{\prime}\rVert^{2}/\sigma^{2}})$
with bandwidth $\sigma=0.2$ and a regularization parameter for the kfda
methods of $\lambda=10^{-2}$ (in the appendix we show the effect of the
regularization in Fig. 4. Note that for $\lambda\to\infty$, kfda and mmd
methods coincide). We showcase the effect of varying splitting ratios $r$ when
the kernel is fixed a-priori (thus we can apply mmd-boot and kfda-boot). With
fixed kernel, opt-mmd-boot essentially discards the training data. We estimate
the test power (rejection rate) with fixed overall sample size $n=m=100$. We
observe that the witness methods achieve highest power for a 50/50 split,
given a fixed kernel and fixed regularization. We also observe that the boot
approaches outperform the witness methods in this case.

However, in practice, it is unlikely that we can pick a powerful kernel and
regularization _a priori_. Therefore, for the middle panel of Figure 1, we
optimize the kernel function over a class of kernels $\mathcal{K}_{10}$
consisting of ten Gaussian kernels with bandwidths on a logarithmic range from
$10^{-3}$ to $10^{1}$. Additionally, for kfda-witness we cross-validate over
five candidate regularizations on a log range from $10^{-4}$ to $10^{3}$. In
this case, the witness methods attain the highest power at a splitting ratio
$r>1/2$, and opt-mmd-witness outperforms opt-mmd-boot for the majority of
splitting ratios and also globally. For the right panel, we use the same
setting, but fix the splitting ratio at $r=1/2$ and vary the sample size. As
we expect, all tests are consistent and we observe that both WiTS test
approaches outperform opt-mmd-boot at a 50/50 split.

#### Benchmark Experiments.

Liu et al. (2020) benchmarked several deep classification two-sample tests
(c2st-l, c2st-c) against MMD with an optimized deep kernel (mmd-d, mmd-o) and
the optimized tests (me, scf) of Jitkrittum et al. (2016). We implement opt-
mmd-witness on top of their proposed method mmd-d, which optimizes a deep
kernel (Liu et al., 2020, Section 5). Therefore after the kernel optimization,
we use the training data to define the MMD witness function (Eq. (4)) and then
proceed with WitnessTest from Algorithm 1. We also run kfda-witness with grid
search over the same kernels and regularization as for the previous
experiments. We run the experiments on two benchmarks. First, an adopted Blobs
problem, with multiple different covariances (Liu et al., 2020, Figure 1) (see
Figure 5 in the appendix), introduced to show the limitations of MMD with
translation-invariant kernels. Second, the Higgs dataset (Baldi et al., 2014)
where "we compare the jet $\phi$-momenta distribution ($d=4$) of the
background process, $P$, which lacks Higgs bosons, to the corresponding
distribution $Q$ for the process that produces Higgs bosons" (cited from Liu
et al. (2020)). For the Higgs dataset we consider sample sizes larger than a
thousand per class. To speed up the computation of the kfda-witness, we
approximate the solution with $M=500$ Nyström centers, see Appendix C, which
underlines the scalability of our approach. For both datasets we observe
higher power of the WiTS tests we propose, see Figure 2. We emphasize that we
used the implementation of Liu et al. (2020), without changing the deep
architecture or any hyperparameters.

##  7 CONCLUSION

We introduced a principled approach to learn optimal witness functions for
two-sample testing. The approach consists of two-stages: First, we learn a
witness on a subset of the observations by maximizing a test-power criterion.
In the second stage, we simply test whether the witness function attains the
same mean on the test samples, and efficiently simulate the null distribution
via permutations. We further showed how to adopt recent tests based on
optimized Maximum Mean Discrepancy into a witness two-sample test. Liu et al.
(2020) advocated optimizing a (deep) kernel in the training stage. Our
experiments show, however, that explicitly learning a one-dimensional witness
can perform better than learning a high-dimensional representation (a kernel
function) in the training stage.

Our results extend beyond kernel methods since we derive a principled
objective to train a one-dimensional function optimal for two-sample testing.
This objective and the proposed testing procedure can be applied with any
function class. The proposed framework thus not only allows domain experts to
perform two-sample tests with the models most suitable to the data at hand,
but can also easily incorporate model selection techniques developed for
classification and regression tasks to optimize for the best parameter
settings.

#### Acknowledgments

This work was in part supported by the German Federal Ministry of Education
and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039B) and the
Machine Learning Cluster of Excellence number 2064/1 – Project 390727645.

## References

  * Balasubramanian et al. (2021) K. Balasubramanian, T. Li, and M. Yuan.  On the optimality of kernel-embedding based goodness-of-fit tests.  _Journal of Machine Learning Research_ , 22(1):1–45, 2021.
  * Baldi et al. (2014) P. Baldi, P. Sadowski, and D. Whiteson.  Searching for exotic particles in high-energy physics with deep learning.  _Nature communications_ , 5(1):1–9, 2014.
  * Borgwardt et al. (2006) K. M. Borgwardt, A. Gretton, M. J. Rasch, H.-P. Kriegel, B. Schölkopf, and A. J. Smola.  Integrating structured biological data by Kernel Maximum Mean Discrepancy.  _Bioinformatics_ , 22(14):49–57, 2006.
  * Cai et al. (2020) H. Cai, B. Goggin, and Q. Jiang.  Two-sample test based on classification probability.  _Statistical Analysis and Data Mining: The ASA Data Science Journal_ , 13(1):5–13, 2020.
  * Chatalic et al. (2022) A. Chatalic, N. Schreuder, A. Rudi, and L. Rosasco.  Nystr$\backslash$" om kernel mean embeddings.  _arXiv:2201.13055_ , 2022.
  * Chen and Friedman (2017) H. Chen and J. H. Friedman.  A new graph-based two-sample test for multivariate and object data.  _Journal of the American Statistical Association_ , 112(517):397–409, 2017.
  * Cheng and Cloninger (2019) X. Cheng and A. Cloninger.  Classification logit two-sample testing by neural networks.  _arXiv:1909.11298_ , 2019.
  * Chwialkowski et al. (2015) K. Chwialkowski, A. Ramdas, D. Sejdinovic, and A. Gretton.  Fast two-sample testing with analytic representations of probability measures.  In _NeurIPS_ , 2015.
  * Friedman (2003) J. H. Friedman.  On multivariate goodness of fit and two sample testing.  _Stanford Linear Accelerator Center–PUB–10325_ , 2003.
  * Friedman and Rafsky (1979) J. H. Friedman and L. C. Rafsky.  Multivariate Generalizations of the Wald-Wolfowitz and Smirnov Two-Sample Tests.  _The Annals of Statistics_ , 7(4):697 – 717, 1979\.
  * Fromont et al. (2012) M. Fromont, B. Laurent, M. Lerasle, and P. Reynaud-Bouret.  Kernels based tests with non-asymptotic bootstrap approaches for two-sample problems.  In _COLT_ , 2012.
  * Fromont et al. (2013) M. Fromont, B. Laurent, and P. Reynaud-Bouret.  The two-sample problem for Poisson processes: Adaptive tests with a nonasymptotic wild bootstrap approach.  _The Annals of Statistics_ , 41(3):1431 – 1461, 2013.
  * Fukumizu et al. (2005) K. Fukumizu, F. R. Bach, and A. Gretton.  Statistical convergence of kernel CCA.  In _NeurIPS_ , 2005.
  * Gretton et al. (2012a) A. Gretton, K. M. Borgwardt, M. J. Rasch, B. Schölkopf, and A. Smola.  A kernel two-sample test.  _Journal of Machine Learning Research_ , 13:723–773, 2012a.
  * Gretton et al. (2012b) A. Gretton, D. Sejdinovic, H. Strathmann, S. Balakrishnan, M. Pontil, K. Fukumizu, and B. K. Sriperumbudur.  Optimal kernel choice for large-scale two-sample tests.  In _NeurIPS_ , 2012b.
  * Harchaoui et al. (2008a) Z. Harchaoui, F. Bach, and E. Moulines.  Testing for homogeneity with kernel fisher discriminant analysis.  _arXiv:0804.1026_ , 2008a.
  * Harchaoui et al. (2008b) Z. Harchaoui, F. R. Bach, and E. Moulines.  Testing for homogeneity with kernel Fisher discriminant analysis.  In _NeurIPS_ , 2008b.
  * Hotelling (1931) H. Hotelling.  The generalization of student’s ratio.  _The Annals of Mathematical Statistics_ , 2(3):360–378, 1931.
  * Jitkrittum et al. (2016) W. Jitkrittum, Z. Szabó, K. P. Chwialkowski, and A. Gretton.  Interpretable distribution features with maximum testing power.  In _NeurIPS_ , 2016.
  * Kim et al. (2021) I. Kim, A. Ramdas, A. Singh, and L. Wasserman.  Classification accuracy as a proxy for two-sample testing.  _The Annals of Statistics_ , 49(1):411 – 434, 2021.
  * Kirchler et al. (2020) M. Kirchler, S. Khorasani, M. Kloft, and C. Lippert.  Two-sample testing using deep learning.  In _AISTATS_ , 2020.
  * Kübler et al. (2020) J. M. Kübler, W. Jitkrittum, B. Schölkopf, and K. Muandet.  Learning kernel tests without data splitting.  In _NeurIPS_ , 2020.
  * Lehmann and Romano (2005) E. L. Lehmann and J. P. Romano.  _Testing statistical hypotheses_.  Springer Texts in Statistics. Springer, third edition, 2005.
  * Li and Yuan (2019) T. Li and M. Yuan.  On the optimality of gaussian kernel based nonparametric tests against smooth alternatives.  _arXiv:1909.03302_ , 2019.
  * Liu et al. (2020) F. Liu, W. Xu, J. Lu, G. Zhang, A. Gretton, and D. J. Sutherland.  Learning deep kernels for non-parametric two-sample tests.  In _ICML_ , 2020.
  * Lopez-Paz and Oquab (2017) D. Lopez-Paz and M. Oquab.  Revisiting classifier two-sample tests.  In _ICLR_ , 2017.
  * Meanti et al. (2020) G. Meanti, L. Carratino, L. Rosasco, and A. Rudi.  Kernel methods through the roof: Handling billions of points efficiently.  In _NeurIPS_ , 2020.
  * Mika (2003) S. Mika.  _Kernel Fisher Discriminants_.  Doctoral thesis, Technische Universität Berlin, Berlin, 2003.
  * Mika et al. (1999) S. Mika, G. Rätsch, J. Weston, B. Schölkopf, and K. R. Müller.  Fisher discriminant analysis with kernels.  In _Neural Networks for Signal Processing IX_ , pages 41–48, 1999\.
  * Muandet et al. (2017) K. Muandet, K. Fukumizu, B. Sriperumbudur, and B. Schölkopf.  _Kernel Mean Embedding of Distributions: A Review and Beyond_ , volume 10 of _Foundations and Trends in Machine Learning_.  2017\.
  * Rudi et al. (2017) A. Rudi, L. Carratino, and L. Rosasco.  Falkon: An optimal large scale kernel method.  In _NeurIPS_ , 2017.
  * Schölkopf and Smola (2002) B. Schölkopf and A. Smola.  _Learning with Kernels: Support Vector Machines, Regularization, Optimization, and Beyond_.  MIT Press, Cambridge, MA, USA, 2002.
  * Schölkopf et al. (2001) B. Schölkopf, R. Herbrich, and A. J. Smola.  A generalized representer theorem.  In _COLT_ , 2001.
  * Schrab et al. (2021) A. Schrab, I. Kim, M. Albert, B. Laurent, B. Guedj, and A. Gretton.  Mmd aggregated two-sample test.  _arXiv:2110.15073_ , 2021.
  * Serfling (1980) R. J. Serfling.  _Approximation theorems of mathematical statistics_.  John Wiley & Sons, 1980.
  * Sriperumbudur et al. (2010) B. K. Sriperumbudur, A. Gretton, K. Fukumizu, B. Schölkopf, and G. R. Lanckriet.  Hilbert Space Embeddings and Metrics on Probability Measures.  _Journal of Machine Learning Research_ , 11:1517–1561, 2010\.
  * Sutherland et al. (2017) D. J. Sutherland, H.-Y. Tung, H. Strathmann, S. De, A. Ramdas, A. Smola, and A. Gretton.  Generative models and model criticism via optimized maximum mean discrepancy.  In _ICLR_ , 2017.
  * Williams and Seeger (2000) C. K. I. Williams and M. W. Seeger.  Using the Nyström method to speed up kernel machines.  In _NeurIPS_ , 2000.

##  Appendix A PROOFS

###  A.1 Proof of Theorem 1

###### Proof.

Theorem 1 follows by the application of the CLT; see, e.g., Theorem A, Chapter
1.9.1 in Serfling (1980). The CLT implies
$\sqrt{n+m}(\hat{h}^{n}_{P}-\bar{h}_{P})=\sqrt{n/c}(\hat{h}^{n}_{P}-\bar{h}_{P})\overset{d}{\to}\mathcal{N}(0,\sigma^{2}_{P}/c)$,
analogously for $Q$ and the variances add up. Since
$\hat{\sigma}^{2}_{c}(h)\overset{p}{\to}\sigma_{c}:=\sigma^{2}_{P}/c+\sigma^{2}_{Q}/(1-c)$,
the result follows from Slutsky’s theorem. ∎

###  A.2 Proof of Proposition 1

###### Proof.

Since we assume $\sigma_{c}(h)>0$, it follows that

$\displaystyle\lim_{n_{\text{te}}+m_{\text{te}}\to\infty}\Phi\left(\Phi^{-1}(1-\alpha)-\sqrt{n_{\text{te}}+m_{\text{te}}}\frac{\bar{h}_{P}-\bar{h}_{Q}}{\sigma_{c}(h)}\right)=0,$
(14)

i.e., the asymptotic rate of type-II errors goes to zero, if and only if
$\bar{h}_{P}>\bar{h}_{Q}$. ∎

###  A.3 Derivation of Equation (12)

We use the following definitions: Let
$Z=\\{x_{1},\dots,x_{n_{\text{tr}}},y_{1},\dots,y_{m_{\text{tr}}}\\}$ denote
the pooled training sample and $K$ denote the kernel matrix such that
$K_{ij}=k(z_{i},z_{j})$ for $i,j\in[n_{\text{tr}}+m_{\text{tr}}]$. Let us
define $G\in\mathcal{H}^{n_{\text{tr}}+m_{\text{tr}}}$ such that
$G_{i}=k(z_{i},\cdot)$. And we write $K=G^{\top}G$. Further we define
$v_{1}=(\frac{1}{n_{\text{tr}}},\dots,\frac{1}{n_{\text{tr}}},0,\dots,0)^{\top}\in\mathbb{R}^{n_{\text{tr}}+m_{\text{tr}}}$,
$v_{2}=(0,\dots,0,\frac{1}{m_{\text{tr}}},\dots,\frac{1}{m_{\text{tr}}})^{\top}\in\mathbb{R}^{n_{\text{tr}}+m_{\text{tr}}}$,
and $\delta=v_{1}-v_{2}$. For $l=n_{\text{tr}},m_{\text{tr}}$ we define the
idempotent centering operator
$P_{l}=\mathop{I}_{l}-l^{-1}\bm{1}_{l}\bm{1}_{l}^{\top}$, where $\mathop{I}$
denotes the identity operator and $\bm{1}_{l}$ the $l$ dimensional vector with
all ones. With this we define the
$(n_{\text{tr}}+m_{\text{tr}})\times(n_{\text{tr}}+m_{\text{tr}})$ matrix
$N_{c}=\begin{pmatrix}\frac{1}{c}P_{n_{\text{tr}}}&0\\\
0&\frac{1}{1-c}P_{m_{\text{tr}}}\end{pmatrix}.$ With the preceding
definitions, we obtain $\hat{\mu}_{P}-\hat{\mu}_{Q}=G\delta$,
$\hat{\Sigma}=\frac{1}{n_{\text{tr}}+m_{\text{tr}}}GN_{c}G^{\top}$.

Starting from (10) we estimate the KFDA witness based on the empirical
estimates of $\mu_{P},\mu_{Q},\Sigma$, i.e.,

$\displaystyle\hat{h}_{\lambda}=\mathop{\mathrm{argmax}}\limits_{\begin{subarray}{c}f\in\mathcal{H}\end{subarray}}\frac{\braket{\mu_{\mathbb{X}_{\text{tr}}}-\mu_{\mathbb{Y}_{\text{tr}}},f}}{\braket{f,(\hat{\Sigma}+\lambda\mathop{I})f}^{\frac{1}{2}}}.$
(15)

We first show a _representer Theorem_ for KFDA (Mika, 2003, Sec. 3.4.3).
Therefore, we decompose possible candidate functions
$f=f_{1}+f_{2}\in\mathcal{H}$ into a part $f_{1}$ that lies in the span of the
training data
$\mathcal{S}_{\text{tr}}=\text{span}(\\{k(z_{i},\cdot)|i\in[n_{\text{tr}}+m_{\text{tr}}]\\})$
and $f_{2}$ which lies in the span’s orthogonal complement. Thus, by
definition, we have $\braket{f_{2},k(z_{i},\cdot)}=0$ for all
$i\in[n_{\text{tr}}+m_{\text{tr}}]$. Since $\mu_{\mathbb{X}_{\text{tr}}}$ and
$\mu_{\mathbb{Y}_{\text{tr}}}$ are within $\mathcal{S}_{\text{tr}}$, we have
$\braket{\mu_{\mathbb{X}_{\text{tr}}}-\mu_{\mathbb{Y}_{\text{tr}}},f}=\braket{\mu_{\mathbb{X}_{\text{tr}}}-\mu_{\mathbb{Y}_{\text{tr}}},f_{1}}$.
Similarly, since $\hat{\Sigma}$ is only defined via the training samples in
$Z$, $\hat{\Sigma}$ maps functions from $\mathcal{S}_{\text{tr}}$ to
$\mathcal{S}_{\text{tr}}$ and we have $\hat{\Sigma}f_{2}=0$. Thus for the
denominator of (15) we get

$\displaystyle\braket{f,(\hat{\Sigma}+\lambda\mathop{I})f}=\braket{f_{1},(\hat{\Sigma}+\lambda\mathop{I})f_{1}}+\lambda\|f_{2}\|^{2}\geq\braket{f_{1},(\hat{\Sigma}+\lambda\mathop{I})f_{1}}.$
(16)

We have shown that the nominator of (15) stays constant, if we add a function
$f_{2}$ that is not is not in $\mathcal{S}_{\text{tr}}$ and the denominator
can only grow. This implies that the maximum in (15) is attained for a
function in $\mathcal{S}_{\text{tr}}$ and we can expand it as
$\hat{h}_{\lambda}(\cdot)=\sum_{i=1}^{n_{\text{tr}}+m_{\text{tr}}}\hat{\alpha}_{i}k(z_{i},\cdot)$.
Hence the solution is

$\displaystyle\hat{\alpha}$
$\displaystyle=\mathop{\mathrm{argmax}}\limits_{\alpha\in\mathbb{R}^{n_{\text{tr}}+m_{\text{tr}}}}\frac{\braket{\mu_{\mathbb{X}_{\text{tr}}}-\mu_{\mathbb{Y}_{\text{tr}}},\sum_{i=1}^{n_{\text{tr}}+m_{\text{tr}}}\alpha_{i}k(z_{i},\cdot)}}{\braket{\sum_{i=1}^{n_{\text{tr}}+m_{\text{tr}}}\alpha_{i}k(z_{i},\cdot),(\hat{\Sigma}+\lambda\mathop{I})\sum_{i=1}^{n_{\text{tr}}+m_{\text{tr}}}\alpha_{i}k(z_{i},\cdot)}^{\frac{1}{2}}}$
(17)
$\displaystyle=\mathop{\mathrm{argmax}}\limits_{\alpha\in\mathbb{R}^{n_{\text{tr}}+m_{\text{tr}}}}\frac{\delta^{\top}K\alpha}{\left(\alpha^{\top}\left(\frac{KN_{c}K}{n_{\text{tr}}+m_{\text{tr}}}+\lambda
K\right)\alpha\right)^{\frac{1}{2}}}.$ (18)

The solution to this is (Mika, 2003, Sec. 3.2)555For a sanity check, simply
compute the gradient of (17) and set it to zero.

$\displaystyle\left(\frac{KN_{c}K}{n_{\text{tr}}+m_{\text{tr}}}+\lambda
K\right)\hat{\alpha}=K\delta\qquad\Longleftrightarrow\qquad\hat{\alpha}=\left(\frac{KN_{c}K}{n_{\text{tr}}+m_{\text{tr}}}+\lambda
K\right)^{-1}K\delta.$ (19)

###  A.4 Convergence of $\hat{h}_{\lambda}$

We will show that $\hat{h}_{\lambda}\to
h_{\lambda}=(\Sigma+\lambda\mathop{I})^{-1}(\mu_{P}-\mu_{Q})$ in probability.

###### Proof.

First, we observe that

$\displaystyle\hat{h}_{\lambda}-h_{\lambda}$
$\displaystyle=(\hat{\Sigma}+\lambda
I)^{-1}(\mu_{\mathbb{X}_{\text{tr}}}-\mu_{\mathbb{Y}_{\text{tr}}})-({\Sigma}+\lambda
I)^{-1}(\mu_{P}-\mu_{Q})$ $\displaystyle=(\hat{\Sigma}+\lambda
I)^{-1}(\mu_{\mathbb{X}_{\text{tr}}}-\mu_{\mathbb{Y}_{\text{tr}}})-(\hat{\Sigma}+\lambda
I)^{-1}(\mu_{P}-\mu_{Q})$ $\displaystyle\qquad+(\hat{\Sigma}+\lambda
I)^{-1}(\mu_{P}-\mu_{Q})-({\Sigma}+\lambda I)^{-1}(\mu_{P}-\mu_{Q})$
$\displaystyle=(\hat{\Sigma}+\lambda
I)^{-1}\left[(\mu_{\mathbb{X}_{\text{tr}}}-\mu_{\mathbb{Y}_{\text{tr}}})-(\mu_{P}-\mu_{Q})\right]+\left[(\hat{\Sigma}+\lambda
I)^{-1}-({\Sigma}+\lambda I)^{-1}\right](\mu_{P}-\mu_{Q}).$

Thus it follows that

$\displaystyle\lVert\hat{h}_{\lambda}-h_{\lambda}\rVert_{\mathcal{H}}$
$\displaystyle\leq\lVert(\hat{\Sigma}+\lambda\mathop{I})^{-1}[(\mu_{\mathbb{X}_{\text{tr}}}-\mu_{\mathbb{Y}_{\text{tr}}})-(\mu_{P}-\mu_{Q})]\rVert_{\mathcal{H}}$
$\displaystyle\qquad+\lVert[(\hat{\Sigma}+\lambda\mathop{I})^{-1}-({\Sigma}+\lambda\mathop{I})^{-1}](\mu_{P}-\mu_{Q})\rVert_{\mathcal{H}}$
$\displaystyle=(A)+(B).$

#### Probabilistic bound on $(A)$.

By the triangle inequality,

$\displaystyle\lVert(\hat{\Sigma}+\lambda\mathop{I})^{-1}[(\mu_{\mathbb{X}_{\text{tr}}}-\mu_{\mathbb{Y}_{\text{tr}}})-(\mu_{P}-\mu_{Q})]\rVert_{\mathcal{H}}$
$\displaystyle\leq\lVert(\hat{\Sigma}+\lambda\mathop{I})^{-1}\rVert\lVert(\mu_{\mathbb{X}_{\text{tr}}}-\mu_{\mathbb{Y}_{\text{tr}}})-(\mu_{P}-\mu_{Q})\rVert_{\mathcal{H}}$
$\displaystyle\leq\lVert(\hat{\Sigma}+\lambda\mathop{I})^{-1}\rVert(\lVert\mu_{\mathbb{X}_{\text{tr}}}-\mu_{P}\rVert_{\mathcal{H}}+\lVert\mu_{Q}-\mu_{\mathbb{Y}_{\text{tr}}}\rVert_{\mathcal{H}}).$

By the spectral theorem, $\|(\hat{\Sigma}+\lambda
I)^{-1}\|=\sup_{\hat{l}\in(\hat{l}_{k})_{k=1}^{\infty}}\frac{1}{\hat{l}+\lambda}\leq
1/\lambda$ where $(\hat{l}_{k})_{k=0}^{\infty}$ are the eigenvalues of
$\hat{\Sigma}$ and by definition non-negative. Then, the
$\sqrt{n}$-convergence of $(A)$ follows from the $\sqrt{n}$-convergence of the
kernel mean embeddings
$\left\lVert\mu_{\mathbb{X}_{\text{tr}}}-\mu_{P}\right\rVert_{\mathcal{H}}=\mathcal{O}_{p}(n_{\text{tr}}^{-1/2})$
and
$\left\lVert\mu_{Q}-\mu_{\mathbb{Y}_{\text{tr}}}\right\rVert_{\mathcal{H}}=\mathcal{O}_{p}(m_{\text{tr}}^{-1/2})$;
see, e.g., Muandet et al. (2017, Theorem 3.4). That is,
$(A)=\mathcal{O}_{p}(\min(n_{\text{tr}},m_{\text{tr}})^{-1/2})$.

#### Probabilistic bound on $(B)$.

Using the identity $C^{-1}-D^{-1}=C^{-1}(D-C)D^{-1}$, we can rewrite $(B)$ as

$\displaystyle\lVert[(\hat{\Sigma}$
$\displaystyle+\lambda\mathop{I})^{-1}-({\Sigma}+\lambda\mathop{I})^{-1}](\mu_{P}-\mu_{Q})\rVert_{\mathcal{H}}$
$\displaystyle=\lVert(\hat{\Sigma}+\lambda\mathop{I})^{-1}(\hat{\Sigma}-\Sigma)({\Sigma}+\lambda\mathop{I})^{-1}(\mu_{P}-\mu_{Q})\rVert_{\mathcal{H}}$
$\displaystyle\leq\lVert(\hat{\Sigma}+\lambda\mathop{I})^{-1}\rVert\lVert\hat{\Sigma}-\Sigma\rVert\lVert({\Sigma}+\lambda\mathop{I})^{-1}(\mu_{P}-\mu_{Q})\rVert_{\mathcal{H}}$
$\displaystyle\leq\lVert(\hat{\Sigma}+\lambda\mathop{I})^{-1}\rVert\lVert\hat{\Sigma}-\Sigma\rVert_{\text{HS}}\lVert({\Sigma}+\lambda\mathop{I})^{-1}(\mu_{P}-\mu_{Q})\rVert_{\mathcal{H}},$

where we used that the operator norm is upper bounded by the Hilbert-Schmidt
norm. Let $n:=n_{\text{tr}}+m_{\text{tr}}$. Then, since
$\lVert(\hat{\Sigma}+\lambda\mathop{I})^{-1}\rVert\leq 1/\lambda$, the
$\sqrt{n}$-convergence of $(B)$ follows from the $\sqrt{n}$-convergence of the
covariance operator, i.e.,
$\lVert\hat{\Sigma}-\Sigma\rVert_{\text{HS}}=\mathcal{O}_{p}(n^{-1/2})$
(Fukumizu et al., 2005, Lemma 4). That is,
$(B)=\mathcal{O}_{p}((n_{\text{tr}}+m_{\text{tr}})^{-1/2})$.

Combining the rates of $(A)$ and $(B)$ yields the overall rate of convergence:
$\lVert\hat{h}_{\lambda}-h_{\lambda}\rVert_{\mathcal{H}}=\mathcal{O}_{p}(\min(n_{\text{tr}},m_{\text{tr}})^{-1/2})$.
∎

###  A.5 Witness objective vs. kernel optimization objective in MMD tests

In MMD-based two sample tests, the most common estimate of the MMD is the
U-statistic estimate, defined as (Gretton et al., 2012a)

$\displaystyle\widehat{\text{MMD}^{2}_{u}}=\frac{1}{n(n+1)}\sum_{i\neq
j}H_{ij},$ (20)

with
$H_{ij}=\braket{k(x_{i},\cdot)-k(y_{i},\cdot),k(x_{j},\cdot)-k(y_{j},\cdot)}$.
The objective function used in Sutherland et al. (2017); Liu et al. (2020)
bases on the asymptotic variance of the estimator under the alternative
hypothesis. If the population value of $\text{MMD}^{2}$ is positive, then the
distribution of the estimate is asymptotically normal (Serfling, 1980, Section
5.5.1),
$\sqrt{n}\left(\widehat{\text{MMD}_{u}^{2}}-\text{MMD}^{2}\right)\overset{d}{\to}\mathcal{N}(0,\sigma^{2}_{H_{1}}),$
with
$\sigma^{2}_{H_{1}}=4({\mathbb{E}\left[{H_{12}H_{13}}\right]}-{\mathbb{E}\left[{H_{12}}\right]}^{2})$
(Liu et al., 2020). This can be used to derive an asymptotic test power
criterion, which is given as the signal-to-noise ratio
$J=\frac{\text{MMD}^{2}}{\sigma_{H_{1}}}$ (Sutherland et al., 2017, Sec. 2.1).

We show, that the power criterion $J=\frac{\text{MMD}^{2}}{\sigma_{H_{1}}}$
corresponds to the SNR criterion we derived in (8). It is an easy exercise to
show that

$\displaystyle\sigma^{2}_{H_{1}}=4$ $\displaystyle\left({\mathbb{E}_{X\sim
P}\left[{\braket{\mu_{P}-\mu_{Q},k(X,\cdot)}^{2}}\right]}+{\mathbb{E}_{Y\sim
Q}\left[{\braket{\mu_{P}-\mu_{Q},k(Y,\cdot)}^{2}}\right]}\right.$
$\displaystyle\left.\quad-\braket{\mu_{P}-\mu_{Q},\mu_{P}}^{2}-\braket{\mu_{P}-\mu_{Q},\mu_{Q}}^{2})\right).$

Recalling the definition of the covariance operator
$\Sigma_{P}={\mathbb{E}\left[{k(X,\cdot)\otimes
k(X,\cdot)}\right]}-\mu_{P}\otimes\mu_{P}$, we obtain

$\displaystyle\sigma^{2}_{H_{1}}$
$\displaystyle=4\braket{\mu_{P}-\mu_{Q},(\Sigma_{P}+\Sigma_{Q})(\mu_{P}-\mu_{Q})}=2\braket{\mu_{P}-\mu_{Q},(2\Sigma_{P}+2\Sigma_{Q})(\mu_{P}-\mu_{Q})}$
$\displaystyle=2\braket{\mu_{P}-\mu_{Q},\Sigma(\mu_{P}-\mu_{Q})},$

where we used $\Sigma=\Sigma_{P}/c+\Sigma_{Q}/(1-c)$ and $c=1/2$ for balanced
samples.

Using $h_{k}^{P,Q}=\mu_{P}-\mu_{Q}$, we have

$\displaystyle J(P,Q|k)$
$\displaystyle=\frac{\text{MMD}^{2}}{\sigma_{H_{1}}}=\frac{\braket{\mu_{P}-\mu_{Q},\mu_{P}-\mu_{Q}}}{\sqrt{2}\braket{\mu_{P}-\mu_{Q},\Sigma(\mu_{P}-\mu_{Q})}^{\frac{1}{2}}}=\frac{\braket{\mu_{P}-\mu_{Q},h_{k}^{P,Q}}}{\sqrt{2}\braket{h_{k}^{P,Q},\Sigma
h_{k}^{P,Q}}^{\frac{1}{2}}}$ (21)
$\displaystyle=\frac{1}{\sqrt{2}}\text{SNR}(h_{k}^{P,Q}).$ (22)

###  A.6 MMD of nonparametrically optimized kernel corresponds to KFDA

Consider a fixed kernel $k$ and denote by $\mathcal{A}$ the set of bounded
positive operators on $\mathcal{H}_{k}$. For the nonparametric class of
kernels
$\mathcal{K}=\\{{k}_{A}|k_{A}(x,y)=\braket{Ak(x,\cdot),Ak(y,\cdot)},A\in\mathcal{A}\\}$
using opt-mmd-witness leads to exactly the same witness function as using
kfda-witness.

###### Proof.

Writing inner products in the original RKHS with kernel $k$ for kernel $k_{A}$
we have the regularized $J$ criterion

$J_{A}^{\lambda}=\frac{\braket{A(\mu_{P}-\mu_{Q}),A(\mu_{P}-\mu_{Q}})}{\braket{A(\mu_{P}-\mu_{Q}),A(\Sigma+\lambda\mathop{I})AA(\mu_{P}-\mu_{Q})}^{\frac{1}{2}}}.$

We define $\delta_{A}:=A^{2}(\mu_{P}-\mu_{Q})$ and obtain

$\displaystyle
J_{A}^{\lambda}=\frac{\braket{\mu_{P}-\mu_{Q},\delta_{A}})}{\braket{\delta_{A},(\Sigma+\lambda\mathop{I})\delta_{A}}^{\frac{1}{2}}},$
(23)

which looks almost like (10). The solution to (10) is (11) which implies that
$\tilde{A}_{\lambda}=(\Sigma+\lambda\mathop{I})^{-\frac{1}{2}}$ defines the
optimal kernel

$\displaystyle\tilde{k}_{\lambda}(x,x^{\prime})$
$\displaystyle:=\braket{(\Sigma+\lambda\mathop{1})^{-\frac{1}{2}}k(x,\cdot),(\Sigma+\lambda\mathop{1})^{-\frac{1}{2}}k(x^{\prime},\cdot)}_{\mathcal{H}}$
$\displaystyle=\braket{k(x,\cdot),(\Sigma+\lambda\mathop{1})^{-1}k(x^{\prime},\cdot)}_{\mathcal{H}}.$

Based on the empirical estimates the MMD witness of the optimized kernel would
be (expressed in terms of the original kernel $k$)

$\displaystyle
h_{\tilde{k}_{\lambda}}^{\mathbb{Z}_{\text{tr}}}=(\hat{\Sigma}+\lambda\mathop{1})^{-1}(\mu_{\mathbb{X}_{\text{tr}}}-\mu_{\mathbb{Y}_{\text{tr}}})=\hat{h}_{\lambda},$
(24)

i.e., the witness of opt-mmd-witness coincides with the kfda-witness in the
original RKHS. ∎

##  Appendix B FURTHER EXPERIMENTS AND DETAILS

Figure 3: Rejection Rates for true null hypothesis (Type I error) at
$\alpha=0.05$. Left: Standard Blobs dataset (500 iterations). Middle: Blobs
dataset of Liu et al. (2020), kfda-witness is only average over 100 trials the
others over $10\times 100$, therefore kfda-witness has higher variance. Right:
Higgs dataset

Figure 4: Effect of regularization on KFDA. We consider the same setting as in
the left panel of Fig. 1 (fixed kernel and fixed regularization and $n=m=100$)
but for different regularization. Left $(\lambda=10^{3})$: For large
regularization KFDA converges to MMD. Middle ($\lambda=10^{-2}$): For a good
regularization the KFDA approaches clearly outperform the corresponding MMD
approaches. Right $(\lambda=10^{-4}$): If the regularization is too small for
a given sample size (here $n=100$) , then KFDA overfits in the training phase,
which leads to a reduction in test power.

Figure 5: Left: Draws from Blobs dataset for the instructive experiments. The
distributions are mixtures of nine Gaussians, with anisotropic covariance (but
the same covariance matrix across blobs). The covariance matrix of $Q$ is
rotated by $\theta=\pi/4$ relative to the covariance matrix of $P$. To
simulate the null hypothesis we use $\theta=0$, which corresponds to drawing
both samples from $P$. Right: Blobs dataset used for Figure 2 as suggested by
Liu et al. (2020, Figure 1). In this case, $P$ has isotropic Gaussian, the
blobs in $Q$ are anisotropic and have different covariance matrices. To
simulate the null hypothesis, we draw both samples from $P$.

This section provides supplementary information on our experiments. We provide
code upon personal request.

#### Datasets.

We used two different versions of the Blobs dataset. We show random draws for
both cases in Figure 5. For the benchmark experiments we also used the Higgs
dataset (Baldi et al., 2014), which is part of the _UCI Machine Learning
Repository_ (https://archive.ics.uci.edu/ml/datasets/HIGGS). We used a version
that is ready for Python usage provided by Liu et al. (2020)
(https://drive.google.com/open?id=1sHIIFCoHbauk6Mkb6e8a_tp1qnvuUOCc). To
ensure the comparability we follow the implementation of Liu et al. (2020) and
draw samples from the Higgs dataset _without replacement._

#### Effect of regularization of kfda-witness.

In the left panel of Figure 1, we chose a fixed regularization
$\lambda=10^{-2}$ for the KFDA methods. In Figure 4, we show the effect of
choosing a bad regularization. If the regularization is too large (left), then
KFDA coincides with MMD. On the other hand, if the regularization is too small
(right), then the effect of inaccurately estimating the covariance operator
might as well lead to a reduced test power. For good performance it is thus
important to chose a suitable regularization. This can be automated by
including a model selection procedure, such as cross-validation, in the
training stage.

#### Estimation of Rejection Rates.

For the instructive experiments in Figure 1 we estimate the rejection rates by
repeating the whole two-stage procedure 1000 times. For the benchmark
experiments we use 100 iterations of the two-stage procedure for kfda-witness.
For all the other methods in the benchmark experiments, we follow the
implementation of Liu et al. (2020) and estimate the rejection rates by
running the first stage ten times and estimating the rejection rate over 100
independent test sets for each run of the first stage. The reason for this is,
that the first stage is quite slow (training a neural network).

#### Type-I errors.

We report Type-I errors for all three different datasets in Figure 3.

##  Appendix C APPROXIMATE COMPUTATION OF THE KFDA WITNESS

Algorithm 2 Pseudocode for the FdaFalkon algorithm. Adopted for KFDA from
(Meanti et al., 2020)

1:function FdaFalkon($Z,\bm{y},k,\lambda,m$, t)

2: $Z_{m},\bm{y}_{m}\leftarrow$ RandomSubsample($(Z,\bm{y}),m$)

3: $T,A\leftarrow$ Preconditioner($Z_{m},\bm{y}_{m},\lambda$)

4: function LinOp($\bm{\beta}$)

5: $\bm{v}\leftarrow A^{-1}\bm{\beta}$

6: $\bm{c}\leftarrow k(Z_{m},Z)NN^{\top}k(Z,Z_{m})T^{-1}\bm{v}$

7: return $A^{-\top}(T^{-\top}\bm{c}+\lambda n\bm{v})$

8: $R\leftarrow A^{-\top}T^{-\top}k(Z_{m},Z)\bm{y}$

9: $\bm{\beta}\leftarrow$ ConjugateGradient($\textsc{LinOp},R,t$)

10: return $T^{-1}A^{-1}\bm{\beta}$, $Z_{m}$

13:function Preconditioner($Z_{m},\bm{y}_{m},\lambda$)

14: $K_{mm}\leftarrow k(Z_{m},Z_{m})$

15: $T\leftarrow\mathrm{chol}(K_{mm})$

16: $K_{mm}\leftarrow\frac{1}{m}TN_{m}N_{m}T^{\top}+\lambda\bm{I}$

17: $A\leftarrow\mathrm{chol}(K_{mm})$

18: return $T,A$

19:function kfdaWitness($\mathbb{Z}_{\text{tr}},k,\lambda$)

20: $Z\leftarrow$ Concatenate($\mathbb{Z}_{\text{tr}}$)

21:
$\bm{y}=[1]*\textsc{len}(\mathbb{X}_{\text{tr}})+[-1]*\textsc{len}(\mathbb{Y}_{\text{tr}})$

22: $m=\textsc{len}(Z)$ $\triangleright$ # Nyström centers

23: $\alpha,Z\leftarrow$ FdaFalkon($Z,\bm{y},k,\lambda,m$)

24: return $h_{\lambda}=\sum_{i=1}^{m}\alpha_{i}k(z_{i},\cdot)$

In this section we will use $n$ instead of $n_{\text{tr}}$ and $m$ instead of
$m_{\text{tr}}$ to keep the notation more concise. In A.3, we showed that the
exact solution for the estimate of the KFDA witness is given by

$\displaystyle\hat{h}_{\lambda}(\cdot)$
$\displaystyle=\sum_{i=1}^{n+m}\hat{\alpha}_{i}k(z_{i},\cdot),$ (25)
$\displaystyle\hat{\alpha}$ $\displaystyle=\left(\frac{KN_{c}K}{n+m}+\lambda
K\right)^{-1}K\delta.$ (26)

######  Remark 1.

The problem with computing the KFDA witness is that a naive implementation
scales cubically with the pooled sample size. In this section, we thus derive
an approach that builds on recent results, that show that one can essentially
get optimal convergence guarantees while only using $\mathcal{O}((n+m)^{3/2})$
time. Therefore two steps are needed. First, the solution is approximated with
$M=\mathcal{O}((n+m)^{\frac{1}{2}})$ Nystrom centers. Second the solution with
for the Nystrom centers is found via conjugate gradient, where a
preconditioner is computed again with only $M$ datapoints.

We take an approach similar to Rudi et al. (2017); Meanti et al. (2020). We
will thus explicitly assume that the function $h$ has the parametric form

$\displaystyle
h_{\tilde{\alpha}}(x)=\sum_{m=1}^{M}\tilde{\alpha}_{i}k(x,\tilde{z}_{i}),$
(27)

with
$M=\\{\tilde{z}_{1},\dots,\tilde{z}_{M}\\}\subseteq\\{x_{1},\dots,x_{n},y_{1},\dots,y_{m}\\}$
(we overload notation and use $M$ to denote the set itself as well as its
size). We take the notation introduced in Section 4 and constrain to the case
$c=\frac{1}{2}$. In this case we can use $N=\begin{pmatrix}P_{n}&0\\\
0&P_{m}\end{pmatrix}=\frac{N_{c}}{2}$, instead of $N_{c}$. Note that this only
affects the scaling of the solution (if we also scale $\lambda$ accordingly),
which is unimportant for WiTS tests. Using $N$ instead of $N_{c}$ has the
advantage that $N$ itself is idempotent $N=NN^{\top}$, which makes the
following easier. Nevertheless, it is straightforward to use the below
algorithm for any $c\in(0,1)$, simply by using
$N_{c}=\begin{pmatrix}\frac{1}{\sqrt{c}}P_{n}&0\\\
0&\frac{1}{\sqrt{1-c}}P_{m}\end{pmatrix}\begin{pmatrix}\frac{1}{\sqrt{c}}P_{n}&0\\\
0&\frac{1}{\sqrt{1-c}}P_{m}\end{pmatrix}$.

In the following we denote with $K_{ZM}$ the $(n+m)\times M$ matrix of entries
$k(z_{i},\tilde{z}_{j})$ and $K_{MZ}$ its transpose. We can then rewrite the
terms in our objective

$\displaystyle\braket{\hat{\mu}_{P}-\hat{\mu}_{Q},h_{\tilde{\alpha}}}=\delta^{\top}K_{ZM}\tilde{\alpha},$
(28) $\displaystyle\begin{aligned}
&\braket{h_{\tilde{\alpha}},(\hat{\Sigma}+\lambda\mathop{1})h_{\tilde{\alpha}}}\\\
&\;=\tilde{\alpha}^{\top}\left(\frac{1}{n+m}K_{MZ}NN^{\top}K_{ZM}+\lambda
K_{MM}\right)\tilde{\alpha}.\end{aligned}$ (29)

Let us define $R_{MZ}:=K_{MZ}N$. This is a $M\times(n+m)$ matrix. Note that
$N$ is the sum of the identity and two 1-sparse matrices, hence computing
$R_{MZ}$ requires only $\mathcal{O}((n+m)\cdot M)$ operations.

With our considerations from above we can write the optimal coefficients as

$\displaystyle\tilde{\alpha}^{*}=\left(R_{MZ}R_{MZ}^{\top}+(n+m)\lambda
K_{MM}\right)^{-1}K_{MZ}\delta,$ (30)
$\displaystyle\Leftrightarrow\left(R_{MZ}R_{MZ}^{\top}+(n+m)\lambda
K_{MM}\right)\tilde{\alpha}^{*}=K_{MZ}\delta$ (31)

Computing $R_{MZ}R_{MZ}^{\top}$ explicitly costs $\mathcal{O}((n+m)M^{2})$
operations and would thus dominate the cost of our previous operations.
However, (31) is now exactly in the same form as Eq. (8) in Rudi et al.
(2017). Thus from this point onwards we can build on their results to
efficiently find a solution.

The key idea of Rudi et al. (2017) is to find an efficient way to precondition
the system of linear equations in (31). In analogy, we propose to use the
following preconditioner

$\displaystyle
BB^{\top}=\left(\frac{n+m}{M}R_{MM}R_{MM}^{T}+\lambda(n+m)K_{MM}\right)^{-1},$
(32)

where $R_{MM}:=K_{MM}N_{M}$ and $N_{M}$ is defined in analogy to $N$ but only
with the $M$ Nyström centers. The preconditioner (32) thus corresponds to the
ideal preconditioner of the problem without Nyström approximation but only $M$
points to start with.

Using this preconditioner we use $t$ conjugate gradient steps to solve

$\displaystyle B^{\top}\left(R_{MZ}R_{MZ}^{\top}+(n+m)\lambda
K_{MM}\right)B\beta=B^{\top}K_{MZ}\delta.$ (33)

If $\hat{\beta}$ is the approximate solution after $t$ steps, we obtain an
approximate solution as

$\displaystyle\hat{\alpha}=B\hat{\beta}.$ (34)

The algorithm is described in Algorithm 2 and has overall complexity of
$\mathcal{O}((n_{\text{tr}}+m_{\text{tr}})Mt+M^{3})$ in time and
$\mathcal{O}(M^{2})$.