Neural Tangent Kernel Maximum Mean Discrepancy

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.


Reviews and author responses:



In 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.   

Pros.
1.  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.
2. The paper is well organized and well written. 


Cons.

1.  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 ? 
2. In the experiments, ReLU activation function is used.  Does the ReLU-NTK satisfy the assumption?
3. The analysis of the SGD  training actually is a one-pass of the whole training set. Does the result hold for actual SGD?   
4. 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.


Minor:
Typos.  definiation of v_{pp} in Eq.(13).


Cheng et al. Kernel MMD Two-Sample Tests for Manifold Data, 2021


=============== POST-REBUTTAL ==================

Thanks 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. 

 The limitations have been stated in section 5.

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.

The 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.  ## Contribution

The 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.

The 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.

While 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.

While 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:
 - 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;
 - 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).

Finally, 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.

[1] A. Gretton et al. A Kernel Two-Sample Test. JMLR 13(25):723−773, 2012.\
[2] W. Jitkrittum et al. Interpretable Distribution Features with Maximum Testing Power. NIPS 2016.


## Writing

The paper is well written, clear and easy to read.

I would only ask the authors to improve the readability of Figure 4 (right) for color-blind people or grayscale printers.


## Other Remarks & Questions

### Biased MMD estimator

Could the authors explain why the considered squared MMD estimator is the biased one, whereas an unbiased estimator could be used?

### Training with Batch 1

Is 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?

### Infinite-Width Neural Networks

The 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.

[3] R. Novak et al. Neural Tangents: Fast and Easy Infinite Neural Networks in Python. ICLR 2020.

### Typos

 - Lines 19-20: "a good kernel function high dimensional data" should be "a good kernel function for high dimensional data".
 - Line 105: "witnesses function" should be "witness function".
 - There is an extra period line 137.
 - The sentence in lines 218-219 seems grammatically incorrect.
 - Line 330: the window size can also be 200, as in Figure 4.


## Overall Sentiment

By 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.

***

## Post-Rebuttal Update

I 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". 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.

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.  Pros:

1. 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.

2. 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.

3. Online human activity change-point detection is a very interesting application for the field.

Cons:

1. 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].

2. 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. 

3. It would be better to introduce Eq. (9) after Lemma 2.2.

In general, I would be happy to accept this paper if more experiments can be used to show the test power of the proposed test.

-----------------------------------------
The authors have addressed my concerns well. I tend to accept this paper now.
 The authors have adequately addressed the limitations and potential negative societal impact of their work.

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.  In general, this is a well-written submission with nice theoretical characterizations and empirical results. 

I 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).

Another 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. It would be great if the authors could discuss the validity of the assumptions involved. 

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