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/ (1 (F) (yddd
1 XX
1 ny
p p:=
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 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, e.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 (()) = - 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 (4) 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) (0-())))))))))))))))())) (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 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
(because global constant normalization of the kernel
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'ddd
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 1A,
(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(e (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 (322ppoveess 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( (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 || 2 = 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 O 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///yhheeggg - 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 of 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
180
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 indler
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 Pdata= 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-183021
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:1906.044 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 8139-8144 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 120.: 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-Zanne
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 1717-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": "sMf7z8173di", "reviewer": "Reviewer_AatE", "summary": "\n\nIn this paper, the authors propose an NTK-MMD to perform NTK based two-sample test.   A \"short-time trained\" neural network is used to approximate the NTK-MMD statistic. ", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "\n\nPros.\n1.  The idea of building a connection between kernel MMD and neural networks based on NTK is interesting.  It can reduce the time complexity of kernel MMD and take advantage of network training.\n2. The paper is well organized and well written. \n\n\nCons.\n\n1.  It seems that the proofs of Theorem 3.1-3.3 follow from the techniques in (Cheng et al. 2021), which does not rely on NTK too much.  I am not sure what is the role of NTK here.  What is the key property of NTK needed to prove Theorem 3.1-3.3 ? \n2. In the experiments, ReLU activation function is used.  Does the ReLU-NTK satisfy the assumption?\n3. The analysis of the SGD  training actually is a one-pass of the whole training set. Does the result hold for actual SGD?   \n4. The paper focuses on the short-time training regime for approximation. I am not sure what is the exact range of the short-time period.  The authors claim a O(t) approximation error. It is better to empirically evaluate the approximation error to show the range of short-time periods.\n\n\nMinor:\nTypos.  definiation of v_{pp} in Eq.(13).\n\n\nCheng et al. Kernel MMD Two-Sample Tests for Manifold Data, 2021\n\n\n=============== POST-REBUTTAL ==================\n\nThanks for the authors providing a detailed response.  Most of my concerns have been addressed.      Now I understand how NTK is involved in bridging the empirical computed test statistic and exact NTK-MMD.   Considering this,   I decided to increase my score. \n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations have been stated in section 5.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "dn9YNsFgbka", "reviewer": "Reviewer_QV5c", "summary": "This paper proposes to instantiate the Maximum Mean Discrepancy (MMD) with a specific type of kernels: Neural Tangent Kernels (NTKs), which are kernels associated to neural architectures. By linking the NTK MMD and its witness function to the training of a neural network over the integral probability metric loss, the authors find that they can both be approximated using neural networks, thereby alleviating the high computational cost of computing kernel matrices.\n\nThe authors then investigate the testing power (to distinguish between two sampled distributions) of this statistic with theoretical guarantees and experimental studies. In particular, a specific bootstrap method is proposed in order to choose the statistical test threshold without retraining the neural network witness.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "## Contribution\n\nThe proposed idea of leveraging NTKs for the MMD statistic is relevant and valuable to the community. This contribution is interesting in several ways. As well motivated in the paper, this choice may enable computational power and memory savings thanks to the approximation of the NTK MMD by a simple neural network training, and NTKs might be interesting kernel choices for complex high-dimensional data due to their direct link with neural networks. Furthermore, this paper contributes to the study of the capabilities of NTKs, which, despite extensive theoretical studies and some successful applications since their recent introduction, have never been studied under the context of MMD.\n\nThe introduction of the NTK MMD and its approximation using neural network training is well explained, simple and sound. The quantification of the approximation that is found to be linear in the training time supports the leveraged approximation, all the more so as similar results for SGD-based optimization are provided as well.\n\nWhile some other applications of MMD might have been considered, the authors choose to study the power of the NTK MMD statistic to distinguish between two sampled distributions, which is a standard and reasonable choice. The derived theoretical results regarding Type I and II errors are, to my knowledge, novel and valuable, more particularly regarding the proposed test-only bootstrap scheme which avoids to retrain the neural network for threshold calibration. Even though, to the best of my understanding, these results are not specific to NTKs from the proof in the appendix, the progress seems significant enough, especially since the properties of NTKs largely remain misunderstood.\n\nWhile diverse and presenting satisfying results, the experiments unfortunately do not completely manage to show the advantages of the NTK MMD statistic. The performance advantage of the NTK MMD statistic with respect to the Gaussian baseline does not seem significant, and the potential computational advantages of the approximate NTK MMD are not illustrated nor mentioned. Tackling these issues would strengthen the relevance of NTK MMD and further support the paper's motivation. I would in particular suggest the authors to investigate the following two potential improvement leads:\n - changing the neural network architecture (and thus its respective NTK): the considered architectures only feature one hidden layer and use the softplus activation, whereas standard architectures are deeper and use ReLU-like activations, which may improve networks and NTKs performance;\n - comparing NTK MMD with RBF MMD on complex high-dimensional data, to show the relevance of NTKs and the potential computational advantages of the approximate NTK MMD (e.g. on the MNIST experiment in the paper for which the authors consider no baseline).\n\nFinally, since linear-time MMD statistics have already been proposed in the literature (e.g. in [1, Section 6] and [2]), the authors might consider comparing, or at least discussing, the computational or efficiency advantages of the proposed NTK MMD with respect to this line of research.\n\n[1] A. Gretton et al. A Kernel Two-Sample Test. JMLR 13(25):723−773, 2012.\\\n[2] W. Jitkrittum et al. Interpretable Distribution Features with Maximum Testing Power. NIPS 2016.\n\n\n## Writing\n\nThe paper is well written, clear and easy to read.\n\nI would only ask the authors to improve the readability of Figure 4 (right) for color-blind people or grayscale printers.\n\n\n## Other Remarks & Questions\n\n### Biased MMD estimator\n\nCould the authors explain why the considered squared MMD estimator is the biased one, whereas an unbiased estimator could be used?\n\n### Training with Batch 1\n\nIs there a reason why the authors choose to train neural networks with a batch size of only 1, when a larger batch size might have been possible?\n\n### Infinite-Width Neural Networks\n\nThe authors do not mention the canonical context when dealing with NTKs, which is the infinite-width limit of the trained neural networks. In this setting, the NTK remains constant during training under mild conditions, and can be analytically computed for some architectures (e.g. thanks to [3]). I would guess that increasing the width of the hidden layers may bring the approximate NTK MMD closer to its exact expression, but I would be interested to hear the author's thoughts on this matter.\n\n[3] R. Novak et al. Neural Tangents: Fast and Easy Infinite Neural Networks in Python. ICLR 2020.\n\n### Typos\n\n - Lines 19-20: \"a good kernel function high dimensional data\" should be \"a good kernel function for high dimensional data\".\n - Line 105: \"witnesses function\" should be \"witness function\".\n - There is an extra period line 137.\n - The sentence in lines 218-219 seems grammatically incorrect.\n - Line 330: the window size can also be 200, as in Figure 4.\n\n\n## Overall Sentiment\n\nBy considering the interesting and valuable contribution of the paper as well as the limitations of its experiments, I vote for a \"Marginally above the acceptance threshold\". However, I believe that the described limitations can be tackled during the rebuttal, and would be pleased to increase my score in the case of a positive feedback.\n\n***\n\n## Post-Rebuttal Update\n\nI acknowledge the authors' response and thank them for their answers. The response is thorough, tackles the raised concerns and provides the suggested new material. I believe that this will make the submission substantially stronger. Therefore, as far as I am concerned, I choose to increase my score to a \"Good paper, accept\".", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations are adequately addressed, but the societal impact is not discussed by the authors. I would suggest to briefly evaluate such impact, for examples in GAN applications.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "8"}, "extra_text": {}}, {"review_id": "UG-rj_Yq-Vg", "reviewer": "Reviewer_jdVg", "summary": "Two-sample testing is a fundamental problem in statistics and machine learning. Recently, the network-based MMD test has shown its great power in complex datasets and real-world detections. To track this research line, this paper presents a novel NTK-based MMD statistic by identifying a new connection between NTK and the MMD statistic. This connection enables developing 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, this connection allows us to understand the NTK test statistic properties by adapting existing theories for kernel MMD.", "questions": "", "limitations": "", "rating": 6, "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": "Pros:\n\n1. Combining NTK and MMD is a very interesting topic. The proposed test statistic can be computed by a short-time training of a 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 linear in training time, which leads to the construction of the NTK-MMD test statistic.\n\n2. The statistical properties of the NTK-MMD, including the Type-I error and the testing power, establish the conditions under which the test is powerful.\n\n3. Online human activity change-point detection is a very interesting application for the field.\n\nCons:\n\n1. My major concerns include the weak baselines. The selected baselines are not powerful. MMD with Gaussian kernel (median distance as bandwidth) works poorly in some kinds of data. Thus, I do not know the actual test power of the proposed test. The MMD test with the optimized kernel, ME test, SCF test, C2ST tests should be baselines. This can help us understand the test power of the proposed statistic well. As I know, [34] provides some benchmark datasets. One benchmark dataset can be used to validate the test power of the proposed method. Mean-shift and covariance-shift are not representative to validate the test power of tests. For example, gmmd can detect the mean shift but not the shift in the HDGM dataset introduced in [34].\n\n2. It is not clear why the NTK-MMD is computationally efficient and memory-efficient for the readers not familiar with NTK. More explanations are required. \n\n3. It would be better to introduce Eq. (9) after Lemma 2.2.\n\nIn general, I would be happy to accept this paper if more experiments can be used to show the test power of the proposed test.\n\n-----------------------------------------\nThe authors have addressed my concerns well. I tend to accept this paper now.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have adequately addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"ethical_concerns": "No.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7hrs"}, "extra_text": {}}, {"review_id": "Ed2hbOC14Ll", "reviewer": "Reviewer_VLkd", "summary": "The paper proposes to use neural tangent kernel (NTK) to approximate a neural-network maximum mean discrepancy (MMD) metric where the kernel is defined by a neural network. The paper characterizes various properties of the NTK MMD metric, including approximation bound to the finite-width neural network MMD, NTK MMD's test power, and threshold calibration. Numerical results on simulation data confirmed that NTK MMD has similar test power compared to 2 baselines. Results on real data demonstrate NTK MMD's real-world applications including change point detection.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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": "In general, this is a well-written submission with nice theoretical characterizations and empirical results. \n\nI have a slight confusion about the extend to which the network needs to be trained (perhaps there is something obvious I missed). According to the approximation error, the longer the network is trained, the worse the approximation. Intuitively, the less you train the network, the better the approximation. In the extreme, for example, maybe using just 1 or 5 data points would give better approximation than using an entire dataset? There seems to be a trade-off in the number of data points used to train the network, the approximation error, and the empirical performance. It would be great if the authors could clarify (or point to the sections in the paper that involves this discussion).\n\nAnother minor concern: the assumption in line 115 that $K_0\\approx \\hat{K}_t$ seems to be very important. It would be great if the authors could comment the conditions under which this assumption will hold (in addition to a small step size). I was under the impression that various factors, including the network width/depth, the complexity of the data, the network weight initialization scheme, among others, could impact the validity of this assumption.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "It would be great if the authors could discuss the validity of the assumptions involved. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes to use NTK to MMD, which connects MMD and neural networks.  Based on this link, the paper proposes a way of solving the high computational cost of kernel methods by approximation using neural networks.  The experimental comparisons could be improved, but generally acceptable.  Overall, the connection between NTK/NN and MMD is a new direction and the work is worth presented in NeurIPS.
Decision: Accept (Poster)