Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper studies an in-depth analysis of runtime and memory requirements for computing the finite-width NTK. The authors analyze computing costs of Jacobian-vector products (and vice versa) for both fully-connected and convolutional neural networks. They also improve the NTK computation cost by leveraging the structure of neural networks resulting in drastic speedup and memory saving. Finally, they make all their implementations open-source based on the JAX library.
main_review: - This paper is well-written and easy to understand. The novelty of this work is fairly weak. All results come from simple computations of automatic differentiation, and the reduced cost of NTK follows from amortizing simple linear operations, which is not surprising at all. However, these weaknesses might be covered by their open-source implementation, as it can be significantly important and useful to future works.

- The infinite-width NTK can be exactly computed without Jacobians (Arora et al., 2019) and it can be much efficient than Jacobian-based approaches. Does the finite-width NTK have concrete benefits compared to the infinite-width NTK?

- Minor Issue:
  - It would be great if more details of Automatic Differentiation (AD) operations to derive memory costs of JVP/VJP are provided.
  - As assumed $O=\mathcal{O}(LW)$, the time and memory costs in Section 3.2~3.4 can be reduced without $O$.
  - The method names in Figure 1 are too small to identify. It would be great to make larger  font size.


summary_of_the_review: The review score is all about the battle of lack of novelty versus the impact of implementations. I judge this work to be under the bar of acceptance for now but am willing to raise it depending on the author's feedback on the importance of computing finite width NTK.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper studies the practical compute and memory requirements to computing the Neural Tangent Kernel (NTK), introducing two new approaches to doing so for standard NN primitives (structured derivative and NTK-vector product) which each have advantages (in terms of variables like batch size, output dim) over the naive jacobian contraction method. The authors provide experiments demonstrating the advantages of their approaches across a range of architectures and hardware. The authors provide open-source code which seems to be integrated neatly with the JAX and Neural Tangents frameworks.
main_review: Strengths:
- Computing the NTK is a popular area at the moment and the open-source code should make some experiments possible that were not previously for some researchers.
- The approaches proposed exploits nice tricks relating to autodifferentiation (like jvps and vjps) and structure in the computation graphs, which seems tailor made for JAX.
- The experiments are convincing are back up the author's analysis of the advantages of their methods
- The paper is generally very well written and clear, I especially like the colour-coded variables like N,O,P etc

Weaknesses (most of these are not huge or beyond the current scope of the paper):
- The actual technical contribution is not huge, and the paper is more about implementing known ideas from autodiff, linear algebra, and particular aspects of NN structure.
- While improvements on previous implementations is great (and i suspect is now optimal?), it seems these improvements are still far more expensive that standard sgd training (quadratic in batch size, and output size vs linear), which will hamper the impact and usage of this work. For example, 1 second to compute the NTK of a resnet 18 on a v100 for a single pair of imagenet inputs is not ideal. Also is average pooling applied here? I recall average pooling making computational requirements very tricky for CNNs. This may seem a bit unfair a criticism but I feel it is valid given that this is a code paper.
- While it is true to JVP and vjp are on the same order of complexity as forward passes, they are strictly more expensive (i believe a jvp is on the order of 3 forward passes), so while the orders you provide e.g. in Table 1,2 are true, in practice the methods will translate to be slightly slower that simply taking the values in Table 1,2 as gospel. I think it would be good to mention this.
- I think JAX is nice, but wondering if the authors also be implementing their approaches in other frameworks like PyTorch? A lot of non-google/deepmind ppl I know prefer PyTorch to a framework like JAX, and so I believe the impact of the proposed approaches would be greater if so. I realise there are properties with the way that autograd is written in pytorch that makes this more difficult, but I feel like an implementation in just JAX serves to benefit a subset of the research community not the whole, particularly one with a commercial motive.

Minor points:
- The legend in fig 1 left blocks the graph
- When you write mathematically in the NTK in the abstract and eq 1, it's really not clear what f, theta, x are? Nor what dimensions that have. I know it's standard notation but some may not have seen this. You can probably get away with it in the abstract for space, but in eq 1 you should spell these out imo.

summary_of_the_review: Nice code which should be used by some researchers. Using methods/ideas that are established. Concerns that calculating NTK is still too expensive (which it is by nature, not fault of this paper) to prohibit some researchers and also that uptake in using the code will be only by a subset of the community.
review_title: null

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: Experiments are convincing. The provided code is helpful for researchers who need a fast computation of NTK. Though the ideological (mathematical part) is very simple.
main_review: The paper describes three methods for computing the neural tangent kernel for a given a batch of vectors and a feed-forward neural network. The first method, Jacobian contraction is quite straightforward and is not suggested for use. The second one,  NTK-vector products is based on reducing the task to the computation of Jacobian-vector and vector-Jacobian products (which themselves can be computed in a fashion similar to forward pass). The method is suggested for use for networks whose width is smaller than the dimension of the output space (probably for contractive autoencoders). The third method, Structured derivatives, is based on a simple identity (9) that simplifies the computation of NTK over weights of the same layer. It turns out that this method is preferable for shallow networks or the typical sitution when  width is larger than the dimension of output space.

Experiments are convincing. The provided code is helpful for researchers who need a fast computation of NTK. Though the ideological (mathematical part) is very simple.
summary_of_the_review: I would not recommend the paper for publication at such a top-tier conference as ICLR. The simplicity could be justified, in principle, by the novelty. But this paper is too simple, and the novelty is moderate.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This work aims to solve the computation problem of the finite-width neural tangent kernel (NTK), which is a central object in deep learning. The authors analyze the computation and memory requirements for finite-width NTK and propose two novel algorithms that can improve efficiency. Open-source has been provided by the authors.
main_review: Strengths

(1) This work provides a comprehensive study on various ways to compute the finite-width NTK. In particular, the time cost and memory cost are given for each case, providing a detailed reference for researchers in related fields.

(2) The author makes the article easy to understand through careful structural planning.

Weakness:

(1) I am confused with the setting that ${\rm O} = O( {\rm L W})$, for which the authors demonstrated that the number of logits is dominated by the product of width and depth.  If I understand correctly, $\rm O$ is the number of neurons in the last layer of the neural network. Thus it should be independent with ${\rm W}$ and ${\rm L}$. 

(2) The dashed lines in Figure 1 are hard to distingish, I suggest the authros could make it clearer. 
summary_of_the_review: Overall, this article is clearly written and well structured. If the author can solve the problem of unclear dotted line in Figure 1, it will be more perfect.

Besides, this work builds on JAX. However, the most popular deep learning framework used is PyTorch and Tensorflow. It would be better if the authors discuss why they choose JAX instead of other frameworks.
review_title: null

Meta-review: This is an interesting and carefully-presented work which discusses how to implement finite-width NTKs more efficiently.  Overall, the reviews were slightly tending positive, though with a variety of concerns, including some concern that the contribution is not sufficiently substantial.  In my own perusal of the paper, personally I feel it could be made more compelling if (a) more speedups could be considered, including ones with various tradeoffs, for instance via randomized linear algebra, (b) explicit consequences on various prediction tasks, rather than plotting wall-clock times (e.g., as this paper cites many works which tried to use finite-width NTK, and as this paper claims massive speedups, then it will be able to repeat some of those experiments at much larger sizes, which should lead to interesting and valuable larger-scale experiments which ideally have some new phenomena, but are even interesting if they simply confirm the smaller-scale phenomena).  As a separate concern, I second the comments of one reviewer, that part of this paper's contribution is to a single software package, which is moreover listed in the abstract (and not just part of the standard code release, e.g., as a footnote); this feels a little strange, like an announcement of a code release, and further limits the impact to general machine learning researchers (for instance, I feel completing some of my preceding suggestions could result in, say, researchers who use other software feeling eager to re-implement this).  Overall, I urge the authors to continue with their interesting work and aim to resolve these concerns and those of the reviewers.
Decision: Reject