Neural Tangent Kernel Maximum Mean Discrepancy
Authors: Xiuyuan Cheng
Yao Xie
Venue: NeurIPS 2021 Poster (2021)
Keywords: Maximum Mean Discrepancy (MMD)
Neural Tangent Kernel (NTK)
two-sample test
kernel methods
change-point detection

Abstract: We present a novel neural network Maximum Mean Discrepancy (MMD) statistic by identifying a new connection between neural tangent kernel (NTK) and MMD. This connection enables us to develop a computationally efficient and memory-efficient approach to compute the MMD statistic and perform NTK based two-sample tests towards addressing the long-standing challenge of memory and computational complexity of the MMD statistic, which is essential for online implementation to assimilating new samples. Theoretically, such a connection allows us to understand the NTK test statistic properties, such as the Type-I error and testing power for performing the two-sample test, by adapting existing theories for kernel MMD. Numerical experiments on synthetic and real-world datasets validate the theory and demonstrate the effectiveness of the proposed NTK-MMD statistic.

Neural Tangent Kernel Maximum Mean Discrepancy
Xiuyuan Cheng Department Mathematics Duke University xiuyuan... cheng@duke edu
Yao Xie
H. Milton Stewart School of Industrial and Systems Engineering Georgia Institute of Technology yao xie@isye gatech. edu
Abstract
We present novel neural network Maximum Mean Discrepancy (MMD) statis- tic by identifying a new connection between neural tangent kemel (NTK) and MMD. This connection enables memory-efficient approach to compute the MMD statistic and perform NTK based two-sample tests towards addressing the long-standing challenge of memory and computational complexity of the MMD statistic. which is essential for online im- plementation assimilating new samples. Theoretically, such connection allows us to understand the NTK test statistic properties, such as the Type-I error and testing power for performing the two-sample test, by adapting existing theories for kernel MMD. Numerical experiments on synthetic and real-world datasets validate the theory and demonstrate the effectiveness of the proposed NTK-MMD statistic.
develop computationally efficient and
Introduction
Maximum Mean Discrepancy (MMD) statistic is popular method in machine learning and statistics. In particular, kernel MMD
has been applied to evaluating and training neural network Though a widely used non-parametric test [23] kernel MMD encounters several challenges in practice. The roadblocks for large-scale implementation of kernel MMD involve heavy memory requirement (due to the computation and storage of the Gram matrix, which grows quadratically with the data size) and the choice of a good kernel function for high dimensional data. While Gaussian RBF kernel was shown to provide metric between pairs of probability distributions with infinite data samples, applying isotropic Gaussian kernel to data in applications, such image data and discrete events data, may invoke issues in terms of kernel
23
generative models 34
44
31
3
30
expressiveness 25 29 36 and sampling complexity [8]]
A potential path forward in developing more computationally and memory-efficies testing statistics to leverage deep neural networks' representation and optimization advantage. For example, idea of training classification neural network for testing problems has been revisited recently in and the connection between classification and two-sample testing dates back to carlier works 43 40), However, in applying deep models to testing problems, the test consistency analysis is usually incomplete due to the lack of optimization guarantee of the trained network. For one thing, assuming perfect training of deep network achieve global minimizer is too strong assumption to fulfill in practice.
37 10
19
A recent focus of neural network optimization research is the so-called lazy training regime of over-parametrized neural networks [13, where the neural network training dynamics exhibit certain linearized property and provable learning guarantee can be obtained 33 17 15. I In this regime, the training time is sufficiently short, and networks are sufficiently parametrized such that network parameters stay close to the randomized initial values over the training process. In particular, the Neural Tangent Kernel (NTK) theory, as firstly described by (28), shows that the network optimization can be well approximated by the Reproducing Kernel Hilbert Space (RKHS) formulation. The NTK
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

theory has been developed for general neural network architectures, including deep fully connected
networks 48
convolutional networks 6 [35, graph neural networks [T6], and residual networks The RKHS approach by NTK has been shown theoretically and empirically to characterize
45
27 7
the wide neural network training dynamic in the early stage.
The current work stems from a simple observation that short-time training of a network is approxi- mately equivalent to computing the wimess function of a kernel MMD with NTK at time zero, when the training objective equals the difference between sample averages of the network function on two samples. The proposed test statistic, called NTK-MMM approximates the classical kernel MMD with NTK, and the eπror at training time t can be bounded to be (tt) under the linearization of the NTK theory. The theoretical benefit of translating the network- based statistic into a kernel MMD is that the testing power of the latter can be analyzed based on previous works. Algorithm- wwie,, the network-based test. statistic can computed on the fly: thanks the form of linear accumulation of the training objective, the training allows small-batch, e.g.. batch size 1 and 1 epoch of training (1 pass of the samples), under the NTK approximation. (The "NTK approximation" in this paper refers to approximating the time-t NTK by the time-zero kernel, both at finite width.) To calibrate the testing threshold needed to prevent false alarm, we introduce an asymmetric MMD using training-testing split and theoretically prove the testing power, where the threshold is estimated from bootstrapping on the test split only and thus avoids retraining network.
Our main contributions include the following: (i) We introduce a neural network-based test statistic called NTK-MMD, which can be computed by short-time training of neural network, particularly online learning using one-pass of the training samples and batch size one. The NTK approximation error of the MMD statistic is shown to be (t). that is, linear training time, and the result extends to Stochastic Gradient Descent training: (ii) We characterize the statistical properties of the NTK-MMD, including the Type-I error and the testing power, which establish the conditions under which the test is powerful; we further introduce data split scheme such that the test threshold can be estimated without network retraining with provable testing power guarantee; (iii) The efficiency of the proposed NTK-MMD test is demonstrated on simulated and real-world datasets.
At the same time, we are aware of the limitations of NTK in explaining deep network optimization. expressiveness power, and so on. We discuss limitations and extensions in the last section. In particular, this paper focuses on demonstrating the power of NTK-MMD statistics for the two-sample test, while the proposed computationally and memory-effiiien NTK-MMD statistics can also be used for other applications of MMD statisties 24 and hypothesis tests.
2 Method
2.1 Preliminary: Kernel MMD
We start by reviewing a few preliminaries. Consider data in X C Rd, sampled from two unknown distributions with densities p and q. Given two data sets
1,.... nx}. == {y, ~ q. ii.d., j 1, ,ny.. we would like to test whether or not they follow the same distribution. This is equivalent to perform the following hypothesis test Ho p q versus H, :p + g. The classical kernel MMD considers test functions in the RKHS of positive semi-definite kernel K(x, v), which can be, for instance, the Gaussian RBF kernel. The (squared and biased) empirical kernel MMD statistic is given by 23
X {z, ~ p. i.i.d.,
(1)
J/ ( (F) (yddd p p:=
1 XX
1 ny
q
5y (2)
nx
ny
The null hypothesis is rejected if MMD > tihres. where tehres is the user-specified test threshold (usually, chosen to control the false alarm up to certain level). The (empirical) witness function of the MMD statistic, (x:) fx (xx, b (d)d,,, indicates where the two densities differ. The Type-I error of the test is defined as P[MMD > [shres] under Ho. and the Type-Il error as & < /4hres] under H1; the power is defined as one minus the Type-Il error. For an alternative distribution q of p. the test errors depend on q. the sample sizes nx and ny, as well as the kernel function K (x,, y). Theoretically, the test power of kernel MMD has been analyzed in (3]. investigated for high dimensional Gaussian data in (38), and for manifold data in [2].
2.2 NTK-MMD statistic
As the proposed NTK-MMD framework can be used on different network architectures, we write the neural network mapping abstractly as f(x; 0), which maps from input X' to R, and 0 is the
2

---

network parameters. Use X UUY& as the training dataset, and let p and q be as in we choose
particular training objective function as
T(0)) -100=-/aa 1(: -
(x)dd -(W))-- 1 ax (an 0)
1 nY
nx
The choice of this objective function is critical in establishing the connection between NTK and MMD. Optimizing this objective will lead divergence of the network function if train for long time. However, if the training is only for a short time, the network function remains finite. Here we mainly focus on short-time training of the network, and particularly the online training setting where the number of epochs is 1, that is, only 1-pass of the training samples is used. We will also show that the method allows using minimal batch size for online learning without affecting the NTK approximation of the MMD statistic, c.f. Remark 2.2
Following the convention in NTK literature, below we formulate in terms of continuous-time Gradient Descent (GD) dynamic of the network training. The extension to discrete-time Stochastic Gradient Descent (SGD) holds with a small learning rate (Remark[.22) The network parameter ((t) evolves according (l)) = - and we define u(x; () == f(x 0())). which the network mapping function at time L. Suppose the network is trained for short time <>> O, we define
((x) = (u((z,)) O)).
(4)
and the test statistic, which depends on time t, is
Thet (t) == - (z)dr== (((O())
(5)
The function 0 is the difference of the network mapping after a short-time training from the initial one, and we call it the witness function of network NTK-MMD statistic. As revealed by (5). (without calibrating the test threshold) the test statistics T is nothing but the decrease in the training objective, and comes as a by-product of network training at no additional computational cost. We show in next subsection that at small f, the statistic Tues (t) provably approximates the classical MMD statistic with the NTK, i.c. TNTT MMD K (X, Y) where K is the NTK time = 0 as in (8. Algorithmically, we will perform two-sample test using Thee (L) by comparing with threshold lchres-
2.3 NTK approximation of MMD statistic
In the continuous time training dynamic of the network, we consider the NTK 28 kernel function defined for / as
(x(x,x')= (Voi(x; 0()), (x(;
(6)
The following lemma follows directly by construction, and the proof is in Appendix A. Lemma 2.1. The network function u(x, 1) satisfies that for t > O,
u(x, t) u(x 0) J. R.(x, (())a"((((((())))))))))( (ol'''dd'dss
It has been shown (in 5 among others) that for the short-time training (lazy training regime), the kernel (6) can be well-approximated by the kernel at time = O, namely
x') == (Vo 0(0)), (ef(x'(0()))
which is only determined by the network weight initialization ((0). Assuming Ko ~ K, in Lemma 2.1 the proposed test statistic Ihee(() (t) as in (5) can be viewed as
Pnee (t) ~ TN == fxx Kulz, ('(z,2')06 (())) ()
which is the kernel MMD statistic with NTK. (See Remark A. for a discussion on biased/unbiased MMD estimator.) In below, we show Proposition 2.1 that the approximation fnet zs Tyrk has O(t)
error, and we experimentally verify the similarity of the two statistics in Subsection 4.2 Throughout the paper, we compute Thet by neural network training, and we call TNTK the exact NTK-MMD
3

---

which is for theoretical analysis. The theoretical benefit of translating Tuet into TNTK lies in that testing power analysis of TNK follows existing methods which is detailed in Section 3
Suppose neural network parameter is in RM and 0 E O, where θ is a domain in R which contains the Euclidean ball (() ro), where we assume To is an O(1) constant. For vector valued function g: (x, Q) - Id and U C θ we denote the infinity norm as | == supxEx llg 0)ll- When 9 maps to a matrix, the notation denotes (the infinity norm over (X x U) of) the operator norm. The test statistic approximation error in Proposition 2.1 directly follows the following lemma concerning the uniform approximation of the kernels. All proofs in Appendix A.I
Lemma 2.2 (NTK kernel approximation), Suppose f is C2 on (X. 6) and Vefl S L for some positive constant Lj. Then for any 0 < r < To. when 0 < &fr := r/ (2L
(1) f(t) stays inside the Euclidean ball B, == r)
(2) Define Cj,- == H|ID&[lIx, [Vofll B,, we have that
sup ( - Ko(x,x') S C,rtt. a..'E
(10)
Remark 2.1 (Boundedness of [[ When p are unbounded density (gaussian), and activation function is relu softplus, the uniform boundednesss of o/ may fail. However, for sub-exponential densities, apply standard truncation argument, and when we restrict to compactly supported distributions. In practice. we standardize the data to be compact domain in R
Proposition 2.1 (Test statistic approximation). The condition on f(x, 0) is the same as in Lemma 2.2 and for 0 < < To, the constants fjr and Crr are as therein. Then, when 0 < t < ff we have that
Fnee (t) Trkl S 2C
Remark 2.2 (SGD and online training). The above error bound analysis based on Taylor expansion can extend to discrete-time GD dynamic by showing that the time discretization introduces higher-order eгror when t is small. In the SGD setting, c.g., the online learning of I epoch, batch size one, and learning rate a, the network parameters are updated after scanning each training sample on the fly. Let 02 be the network after scanning k many samples, we show in Appendix (A.2 that the difference |10g- Bull can be bounded by O(ak/ /n), and the trained network witness function after epoch approximates the witness function with the pero-time NTK kernel up to an O(a) error. The learning rate a has the role of training time t. The fact that batch size will not affect the NTK approximation of the network training is a result of that the loss (3) is a linear accumulation over samples, which may not hold for other loss types. The compatibility with online learning and training with very small batch size of NTK-MMD statistic makes it convenient for deep network training. especially under memory constraints.
2.4 Computational and memory efficiency
The update of network parameters in NTK-MMD training can be viewed as an implicit computation of the inner-product between high dimensional kernel feature maps (V ((t)), Vof ((t()) (by chain rule, c.f. (24) (25) in Appendix [..]]. The network witness function 9 defined in 4 is parametrized and stored in trained network parameters. This allows the (approximate) evaluation of kernel on a test sample x'' without computing the gradient (x; 0) explicitly. It also means that the NTK network witness function can be evaluated on any new x' without revisiting the training set. In contrast, traditional kernel MMD computes kernel witness function (defined as K K(x, y) (f-())))) (y)dy 23 on new point x' by pairwise computation between x' and samples in datasets X and Y.
can be computed via batch- satch-size-one training over one- pass of the training set (c.f. Remark 2.2 and experimentally verified in Table A.2) The gradient field evaluation (back propagation) is only
conducted on the training set but not the testing test, and the bootstrap calibration of the test threshold can be computed from test set only (c.f. Section 3.2)). Thus, by using small learning rate (allowed by one can incorporate large number of training samples via more training iterations without worsening the approximation error to exact NTK-MMD, which will improve testing power. This "separation" of training and testing, in memory and computation, of NTK-MMI allows scalable online learning as well as efficient deployment of the network function on potentially large test sets.
floating point precision, c.f. Remark
3 Theoretical properties of NTK-MMD
In this section, we prove the testing power (at a controlled test level) of the NTK-MMD statistic Tvtk in (9) with large enough finite samples. We also introduce asymmetric version of the

---

MMD statistic using shining-testingg dataset splitting, which enables the bootstrap estimation of the threshold of the test thrres without retraining of the neural network.
3.1 NTK-MMD without data splitting
We write the NTK kernel Ko(æ, x') as K (2x, 2") omitting the subscript, and assume that ((x, x') is uniformly bounded, that is, su K (ar, 2) < B < oc for some positive constant B. Without (because global constant normalization of the kernel
loss of generality, we assume that B introduces global constant multiplied to the test statistic, and does not change the testing). By that the kernel is PSD, we thus have that
sup [K x') 1
(11)
ot, FEX
We omit the NTK subscript and denote the exact WTK-MMD statistic in ( as T. The corresponding population statistic is the squared MMD of kernel K K(x,
6k = MMD (p. = Ja K(z, w) ((x,
()) (y)dddyy.
(12)
By the uniform boundedness I the kernel K(x, is in L3 (x (p+ (x) ddz'dddd)))
We define the squared integrals of the kernel
Ppp == Ex-p.spp K (zz, v)² V'gq :- Ee-nara K (I,,)²
Ugg == Ex K (ar v)²
(13)
In addition, we assume that as n == nx + ny increases, n nx/n stay bounded and approaches Px 1). Equivalently, there is some constant 0 < 1 such that for large enough n,
en + 1 S nx, ny S n, i i= 1,2.
(14)
Without loss of generality, we assume that (14) always holds for the considered. Theorem 3.1 (Test power of TNIK). Suppose and 14) hold, and
(i) Under Hi, p & q, the squared population kernel MMD 8g as in 12 is stricily positive,
(ūi The three integrals as in 13 Ppps Vpg Vap are all bounded by a constant v $ 1.
Define A == v (oog(( 4/ (leevo)), and let the threshold for the test be lehres = 4/(cn) + 4d1 v v/on.
Then, if for some Xa > is large enough such that
1
{
1 9v
V (8()1 + X2))
max
max
(15)
0%
then under Ho. PP[
[thres] S Aeveli and under 11, PP S (ehrrs)] S 3e x2/88
The proof uses the U-statistic concentration analysis, and is left to Appendix B As revealed by the proof, the diagonal entries in the kernel matrix contribute to the O(1/n) term, and thus switching from the biased estimator of MMD (9) to the unbiased estimator gives similar theoretical results.
Remark 3.1 (Choice of h.res). The choice of fehres in the above theorem is theoretical one and may not be optimal, due to the use of concentration inequality and the relaxation of the bounds by using constants v and c. By definition, the optimal value of 4thres is the (1 11 -Oeeval (elleval-quantilee of the distribution of 7 under Ho. The asymptotic choice may be obtained analytically according to the limiting distribution of the MMD statistic, c.f. Remark B. The threshold 4thres is also computed by bootstrap strategy in practice [4] (called "full- ftull-bootsrrap" in next subsection). The bootstrap approach permutes the labels in data sets X and Y, and since in Tuet the witness function ((x) is computed by neural network training, this will incur retraining of the network. A solution to avoid retraining by adopting test set for bootstrap estimation of three is introduced in next subsection.
3.2 Threshold calibration by data splitting
As shown in Theorem 3.I and Remark [3.1] in the theoretical characterization of test power (at a required test level) the test threshold plays critical role. In practice, we need more precise threshold to exactly control the false alarm under the null hypothesis. In this section, we discuss how to set the threshold in two settings: fixed-sample and pilot-data. Nevertheless, we would like to mention that there exist applications where the threshold is not needed, and the symmetric MMD (without training/test split) can be used as a measurement of distribution divergence.
5

---

Fixed-sample setting We first consider the setting where we have a fixed number of samples from p and 9. To obtain precise threshold to control the false alarm, we need to split data into two non-overlapping parts: one part for training neural networks (compute the witness function) and one part data for bootstrapping and calibrating the threshold. We want to highlight that here we develop a scheme for threshold calibration such that no retraining of the witness function is necessary.
We randomly split the datasets X and Y into training and testing sets, X X(1) U X(2) and Y = Y()) U Yea), and compute an asymmetric version of kernel MMD (the subscript is for "asymmetric")
Ta = S. K(x, x(xxx")60 (21) ((2)) (x) ddrde'
(16)
where P(o) and (()) are the empirical measures of datasets X() and (()) respectively, i 1, Define (x.(()) (x) and ny, = |Yoo) = 1, 2. Similarly as in Section 2. the MMD statistic (16) with
K(2, v) = Ku (x, y), the time NTK, can be approximated by
net (1) = J( (1))(() (x)de, 9(1)) (ū(o, ú(2;, 0)),
(17)
for small time t, where ū(æ,t) is the network function trained by minimizing L(
- Sx fx; (ō() (z)dr the training set D,, (X(1), y(} with binary labels {1.
Same as in Lemma 2.2 Proposition 2.1 the difference [Ta - f aet (t) can be bounded to be (tt).
We theoretically analyze the testing power of Ta where K(x, = Ko y) in below.
The benefit of splitting the test set lies in that once the witness function ((1)(2)) is trained from Der. one can do a test-only bootstrap which is to compute
Tull = J (,)(=)00(2) (x)
(18)
where P/22) and (22 are empirical measure of samples in Dic = (x(2), y(2) by randomly permuting the 12,122) + my.12) many binary class labels. Since permuting test labels does not affect (11) the test-only bootstrap does not require retraining of the neural network nor revisiting the training samples. Alternatively, one can permute the binary class labels in both Dir and Die. and this will require to retain the neural network to obtain the new witness function ((1) given the new class labels of D... We call such a bootstrap the theffll-bbol full-bootstrap. The full-bootstrap can be applied to the symmetric MMD statistic without test splitting as well, namely the setting of Theorem 3.1 to obtain an estimate of optimal htres in practice.
We give two theoretical results the testing power guarantee of the asymmetric NTK-MMD statistic (16): For rest-only bootstrap, Theorem (.22ppoveess testing power by restricting to good events over the randomness of Dry: For full bootstrup, the guarantee is provided in Theorem 3.3 which is the counterpart of Theorem 3.1
All proofs are in Appendix B
We assume the balance-ness of the two samples as well as the training and testing splitting, that is, 7xx(1) /xx + Px,(1), my,,m/ny + (Y,(1) nx/n + px. and the three constants are all in (0, 1). With nx + ny, we assume that for constant 0 < Co 1,
Cu S (x),))) Ry,() $ n, i= 1,2.
(19)
We denote by (I1) the randomness over Dtr, and ((2) that over Dre. Theorem 3.2 (Test power of T., test-only bootstrap). Suppose that and (ii) in Theorem 3.1 hold, and 0 < < 1 is a small number. Define A == /4log(4 (log(4//)eve) (eeeel)
19p and the conditions (i)
A(1)= / log(8/7 8/7), and set the threshold as luures = 4(V1.1A) + Aa)))//// (V/(c))) Ifn is large
enough such that n ( /2/(cc)), and for some >x2), 0.
max
max(A A )
12
16v 8%
(2x()) + VI.((A2)) + X12)2))
(20)
then, under both Ho and H₃ there is good event over the randomness of Der which happens w.p.> 1 7, under which, conditioning on Drr, P(2) (Ta tehres] Qlevel under Ho, and IP P(z) (T
tshree] < 4e -a,,a/ under H1,

---

Remark 3.2 (Sampling complexity). Compared to the full bollstoott result Theorem 3.3 the additional requirement on n is that n. needs to be greater than (Agr//2²² up absolute constant, and thus when v/02 V 2 the (v/6g) J-term still dominates the needed lower bound of n, same as in Theorems 3.1 and 3.3 (Here we treat )(1). A12),1 and X (1) constants. Because the constant y controls the good event probability over the randomness of Dir. thus if can be chosen to be of the same order as Olevel, then Aς1) has the same order as A(2). The result shows that with test split and test-only bootstrap (avoiding retraining), the test power has the same order of needed sampling complexity, n 2 v/0k.
k. as full bootstrap, with high probability and for large enough
Theorem 3.3 (Test power of Te, full bootstrap). Suppose that (II). (9) and the conditions (i) and (ii) in Theorem 3.7 hold. Define ) == /8 log 4/ ovvee)) and let the threshold for the test be
tehress 4 can Then, for some x2 > O, n is large enough such that
n > max
1 Ca
1 90
max N2)?
V (4(21 + +x2))²
(21)
82 K
then under Ho. P[f > fthres] S Meveli and under H1, P(f S tchres] I 4e 4e-eex x2/88 -22/88
Pilot data setting, This section considers the setting where we may have many samples for one distribution, c.g., the p. For instance, change-point detection, where we are interested in detecting a shift in the underlying data distribution, there can be a large pool of pilot data before the change happens, collected historically and representing the normal status. We may have fewer data samples for the distribution q. For such a case, we can use data from the reference pool represent distribution p to train the model and calibrate the threshold, e.g., using bootstrap. Since such "training" is done offline, we can afford the higher computational cost associated with training the model multiple times. In short, our strategy is to pre-compute the detector (retrain multiple times) and then use boostrap to obtain the threshold thres for detector: (i) compute the symmetric MMD T as in (9) on {p, q). where q is the new coming test samples (e.g. in thange-point detection), and ý is from the pool; (ii) pre-compute the symmetrie MMD Taull on {p2. 122) } from the pool of samples, with retraining, and obtain the "true" threshold for T'. Retraining of the network is expensive, but this pre-computation and not counted in the online computation.
4
Numerical experiments
The section presents several experiments to examine the proposed method and validate the theory.
4.1
Gaussian mean and covariance shifts
Set-up. Consider Gaussian mean shift and covariance shift in R 100, where nx = ny 200; p is the distribution of (0, Ia), = 100: (i) Mean-shift: q is the distribution of N(H,, Ia), where || = which varies from 0 to 0.8 and (ii) Covariance-shift g is the distribution of (0. Ia + pE), where E an d-by-d all-ones matrix, and changes from 0 to 0.16. We split training and test sets into halves, and compute the asymmetric network approximated NTK-MMD statistic Ta.net [7]) and estimate the test threshold by the quantile of (18); Ho is rejected if fa.net > fthres- We use 2-layer network (1 hidden layer) with soft-plus activation. The online training is of epoch (1 pass over the training set) with batch-size = 1. The bootstrap estimate of test threshold uses Rbooa = 400 permutations. The testing power is approximated by 7ym = 500 Monte Carlo replicas, and we compare with the benchmarks by (i) Hotelling's T-test, and (ii) Gaussian kernel MMD test (median distance bandwidth) (23). The median distance bandwidth is a reasonable choice for detecting high dimensional Gaussian mean shift (38). Both Gaussian kernel MMD and Hotelling's ' Test have access to all the samples D., U Dre. More experimental details are in Appendix C.T
The results are shown in the left two plots in Figure I The NTK-MMD test gives comparable but slightly worse power than the other two benchmarks on the mean shift. On the covariance shift test, the network MMD test gives equally good power as the Gaussian MMD. For the Gaussian covariance shift case, we also compute the testing power when only part of the training samples are used in the online training, and the results are shown in the right two plots in Figure [] increases as the neural network scans more training samples, and when the covariance shifts are larger the transition takes place with smaller training sample size.
Results.
Testing power
'Code available hattps://githee httt::://gttuub..o///yhheeegg - MMD/
7

---

25082 shilt, d-100
covariance shift, d-100
comrance testing power
covatiance skift Solting
net
net Hozling grand
0.8 Hoxilingg
0.8
0.8
0.05
0.0
0.6
0.6
0.6
08
0.1
0.4
0.4
0.4
0.4
0.2
02
0.2
0.2
0.4
0.6
0.0
0 0.05 0.1
0:15
50
100
150
200
100 200 300 400 500
trolning semple size
training sample sioe
Figure 1: (Left two plots) Estimated testing power on Gaussian mean shift (change size is 6) and Gaussian covariance shift (change size is p) in 10000 where datasets X and Y have 200 samples respectively, and the training and testing splitting is half-half, i.c. Ber = Aiee = 200. Test power is estimated from Rrun 500. (Right two plots) Estimated testing power training set (batch size =1) and over varying values of p. also plotted as color field.
function of the number of samples processed in the 1ppsss of the
covtriance shifft. d-100
cosuriance shift, x6-2
trorani-mlly Loxdstzap
fall boxistrup
0.8
-+-8k1
18.1
1892
1062
30
40
0.4
0.4
20
0.2
0.2
0
0.25
0.5
0.05
-0.06
Figure 2: (Left two plots) Estimated testing power from frum 500 of the covariance shift test in Figure in R and IR?. nx = ny = 200, using three statistics: Time (net), TYTK with test set only bootstrap (ntk1) and with full bootstrap (ntk2) the training and testing splitting is half-half. (Right two plots) Test statistics T', (red cross), the empirical distribution of muml using the test-only bootstrap and the full bootstrap (blue bars), and the estimated threshold (green circle). Computed from NTK kernel at = and Thooot 400.
In addition, we show in Appendix C.2 that NTK-MMD gives similear performance with varying network architectures, activation functions (like relu), and SGD configurations, and possibly better testing power with a larger network depth and width (Tables A. and A.2). We also compare with linear-time kernel MMD in Appendix C.7] As shown in Table NTK-MMD outperforms linear-
time gMMD as in [23, Section 6], and underperforms the full gMMD which however requires ((n?) computation and storage.
4.2
Comparison of î and 7NTK
Set-up. Since we use a 2-layer fully-connected network, the finite-width NTK at ( = 0 (using initialized neural network parameters) can be analytically computed, which gives an -y---r asymmetric kernel matrix K. The expression of K and more details are provided in Appendix C.3 This allows computing the exact NTK-MMD (16), as well as the (i) full bootstrap and (i) the test-only bootstrap of the MMD statistic under Ho by (i) permuting both rows and columns simultancously and (ii) only permuting rows of the matrix K.
Results. To verify the O(t) discrepancy as in Proposition 2.1 we first compute the numerical
values of TKK and Tnet (t) for different values of t (which corresponds to different learning rate as explained in Remark 2.2 and Appendix (A.2) and the relative approximation eиror defined as err = = [t)
TNTK| /|Tk] The results are shown in Figure A. The fitted scaling of the error for softplus activation is about 0.96, which agrees with the theoretical O(t) error. Switching to relu,
the order is not close to 1 (instead 0.62) but Thce (t) still gives good approximation of TwTK the relative error achieves about 10-3 The comparison of the testing power of network approximate NTK statistic and the exact NTK statistic tests are shown in Figure 2 In the high dimensional Gaussian covariance shift test (d = 100), the powers of the three tests are similar. When reducing dimension to d 2, the full-bootstrap NTK tests show slightly different testing power than the other two. The network approximate NTK and NTK with test-only bootstrap always show almost the same testing power, consistent with the theory in Subsection 2.3 In the experiment on R data, the estimated threshold by boll-bootstrap is smaller than by test-only bootstrap (right two plots), which explains the possibly better-testing power.
4.3 Comparison to neural network classification two-sample tests Set-up. We experimentally compare NTK-MMD and state-of-the-arr classification two-sample test (C2ST) basclines, which are neural network based tests. Following (36), we compare with C2ST-S,
8

---

3m
10
10
10
sumber
Junety shange
001
000
000
Figure 3: NTK-MMD statistic to detect distribution ahundance change in MNIST digit image data; ne. 2000 and nle is about 4000. (Top) Most left: Change of statistic over the number of samples in the online training
(batch size =1) of 2-layer convolutional network. From 2nd-5th columns: The MMD statistic T., compared with the empirical distribution under Ho via test-only bootstrap. at four times along the online training (red cireles on the left plot). (Bottom) Most left: The change of distribution of MNIST dataset embedded 2D by tSNE [46). From 2nd-5th columns: The witness function å plotted as a color field over the samples, at the four times corresponding to the upper panel plots.
Ber
2000
4000
6000
8000
MEA
~ 10.0
30.0
58.0
75.0
SCF*
5.0
~ 6.0
~ 10.0
15.0
C2ST-S (Adam)
(61.6)
14.0 (95.8)
39.1 (100.0)
61.2 (100.0)
C2ST-L (Adam)
14.1 (87.8)
38.4 (100.0)
76.4 (100.0)
92.9 (100.0)
C2ST-S (SGD)
(13.9)
10.6 (50.0)
10.8 (94.4)
14.8 (99.6)
C2ST-L (SGD)
6.7 (22.2)
12.8 (81.6)
22.1 (100.0)
34.6 (100.0)
NTK-MMD
7.1
9.6
13.7
17.9
Table 1: Test power on Gaussian mixture data Example I, dimension d = 10. ("recovered from Figure 3 in (36]
means about.) For C2ST's, the number outside brackets for epoch and brackets for epoch 10.
the classification accuracy test [37], and C2ST-L, the classification logit test 10 details are given in Appendix C.4 The data distributions are:
Experimental
Example Gaussian mixture. fixed dimension d = 10 and increasing fur, which the same
setting as Figure 3 (left 2 plots) in [36]. Numbers in Table I show testing power (in %).
Example 2: Modified Gaussian mixture (from Example 1), the covariance shift is T + 0.1 in both mixtures, where is all-one matrix with zeros on the diagonal. Dimension d = 10, and number of training samples Bhr increases. The test power is shown in Table 2
Results. On Example 1, NTK-MMD performs similar to SCF test in most cases, better than C2ST-S (SGD 1-epoch), and is worse than the other baselines. On Example 2, NTK-MMD outperforms C2ST baselines in several cases, c.g., constantly better than C2ST-S (SGD and Adam, 1-epoch) be sensitive to training
and comparable to C2ST-L (SGD 1-epoch). Note that C2ST baselines hyperparameters, such as the choice of optimization algorithm (SGD or Adam) and number epochs. As far as the authors are aware of, there is theoretical training guarantee of C2ST tests. In contrast, NTK-MMD has theoretical training guarantees due to the provable approximation to a kernel MMD, The weakness of NTK-MMD. though, is that the NTK kernel may not be discriminative distinguish certain distribution departures, like in Example 1. The expressiveness power of NTK-MMD may be theoretically analyzed, for example, in the infinite-width limit using the analytical formula, as the infinite-width NTK has been shown to be universal for data on hyperspheres (28). Overall, the results suggest that the performances of the three neural network tests depend on the data distributions, which is anticipated for any hypothesis test. Further theoretical investigations are postponed here.
Yhr C2ST-S (Adam)
500
1000
1500
2000
21.8 (28.1)
62.2 (53.8)
79.4 (74.0)
94.6 (85.2)
C2ST-L (Adam)
48.5 (49.4)
92.8 (82.6)
99.5 (96.3)
100,0 (98.8)
C2ST-S (SGD)
7.4 (28.3)
22.7 (79.7)
35.3 (92.4)
54.9 (96.8)
C2ST-L (SGD)
(52.2)
56.8 (97.6)
81.4 (99.9)
97.3 (100.0)
NTK-MMD
34.3
68.9
88.8
95.9
Table 2: Test power on Gaussian mixture data Example 2. dimension d = 10.

---

bending
threwing
Hotelling
MMD
NTK
200
30
detertion statieil
1.0
detection statiztc
0.8
DADeg
detection statlste
100
0.4
10
0.
491
498
537
571 592
600
time
400
600
800
400
E00
800
400
600
800
time indore
Simoe incler
time indor
Figure 4: (Left) Example data sequence the human activity dataset: before and after change-point. (Right) Detection statistics computed by Hotelling's statistic, Gaussian MMD, and NTK-MMD, on human action trajectory dataset, using window size 100 (blue), 150 (red), and 200 (yellow), respectively. The change point is time index 550 (dotted black).
4.4 MNIST distribution abundance change
Dataset. We take the original MNIST dataset, which contains 28 28 gray-scale images, and construct two densities p and q by subsampling from the 70000 images in 10 classes, following [[2]: p is uniformly subsampled from the MNIST dataset. p Pdatas and has a change of abundance 0.85paata + 0. 00.5hoohrrt, where Poohort is the distribution of a subset of the class of digit "I" having about 1900 samples. The Peohort is illustrated in the left bottom plot in Figure 3 The two samples X and Y have nxx = 3000, my = 2981, and we randomly split X and Y make the training set DDr = (X())) (())), nx,(1) = (y,(1) 1000, and the rest is the test set Die.
Results. Using a 2-layer convolutional nerual network, we compute the network MMD statistic To net T77 and the test-only bootstrap (18). The online training uses batch size =1 and one epoch, and more experimental details are in Appendix
The results are shown in Figure 3 The NTK-MMD statistic already shows testing power after being trained on 50 samples, and in the later stage of training, the NTK witness function 9(1) identifies the region of the abundance change.
4.5 Online human activity change-point detection
Set-up. We present an illustrative example using NTK-MMD test statistic to perform online change- point detection: detecting human activity transition. We consider a real-world dataset, the Microsoft Research Cambridge-12 12 (MSRC-12) Kinect gesture dataset [18]. The data sequence records human subject repetitively bending the body/picking up and throwing a ball before/after the change happens. After preprocessing, the sequential dataset contains 1192 frames (samples) and 54 attributes (data samples are IR55), with change of action from "bending" to "throwing" at time index 550. More description of the dataset and experimental details is provided in Appendix C.6 Example samples before and after the change point are shown in the left of Figure 4 The algorithm is based on a sliding window which moves forward with time, and we compute the detection statistic every ten frames; such a procedure can be viewed as the Shewhart Chart in the literature [47]; scanning MMD statistic has been used in [2] The window size is chosen to be 100, 150, and 200, respectively. We take block of data (same size as the window) before the time index 300 (to use as the pilot samples) and compare with the distribution of data from the sliding window to compute the detection statistic. If there is a change-poine point, the detection statistic will show a large value.
Results. The other two detection statistics are computed by (i) Gaussian MMD (with bandwidth chosen to median distance) and (ii) Hotelling's T statistics. The results are shown in Figure + where both the Gaussian MMD and the NTK-MMD statistics can detect the change: the detection statistic value remains low before the change and remains high after the change point, and both are better than the Hotelling statistic.
5 Discussion
The current work can naturally be extended in several aspects. First, the analysis of NTK approxi- mation error may be extended, c.g., to other activation functions, and under the infinite-width limit. Second, considering other training objectives may allow us to compare NTK-MMD to other neural network classification tests. At the same time, the limitation of lazy-regime training has been studied in 21 22 39,, which indicates that NTK theory cannot fully characterize the modeling ability of deep networks. It has also been shown that the expressiveness of the NTK kernel may be restricted to certain limited type of kernels 8 20 D. This motivates extensions of NTK for studying deep
network training [26 41 Finally, the application may extend to various hypothesis testing tasks as well as deep generative models. We thus view the current work as a first step towards understanding the role and potential of trained neural networks in testing problems and applications.
10

---

Acknowledgement
The authors thank Alexander Cloninger and Galen Reeves for helpful discussion on the initial version of the paper, and the anonymous reviewers for helpful feedback. The work is supported by NSF DMS-21 34037.. XC is partially supported by NSF, NIH and the Alfred P. Sloan Foundation. YX is partially supported by NSF (CAREER Award 1650913, DMS- NSSS-938106, and DMS-18302)
References
[1] Zeyuan Aluen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterizatoon In International Conference on Machine Learning, pages 242-252. PMLR, 2019.
[2] Niall H Anderson, Peter Hall, and D Michael Titterington. Two-sample test statistics for measuring discrepancies between two multivariate probability density functions using kernel- based density estimates. Journal of Multivariate Analysis, 50(1):41 1994.
[3] Michael Arbel, Anna Korba, Adil Salim, and Arthur Gretton. Maximum mean discrepancy gradient flow. arXiv preprint 1906v: 04370, 2019.
[4] Miguel A Arcones and Evarist Gine. On the bootstrap of u and v statistics. The Annals of Statistics, pages 655-674, 1992.
[5] Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine- Fine-grained analysis of op- timization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning, pages 322-332. PMLR, 2019.
[6] Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with infinitely wide neural net. In Advances in Neural Information Processing Systems, pages 8 8139-810 2019.
[7] Yuval Belfer, Amnon Geifman, Meirav Galun, and Ronen Basri. Spectral analysis of the neural tangent kernel for deep residual networks. arXiv preprint arXiv:2104. 00.00003,.. 2021.
[8] Alberto Bietti and Julien Mairal. On the inductive bias neural tangent kernels. arXiv preprint 1rXiv:1905.121 1201: 2019.
[9] Lin Chen and Sheng Xu. Deep neural tangent kernel and laplace kernel have the same rkhs. arXiv preprint arXiv:2009.10683, 2020.
[10] Xiuyuan Cheng and Alexander Cloninger. Classification logit two-sample testing by neural networks. arXiv preprint arXiv: /909. /298, 2019.
[11] Xiuyuan Cheng, Alexander Cloninger, and Ronald R Coifman. Two-sample statistics based on anisotropic kernels. Information and Inference: A Journal of the IMA, 93))6777-7 2020.
[12] Xiuyuan Cheng and Yao Xie. Kernel two-sample tests for manifold data. arXiv preprint arXiv:2105.03425 2021.
[13] Lénaïc Chizat, Edouard Oyallon, and Francis Bach. On lazy training in differentiable program- ming. Advances Neural Information Processing Systems, 32:2937-2947, 2019.
[14] Kacper P Chwialkowski, Aaditya Ramdas, Dino Sejdinovic, and Arthur Gretton. Fast two- sample testing with analytic representations of probability measures. In Advances in Neural Information Processing Systems, pages 1981-1989, 2015.
[15) Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning, pages 1675-1685, PMLR, 2019.
[16] Simon S Du, Kangcheng Hou, Barnabás Póczos, Ruslan Salakhutdinov, Ruosong Wang, and Keyulu Xu. Graph neural tangent kernel: Fusing graph neural networks with graph kernels. arXiv preprint 1905.1312..

---

[17] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations. 2018.
[18] Simon Fothergill, Helena Mentis, Pushmeet Kohli, and Sebastian Nowozin. Instructing people for training gestural interactive systems. In Proceedings of the SIGCHI Conference on Human Factors in Computing Systems, pages 1737-1746, 2012.
[19] Jerome Friedman. On multivariate goodness-of-fin and two-sample testing, Technical report, Stanford Linear Accelerator Center, Menlo Park, CA (US), 2004.
[20] Amnon Geifman, Abhay Yadav, Yoni Kasten, Meirav Galun, David Jacobs, and Ronen Basri. On the similarity between the laplace and neural tangent kernels. arXiv preprint arXiv:2007.01580,
2020.
[21] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Limitations of lazy training of two-layers neural networks. arXiv preprint arXiv: 190,09990,,89,, 2019.
[22] Behrooz Ghorbani, Song Mei, Theodor Misiakiewicz, and Andrea Montanari. When do neural networks outperform kernel methods? arXiv preprint arXiv:2006. /3409, 2020.
[23] Arthur Gretton, Karsten M Borgwardt, Malte J Rasch, Bernhard Schölkopf, and Alexander Smola. A kernel two-sample test. Journal of Machine Learning Research, 13(M27):723-773 2012.
[24] Arthur Gretton, Philipp Hennig, Carl Edward Rasmussen, and Bernhard Schölkopf. New directions for learning with kernels and gaussian processes (dagstuhl seminar 16481). In Dagstuhl Reports, volume 6. Schloss Pagstubl-Leibniz-Zrnnn
fuer Informatik, 2017.
[25] Arthur Gretton, Dino Sejdinovic, Heiko Strathmann, Sivaraman Balakrishnan, Massimiliano Pontil, Kenji Fukumizu, and Bharath K Sriperumbudur. Optimal kernel choice for large-scale two-sample tests. In Advances in neural information processing systems, pages 1205-1213, 2012.
[26] Jiaoyang Huang and Horng-Tzer Yau. Dynamics of deep neural networks and neural tangent hierarchy. In International Conference on Machine Learning, pages 4542 4551. PMLR, 2020.
[27] Kaixuan Huang, Yuging Wang, Molei Tao, and Tuo Zhao. Why do deep residual networks gen- eralize better than deep feedforward networks? neural tangent kernel perspective. Advances in Neural Information Processing Systems, 33, 2020.
[28] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: convergence and generalization in neural networks. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pages 8580-8589, 2018.
[29] Wittawat Jitkrittum, Zoltán Szabó, Kacper P Chwialkowski, and Arthur Gretton. Interpretable distribution features with maximum testing power. In Advances in Neural Information Process- ing Systems, pages 181-189, 2016.
[30] Jonas M Kübler, Wittawat Jitkrittum, Bernhard Schölkopf, and Krikamol Muandet. Learning kernel tests without data splitting. Advances in Neural Information Processing Systems, 2020.
[31] Chun-Liang Li, Wei-Cheng Chang, Yu Cheng, Yiming Yang, and Barnabás Póczos. Mmd gan: Towards deeper understanding of moment matching network. In Advances in Neural Information Processing Systems, pages 2203-2213, 2017.
[32] Shuang Li, Yao Xie, Hanjun Dai, and Le Song. Scan b-statistic for kernel change- point detection. Sequential Analysis, 38(4):503-54) 14450000544,,,0 2019.
[33] Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent structured data. In Proceedings of the 32nd International Conference Neural Information Processing Systems, pages 8168-8177, 2018
[34] Yujia Li, Kevin Swersky, and Rich Zemel. Generative moment matching networks. In Interna- tional Conference on Machine Learning, pages 1718-1727, 2015.
12

---

[35] Zhiyuan Li, Ruosong Wang, Dingli Yu, Simon S Du, Wei Hu, Ruslan Salakhutdinov, and Sanjeev Arora. Enhanced convolutional neural tangent kemels. arXiv preprint arXiv:1911.00809, 2019.
[36] Feng Liu, Wenkai Xu, Jie Lu, Guangquan Zhang, Arthur Gretton, and DJ Sutherland. Learning deep kernels for non-parametric two-sample tests. arXiv preprint arXiv:2002.09116, 2020.
[37] David Lopez-Paz and Maxime Oquab. Revisiting classifier two-sample tests. In International Conference on Learning Representations, 2017.
[38] Aaditya Ramdas, Sashank Jakkam Reddi, Barnabás Póczos, Aarti Singh, and Larry Wasserman. On the decreasing power of kernel and distance based nonparametric hypothesis tests in high dimensions. In Twenty-Nimth AAAI Conference on Artificial Intelligence, 2015.
[39] Maria Refinetti, Sebastian Goldt, Florent Krzakala, and Lenka Zdeborová, Classifying high- dimensional gaussian mixtures: Where kernel methods fail and neural networks succeed. arXiv preprint arXiv: 17442 2021.
[40] Mark D Reid and Robert C Williamson. Journal of Machine Learning Research, 12(Mar):731-817, 2011.
Information, divergence and risk for binary experiments.
[41] Daniel A Roberts, Sho Yaida, and Boris Hanin. The principles of deep learning theory. arXiv preprint arXiv:2106. 10165, 2021.
[42] Robert Serfling. Approximation theorems of mathematical statistics, 1981.
[43] Bharath Sriperumbudur, Kenji Fukumizu, Arthur Gretton, Bernhard Schölkopf. and Gert RG Lanckriet. On integral probability metrics, d--divergences and binary classification. arXiv preprint arXiv:0901.2696 2009.
[44] Dougal J Sutherland, Hsiao-Yi Yu Tung, Heiko Strathmann, Soumyajit De, Aaditya Ramdas, Alex Smola, and Arthur Gretton. Generative models and model criticism via optimized maximum mean discrepancy. arXiv preprint arXiv: 04488, 2016.
[45] Tom Tirer, Joan Bruna, and Raja Giryes. Kernel-based smoothness analysis of residual networks arXiv preprini arXiv:2009. 10008, 2020.
[46] Laurens Van Der Maaten. Accelerating t-sne using tree-based algorithms. The Journal of Machine Learning Research, 15(1):3221-322 2014.
[47] Liyan Xie, Shaofeng Zou, Yao Xie, and Venugopal V Veeravalli. Sequential (quickest) change detection: Classical results and new directions. IEEE Journal on Selected Areas in Information Theory, 2021.
[48] Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over- parameterized deep relu networks. Machine Learning, 109(3)4467-492, 2020.
13


Reviews:
[{"review_id": "h09zAv9tjcj", "reviewer": "Reviewer_Kznb", "summary": "The paper addresses two-sample testing with a neural-network-based maximum mean discrepancy. \nThe paper essentially considers two types of test statistics: (a) MMD defined by a neural tangent kernel (NTK), and (b) mean differences based on neural network (NN) features. \nFor two-sample testing, the paper seems to consider the second kind only; the test is specifically a mean difference test, where data features are given by a NN trained on a held-out dataset. ", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Overview\nI appreciate the authors’ effort to develop a fast, high-power two-sample test with NN features. \nThe paper, however, needs improvement. In particular, the proposed tests lack empirical and theoretical justifications at multiple levels. \n\nAs described in Summary, the authors consider the mean-difference test (the test stat is defined in Eq. (17)) in their two-sample test experiments. \nWhile the test is promising and could be improved with a more elaborate NN, the authors failed to show its utility in their experiments; We can see that the performance is comparable to the MMD test (Gretton et al., 2012) with the Gaussian kernel. Therefore, leaving aside its issues (discussed below) as a two-sample test, the paper is not practically compelling and not ready for publication.  \n\nFor the other statistic, defined by the NTK of a trained network, the authors did not show its merits relative to a variety of existing approaches, making it difficult to value the paper's contribution. What makes the paper not appealing is the following: \n* There is no comparison to other methods in the MNIST experiment\n* Changepoint detection; apart from the lack of competitors, the behaviour of the test looks almost identical to the Gaussian kernel MMD. \n\n# Detailed comments \n\n## On the mean difference test (the test-only bootstrap)\nIn addition to the aforementioned weakness, the test lacks theoretical comparison. \nExisting tests, such as the original MMD test and a deep kernel version  (Liu et al., 2020), are at least asymptotically guaranteed to distinguish any pair of distributions. \nOn the other hand, the proposed test lacks such property; the defined maximum mean discrepancy (or integral probability metric) is not guaranteed to separate two distributions (i.e., it is only zero if and only two distributions agree), and optimisation might get stuck in bad local optima. \nGiven that the empirical performance is not well demonstrated, it is desirable, while challenging, that the paper discusses when and how the test works better (e.g., what kind of alternatives the test is adept at detecting).\nAlso, if the paper's focus is on improving the time complexity of quadratic time MMD tests, it would be helpful to contrast the proposed approach with existing linear time tests such as ChwialkowskiJ et al., (2015); Jitkrittum et al., (2016). Currently, the paper lacks such obvious baselines. \n\nNote that the power analysis (Theorem 3.2) addresses a different test; the analysed test uses a threshold defined by a concentration inequality, whereas the proposed test derives its threshold with permutation bootstrap. It would be helpful to state the relation between them. \n\n## On the NTK test \nThe concentration analysis in Theorem 3.1 defines a test based on the NTK statistic. The analysis has the following issues: \n* Constant normalisation of the kernel. The constant B is typically unknown for the NTK;  the subsequent concentration analysis implicitly depends on this constant, and the test may not be easy to implement in practice. Note that if we normalise the kernel $K$ as $K(x, y)/(\\sqrt{K(x, x)}\\sqrt{K(x, x)})$, we would need a different proof for the approximation (Eq. (10)) of the trained NTK with the NTK at initialisation. \n* The NTK Approximation in Theorem 3.1 is not experimentally justified enough. The experiment considers a single null scenario, which makes the justification dubious. \n\nAdditionally, if the NTK test behaves like that of the NTK at initialisation without features learned,  the test’s benefit is unclear (how can it achieve higher power?)\n\n\n## On the full-bootstrap test\nThis test was not compared to existing methods. Also, the proposed bootstrap method is not theoretically justified with proof as in the foregoing test-only bootstrap test. Its behaviour under the null is also not thoroughly investigated in experiments. \n\n## Changepoint detection\nIt is unclear that the changepoint detection experiment highlights the proposed statistic. At the end of the day, we must determine if a change occurs, which requires a threshold. Merely tracking a statistic's value is not useful for this purpose. \n\n\n===== post rebuttal =====\n\nI thank the authors for their rebuttal. I misunderstood the proposed statistics at the time of writing my review, and the rebuttal clarified this. My evaluation on this work remains the same as the following questions are not satisfactorily answered: \n\n1. As two-sample tests, are the proposed tests valid and useful?\n2. Is the proposed statistic Eq. (5) useful for downstream applications?\n\nOn the first question, I doubt the validity of the test based on $\\hat{T}_{net}$. The approximation by the NTK-MMD (Proposition 2) involves an unknown constant. If we determine the threshold by the NTK-MMD, we need to take into account the deviation caused by the approximation, which requires the constant. Even when the constant in Proposition 2 is available, it depends on the network architecture and initialisation and could be large, which could result in a conservative test. As the authors did not actually use the test in their experiments (and did not report the type-I errors of the test in the added experiments), the validity of the test is dubious. \n\nSecond, of the two bootstrap tests, I will focus on the test-only bootstrap as it is the test used in their experiments. The type-I guarantee for this test is okay unlike the first test discussed above. The test is not practically compelling:\n\n* The reported results in Additional Table 3 do not show significant improvements by the test (see the linear time baseline ME).\n* Additional Table 4 does show the superiority over the linear time MMD, but somehow it does not show other linear time baselines (including the ME test), which makes the result arbitrary.\n\n\n Is the proposed statistic Eq. (5) useful for downstream applications?\nMy understanding is that the authors base their argument on\n\n1.  fast computation (linear-time)\n2. online training (can be used for stream data)\n3. learning guarantee (we can obtain globally optimal network params).\n\nThe first two points are definitely true; however, they apply to any neural network IPM like Eq. (3). It is unclear why the witness function of the form in Eq. (4) is desirable. For the third point, my impression is that this kind of argument is only applicable to shallow networks; even if we can obtain an optimal network within such a class, it is unclear how useful the witness function is (justification is not provided in the paper except for the approximation by the NTK-MMD). Apart from the optimisation consideration, overall, the paper lacks comparison in downstream tasks and fail to showcase the proposed method. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "24 hours"}, "extra_text": {}}, {"review_id": "ZnZamTCqokL", "reviewer": "Reviewer_aee6", "summary": "This paper tried to apply the Neural Tangent Kernel (NTK) as the kernel function of MMD for a two-sample test task. Setting the MMD (without square; using hypothesis space of NN as the function space of MMD) as the training objective of the NN, the paper utilized the NTK which is trained after 1 epoch as the kernel function used in MMD. Theoretically, this paper provided the approximation error analysis for the NTK in the initialization state and the NTK after a short training and gave detailed bounds of NTK-MMD's testing power by deriving the sample complexity for a given level of Type-I error. Finally, the paper conducted a lot of experiments (on artificial datasets and real-world datasets) to verify the effectiveness of NTK-MMD.", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "While the content is quite dense and solid, the paper is not that strongly motivated and lack of insights in both the theoretical and experimental parts. Some concerns are as follows.\n\n1. Although the paper provided theoretical properties of NTK-MMD in Section.3, the derived bounds about testing power seem to have no specific relation with NTK. Besides, the paper did not compare the bounds of NTK with bounds of other common kernel functions and thus gave little insight. For example, is this testing power bound of NTK better than RBF/Laplacian kernel, or what's the contrast between them? If the bound holds for general PSD kernels, then why introduce NTK?\n\n2. In the experiments, the paper only uses common finite neural networks, which is rather narrow compared to the needed infinite width that NTK/lazy training theory applies. In this narrow case, the parameters should change a lot after training, is it right?\n\nMinor issue:\nPage 4. Remark 2.2. (SGD and online training).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "NAN", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8 hours"}, "extra_text": {}}, {"review_id": "FlScvV5S48P", "reviewer": "Reviewer_tRT5", "summary": "The paper introduces a new neural network Maximum Mean Discrepancy statistic to achieve both computational and memory efficiency. The authors further characterize the statistical properties of the statistic including Type-I and Type-II errors. They also perform both simulation and real data analysis to address the efficiency of the proposed testing procedure.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The main contribution of this paper lies in developing a novel neural network MMD statistic. It introduces the advantage of deep neural networks' representation and optimization to kernel MMD testing procedure, which solves the computational and memory difficulties of using MMD statistic. The observation that training of a network in a short time nearly means computing the witness function of a kernel MMD at time zero given some specific choice of objective function is insightful.\n\nMy minor concerns come from theoretical perspective. The analysis here seems to follow Cheng and Xie (2021). The authors may discuss more on the theoretical novelty. In addition, the paper only discusses the balanced sample setting. Some discussion on unbalanced data is appreciated, even from practical side.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors adequately addressed the limitations and potential negative societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "40S9-w9GqM", "reviewer": "Reviewer_e8cj", "summary": "The paper presents a theoretical method to analyze the neural network-based two-sample test methods. In particular, it does so by identifying a correspondence between NTK and MMD", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think the idea is novel and the theory / empirical validation is solid, but I am concerned with the following two problems that I deem crucial.\n\n(a) The NTK-MMD correspondence only applies to the loss function given in Eq.3, which is both artificial (not used in practice) and ill-defined (no global convergence). A study based on such kind of a loss function cannot convince me unless it is shown that this loss function can be useful in practice.\n\n(b) The paper is also not sufficiently motivated in my opinion. In particular, I do not know how important neural network-based methods are for the problem of two-sample testing. For example, is it a very popular approach? Does it empirically achieve better performance than the traditional two-sample test methods? The paper does not say this, and I remain unconvinced about the importance of the neural network-based two-sample testing approach\n\nIn summary, I find the motivation and the setting of this paper contrived and of uncertain machine learning importance. I might consider raising the score if the authors convince me that Eq.3 is not as artificial as it seems by either showing that (1) it is useful (achieve SOTA) or that (2) standard two-sample testing loss functions can be approximated by this one (so that eq3 has some universal character to some extent) or (3) both. Problem (b) is easier to fix I guess, but the authors need to cite some relevant works to convince me (and other readers) that neural network-based approaches are indeed an important approach", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Sufficiently addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1 hour"}, "extra_text": {}}]

Author rebuttal:

Meta-review: In this paper, the authors proposed a theoretical method to analyze the neural network-based two-sample test methods. Though the paper is dense and solid, the motivation is not very clear. Some reviewers pointed out that theoretical insight was insufficient and numerical experiments were not convincing to show the effectiveness of the proposed method. Following the careful discussion, the reviewers agree that the paper is not yet ready for publication. 
Decision: Accept (Poster)