Venue: ICLR 2020 Conference

Reviews:

final_score: rating: 3: Weak Reject
review_content: Summary: A Jax based neural tangents kernel library is introduced, with native GPU, TPU, and XLA support. Due to the correspondences with infinite neural network kernels (NNGPs), these kernels are also able to be computed for (essentially) free. Layers which do not emit an analytical form (e.g. Tanh or Softplus) can be implemented using Monte Carlo estimates. Several engineering-based experiments are performed demonstrating the potential scalability of their library.

While I really enjoyed reading the paper and believe that this library could be extremely practically useful, I vote to reject this paper because I do not feel that it has sufficient novelty to be a paper on its own in light of Lee et al, 2019. 

Edit: post rebuttal, I'm bumping my score to a weak reject but would have minimal qualms if this paper were to be accepted. I find the further experiments performed by the authors of very good quality overall, but I'm still not particularly satisfied by their `argument that the codebase itself is distinct enough from separate related work. It's unfortunately a bit hard, from a machine learning researcher side, to review the quality of a codebase in and of itself. 

Significance: Having played around with the code a bit, I find that the library itself is of very high quality and is pretty straightforward to use. I could definitely see myself using this library in the future for research work.

However, my primary concern with this paper is that it’s not sufficiently distinct from the previous work of Lee et al, 2019. After all, most of the experiments in that paper would have required the type of implementation that is described in greater detail in this paper. 

To be able to vote to accept this paper, I will have to see an experiment that is practically performed with the current library in order to distinguish it from previous work (specifically Lee et al, 2019). In recent work, Arora et al, 2019 (Note: I do not consider this reference in my review other than to be mentioned as an example of an experiment that could be run with your library) run neural tangent kernels on tabular data using kernel SVMs. One other potential example would be a kernel SVM in this manner on CIFAR-10. An alternative example would be to exploit the Gaussian process representation and test out both NTKs and NNGPs in comparison to standard kernels for GPs and NNs on UCI regression tasks.

Originality: Again, a very efficient and easy to use implementation of neural tangent kernels would be a great boost to the community. This is doubly so as Jax is easy and pretty straightforward to use and is quite numpy like. 

Again, I am very concerned with originality in comparison to Lee et al, 2019. Even checking out the link to their codebase provides a github repo that is quite similar to this one. Given that ICLR is a venue of similar domain to NeurIPS, it’s not clear to me why this paper ought to be anything other than a separate supporting tech report. If this paper had been submitted to something like SysML (edit: or JMLR MLOSS), I would see the distinctness instead.

Clarity: I find the paper to be extremely well-written and easy to follow. The addition of code snippets throughout is very well done, even if it’s a bit overkill. I don’t know what adding a half page long description of an infinitely wide WideResNet adds to the paper when that space could be better used by another experiment. 

Quality: I find the experiments performed to be very well constructed. Below are a few mostly minor comments on the experiments:

In Figure 1 on the right, I would have liked to have seen the posterior predictive for a NNGP with the same kernel as well. 

In Figure 2, why is the NNGP slower to converge to the analytic values here? Obviously, the rates of convergence are the same, but the constants seem different.

In Figure 3 (and throughout the experiments), does “full Bayesian inference for CIFAR10” mean that you treated the classification labels as regression targets? If so, how was classification error measured.

In Section 3.1, you mention that the library “leverages block-diagonal structure” to invert the CIFAR10 covariance matrix for each class (still 50k x 50k). Possibly this is because I haven’t had the chance to use TPUs, but I’m currently struggling to see how one could form and invert (via Choleskys) matrices of this size (50k x 50k) on a standard GPU (or CPU). Could the authors please clarify how they did this (whether through iterative methods, another structure exploiting trick, lots of memory, etc.)?

second edit: I was also unable to respond to the final comment about the UCI experiments in its own comment, but thank you for providing the estimated depths. These results definitely show the potential software promise of the codebase and open some interesting new research questions as a result.

References:

Arora, S., et al., Harnessing the Power of Infinitely Wide Deep Nets on Small-data Tasks, https://arxiv.org/abs/1910.01663

Lee, J., et al., Wide Neural Networks of any Depth Evolve as Linear Models Under Gradient Descent, NeurIPS, 2019, https://arxiv.org/abs/1902.06720

review_title: Official Blind Review #3

final_score: rating: 8: Accept
review_content: POST-REBUTTAL COMMENTS

I appreciate the response from the authors. 

I particularly like the comparison table in the response to the other reviewer and ought to be highlighted in the paper.

If I were to start this line of research, I would be inclined to expand on the codebase. The contribution is significant. Hence, I am bumping up my score to accept.


PRIOR FEEDBACK

The contribution of this work lies in providing a library for working with the existing variants of infinite-width neural networks and avoiding the need to derive the NNGP and NT kernels for each architecture by hand. The authors have firstly shown performance comparisons between inferences between finite vs. infinitely wide neural networks. The authors then go into some implementation details with their library. The authors have provided the code and cookbook in the links provided in the abstract. On the overall, I like this effort which is timely.

Some additional suggestions below:

I would like to see an additional metric for performance comparison of probabilistic models, which is often used in the GP literature: mean negative log probability.

It would also be interesting to see how the posterior variance (e.g., Fig. 1 right) evolves over the entire space during training. 

I would have preferred a more detailed discussion about the implementation on transforming tensor ops to kernel ops in Section 3.

For the summary of contributions, can you give the corresponding section number to refer to when you demonstrate each feature? For example, is the 4th feature (i.e., exploring weight space perspective) demonstrated in the paper?

Can the authors elaborate on the ease of expanding their library for the new developments in this field?


Minor issues:

Page 1: Gaussian Procesesses?
Page 4: it’s infinite?
Fig. 4: I would have preferred the indices to be placed as subscripts instead of superscripts.
Page 8: it’s order of dimensions?
review_title: Official Blind Review #1

final_score: rating: 6: Weak Accept
review_content: This work develops a library for working with a class of infinitely wide neural networks, in particular those corresponding to neural tangent kernels (NTKs) and neural network Gaussian processes (NNGPs). The theory for these two kernels was well developed in a series of recent papers, and this library provides an automatic way to transform any appropriate neural net architecture into its corresponding NTK and NNGP.

Infinitely wide neural networks have been a popular subject of theoretical research and have been observed to have highly nontrivial performance on a variety of tasks (e.g. CIFAR-10 classification). It's really nice to see the development of such a library, which I believe could benefit the deep learning community a lot, especially for theoretical research on NTK.

I appreciate this work a lot. Currently I can only give weak accept instead of accept for a couple of reasons:
1. The theory and formulae of NTKs and NNGPs were well developed. This work mostly consists of implementing and modularizing them. The research contribution is relatively low.
2. As commented in the paper and if I understand correctly, the current library cannot scale to large datasets for CNNs with pooling. This would make the computation much more expensive (and probably infeasible without additional techniques and huge computing power) as mentioned in [Novak et al. 2019] and [Arora et al. 2019]. However pooling seems extremely useful for NTKs and NNGPs on image datasets. I think this makes this work somewhat less exciting than it may sound.
review_title: Official Blind Review #2

Meta-review: This paper presents a software library for dealing with neural networks either in the (usual) finite limit or in the infinite limit. The latter is obtained by using the Neural Tangent Kernel theory. 

There is variance in the reviewers' scores, however there has also been quite a lot of discussion, which has been facilitated by the authors' elaborate rebuttal. The main points in favor and against are clear: on the positive side, the library is demonstrated well (especially after rebuttal) and is equipped with desirable properties such as usage of GPU/TPU, scalability etc. On the other hand, a lot of the key insights build heavily on prior work of Lee et al, 2019. However, judging novelty when it comes to a software paper is more tricky to do, especially given that not many such papers appear in ICLR and therefore calibration is difficult. This has been discussed among reviewers. 

It would help if some further theoretical insights were included in this paper; these insights could come by working backwards from the implementation (i.e. what more can we learn about infinite width networks now that we can experiment easily with them?).

Overall, this paper should still be of interest to the ICLR community.

Decision: Accept (Spotlight)