Maximum Mean Discrepancy Meets Neural Networks: The
  Radon-Kolmogorov-Smirnov Test
Maximum mean discrepancy (MMD) refers to a general class of nonparametric
two-sample tests that are based on maximizing the mean difference over samples
from one distribution $P$ versus another $Q$, over all choices of data
transformations $f$ living in some function space $\mathcal{F}$. Inspired by
recent work that connects what are known as functions of $\textit{Radon bounded
variation}$ (RBV) and neural networks (Parhi and Nowak, 2021, 2023), we study
the MMD defined by taking $\mathcal{F}$ to be the unit ball in the RBV space of
a given smoothness order $k \geq 0$. This test, which we refer to as the
$\textit{Radon-Kolmogorov-Smirnov}$ (RKS) test, can be viewed as a
generalization of the well-known and classical Kolmogorov-Smirnov (KS) test to
multiple dimensions and higher orders of smoothness. It is also intimately
connected to neural networks: we prove that the witness in the RKS test -- the
function $f$ achieving the maximum mean difference -- is always a ridge spline
of degree $k$, i.e., a single neuron in a neural network. This allows us to
leverage the power of modern deep learning toolkits to (approximately) optimize
the criterion that underlies the RKS test. We prove that the RKS test has
asymptotically full power at distinguishing any distinct pair $P \not= Q$ of
distributions, derive its asymptotic null distribution, and carry out extensive
experiments to elucidate the strengths and weakenesses of the RKS test versus
the more traditional kernel MMD test.
Seunghoon Paik, Michael Celentano, Alden Green, Ryan J. Tibshirani
# Maximum Mean Discrepancy Meets Neural Networks:

###### Abstract

Maximum mean discrepancy (MMD) refers to a general class of nonparametric two-sample tests that are based on maximizing the mean difference between samples from one distribution \(P\) versus another \(Q\), over all choices of data transformations \(f\) living in some function space \(\mathcal{F}\). Inspired by recent work that connects what are known as functions of _Radon bounded variation_ (RBV) and neural networks (Parhi and Nowak, 2021, 2023), we study the MMD defined by taking \(\mathcal{F}\) to be the unit ball in the RBV space of a given smoothness degree \(k\geq 0\). This test, which we refer to as the _Radon-Kolmogorov-Smirnov_ (RKS) test, can be viewed as a generalization of the well-known and classical Kolmogorov-Smirnov (KS) test to multiple dimensions and higher orders of smoothness. It is also intimately connected to neural networks: we prove that the witness in the RKS test--the function \(f\) achieving the maximum mean difference--is always a ridge spline of degree \(k\), i.e., a single neuron in a neural network. We can thus leverage the power of modern neural network optimization toolkits to (approximately) maximize the criterion that underlies the RKS test. We prove that the RKS test has asymptotically full power at distinguishing any distinct pair \(P\neq Q\) of distributions, derive its asymptotic null distribution, and carry out experiments to elucidate the strengths and weaknesses of the RKS test versus the more traditional kernel MMD test.

## 1 Introduction

In this paper, we consider the fundamental problem of nonparametric two-sample testing, where we observe independent samples \(x_{i}\sim P\), \(i=1,\ldots,m\) and \(y_{i}\sim Q\), \(i=1,\ldots,n\), all samples assumed to be in \(\mathbb{R}^{d}\), and we use the data to test the hypothesis

\[H_{0}:P=Q,\quad\text{versus}\quad H_{1}:P\neq Q.\]

Though there are many forms of nonparametric two-sample tests, we will focus on a class of tests based on _maximum mean discrepancy_ (MMD), which measures distance between two sets of samples by taking the maximum difference in sample averages over a function class \(\mathcal{F}\):

\[\rho(P_{m},Q_{n};\mathcal{F})=\sup_{f\in\mathcal{F}}\,|P_{m}(f)-Q_{n}(f)|. \tag{1}\]

Here \(P_{m}=\frac{1}{m}\sum_{i=1}^{m}\delta_{x_{i}}\) is the empirical distribution and \(P_{m}(f)=\frac{1}{m}\sum_{i=1}^{m}f(x_{i})\) the corresponding empirical expectation operator based on \(x_{i}\), \(i=1,\ldots,m\); and likewise for \(Q_{n}\) and \(Q_{n}(f)\).

The discrepancy in (1) is quite general and certain choices of \(\mathcal{F}\) recover well-known distances that give rise to two-sample tests. A noteworthy univariate (\(d=1\)) example is the _Kolmogorov-Smirnov_ (KS) distance (Kolmogorov, 1933; Smirnov, 1948). Though it is more typically defined in terms of cumulative distribution functions (CDFs) or ranks, the KS distance can be seen as the MMD that results when \(\mathcal{F}\) is taken to be the class of functions with at most unit _total variation_ (TV).

Some attempts have been made to define a multivariate KS distance based on multivariate CDFs (Bickel, 1969) or ranks (Friedman and Rafsky, 1979), but the resulting tests have a few limitations, such as being expensive to compute or exhibiting poor power in practice. In fact, even the univariate KS test is known to be insensitive to certain kinds of alternatives, such as those \(P\neq Q\) which differ "in the tails" (Bryson, 1974).
This led Wang et al. (2014) to recently propose a higher-order generalization of the KS test, which uses the MMD based on a class of functions whose _derivatives_ are of bounded TV.

In this work we propose and study a new class of MMDs based on a kind of multivariate total variation known as _Radon total variation_ (RTV). The definition of RTV is necessarily technical, and we delay it until Section 2. For now, we remark that if \(d=1\) then RTV reduces to the usual notion of total variation, which means that our RTV-based MMD--the central focus of this paper--directly generalizes the KS distance to multiple dimensions. We refer to this RTV-based MMD as the _Radon-Kolmogorov-Smirnov_ (RKS) distance, and the resulting two-sample test as the RKS test. We also will consider the MMD defined using a class of functions whose derivatives are of bounded RTV, which in turn generalizes the higher-order KS distance of Wang et al. (2014) to multiple dimensions. Sadhanala et al. (2019) recently showed that the higher-order KS test can be more sensitive to departures in the tails, and we will give empirical evidence that our higher-order RKS test can be similarly sensitive to tail behavior.

Our test statistic also has a simple but revealing characterization in terms of _ridge splines_. A ridge spline of integer degree \(k\geq 0\) is defined for a direction \(w\in\mathbb{S}^{d-1}\) and an offset \(b\in\mathbb{R}\) as the truncated polynomial \(f(x)=(w^{\top}x-b)_{+}^{k}\). Here we use \(\mathbb{S}^{d-1}\) to denote the unit \(\ell_{2}\) sphere in in \(\mathbb{R}^{d}\) (the set of vectors in with unit \(\ell_{2}\) norm), and we use the abbreviation \(t_{+}=\max\{t,0\}\). We prove (Theorem 3) that a \(k^{\text{th}}\) degree ridge spline _witnesses_ the \(k^{\text{th}}\) degree RKS distance: there is a \(k^{\text{th}}\) degree ridge spline \(f\) with unit Radon total variation for which \(P_{m}(f)-Q_{n}(f)\) equals the RKS distance. Thus the RKS test statistic can be equivalently written as

\[T_{d,k}=\max_{(w,b)\in\mathbb{S}^{d-1}\times[0,\infty)}\,\bigg{|}\frac{1}{m} \sum_{i=1}^{m}(w^{\top}x_{i}-b)_{+}^{k}-\frac{1}{n}\sum_{i=1}^{n}(w^{\top}y_{i }-b)_{+}^{k}\bigg{|}. \tag{2}\]

Ridge splines--and in particular the first-degree ridge spline \(f(x)=(w^{\top}x-b)_{+}\), which is also called a rectified linear (ReLU) unit--are fundamental building blocks of modern neural networks. From this perspective, the representation (2) says that the RKS distance is achieved by a single neuron in a two-layer neural network. In fact, the connections between Radon total variation and neural networks run deeper. The RKS distance can be viewed as finding the maximum discrepancy between sample means over all two-layer neural networks (of arbitrarily large width) subject to a weight decay constraint (Section 2). Although we focus throughout on two-layer neural networks, the MMD defined over deeper networks can be viewed the RKS distance in the representation learned by the hidden layers.

The representation (2) also hints at several interesting statistical properties. Let \(X\sim P_{m}\) and \(Y\sim Q_{n}\) be random variables drawn from the empirical distributions of \(\{x_{i}\}_{i=1}^{m}\) and \(\{y_{i}\}_{i=1}^{n}\), respectively. The \(0^{\text{th}}\) degree RKS distance finds the direction \(w\in\mathbb{S}^{d-1}\) along which the usual Kolmogorov-Smirnov distance--maximum gap in CDFs--between the univariate random variables \(w^{\top}X\) and \(w^{\top}Y\) is as large as possible. For \(k\geq 1\), the \(k^{\text{th}}\) degree statistic finds the direction \(w\in\mathbb{S}^{d-1}\) along which the higher-order KS distance--maximum gap in truncated \(k^{\text{th}}\) moments--between \(w^{\top}X\) and \(w^{\top}Y\) is maximized.

Thus, speaking informally, we expect the RKS test to be particularly sensitive to the case when the laws of \(w^{\top}X\) and \(w^{\top}Y\) differ primarily in just a few directions \(w\) (because we are taking a maximum over such directions in (2)). In other words, we expect the RKS test to be sensitive to _anisotropic_ differences between \(P\) and \(Q\). On the other hand, taking a larger value of \(k\) results in higher-order truncated sample moments, which can be much more sensitive to small differences in the tails. Figure 1 gives an illustrative example. We
consider two normal distributions with equal means and covariances differing in just one direction (along the x-axis). The figure displays the witnesses (ridge splines) for the RKS test when \(k=0,1,2\). For larger \(k\), the witness is more aligned with the direction in which \(P\) and \(Q\) vary, and also places more weight on the tails.

Summary of contributions.Our main contributions are as follows.

* We propose a new class of two-sample tests using a \(k^{\text{th}}\) degree Radon total variation MMD. We prove a representer theorem (Theorem 3) which shows that the MMD is witnessed by a ridge spline (establishing the representation in (2)). This connects our proposal to two-sample testing based on neural networks, and helps with optimization.
* Under the null \(P=Q\), we derive the asymptotic distribution of the test statistic (Theorem 5).
* Under the alternative \(P\neq Q\), we give upper bounds on the rate at which the test statistic concentrates around the population-level MMD. Using the fact that the population-level MMD is a metric, we then show that the RKS test is _consistent_ against any fixed \(P\neq Q\): it has asymptotic error (sum of type I and type II errors) tending to zero (Theorem 6).
* We complement our theory with numerical experiments to explore the operating characteristics of the RKS test compared to other popular nonparametric two-sample tests.

Related work.Recently, there has been a lot of interest in multivariate nonparametric two-sample tests using kernel MMDs (Gretton et al., 2012), and energy distances (Baringhaus and Franz, 2004; Szekely and Rizzo, 2005). The latter are in many cases equivalent to a kernel MMD (Sejdinovic et al., 2013). It is worth being clear that the Radon bounded variation space is _not_ an RKHS, and as such, our test cannot be seen as a special case of a kernel MMD. There are also many other kinds of multivariate nonparametric two-sample tests, such as graph-based tests using \(k\)NN graphs (Schilling, 1986; Henze, 1988) or spanning-trees (Friedman and Rafsky, 1979). The RKS test statistic (2) avoids the need to construct a graph over the samples, though we note that one could also use a graph to measure (approximate) multivariate total variation of a different variety (sometimes called _measure-theoretic_ TV, which is not the same as Radon TV), and one could then form a related two-sample test, accordingly.

Some literature refers to maximum mean discrepancy (MMD) with respect to an arbitrary function class as an _integral probability metric_ (IPM) (Muller, 1997). While IPMs look at differences in \(\mathrm{d}P-\mathrm{d}Q\), closely related are \(\phi\)-divergences, which look at differences in \(\mathrm{d}P/\mathrm{d}Q\)(Sriperumbudur et al., 2009).

In machine learning, two-sample tests play a fundamental role in generative adversarial networks (GANs) (Goodfellow et al., 2014), which have proven to be an effective way to generate new (synthetic) draws that adhere to the distribution of a given high-dimensional set of samples. In short, GANs are trained to produce synthetic data which a two-sample test cannot distinguish from "real" samples. However, optimizing a GAN is notoriously unstable. It has been suggested that the stability of GANs can be improved if the underlying two-sample test is based on an MMD, with constraints on, e.g., the Lipschitz constant (Arjovsky et al., 2017), RKHS norm (Li et al., 2017), or Sobolev norm (Mroueh et al., 2018) of the witness. Unlike Radon TV, these function classes do not have an intrinsic connection to neural networks, which in practice are most commonly used to train the discriminator in a GAN.

At a technical level, our paper falls into a line of work exploring neural networks from the functional analytic perspective. The equivalences between weight decay of parameters in a two-layer neural network and first-degree (\(k=1\)) Radon total variation were discovered in Savarese et al. (2019); Ongie et al. (2020), and further formalized and extended by Parhi and Nowak (2021) to the higher-degree case (\(k>1\)). We make use of these equivalences to prove results on the Radon TV MMD. Lastly, we note that the sensitivity of neural networks to variation in only a few directions was observed even earlier, in Bach (2017).

## 2 The Radon-Kolmogorov-Smirnov test

For independent samples \(x_{i}\sim P\), \(i=1,\ldots,m\) and \(y_{i}\sim Q\), \(i=1,\ldots,n\), and a given integer \(k\geq 0\), the RKS test statistic \(T_{d,k}\) is defined as in (2). This statistic searches for a marginal (defined by projection onto the
direction \(w\)) with the largest discrepancy in a truncated \(k^{\text{th}}\) moment. As with any two-sample test, we can calibrate the rejection threshold for the RKS test statistic using a permutation null. Given any user-chosen level \(\alpha\in[0,1]\), we reject \(H_{0}:P=Q\) if

\[\frac{\sum_{j=1}^{B}\mathds{1}\big{\{}T_{d,k}(z_{\pi_{j}})\geq T_{d,k}(z)\big{\}} +1}{B+1}\leq\alpha. \tag{3}\]

Here we use \(z=(x_{1},\ldots,x_{m},y_{1},\ldots,y_{n})\in\mathbb{R}^{m+n}\) to denote the original data sequence, and we use \(T_{d,k}(z)\) to emphasize that the statistic in (2) is computed on \(z\). Further, each \(\pi_{j}\) is a permutation of \(\{1,\ldots,n\}\), drawn uniformly at random, and we use \(T_{d,k}(z_{\pi_{j}})\) to denote the test statistic computed on the permuted sequence \(z_{\pi_{j}}=(z_{\pi(1)},\ldots,z_{\pi(m+n)})\). Standard results on permutation tests imply that the test defined by the rejection threshold (3) has type I error control at the level \(\alpha\).

Finding the value of \(w\) and \(b\) achieving the supremum in (2) is a nonconvex problem. However, it bears a connection to neural network optimization, which is of course among the most familiar and widely-studied nonconvex problems in machine learning. We show in Appendix I that problem (2) can be equivalently cast as finding the maximum discrepancy between sample means over all two-layer neural networks (of arbitrarily large width) under a weight decay-like constraint. In particular, if we denote by \(f_{(a_{j},w_{j},b_{j})_{j=1}^{N}}\) the \(N\)-neuron two-layer neural network defined by

\[f_{(a_{j},w_{j},b_{j})}(x)=a_{j}(w_{j}^{\top}x-b_{j})_{+}^{k}\quad\text{and} \quad f_{(a_{j},w_{j},b_{j})_{j=1}^{N}}(x)=\sum_{j=1}^{N}\,f_{(a_{j},w_{j},b_{ j})}(x), \tag{4}\]

then the RKS test statistic is equivalently

\[T_{d,k} =\max_{f=f_{(a_{j},w_{j},b_{j})_{j=1}^{N}}} P_{m}(f)-Q_{n}(f)\] (5) subject to \[\sum_{j=1}^{N}\,|a_{j}|\|w_{j}\|_{2}^{k}\leq 1,\] \[(a_{j},w_{j},b_{j})\in\mathbb{R}\times\mathbb{R}^{d}\times[0, \infty),\;j=1,\ldots,N.\]

This equivalence holds for any number of neurons \(N\geq 1\), which is why we say that the network can be of "arbitrarily large width".

In Section 4, we use a Lagrangian form of this constrained optimization, allowing us to leverage existing deep learning toolkits to compute the RKS distance. As with neural network optimization in general, we cannot prove that any practical first-order scheme--such as gradient descent or its variants--is able to find the global optimum of (5), or its Lagrange form. However, we find empirically that the test statistic resulting from such first-order schemes has favorable behavior in practice, such as stability (with respect to the choice of learning rate) and strong power (with respect to certain classes of alternatives). Finally, we note that in practice we use the permutation threshold (3) to calibrate the statistic computed by first-order optimization, i.e., the same first-order optimization scheme is applied to the original data sequence and to every permuted sequence. In effect, we can view this as _redefining_ the RKS test to be the output of such a first-order scheme, and as the permutation null controls type I error for any measurable function of data, it also controls type I error for a test statistic computed this way.

### Functions of Radon bounded variation

The RKS distance, like the univariate KS distance, can be viewed as the maximum mean discrepancy with respect to a function class defined by a certain notion of smoothness. Whereas the univariate KS distance is the MMD over functions with total variation is bounded by 1, the RKS distance is the MMD over functions with Radon total variation bounded by 1, subject to a boundary condition. This will be shown a bit later; first, we must cover preliminaries needed to understand Radon total variation.

Heuristically, the \(k^{\text{th}}\) degree Radon total variation of \(f:\mathbb{R}^{d}\to\mathbb{R}\) measures the average smoothness of the "marginals" of \(f\). To be more concrete, consider a function \(f:\mathbb{R}^{d}\to\mathbb{R}\) which is infinitely differentiable, and whose derivatives of any order decay super-polynomially: \(f\in C^{\infty}(\mathbb{R}^{d})\) and \(\sup_{x\in\mathbb{R}^{d}}|x^{\alpha}\partial^{\beta}f(x)|<\infty\). Here
\(\alpha=(\alpha_{1},\ldots,\alpha_{d})\) and \(\beta=(\beta_{1},\ldots,\beta_{d})\) are multi-indices (with nonnegative integer elements), and we use the abbreviations

\[x^{\alpha}=\prod_{i=1}^{d}x_{i}^{\alpha_{i}},\quad\text{and}\quad\partial^{\beta }f=\partial_{x_{1}}^{\beta_{1}}\cdots\partial_{x_{d}}^{\beta_{d}}f.\]

The set of such functions is called the _Schwartz class_, and denoted by \(\mathcal{S}(\mathbb{R}^{d})\). The _Radon transform_ of a function \(f\in\mathcal{S}(\mathbb{R}^{d})\) is itself another function \(\mathcal{R}\{f\}:\mathbb{S}^{d-1}\times\mathbb{R}\to\mathbb{R}\) defined as

\[\mathcal{R}\{f\}(w,b)=\int_{w^{\top}x=b}f(x)\,\mathrm{d}x,\]

where the integral is with respect to the Lebesgue surface measure on the hyperplane \(\{w^{\top}x=b\}\). As a function of \(b\), the Radon transform can be viewed as the marginal of the function \(f\) in the direction \(w\). The \(k^{\text{th}}\) degree Radon total variation of \(f\) is then defined to be (Parhi and Nowak, 2021, 2023):1

Footnote 1: A remark on notation and nomenclature: we use \(\mathrm{RBV}^{k}\) to refer to the space which Parhi and Nowak (2021); Parhi (2022); Parhi and Nowak (2023) instead call \(\mathrm{RBV}^{k+1}\). That is, we index based on the _degree_ of the underlying ridge spline, whereas the aforementioned papers index based on the _order_, traditionally defined as the degree plus one, \(m=k+1\).

\[\|f\|_{\mathrm{RTV}^{k}}=c_{d}\int_{\mathbb{S}^{d-1}\times\mathbb{R}}\big{|} \partial_{b}^{k+1}\Lambda^{d-1}\mathcal{R}\{f\}(w,b)\big{|}\,\mathrm{d}\sigma (w,b),\]

where \(\sigma\) is the Hausdorff measure on \(\mathbb{S}^{d-1}\times\mathbb{R}\), \(c_{d}=1/(2(2\pi)^{d-1})\), and \(\Lambda^{d-1}\) is the "ramp filter" defined by \(\Lambda^{d-1}=(-\partial_{b}^{2})^{\frac{d-1}{2}}\) with fractional derivatives interpreted in terms of Riesz potentials (see Parhi and Nowak (2021, 2023) for details). Equivalently, if we define an operator by \(R_{k}\{f\}(w,b)=c_{d}\partial_{b}^{k+1}(-\partial_{b}^{2})^{\frac{d-1}{2}} \mathcal{R}\{f\}(w,b)\), then we can write \(k^{\text{th}}\) degree Radon TV of \(f\) simply as \(\|f\|_{\mathrm{RTV}^{k}}=\|R_{k}\{f\}\|_{L^{1}}\). The mapping \(f\mapsto\|f\|_{\mathrm{RTV}^{k}}\) is a seminorm, which we call the \(\mathrm{RTV}^{k}\) seminorm.

Parhi and Nowak (2021, 2023) and Parhi (2022) extend the \(\mathrm{RTV}^{k}\) seminorm beyond Schwartz functions using functional analytic tools based on duality. The space of functions with bounded \(\mathrm{RTV}^{k}\) seminorm is referred to as the \(\mathrm{RBV}^{k}\) space (where \(\mathrm{RBV}\) stands for "Radon bounded variation"). After this extension, one can still think of \(\mathrm{RTV}^{k}\) seminorm as measuring average higher-order smoothness of the marginals of \(f\), but it does not require that all derivatives and integrals exist in a classical sense. We refer the reader especially to Parhi (2022) for the complete functional analytic definition of the \(\mathrm{RBV}^{k}\) space. For our purposes, it is only important that \(\mathrm{RBV}^{k}\) functions enjoy the representation established by the following theorem. Note that in order to define the \(\mathrm{RBV}^{k}\) space, Parhi and Nowak (2021); Parhi (2022) must view it as a space of equivalence classes of functions under the equivalence relation \(f\sim g\) if \(f(x)=g(x)\) for Lebesgue almost every \(x\in\mathbb{R}^{d}\). We denote elements of \(\mathrm{RBV}^{k}\) as \(\mathsf{f}\), and the functions they contain as \(f\in\mathsf{f}\).

**Proposition 1** (Adaptation of Theorem 22 in Parhi and Nowak (2021); Theorem 3.8 in Parhi (2022)).: _Fix any \(k\geq 0\). For each \(\mathsf{f}\in\mathrm{RBV}^{k}\), there exists a representative \(f\in\mathsf{f}\) which satisfies for all \(x\in\mathbb{R}^{d}\)_

\[f(x)=\int_{\mathbb{S}^{d-1}\times[0,\infty)}(w^{\mathsf{T}}x-b)_{+}^{k}\, \mathrm{d}\mu(w,b)+q(x),\]

_for some finite signed Borel measure \(\mu\) on \(\mathbb{S}^{d-1}\times[0,\infty)\) satisfying \(\|\mu\|_{\mathrm{TV}}=\|\mathsf{f}\|_{\mathrm{RTV}^{k}}\), and some polynomial \(q\) of degree at most \(k\), where we use \(\|\cdot\|_{\mathrm{TV}}\) for the total variation norm in the sense of measures._

Proposition 1 is the key result connecting the \(\mathrm{RBV}^{k}\) space to two-layer neural networks. It states that any \(\mathrm{RBV}^{k}\) functions is equal almost everywhere to the sum of a (possibly infinite-width) two-layer neural network and a polynomial. Furthermore, the \(\ell_{1}\) norm of coefficients in the last layer of this neural network (formally, the total variation norm \(\|\mu\|_{\mathrm{TV}}\) of the measure \(\mu\) defined over the coefficients) equals the \(\mathrm{RBV}^{k}\) seminorm of the function in question. As stated, Proposition 1 is a slight refinement of results in Parhi and Nowak (2021); Parhi (2022), which we show in Appendix A.

Parhi and Nowak (2021) consider the nonparametric regression problem defined by minimizing, over all functions \(f\), the squared loss incurred by \(f\) with respect to a given finite data set plus a penalty on \(\|f\|_{\mathrm{RTV}^{k}}\). They show this is solved by the sum of a _finite-width_ two-layer neural network and a polynomial, and hence provide a representation theorem for \(\mathrm{RTV}^{k}\)-penalized nonparametric regression analogous to that for kernel
ridge regression (Wahba, 1990). In a coming subsection, we will establish a similar connection in the context of two-sample testing: the MMD under an \(\mathrm{RTV}^{k}\) constraint is realized by a single neuron.

Before moving on, it is worth providing some intuition for the results just discussed. First, the operator \(R_{k}\) should be understood as transforming a function \(f\) into the Radon domain: namely, it specifies \(f\) via the higher-order derivatives of its marginals. Then, the \(\mathrm{RTV}^{k}\) acts as the \(L^{1}\) norm in the Radon domain. Lastly, the reason two-layer neural networks and ridge splines emerge as the solutions to nonparametric regression and maximum mean discrepancy problems under an \(\mathrm{RTV}^{k}\) constraint is that ridge splines are _sparse_ in the Radon domain. In particular, for \(f(x)=(w^{\top}x-b)_{+}^{k}\), we have (Parhi and Nowak, 2021):

\[R_{k}\{f\}(w,b)=\frac{1}{2}\big{(}\delta_{(w,b)}+(-1)^{k+1}\delta_{(-w,-b)} \big{)},\]

where \(\delta_{(\pm w,\pm b)}\) denote point masses at \((\pm w,\pm b)\). Thus, \(\mathrm{RTV}^{k}\)-penalized optimization problems are solved by two-layer neural networks due to the tendency of the \(L^{1}\) norm to encourage sparse solutions.

### Pointwise evaluation of RBV functions

Recall, the \(\mathrm{RBV}^{k}\) space contains equivalence classes of functions whose members differ on sets of Lebesgue measure zero. Therefore, at face value, point evaluation is meaningless for elements of the \(\mathrm{RBV}^{k}\) space. This poses a problem for defining an MMD with respect to RBV functions, since the MMD depends on empirical averages, which require function evaluations over samples. Fortunately, we show that there is a natural way to resolve point evaluation over the \(\mathrm{RBV}^{k}\) space: each equivalence class in \(\mathrm{RBV}^{k}\) has, possibly subject to a mild boundary condition, a unique representative which satisfies a suitable notion of continuity. We can then define point evaluation with respect to this representative.

For the case \(k=0\), we need to introduce a new notion of continuity, which we call _radial cone continuity_. For \(k\geq 1\), we rely on the standard notion of continuity. Radial cone continuity is defined as follows.

**Definition 1**.: _A function \(f\) is said to be radially cone continuous if it is continuous at the origin, and for any \(x\in\mathbb{R}^{d}\setminus\{0\}\), we have \(f(x+\epsilon v)\to f(x)\) whenever (i) \(\epsilon\to 0\) and (ii) \(v\in\mathbb{S}^{d-1}\) and \(v\to x/\|x\|_{2}\)._

Radial cone continuity can be viewed as a generalization of right-continuity to higher dimensions. Indeed a consequence of radial cone continuity is that for every \(v\in\mathbb{R}^{d}\), the restriction of the function \(f\) to the ray emanating from the origin in the direction \(v\) is right-continuous (i.e., \(t\mapsto f(tv)\), \(t>0\) is right-continuous). The concept of radial cone continuity, however, is stronger than right-continuity along rays--it additionally requires continuity along any path which is tangent to and approaches \(x\) from the same direction as the ray that points from the origin to \(x\). The following theorem shows that (subject to a boundary condition for the case \(k=0\)) every equivalence class \(\mathsf{f}\in\mathrm{RBV}^{k}\) contains a unique continuous representative when \(k\geq 1\), and radially cone continuous representative when \(k=0\).

**Theorem 1**.: _For \(k=0\), if \(\mathsf{f}\in\mathrm{RBV}^{0}\) contains a function that is continuous at the origin, then it contains a unique representative that is radially cone continuous. Moreover, for \(k\geq 1\), each \(\mathsf{f}\in\mathrm{RBV}^{k}\) contains a unique representative that is continuous and this representative is in fact \((k-1)\)-times continuously differentiable._

We prove Theorem 1 in Appendix B. The theorem motivates the definition of the function spaces:

\[\mathrm{RBV}^{0}_{c} =\Big{\{}f\,:\,f\in\mathsf{f}\text{ for some }\mathsf{f}\in \mathrm{RBV}^{0},\,f\text{ radially cone continuous}\Big{\}}, \tag{6}\] \[\mathrm{RBV}^{k}_{c} =\Big{\{}f\,:\,f\in\mathsf{f}\text{ for some }\mathsf{f}\in \mathrm{RBV}^{k},\,f\text{ continuous}\Big{\}},\quad k\geq 1.\]

Unlike \(\mathrm{RBV}^{k}\), the space \(\mathrm{RBV}^{k}_{c}\) contains functions rather than equivalence classes of functions, so that point evaluation is well-defined. The Radon total variation of a member \(f\in\mathsf{f}\) is defined in the natural way as the Radon total variation of the function class in \(\mathrm{RBV}^{k}\) to which it belongs.

### The RKS distance is an MMD

We are ready to show that the RKS distance is equivalent to the MMD over functions \(f\) with \(\|f\|_{\mathrm{RTV}^{k}}\leq 1\), subject to a boundary condition. In this sense, it can be viewed as a generalization of the classical univariate
KS distance, which, as we have mentioned, is the MMD over the space of functions with total variation at most 1. Precisely, the RKS distance is the MMD over the function space

\[\mathcal{F}_{k}=\Big{\{}f\in\mathrm{RBV}^{k}_{c}\,:\,\|f\|_{\mathrm{RTV}^{k}} \leq 1,\;\partial^{\alpha}f(0)\text{ exists and equals }0\text{ for all }|\alpha|\leq k\Big{\}}. \tag{7}\]

The boundary condition \(\partial^{\alpha}f(0)=0\) for all multi-indices \(\alpha=(\alpha_{1},\dots,\alpha_{d})\) with \(|\alpha|=\sum_{j=1}^{d}\alpha_{j}\leq k\) is needed to guarantee that the MMD is finite. We can interpret this condition intuitively as requiring the polynomial \(q\) to be zero in the representation in Proposition 1. Otherwise, \(q\) would be able to grow without bound under the constraint \(\|f\|_{\mathrm{RTV}^{k}}\leq 1\), which we could then use to drive the MMD criterion to \(\infty\), provided that \(P\) and \(Q\) had any moment differences (in their first \(k\) moments).

The following representation theorem proves that the RKS distance is the MMD over (7).

**Theorem 2**.: _Fix any \(k\geq 0\). For any \(f\in\mathcal{F}_{k}\), there exists a finite, signed Borel measure \(\mu\) on \(\mathbb{S}^{d-1}\times[0,\infty)\) such that for all \(x\in\mathbb{R}^{d}\),_

\[f(x)=\int_{\mathbb{S}^{d-1}\times[0,\infty)}(w^{\mathsf{T}}x-b)_{+}^{k}\, \mathrm{d}\mu(w,b),\]

_and \(\|\mu\|_{\mathrm{TV}}=\|f\|_{\mathrm{RTV}^{k}}\). When \(k=0\), the measure \(\mu\) may be taken to be supported on \(\mathbb{S}^{d-1}\times(0,\infty)\)._

We prove Theorem 2 in Appendix C. Note that the representation provided by Theorem 2 is like that in Proposition 1 except that the integral is restricted to \(b\geq 0\), and the polynomial term is \(q=0\). The latter is a consequence of the boundary condition imposed in the definition of \(\mathcal{F}_{k}\).

Finally, the next theorem establishes the equivalence between \(T_{d,k}\) as defined in (2) and the MMD over \(\mathcal{F}_{k}\). Its proof is in Appendix D.

**Theorem 3**.: _Fix any \(k\geq 0\). Define_

\[\mathcal{G}_{k}=\Big{\{}(w^{\mathsf{T}}\cdot-b)_{+}^{k}\,:\,(w,b)\in\mathbb{S }^{d-1}\times[0,\infty)\Big{\}},\]

_where by convention we take \(t_{+}^{0}=1\{t\geq 0\}\). Then for any \(P\) and \(Q\) with finite \(k^{\text{th}}\) moments, we have_

\[\rho(P,Q;\mathcal{F}_{k})=\rho(P,Q;\mathcal{G}_{k}).\]

_In particular, this implies that for the empirical distributions \(P_{m}\) and \(Q_{n}\) over any sets of samples \(\{x_{i}\}_{i=1}^{m}\) and \(\{y_{i}\}_{i=1}^{n}\), we have \(\rho(P_{m},Q_{n};\mathcal{F}_{k})=T_{d,k}\), where the latter is as defined in (2)._

### The RKS distance identifies the null

An important property of the RKS distance for the purposes of two-sample testing is its ability to distinguish any two distinct distributions. In particular, we have the following result, whose proof is in Appendix E.

**Theorem 4**.: _For any \(P,Q\) with finite \(k^{\text{th}}\) moments, \(\rho(P,Q;\mathcal{F}_{k})=0\) if and only if \(P=Q\), and is positive otherwise._

Theorem 4 states that the RKS distance, in the population, perfectly distinguishes the null \(H_{0}:P=Q\) from the alternative \(H_{1}:P\neq Q\). Thus, the function space \(\mathcal{F}_{k}\) is rich enough to make the MMD a _metric_ at the population level. As an interesting contrast, we remark that this is not true of the kernel MMD with a polynomial kernel of degree \(k\), as studied in Gretton et al. (2012). Indeed, if \(P\) and \(Q\) have the same moments up to order \(k\) (where "moments" is to be interpreted in the appropriate multivariate sense) then the kernel MMD in the population between \(P\) and \(Q\) is 0. In other words, we can see that the use of _truncated_ moments, as given by averages of ridge splines, is the key to the discrepancy being a metric.

## 3 Asymptotics

Given any probability distribution \(P\) on \(\mathbb{R}^{d}\) with finite \((2k+\Delta)\)-moments for some \(\Delta>0\), let us define the corresponding Gaussian process \(\mathbb{G}_{P}=\{G_{w,b}:(w,b)\in\mathbb{S}^{d-1}\times[0,\infty)\}\) by

\[G_{w,b}\sim\mathcal{N}\Big{(}0,\,\mathbb{E}_{x\sim P}\big{[}(w^{\top}x-b)_{+}^ {2k}\big{]}\Big{)},\quad\mathrm{Cov}(G_{w,b},G_{w^{\prime},b^{\prime}})= \mathbb{E}_{x\sim P}\Big{[}(w^{\top}x-b)_{+}^{k}(w^{\prime\top}x-b^{\prime})_{+ }^{k}\Big{]}. \tag{8}\]
The importance of this Gaussian process is explained in the next result, which shows that its supremum provides the asymptotic null distribution of the RKS test statistic.

**Theorem 5**.: _Fix any \(k\geq 0\). Assume that \(P\) has finite \((2k+\Delta)\)-moments for some \(\Delta>0\). If \(k=0\), then additionally assume that \(P\) has bounded density with respect to Lebesgue measure. When \(P=Q\), we have_

\[\sqrt{\frac{mn}{m+n}}\,T_{d,k}\,\stackrel{{ d}}{{\to}}\,\sup_{(w,b)\in\mathbb{S}^{d-1}\times[0,\infty)}\,|G_{w,b}|,\quad\text{as }m,n\to\infty.\]

Moreover, with a proper rejection threshold, the RKS test can have asymptotically zero type I error and full power against any fixed \(P\neq Q\).

**Theorem 6**.: _Fix any \(k\geq 0\). Assume \(P,Q\) have finite \((2k+\Delta)\)-moments for some \(\Delta>0\). If \(k=0\), then additionally assume that \(P,Q\) have bounded densities with respect to Lebesgue measure. Let \(p=m+n\) and consider any sequence \(t_{p}\) such that \(t_{p}\to 0\) and \(t_{p}\sqrt{p}\to\infty\) as \(m,n\to\infty\). Then the test which rejects when \(T_{d,k}>t_{p}\) rejects with asymptotic probability 0 if \(P=Q\) and asymptotic probability 1 if \(P\neq Q\)._

We prove Theorems 5 and 6 in Appendices F and G. Roughly, Theorem 5 follows from general uniform central limit theorem results in Dudley (2014), after controlling the bracketing number of the function class \(\mathcal{G}_{k}\). Theorem 6 follows from tail bounds that we establish for the concentration of the RKS distance around its population value, combined with the fact that the population distance is a metric (Theorem 4).

## 4 Experiments

We conduct experiments to compare the power of the RKS test with other multivariate nonparametric tests in multiple settings (different pairs of \(P\) and \(Q\)), both to investigate the strengths and weaknesses of the RKS test, and to investigate its behavior as we vary the smoothness degree \(k\).

### Computation of the RKS distance

First, we discuss computation of the RKS distance, as initially defined in (2). This a difficult optimization problem because of the nonconvexity of the criterion and the unit \(\ell_{2}\) norm constraint on \(w\). Though we will not be able to circumvent the challenge posed by nonconvexity entirely, we can ameliorate it by moving to the overparametrized formulation in (5): this now casts the same computation as an optimization over the \(N\)-neuron neural network \(f\) in (4), subject to a constraint on \(\sum_{j=1}^{N}|a_{j}|\|w_{j}\|_{2}^{k}\) (which is sometimes called the "path norm" of \(f\)). We call this problem "overparametrized" since any individual pair \((w_{j},b_{j})\) would be sufficient to obtain the global optimum in (2), and yet we allow the optimization to search over \(N\) such pairs \((w_{j},b_{j})\), \(j=1,\ldots,N\). Perhaps unsurprisingly, we find that overparametrization (larger \(N\)) generally makes optimization easier--more stable with respect to the initialization and choice of learning rate, discussed in more detail in Appendix K.

Problem (5) is still more difficult than we would like it to be, i.e., not within scope for typical first-order optimization techniques, due to the path norm constraint \(\sum_{j=1}^{N}|a_{j}|\|w_{j}\|_{2}^{k}\leq 1\). Fortunately, it is easy to see that any bound here suffices to compute the RKS distance: we can instead use \(\sum_{j=1}^{N}|a_{j}|\|w_{j}\|_{2}^{k}\leq C\), for any \(C>0\), rescale the criterion \(P_{m}(f)-Q_{n}(f)\) by \(1/C\), and then the value of the criterion at its optimum (the MMD value) is unchanged. This motivates us to instead solve a Lagrangian version of (5):

\[\min_{\begin{subarray}{c}f=f_{(a_{j},w_{j},b_{j})_{j=1}^{N}}\\ (a_{j},w_{j},b_{j})\in\mathbb{R}\times\mathbb{R}^{d}\times[0,\infty)\end{subarray}} -\log\Big{(}|P_{m}(f)-Q_{n}(f)|\Big{)}+\lambda\sum_{i=1}^{N}|a_{i}|\|w_{i}\|_ {2}^{k}, \tag{9}\]

which is the focus in all of our experiments henceforth. The use of a log transform for the mean difference in (9) is primarily used because we find that it leads to a greater stability with respect to the choice of learning rate (Appendix J). Importantly, the choice of Lagrangian parameter \(\lambda>0\) in (9) can be arbitrary, by the same rescaling argument as explained above: given any (local) optimizer \(f^{*}\) in (9), we can simply return the mean difference over the rescaled witness \(f^{*}/\|f^{*}\|_{\mathrm{RTV}^{k}}\), where \(\|f^{*}\|_{\mathrm{RTV}^{k}}=\sum_{i=1}^{N}|a_{i}^{*}|\|w_{i}^{*}\|_{2}^{k}\).
For \(k\geq 1\), we apply the Adam optimizer (a variant of gradient descent), as implemented in PyTorch, to (9). In our experiments that follow, we use \(N=10\) neurons and a few hundred iterations of Adam. On the other hand, for \(k=0\), such a first-order scheme is not applicable due to the fact that the gradient of the \(0^{\text{th}}\) degree ridge spline \((w^{\mathsf{T}}x-b)_{+}^{0}=1\{w^{\mathsf{T}}x\geq b\}\) (with respect to \(w,b\)) is almost everywhere zero. Therefore, as a surrogate, we instead approximate the optimum \(w^{*},b^{*}\) in (2) using logistic regression. Appendix I provides further details on the optimization setup, and Appendices J and K present related sensitivity analyses.

### Experimental setup

We compare the RKS tests of various degrees across several experimental settings, both to one another and to other all-purpose nonparametric tests: the energy distance test (Szekely and Rizzo, 2005) as well as kernel MMD (Gretton et al., 2012) with a Gaussian kernel.

We fix the sample sizes to \(m=n=512\) throughout, and consider four choices of dimension: \(d=2,4,8,16\). For each dimension \(d\), we consider five settings for \(P,Q\), which are described in Table 1. In each setting, the parameter \(v\) controls the discrepancy between \(P\) and \(Q\), but its precise meaning depends on the setting. The settings were broadly chosen in order to study the operating characteristics of the RKS test when differences between \(P\) and \(Q\) occur in one direction (settings 1-4), and in all directions (setting 5). Among the settings in which the differences occur in one direction, we also investigate different varieties (settings 1 and 2: mean shift under different geometries, setting 3: tail difference, setting 4: variance difference). Finally, we note that because the RKS test is _rotationally invariant_,2 the fact that the chosen differences in Table 1 are axis-aligned is just a matter of convenience, and the results would not change if these differences instead occurred along arbitrary directions in \(\mathbb{R}^{d}\).

Footnote 2: Meaning, \(\rho(P_{m},Q_{n};\mathcal{F}_{k})\) is invariant to rotations of the underlying samples \(\{x_{i}\}_{i=1}^{m}\) and \(\{y_{i}\}_{i=1}^{n}\) by any arbitrary orthogonal transformation \(U\in\mathbb{R}^{d\times d}\). This is true because the RTV\({}^{k}\) seminorm is itself invariant to rotations of the underlying domain, as can be seen directly from the representation in Proposition 1.

For the RKS tests, we examine smoothness degrees \(k=0,1,2,3\); for kernel MMD, we set the Gaussian kernel bandwidth according to a standard heuristic based on the median of all pairwise squared distances between the input samples (Caputo et al., 2002). (The energy distance test has no tuning parameters.) For each setting, we compute these test statistics under the null where each \(x_{i}\) and \(y_{i}\) are sampled i.i.d. from the mixture \(\frac{m}{m+n}P+\frac{n}{m+n}Q\), and under the alternative where \(x_{i}\) are i.i.d. from \(P\) and \(y_{i}\) from \(Q\). We then repeat this 100 times (draws of samples, and computation of test statistics), and trace out ROC curves--true positive versus false positive rates--as we vary the rejection threshold for each test. Python code to replicate our experimental results is available at [https://github.com/100shpaik/](https://github.com/100shpaik/).

In Appendix L, we also compare the RKS test to kernel MMD with a polynomial kernel. Importantly, the latter is not actually a nonparametric test, for any finite polynomial order, since the population MMD

\begin{table}
\begin{tabular}{|c||c|c|c|} \hline Setting & \(P\) & \(Q\) & \(v\) \\ \hline Pancake mean shift & one axis \(\sim\mathcal{N}(0,1)\) & one axis \(\sim\mathcal{N}(v,1)\) & \\ (“pancake-shift”) & the others \(\sim\mathcal{N}_{d-1}(0,16I)\) & the others \(\sim\mathcal{N}_{d-1}(0,16I)\) & 0.3 \\ \hline Ball mean shift & \multirow{2}{*}{\(\mathcal{N}_{d}(0,I)\)} & one axis \(\sim\mathcal{N}(v,1)\) & \\ (“ball-shift”) & & the others \(\sim\mathcal{N}_{d-1}(0,I)\) & 0.2 \\ \hline \(t\) coordinate & \multirow{2}{*}{\(\mathcal{N}_{d}(0,I)\)} & one axis \(\sim t(v)\) & \\ (“t-coord”) & & the others \(\sim\mathcal{N}_{d-1}(0,I)\) & 3 \\ \hline Variance inflation & \multirow{2}{*}{\(\mathcal{N}_{d}(0,I)\)} & one axis \(\sim\mathcal{N}(0,v)\) & \\ in one direction & & the others \(\sim\mathcal{N}_{d-1}(0,I)\) & 1.4 \\ (“var-one”) & & \\ \hline Variance inflation & \multirow{2}{*}{\(\mathcal{N}_{d}(0,I)\)} & \multirow{2}{*}{\(\mathcal{N}_{d}(0,vI)\)} & \multirow{2}{*}{1.2} \\ in all directions & & \\ (“var-all”) & & & \\ \hline \end{tabular}
\end{table}
Table 1: Experimental settings. Here \(\mathcal{N}(\mu,\Sigma)\) means the normal distribution with mean \(\mu\) and covariance \(\Sigma\), and \(t(v)\) means the \(t\) distribution with \(v\) degrees of freedom.
does not define a metric (it cannot distinguish distributions \(P,Q\) that agree up to some number of moments, interpreted in the multivariate sense, but differ otherwise). We include comparisons to polynomial kernels for completeness nonetheless.

### Experimental results

Figure 2 displays the ROC curves from all tests across all settings (rows) and all dimensions (columns). We see that in the first four settings (first four rows), the RKS tests exhibit generally strong performance: in each case, one of the \(k^{\text{th}}\) degree RKS tests performs either the best or near-best among the methods considered. This supports the intuition discussed earlier, since in each of these settings, \(P\) and \(Q\) differ in only a small number (in fact, one) direction. Conversely, in the last setting, "var-all", the distributions \(P\) and \(Q\) differ by a small amount in all directions, and the Gaussian kernel MMD and energy distance tests outperform the RKS tests, especially for larger \(d\).

Keeping in mind that larger choices of \(k\) correspond to higher-order truncated moments, we can further inspect the first four settings to try to learn from why the different order RKS tests perform better or worse. In the "pancake-shift" setting, the mean shift occurs in a low-variance direction, leading to nearly separable classes, so the smallest choice \(k=0\) performs best, and better than kernel-based tests. In the "var-one" and "t-coord" settings, \(P\) and \(Q\) have the same mean, but different variance or tails, respectively, in one direction. Thus, larger choices of \(k\) perform better: \(k=1\) for "t-coord", and \(k=2,3\) for "var-one". For a mean shift of between isotropic Gaussians, "ball-shift", degrees \(k=0,1,2\) are all reasonably competitive.

## 5 Discussion

We proposed and analyzed a new multivariate, nonparametric two-sample test, which is a maximum mean discrepancy (MMD) test with respect to a class of functions having unit Radon total variation, of a given smoothness degree \(k\geq 0\). We call this the Radon-Kolmogorov-Smirnov (RKS) test, and it generalizes the well-known Kolmogorov-Smirnov (KS) test to multiple dimensions and higher orders of smoothness.

The RKS test, like any statistical test, is not expected to be most powerful in all scenarios. Roughly, we expect the RKS test to be favorable in settings in which the given distributions \(P\) and \(Q\) differ in only a few directions in the ambient space; and higher orders of smoothness \(k\) can be more sensitive to tail differences along these few directions. When \(P\) and \(Q\) differ in many or all directions, the more traditional kernel MMD test with (say) Gaussian kernel will probably perform better.

It is worth revisiting computation of the RKS distance and discussing some limitations of our work and possible future directions. In this paper, we considered _full batch_ gradient descent applied to (9) (we actually used Adam, but the differences for this discussion are unimportant), i.e., in each update the entire data set \(\{x_{i}\}_{i=1}^{m}\cup\{y_{i}\}_{i=1}^{n}\) is used to calculate the gradient. This results in the following complexity: \(T\) iterations of gradient descent costs \(O((n+m)dNT)\) operations, where recall \(N\) is the number of neurons. However, _mini batch_ gradient descent should be also investigated, and is likely to be preferred in large problem sizes. That said, implementing mini batch gradients for (9) is nontrivial due to the nonseparable nature of the criterion across observations (due to the log and the absolute value in the first term), and pursuing this carefully is an important direction for future work. For now, we note that by just subsetting \(P_{m},Q_{n}\) to a mini batch of size \(b\leq n+m\), this results in a \(T\) iteration complexity of \(O(bdNT)\) operations. To compare kernel MMD: this costs \(O((n+m)^{2}d)\) operations, and thus we can see that for large enough problem sizes computation of the RKS statistic via full batch and especially mini batch gradient descent can be more efficient, provided that small or moderate values of \(N,T\) still suffice to obtain good performance.

Finally, another direction worth pursuing is the use of deeper networks, which could be interpreted as learning the feature representation in which there is the largest discrepancy in projected truncated moments. This would help further distinguish the RKS test from kernel MMD, because the latter is "stuck with" the representation provided by the choice of kernel, and it cannot learn the representation based on the data. A deep RKS test could dovetail nicely with existing architectures (and even pre-trained networks) for related tasks in areas such as generative modeling, and out-of-distribution detection.
Figure 2: ROC curves across the experimental settings described in Table 1. Each row represents a different setting, and each column a different dimension. The RKS, kernel MMD (KMMD), and energy distance tests are compared (each coded by a combination of color and line type).
### Acknowledgments

We thank Rahul Parhi and Jonathan Siegel and helpful discussions about Radon bounded variation spaces. This work was supported by NSF SCALE MoDL grant 2134133, ONR MURI grant N00014-20-1-2787, and the Miller Institute for Basic Research in Science at the University of California, Berkeley.
