Fast Finite Width Neural Tangent Kernel
Authors: Roman Novak
Jascha Sohl-Dickstein
Samuel Stern Schoenholz
Venue: ICLR 2022 Submitted (2022)
Keywords: Neural Tangent Kernel
NTK
Finite Width
Fast
Algorithm
JAX
Jacobian
Software

Abstract: The Neural Tangent Kernel (NTK), defined as the outer product of the neural network (NN) Jacobians, $\Theta_\theta(x_1, x_2) = \left[\partial f(\theta, x_1)\big/\partial \theta\right] \left[\partial f(\theta, x_2)\big/\partial \theta\right]^T$, has emerged as a central object of study in deep learning. In the infinite width limit, the NTK can sometimes be computed analytically and is useful for understanding training and generalization of NN architectures. At finite widths, the NTK is also used to better initialize NNs, compare the conditioning across models, perform architecture search, and do meta-learning. Unfortunately, the finite-width NTK is notoriously expensive to compute, which severely limits its practical utility. 

We perform the first in-depth analysis of the compute and memory requirements for NTK computation in finite width networks. 
Leveraging the structure of neural networks, we further propose two novel algorithms that change the exponent of the compute and memory requirements of the finite width NTK, dramatically improving efficiency.

We open-source (https://github.com/iclr2022anon/fast_finite_width_ntk) our two algorithms as general-purpose JAX function transformations that apply to any differentiable computation (convolutions, attention, recurrence, etc.) and introduce no new hyper-parameters.


Under review as a conference paper at ICLR 2022
FAST FINITE WIDTH NEURAL TANGENT KERNEL
Anonymous authors
Paper under double-blind review
ABSTRACT
The Neural Tangent Kernel (NTK), defined as the outer product of the neural (1)
network (NN) Jacobians, Op(21, x2) emerged as a central object of study in deep learning. the infinite width limit, the NTK can sometimes be computed analytically and ís useful for understanding training and generalization of NN architectures. At finite widths, the NTK is also used to better initialize NNs, compare the conditioning across models, perform architecture search, and do meta-learning. Unfortunately, the finite width NTK is notoriously expensive compute, which severely limits its practical utility.
x2)/00 has
We perform the first in-depth analysis of the compute and memory requirements for NTK computation in finite width networks. Leveraging the structure of neu- ral networks, we further propose two novel algorithms that change the exponent of the compute and memory requirements of the finite width NTK, dramatically improving efficiency.
senitt_width_nng our two algo- rithms as general-purpost JAX function transformations that apply to any differ- entiable computation (convolutions, attention, recurrence, etc.) and introduce no new papeerr-aramters...
We open-source
INTRODUCTION
The past few years have seen significant progress towards theoretical foundation for deep learning. Much of this work has focused on understanding the properties of random functions in high dimen- sions. One significant line of work (Neal, 1994; Lee et al., 2018: Matthews et al., 2018; Novak et al., 2019; Garriga-Alonso
al., 2019; Hron et al., 2020; Yang, 2019) established that in the limit of infi- nite width, randomly initialized Neural Networks (NNs) are Gaussian Processes (called the NNGP). Building on this development, Jacot et al. (2018) showed that in function space the dynamics under gradient descent could be computed analytically using the so-called Neural Tangent Kernel (NTK) and Lee et al. (2019) showed that wide neural networks reduce to their linearization in weight space throughout training. A related set of results (Belkin et al., 2019; Spigler al., 2019) showed that the ubiquitous bias-variance decomposition breaks down as high-dimensional models enter the so- called interpolating regime. Together these results describe learning in the infinite width limit and help explain the impressive generalization capabilities of NNs.
Insights from the wide network limit have had significant practical impact. The conditioning of the NTK has been shown to significantly impact trainability and generalization in NNs (Schoenholz al., 2017; Xiao et al., 2018; 2020). This notion inspired initialization schemes like Fixup (Zhang et al., 2019), Metalnit (Dauphin & Schoenholz, 2019), and Normalizer Free networks (Brock et al., 2021a;b) and has enabled efficient neural architecture search (Park et al., 2020; Chen et al., 202 The NTK has additionally given insight into a wide range of phenomena such as: behavior of Gen- erative Adversarial Networks (Franceschi et al., 2021), neural scaling laws (Bahri et al., 2021), and neural irradiance fields (Tancik et al., 2020). Kernel regression using the NTK has further enabled strong performance on small datasets (Arora et al., 2020), and applications such as dataset distilla- tion (Nguyen al., 2020; 2021) and uncertainty prediction (He et al., 2020; Adlam al., 2020).
Despite the significant promise of theory based on the NTK, computing the NTK in practice is chal- lenging. In the infinite width limit, the NTK can sometimes be computed analytically. However, it remains intractable for many architectures, and finite width corrections can be important to describe actual NNs used in practice. The NTK matrix can be computed for finite width networks as the outer

---

Under review as a conference paper at ICLR 2022
product of Jacobians using forward or reverse mode automatic differentiation (AD),
= [a.(0,z)//00) (21,)
x20000))
(I)
Ox0
OxP
PxO
where f is the forward pass NN function producing outputs in R 0 I' are all trainable parame- ters, and 21 and T2 are two inputs to the network. If inputs are batches of sizes N1 and N2, the NTK is an N1 O N₂O matrix.
Unfortunately, evaluating Eq. (1) is often infeasible due to time and memory requirements.
In this paper, perform the first in-depth analysis of the compute and memory requirements for the NTK as in Eq. (1). Noting that forward and reverse mode AD are two extremes of a wide range of AD strategies (Naumann, 2004; 2008), we explore other methods for computing the NTK leveraging the structure of NNs used in practice. We propose two novel methods for computing the NTK that exploit different orderings of the computation. We describe the compute and memory requirements of our techniques in fully-connected (FCN) and convolutional (CNN) settings, and show that one asymptotically more efficient in both settings. We compute the NTK over wide range of NN architectures and demonstrate that these improvements are robust in practice. We open-source implementations of both methods as JAX function transformations.
RELATED WORK
The finite width NTK (denoted as simply NTK throughout this work) has been used extensively in many recent works, but to our knowledge implementation details and compute costs were rarely made public. Below we draw comparison to some of these works, but we stress that it only serves as sanity check to make sure our contribution is valuable relative to the scale of problems that have been attempted (none of these works had efficient NTK computation as their central goal).
In order to compare performance of models based on the NTK and the infinite width NTK, Arora et al. (2019a, Table 2) compute the NTK of up to 20-layer, 128-channel CNN in a binary CIFAR-2 classification setting. In equivalent setting with the same hardware (NVIDIA V100), are able to compute the NTK of a 2048-channel CNN, i.c. a network with at least 256 times more parameters.
To demonstrate the stability of the NTK during training for wide networks, Lee et al. (2019, Figure S6) compute the NTK of up to 3-layer 2idde or I-layer 2 wide FCNs. In the same setting with the same hardware (NVIDIA V100), handle networks with at least 16 times more parameters.
reach widths of at least 2 and 21 respectively, i.c.
To investigate convergence of WideResNet WRN-28-k (Zagoruyko & Komodakis, 2016) to its infinite width limit, Novak et al. (2020, Figure 2) evaluate the NTK of this model with widening factor k up to 32. In matching setting and hardware, we are able reach the widening factor of least 64, i.e. work with models at least 4 times larger.
To meta-learn NN parameters for transfer learning in MAML--like (Finn al., 2017) setting, Zhou et al. (2021, Table 7) replace the inner training loop with NTK-based inference. They use up to 5-layer, 200-channel CNNs on MinilmageNet (Oreshkin al., 2018) with scalar outputs and batch size 25. In same setting we achieve at least 512 channels, i.e. support models at least 6 times larger.
Park al. (2020, $4.1) use the NTK to predict the generalization performance of architectures in the context of Neural Architecture Search (Zoph & Le. 2017, NAS); however, the authors comment on its high computational burden and ultimately use different proxy. In another NAS setting, Chen et al. (2021a, 1.11.1) use the condition number of NTK to predict a model's trainability. Remark- ing its prohibitive cost, Chen et al. (2021b, Table 1) also use the NTK to evaluate the trainability of several ImageNet (Deng et al., 2009) models such as ResNet 50/152 (He et al., 2016), Vision Transformer (Dosovitskiy et al., 2021) and MLP-Mixer (Tolstikhin et al., 2021). However, in all of the above cases the authors only evaluate pseudo-NTKs i.e. an NTK of scalar-valued function,' which impacts the quality of the respective trainabliityggeeeraliitttioo proxy.
Precisely, computing the Jacobian only for a single logit or the sum of all 1000 class logits. The result is not the full NTK, but rather single diagonal block the sum of its 1000 diagonal blocks (finite width NTK is dense matrix, not block-diagonal).
2

---

Under review as a conference paper at ICLR 2022
Method
Time N'LO
Memory
Use when
Jacobian contraction
NOW + N + NLW + LW
Don't
NTK-vector products
No O"W + +N'LWW
NOW Nº NLW + LW2
0 W or
Structured derivatives
N + LOW2
NOW + NºO NLW + LW2
< W or
Table I: Asymptotic time and memory cost of computing the NTK for FCN. Costs for pair of batches of inputs of size N each, and for L-deep, W -wide FCN with O outputs. Resulting NTK has shape NO NO. NTK-vector products allow a reduction of the time complexity, while Structured derivatives reduce both time and memory complexity. Note: presented are asymptotic cost estimates: in practice, all methods incur large constant multipliers (e.g. least 3x for time; see 13.1). However, this generally does not impact the relative performance of different methods. See $3.6 for discussion, Table 7 for CNN, and Table 2 for more generic cost. analysis.
Method
Time
Memony
Use when
Jacohian contraction
O FP + N'OPP
O' + NO Y
+ NY +
< smalll O, exotic primitives
NTK-vector products
O'O
NO
NY+P
FP OP, large O, small
Structured denvvatvvs
NO OP
OO NOY
FP> large O, large
Table 2: Asymptotic time and memory cost estimates of computing the NTK for a generic function. stands for the number of all parameters in the network, Y stands for size of all pre- activations in the network, FP stands for forward pass, and G and J depend on the structure of FP (SB). For example, FCNs usually have cheap FP I OP, it consists of a single matrix multiply with the parameter matrix, and therefore NTK-vector products are recommended. CNNs, notably when the number of output pixels D is large, have costly FP 2 OP, since it amounts to D matrix multiplies with the parameters, and therefore Structured derivatives are preferred. For precise analysis, see Table 1 for FCN and Table for CNN.
In contrast, in this work we can compute the full 1000 x 1000 NTK on the same models (1000 classes), i.. perform a task 1000 times more costly.
Finally, we remark that in all of the above settings, scaling up by increasing width or by working with the true NTK (vs the pseudo-NTK) should lead to improved downstream task performance due to better infinite width/linearization approximation or higher-quality traiabbiityggeeeralizatoon proxy respectively, which makes our work especially relevant modern research.
3 EFFICIENT FINITE WIDTH NTKs IN A SIMPLIFIED SETTING
To gain intuition for the problem, we start by analyzing and improving the cost of computing the NTK for a simple FCN. See $F for an equivalent analysis of CNNs. We summarize the resulting complexities for FCN in Table 1, CNN in Table 7, and a general takeaway in Table 2.
.... E RO; where number of logits. We denote individual weight matrices as with shapes W x W (except for top- layer 6 of shape O x W), where W the width of the network, and write the set of all parameters as = vec [00
Setting. Consider an L-layer FCN f (0, 0 & (Ot-1... de (0"x)...
the
04] E TRWW*+OW, We further define x' = & (1/ as post-activations (with 0 == x), and o/ == x! as pre-activations with y = f (0, x). See Fig. 5 for a visual schematic of these quantities. For simplicity,
assume that inputs 2 also have width W, and O O (LW), i.e. the number of logits is dominated by the product of width and depth. In ŠL we repeat the same derivations without the latter assumption, and arrive qualitatively identical conclusions.
The NTK of f evaluated at two inputs In and 12 is an O x O matrix defined as
0f(0, x1) 0f(0,I2)"
of (0, ( T
Σ 0 E ROxO
00
00
20
20
where we have defined 0 to be the summands. We omit dependence 1T, B2 and f for brevity.
In $3.1 and $3.2 we describe the cost of several fundamental AD operations that we will use as building blocks throughout the text. We borrow the nomenclature introduced by Autograd (Maclau-
3

---

Under review as a conference paper at ICLR 2022
rin et al.) and describe secoobian-vecton products (JVP), vector-Jacobian products (VJP), as well as the cost of computing the Jacobian ( z08
In $3.3, describe the baseline complexity of evaluating the NTK, by computing two Jacobians and contracting them. This approach is used most (likely all) prior works, and scales poorly with the NN width W and output size O.
In $3.4 we present our first contribution, that consists in observing that many intermediate operations on weights performed by NNs possess certain structure, that allow linear algebra simplifica- tions of the NTK expression, leading to a cheaper contraction and smaller memory footprint.
In $3.5 we present second contribution, where rephrase the NTK computation instantiating itself row-by-row by applying the NTK-vector product function to columns of an identity matrix. As will show, this trades off Jacobian contraction for more forward passes, which proves beneficial in many (but not all) settings.
3.1 RACOBIAN-VECTOR PRODUCTS AND VECTOR-JACOBIAN PRODUCTS
We begin by defining Jacobian- products and vector-Jacobian products:
05 (0,0)))
(r. Og E TR. Igl. +oW
0, E RO,
(3)
08
VJP fc ER (0
00 fo E LLW
(4)
The JVP can be understood as pushing forward a tangent vector in weight space to tangent vector in the space of outputs; by contrast the VJP pulls back cotangent vector in the space of outputs to a cotangent vector in weight space. These elementary operations correspond to forward and reverse mode AD respectively and serve basis for typical AD computations such gradients, Jacobians, Hessians, etc.
Time and memory costs of JVP and VJP are asymptotically equivalent to the cost of the forward pass (FP), except for VJP additionally requires storing all intermediate activations. (see $N and Fig. 6).
For the case of FCNs, the time cost of both operations is therefore
[FP] = [cost of all intermediate layers] + [cost of the top layer] = [LW + [Ow] LW?. For single input, the memory cost of computing both the JVP and the VJP are respectively,
[size of all weights] + [size of activations at single layer] = [LW2 + Ow] + [W + + ~ LW?
[size of all weights] + jsize of activations in all layers) [LW2 + Ow] + [LW + LW
Despite the fact that the VJP requires more memory to store intermediate activations, we see that for FCNs both computations are dominated by the cost of storing the weights.
Batched inputs. If x is a batch of inputs of size N, the time cost of JVP and VJP increases linearly NLW2, The memory cost is slightly more nuanced. Since weights can be shared across inputs, the memory cost of the JVP and VJP are respectively,
(size of all weights] + N (size of activations at single layer]
[LW2 + OW] + N [W + O ~ LW2 + NW + NO,
(size of all weights + N [size of activations in all layers] + N jsize of all weight matrices]
|LW2+ + OW + [ + OJ+N + LW2 + OW) NLW2
The cost of the VJP is dominated by the cost of storing the cotangents in weight space. However, for the purposes computing the NTK, we will be contracting Jacobians layerwise and so we will only need store cotangent weight matrix, O /00 at a time. Thus, for the purposes of this work we end up with the following costs:
JVP costs NL 2 time and LW2 + NW + NO memory.
VJP costs NL 2 time and LW² + NLW + NW2 + NOW memory.
"To declutter notation, we omit the O symbol indicate asymptotic complexity in this work.
4

---

Under review as a conference paper at ICLR 2022
3.2 JACOBIAN COMPUTATION
For neural networks, the Jacobian is most often computed by evaluating the VJP on rows of the identity matrix lo, i.e.
(0, /00| (0, /00 lo IR (R(LW*+ow)x0
(5) follows that computing the Jacobian takes O evaluations of the VJP. However, as mentioned in $3.1, we only need to store one 0f/ 00/00' at time, while the weights and intermediate activations are reused across evaluations. Thus, time and memory costs to compute the Jacobian are respectively,
ON ((cost of all intermediate layers] + [cost of the top layer))
-ON ([LW"]= + [OW] ~ NLOW NO" [size of all weights] + N (size of activations in all layers] ON |size of single weight matrix]
[LW + OW] + NN + 0] ON ON [w² OW] LW2 + NLW + NOW? NO?
Therefore, asymptoticclly...
Jacobian costs NLOW2 NO time and LW2 NLW + NOW 2O memory.
3.3 JACOBIAN CONTRACTION
We now analyze the cost of computing the NTK, starting with the direct computation as the product of two Jacobians. Consider a single summand from Eq. (2):
0f (0,z1)
0f (0, 22)
0
(6)
00
00
0x0
Ox(WXW)
(WxW)x0
The time cost of this contraction is O and the memory necessary instantiate each factor and the result is OW2 + O Repeating the above operation for each 0', we arrive at L time cost and unchanged memory, due being able to process summands sequentially.
Batched inputs. If we consider Ii and I2 to be input batches of size N, then the resulting NTK matrix of shape NO x NO, and the time cost becomes N²L0# while memory grows to [NTK matrix size] + faccorr size] = N²O + 2
What remains is to account for the cost of computing and storing individual cotangents 0 0f/00' which is exactly the cost of computing the Jacobian ($3.2). Adding the costs up obtain
Jacobian contraction costs NºLO2W2 time and 2 + NOW? + 2 + LW NLW memory.
3.4 LEVERAGING STRUCTURED DERIVATIVES COMPUTING THE NTK
We can rewrite 0% in Eq. (6) using the chain rule and our pre- and post-activation notation as:
0 = Tas (0, .1) Oy'.
@ (0, x2) Öy
0s (0,x1)
Dy
By
0f (0,
a0
Oyk.
00
Oy
00
Oy'
OXW
(WXXW)
WxO
At face value, rewriting Eq. (6) this way is unhelpful as it appears to have introduced additional costly contractions. However, recall that y' = O'a', and therefore
dy 80
Ty
Aw x 2
Iw x
(8)
00
where is the Kronecker product. Plugging Eq. (8) into Eq. (7) obtain (see $G)
5

---

Under review as a conference paper at ICLR 2022
os (0,x1) Oy'n
0f (0, 32) T dye.
0
11
J22
1xW
Wx1
WxO
and observe that it takes only O2W time and OW + O memory. Accounting for depth, time cost becomes LO²W, while memory does not change since the summands can be processed sequentially.
Batched inputs. The time cost grows quadratically with the bath size N up to LO2 while the memory cost increases to O NOW to store the resulting NTK and oʃ (0, x) /Oy' factors.
Finally, we need to account for the cost of computing the derivatives, 0 /Oy and post-activations, x'. Notice that both xf and 0f/ōy' arises naturally when computing the Jacobian as the primals and cotangents in layer respectively. However, since we do not need to compute the weight space cotangents explicitly (i.e. we cut the backpropagation algorithm short) the memory cost will be,
[size all weights] N Jsize activations in all layers]
= [LW2 + OW] + NLW + O] ~ Lw2 + NLW.
The extra time cost is asymptotically the cost of O forward passes, NLOW? which is the same as the Jacobian. However, as we will see in experiments, in practice we'll often compute the NTK faster than the Jacobian due not computing the weight space cotangents 0f ///00. Altogether,
By leveraging Structured derivatives in NN computations, we have reduced the cost of NTK to N NLOW time and 2²0 + NOW + LW2 NLW memory.
The key insight was to leverage the constant block-diagonal structure of the pre-activation derivatives 0y' /00'. This idea is quite general; discuss in $4 and detail in the appendix, similar structure exists for many common operations such as convolutions, pooling, and arithmetic.
We highlight that these computational improvements do not emerge automatically in AD. While JAX and other libraries leverage structures analogous to Eq. (8) to efficiently compute single evalu- ations of the VJP and JVP, this knowledge is lost once the (structureless) Jacobian (f, x //00 is instantiated, and cannot be taken advantage of in the following contraction with (0 x2) x2)/00) /00' We discuss how we impose this structure compute the NTK for general neural networks in $4.
3.5 NTK VIA VECTOR PRODUCTS
Computing the Jacobian contraction using Jacobian first instantiates the Jacobian using using VJPs and then performs a contraction. Structured derivatives use a similar strategy, but speed-up the contraction and avoid explicitly instantiating the weight space cotangents. Here avoid performing a contraction altogether at the cost of extra VJP/JVP calls; this ends up being beneficial for FCNs.
R OBE R Applying this function to columns of the identity matrix To allows us to compute the NTK, i.e. ObTo = gg. The cost of evaluating the NTK in this fashion is equal to O times the cost of a single NTK-vector product evaluation OVP(u). We now expand @VP(c) Oat as
We introduce the linear function performing the NTK-vector product: OVP:
0s (0,z,) of (0, ae
T 00 c 20
(0, (1. (v) = (VP(1.82.)) [VIP (v)),
(10)
where we have observed that, contracted from right to left, the NTK-vector product can be ex- pressed as composition of JVP and VJP of the underlying function f. The cost of this operation is asymptotically equivalent to the cost of the Jacobian, since it consists of O VJPs followed by O (cheaper) JVPs. Therefore it costs LOW² Ow time and LW2 + Ow2 O2 " memory.
Batched inputs. In the batched setting Eq. (10) is repeated for each pair of inputs, and therefore time increases by factor of N2 to become LOW2 + 2 However, the memory cost grows only linearly in N (except for the cost of storing the NTK of size N²02), since intermediate activations and derivatives necessary compute the JVP and VJP can be computed for each batch zy and x2

---

Under review as a conference paper at ICLR 2022
separately; these quantities are then reused for every pairwise combination resulting in a memory equivalent to the Jacobian, i.c. Nº + (LW2 + NOW? + NO + NLW), resulting in
NTK computation as sequence of NTK-vector products costs NoLOW + 2 time and N²O? NOW + LW2 + NLW memory.
3.6 SUMMARY
Structured derivatives and NTK-vector products allow a reduction in the time cost of NTK compu- tation in different ways, and Structured derivatives also reduce memory requirements. Structured derivatives are beneficial for wide networks, with large W, and NTK-vector products are beneficial for networks with large outputs O. We confirm our predictions with FLOPs measurements in Fig.
We further confirm our methods can provide orders of magnitude speed-ups and memory savings on all major hardware platforms in Fig. (right) and Fig. 3. However, we notice that our wall-clock time measurements often deviate from predictions due to unaccounted constant overheads of various methods, hardware specifics, padding, and the (largely black-box) behavior of the XLA compiler. Notably, in practice, we find Structured derivatives almost always outperform NTK-vector products.
Finally, we evaluate our methods in the wild, and confirm computational benefits on full ImageNet models in Fig. 2 (ResNets, He et al. (2016)) and Fig. 4 (WideResNets, Zagoruyko & Komodakis (2016); Vision Transformers and Transformer-Resse et al. (2021); and MLP-Mixers Tolstikhin et al. (2021)). Computing the full O x O 1000 x 1000 NTK for many of these models on modern accelerators is only possible with Structured derivatives.
hybrids Dosovitskiy et al. (2021): Steiner
STRUCTURED DERIVATIVES FOR GENERIC FUNCTIONS
Here generalize the idea of leveraging structure in subexpressions presented in $3.4. This section (and our implementation) is not specific to NNs and applies to any differentiable function.
Consider two differentiable functions defined on common input domain:
f:: (0",...,0" 0²) Px. + ft (0" (0"....,0*) E go. (i (1,2)).
For NNs, typically (40....,00) 0 correspond to trainable parameters in layers O,..., L, and fo (0".... ..... Q1.)- [@0 ...... H¹. 2z.) 24 being network inputs, O O being the number of outputs (logits, classes). The NTK is defined as
Ao (fr. = Σ 0f.0f² T 00 a0 E 000
(11)
1=0
Assume the following decompositions of fi into computational graphs made of primitives y::
...., ......1.. fi (0) ..... 0²) fs (0) (0"...., ..),.... . (00 (6......0") 0')) (i E (1,2)))
(12)
with ot (0o. ..... 04) E R In common NNs, A would correspond to pre-activations evaluated on inputs x; in layer k; and, without weight sharing, typically K, = Kg L.. However, we do not impose any relationship between the number of parameter variables L and number of primitives K₁ and K2, allowing arbitrary weight sharing. We can then use the chain rule in Eq. (2) to obtain:
L,K,,Ka Σ L,R
Ou d0
)
0 0y dyy 00
L,K, Ka 0ń Oy ki 0y Σ
T
0/2
O, (21,
(13)
By
1,A1.k2 Byp 00 a0
All methods from $3 perform the sum of contractions in Eq. (13) one way another. Jacobian con- traction uses VJPs to implicitly contract each summand "outside-in", i.c. it first computes 0f:/ 00' terms with VJPs followed by their contraction. As discussed in $3.3, this costs NO [FP] + P
NTK-vector products use both JVPs and VJPs contract "Right-to-left" i.e. first compute 0f2/00' as an implicit contraction of 0 0f2/Oy2 with dyz /00' via VJP, followed by an implicit contraction of the result with 0y./ /00" via JVP, followed by another implicit contraction with 0f: θy with another JVP. Per $3.5 this costs O [FP).
7

---

Under review as a conference paper at ICLR 2022
FLOPs (per NTK entry)
Wall-clock time (TPUv3)
O= logits
16 logits
0- 128 logits
o= logits
o= 166 logits
o= 128 logits
10:
10a
U2204 N
Seconds
20-1
batch
10-4
10*
U2004 91 N
Seconds
ELOP
10-1
10*
10'
10-3
10A
420ig = M
10'
Seconds
1000
104
LoUee 72L
10
10a
10"4
22
2
2
2
2
212
2a
25
2
22
2 2
22 2 2
2
2a 212
Width W
Width W
Width W
Width W
Width W
Width W
Jacobian contraction
Structured derivatives jacobian
W ==0 W2
NTK-vector products
Figure 1: FLOPs (Ieft) and wall-clock time (right) of computing the NTK for 10-layer ReLU
FCN. As predicted by Table 1, our methods almost always outperform Jacobian contraction, allow- ing orders of magnitude speed-ups and memory improvements for realistic problem sizes. FLOPs per NTK entry: We confirm several specific predictions: (1) NTK-veetor products are the best per- forming method for N 1, and have cost equivalent to Jacobian for any width W or output size O (top row); (2) NTK-vector products offer an O-fold improvement over Jacobian contraction (left to right columns); (3) NTK-vector products are equivalent to Jacobian contraction for O = 1 (leftmost column); (4) Structured derivatives outperform NTK-vector products iff O < W (O = W are plot- ted as pale vertical lines, which is where Structured derivatives and -TK-vector products intersect); (5) Structured derivatives approach the cost of Jacobian in the limit of large width W (left to right). (6) All methods, as expected, scale quadratically with width W. Wall-clock runtime: In real appli- cations, given wardware-speciff constraints, padding, and delicate interplay with the XLA compiler,
we observe that: (1) WTK-veetor products improve upon Jacobian contraction for O > 1, but the effect is not perfectly robust (see bottom row for small W and Fig. 3, notably GPU platforms); (2) Structured derivatives robustly outperform all other methods, including simply computing the Ja- cobian, as discussed in $3.4; (3) Structured derivatives have lower memory footprint, and reach up to 8x larger widths (bottom right; missing points indicate out-of- out-of-memory memory), i.c. can process models up to 64x larger than other methods, as discussed in $3.4; (4) All methods have a smaller memory footprint than Jacobian (see $3.1) 1). More: Fig. 3 for other hardware platforms, ŠH for details.
NVIDIA V100
TPUv4
CPU
10a1
10a
=
10
1
10:1
10
1-7x 1-7x100
0
1.6 x 10'
Seconds
1.4x100 10*
Jacoblan contraction
10a
SHlB 0000
1.2 103
NTK-vector products
Structured derivatives
10a
101
jacobian
10'
18 34 50
101
152
200
18 34 50
101 152
200
18 34 50
101
152
200
ResNet depth
ResMet depth
ResNet depth
Figure 2: Wall-clock time cost of computing an NTK for several ResNet sizes on a pair of ImageNet inputs. Structured derivatives allow the NTK to be computed faster and for larger models (see bottom row missing points indicate memoory NTK-vector products, as predicted in $3.6 and Table 2, are advantagcous for large O (bottom row), but are suboptimal when the cost of the forward pass is large relative to the number of parameters, e.g. when there is a lot of weight sharing (see Table 7 and Table 2), which is the case for convolutions. See Fig. 4 for more ImageNet models, $F for analysis of CNN NTK computational complexity, and SH for experimental details.
8

---

Under review as a conference paper at ICLR 2022
However, recall from $3.4, while JVPs and VJPs themselves computationally optimal, higher- order computations like their contraction (Jacobian contraction) or composition (NTK-vector prod- ucts) are generally not. The idea of Structured derivatives is to design rules for efficient computation of such contractions, similarly to how JAX and has rules for efficient JVPs and VJPs. From Eq. (13), in the general case this requires hand-made rules for all pairwise combinations of primitives y/ and 92. Due to quadratic scaling in the number of primitives, we restrict the current implementation to rules that operate on individual primitives y. This still provides substantial computational benefit.
Specifically, our rules identify a few simple types of structure (e.g. block diagonal, constant-block diagonal, tiling) in 0y/ a00, that allow us to simplify the contraction in Eq. (13). In practice this amounts to replacing the inner terms Oy k. /00' and dy /00' with their (much) smaller subarrays, and modifying the instructions passed to np. einsum that contracts all terms. In $C we provide specific descriptions of our rules and their impact on the computational complexity of Eq. (13).
In Table and Table 7 we show that our rules are asymptotically better than Jacobian contraction for matrix multiplications and convolutions, and verify that they are practically beneficial in much wider set of operations used by contemporary ImageNet models in Fig. 2 and Fig. 4.
For both Structured derivatives and NTK-vector products a fully general and rigorous comparison of complexities is not feasible since it will rely upon specifics of the model architecture and the pairs of primitives, 1/1 and y2- present in the network. Nonetheless, we can offer heuristies that suggest when each method will be beneficial. The time complexity of Structured derivatives has the form of NO FP] + NOG + [ OP), where G is related to the cost of contraction, and J to the cost of computing 0y ae (exact values depend on the structure present y1 and 32). This is guaranteed to be no worse than Jacobian contraction for FCNs and CNNs. From Table 2, the performance of NTK-vector products relative to Jacobian contraction ultimately depends on the cost of the forward pass through the network, [FP). relative to OP. In practice this amounts to best performance on models without weight sharing like FCNs.
Owing to the nuanced trade-offss between different computational methods in the general case, we release all our implementations as a single function that allows the user to manually select the desired implementation. For convenience. include an automated setting which will perform FLOPs analysis for each method at compilation time and automatically select the most efficient one.
IMPLEMENTATION
Both algorithms are implemented in JAX (Bradbury et al., 2018) as the following function trans- formation ntk_fn| : [f (0, () f(0, + 10 (21, I2, 0)
0,(x1,22)] i.e. our function accepts any function f with the above signature and returns the efficient NTK kernel function oper- ating on inputs 21 and ø2 and parameterized by f. Inputs x, parameters 0, and outputs f(0, x) can be arbitrary PyTrees. We rely on many utilities from JAX and Neural Tangents (Novak al., 2020).
NTK-vector products algorithm is implemented by using JAX core operations such as vjp, jvp. and vmap to map the NTK-vp function to the To matrix and to parallelize the computation over pairwise combinations of N inputs in each batch æ and x2.
Jaxpr interpreter, built top of the default JAX reverse mode AD interpreter. On a high level, the algorithm performs the sum in Eq. (13). Each summand is contraction of 4 factors: 0f./ 0f//011. 0y./00, 0yz/00, f /0yz.
Structured derivatives algorithm is implemented
First, we linearize f to obtain a computational graph constructed out of limited set (54,² see Table 5) of linear primitives y "'....," Then, we can obtain two factors of./ (yy, 0f2/O0 as part of a backward pass almost identical to calling jax. jacobian (0) To contract these terms with /00 and 0y2/ /00, as described above, we query a dictionary of rules which map primitives to structural description ($C.8); for a given pair of primitives, these rules allow to analytically simplify the contraction and avoid explicitly instantiating the derivatives.
3 JAX leverages similar approach implement only 54 transpose rules for linear primitives for reverse mode differentiation instead of 131 VJP rules (Frostig et al., 2021).

---

Under review as a conference paper at ICLR 2022
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. Tensorflow: system for large- scale machine learning. In 12th USENIX Symposium on Operating Systems Design and Imple- mentation (OSDI 16), 2016.
Ben Adlam, Jaehoon Lee, Lechao Jeffrey Pennington, and Jasper Snoek. Exploring the un- certainty properties of neural networks' implicit priors in the infinite-width limit. In International Conference on Learning Representations, 2020.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. In Advances in Neural Information Processing Systems, pp. 8141 1-8800. Curran Associates, Inc., 2019a.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grainee analysis of optimization and generalization for overparameterized two-layer neural networks. arXiv preprint arXiv:1901.0888,, 2019b.
Sanjeev Arora, Simon S. Du, Zhiyuan Li, Ruslan Salakhutdinov, Ruosong Wang, and Dingli Yu. Harnessing the power of infinitely wide deep nets on small-data tasks. In International Conference on Learning Representations, 2020. URL https://openrevieww
Yasaman Bahri, Ethan Dyer, Jared Kaplan, Jaehoon Lee, and Utkarsh Sharma. Explaining neural scaling laws. arXiv preprint arXiv:2 2021.
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine- learning practice and the classical barinnar trade-off. Proceedings of the National Academy of Sciences, 116(32):15849) 15849-15544 2019.
James Bradbury. Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman- JAX: composable transformations of Python+NumPy programs, 2018 URL
httpgggthh.c.ooooo/e//e//// comggooge//jax
Andrew Brock, Soham De, and Samuel I. Smith. Characterizing signal propagation to close the performance gap in unnormalized resnets. arXiv preprint 101.:2200.008992, 2021a.
Andrew Brock, Soham De, Samuel Smith, and Karen Simonyan. pigh-perforrance large-scale image recognition without normalization. arXiv preprint arXiv:2102.06171,
Wuyang Chen, Xinyu Gong, and Zhangyang Wang. Neural architecture search on imagenet in four gpu hours: A theoretically inspired perspective. In International Conference on Learning Representations, 2021a.
Xiangning Chen, Cho-Jui Hsieh, and Boqing Gong.
When vision transformers outperform resnets
without pretraining strong data augmentations, 2021b.
Yann Dauphin and Samuel Schoenholz. Metainit: Initializing learning by learning to initialize. 2019.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hi- erarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pp. 248 248-255 255.. Ieee, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani. Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations, 2021. URL https:
Simon S Du, Kangcheng Hou, Russ R Salakhutdinov, Barnabas Poczos, Ruosong Wang, and Keyulu Xu. Graph neural tangent kernel: Fusing graph neural networks with graph kernels. In Advances in Neural Information Processing Systems. 2019.
10

---

Under review as a conference paper at ICLR 2022
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In Doina Precup and Yee Whye Tch (eds.). Proceedings of the 34th International Conference Machine Learning. volume 70 of Proceedings Machine Learning Research, pp- 1126-1135. PMLR, 06-11 Aug 2017. URL https: httpss:///ooccednggs.oom wttp:/oeeeeiiin............ wtt.mm/...ee//////eeeemsm finnnn html.
Yves Franceschi, Emmanuel de Bézenac, Ibrahim Ayed, Mickaël Chen, Sylvain Lamprier, and Patrick Gallinari. A neural tangent kernel perspective of gans. arXiv preprint arXiv:2106.05566, 2021.
Roy Frostig, Matthew Johnson, Dougal Maclaurin, Adam Paszke, and Alexey Radul. Decompos- ing reverse-mode mode automatic differentiation. arXiv preprint wrXiv:210S.09465 2021.
Adrià Gariiga-Alonso, Laurence Aitchison, and Carl Edward Rasmussen. Deep convolutional net- works as shallow gaussian processes. International Conference on Learning Representations, 2019.
Andreas Griewank and Andrea Walther. Evaluating Derivatives. Society for Industrial and Applied Mathematics, second edition, 2008. doi: 10.11337719778889877777 URL //epus. siam. org//oi/ass//////oooooooogog/o 10. 1137/1 9780898717761..
Roger Grosse. Neural net training dynamics, January 2021. URL https://www.c.. www.ntoooooo.teeo ""rrooss/cour hatese///cccreesss/55//5 2021/readings/L0z Taylor approximations. .ddf.
Boris Hanin and Mihai Nica. Finite depth and width corrections to the neural tangent kernel. In International Conference Learning Represenlations, 2020. URL httpt::///ppeeeeeeeeee net/ forum?id=S]gndl 4KwB.
Bobby He, Balaji Lakshminarayanan, and Yee Whye Teh. Bayesian deep ensembles via the neu- ral tangent kernel. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Floriaa Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, Decem- ber 6-12, 2020, virtual, 2020. URL http://///cpooooein.s....o neuriiinnsnnnuuis.......m.. httppaper/222/hhaa 11266662222611888817772222222222 Abstract. html.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 770-778, 2016.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2020. URL http://gi teetittubbcco ctt. www..////oo///aa.oooooo
Tom Hennigan, Trevor Cai, Tamara Norman, and Igor Babuschkin. Haiku: Sonnet for JAX, 2020.
URL http://gi thub. www homm/deepmindd.. dmdhhmuikuu...
Jiri Hron, Yasaman Bahri, Jascha Sohl-Dicksteii and Roman Novak. Infinite attention: NNGP and
NTK for deep attention networks. In International Conference on Machine Learning, 2020.
Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and gen- eralization in neural networks. In Advances in Neural Information Processing Systems, 2018.
Jaehoon Lee, Yasaman Bahri, Roman Novak, Sam Schoenholz, Jeffrey Pennington, and Jascha Sohl- dickstein. Deep neural networks gaussian processes. In International Conference on Learning Representations, 2018.
Jachoon Lee, Lechao Xiao, Samuel S. Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve linear models under gradient descent. In Advances in Neural Information Processing Systems, 2019.
Jaehoon Lee, Samuel S Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: empirical study. 2020.
Dougal Maclaurin, David Duvenaud, and Ryan P Adams, Autograd: Effortless gradients in numpy.

---

Under review as a conference paper at ICLR 2022
James Martens and Roger Grosse. Optimizing neural networks with kronecker- factored approximate curvature. In International conference on machine learning, pp. 2408-2417. PMLR, 2015.
Alexander G. de G. Matthews, Jiri Hron, Mark Rowland, Richard E. Turner, and Zoubin Ghahra- mani. Gaussian process behaviour in wide deep neural networks. In International Conference on Learning Representations, 2018.
Herman Müntz. Solution directe l'équation séculaire et de quelques problèmes analogues tran- scendants. C. R. Acad. Sci. Paris, 156:43-46, 1913.
Uwe Naumann. Optimal accumulation of jacobian matrices by elimination methods on the dual computational graph. Mathematical Programming, 99(3):399-422) 221,2000 2004.
Uwe Naumann. Optimal jacobian accumulation is np-complete. Mathematical Programming, 112 (2):427-441, 2008
Radford M. Neal. Priors for infinite networks (tech, rep. no. crg-tr-94-1). University of Toronto, 1994.
Timothy Nguyen, Zhourong Chen, and Jaehoon Lee. Dataset meta-learning from kernel ridge- regression. arXiv preprint M0Xiv:2011.00050 2020.
Timothy Nguyen, Roman Novak, Lechao Xiao, and Jaehoon Lee. Dataset distillation with infinitely wide convolutional networks. arXiv preprint arXiv:2107.1303: 2021.
Roman Novak, Lechao Xiao, Jaehoon Lee, Yasaman Bahri, Greg Yang, Jiri Hron, Daniel A. Abo- lafia, Jeffrey Pennington, and Jascha Sohl-Dickstein. Bayesian deep convolutional networks with many channels are gaussian processes. In International Conference on Learning Representations, 2019.
Roman Novak, Lechao Xiao, Jiri Hron, Jachoon Lee, Alexander A. Alemi, Jascha Sohl-Dickstein and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In International Conference on Learning Reprexentations, 2020. URL https: com/ googlee/nerall-angeen
Boris N. Oreshkin, Pau Rodríguez López, and Alexandre Lacoste. Tadam: Task dependent adaptive metric for improved few-shot learning. In NeurlPS, 2018
Daniel S Park, Jachoon Lee, Daiyi Peng. Yuan Cao, and Jascha Sohl- Sohl-Dicksteen. Towards nngp- guided neural architecture search. arXiv preprint arXiv:201 1.06006, 2020.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Ed- ward Yang. Zachary Vito,o Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Buc, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems 32, pp. 8024-8035. Curran Associates, Inc., 2019. URL http::/papprs..nnurips.. ctw.:rssssriiiriccccpppppppppppp cc/paper/ wwwttorchann-imeraiiiii style- www deep
pdf.
Jeffrey Pennington and Yasaman Bahri. Geometry neural network loss surfaces via random matrix theory. In Doina Precup and Yee Whye Teh (eds.). Proceedings of the 34th Interna- tional Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Re- search, pp. 2798-2806. PMLR, 06-11 Aug 2017. URL https://proceedings.. .ml.eng........................m prrss// v70/pennington177 htmml.
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. International Conference on Learning Representations, 2017.
Stefano Spigler, Mario Geiger, Stéphane d' Ascoli, Levent Sagun, Giulio Biroli, and Matthieu Wyart. A jamming transition from under-to to over- parapemtettiiiiioonn affects generalization in deep learning. Journal of Physics A: Mathematical and Theoretical, 52(47):474001, 2019.
12

---

Under review as a conference paper at ICLR 2022
Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv preprint arXiv:2106.1022
2021.
Matthew Tancik, Pratul Srinivasan, Ben Mildenhall, Sara Fridovich-Kei,, Nithin Raghavan, Utkarsh Singhal, Ravi Ramamoorthi, Jonathan T. Barron, and Ren Ng. Fourier features let net- works learn high frequency functions in low dimensional domains. NeurIPS, 2020.
Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Un- terthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, and Alexey Dosovitskiy. Mip-mixer: An all-mlp architecture for vision, 2021.
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel Schoenholz, and Jeffrey Pennington. Dynamical isometry and mean field theory of CNNs: How to train 10,000-layer vanilla convo- lutional neural networks. In International Conference on Machine Learning, 2018.
Lechao Xiao, Jeffrey Pennington, and Samuel S Schoenholz. Disentangling trainability and gener- alization in deep learning. International Conference on Machine Learning. 2020.
Sho Yaida. Non-Gaussian processes and neural networks at finite widths. In Mathematical and Scientific Machine Learning Conference, 2020.
Greg Yang. Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation. arXiv preprint 1:Xi::1902.04760, 2019.
Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006. 14548, 2020.
Greg Yang. Jeffrey Pennington, Vinay Rao, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. A mean field theory of batch normalization. In International Conference on Learning Representa- tions, 2019.
Sergey Zagoruyko and Nikos Komodakis. Conference, 2016.
Wide residual networks. In British Machine Vision
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. arXiv preprint arXiv: 009311, 2019.
Yufan Zhou, Zhenyi Wang, Jiayi Xian, Changyou Chen, and Jinhui Xu. Meta-learning with neural tangent kernels. In International Conference Learning Representations, 2021. URL https:
Barret Zoph and Quoc V. Le. Neural architecture search with reinforcement learning. 2017. URL fttps://aaxiiv.or//btt 01578.
13


Reviews:
[{"review_id": "_uymrH5KAs6", "reviewer": "Reviewer_Ye97", "summary": "", "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": "Strengths\n\n(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.\n\n(2) The author makes the article easy to understand through careful structural planning.\n\nWeakness:\n\n(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}$. \n\n(2) The dashed lines in Figure 1 are hard to distingish, I suggest the authros could make it clearer. ", "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.", "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.\n\nBesides, 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.\n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "WPVoHUX2aO8", "reviewer": "Reviewer_uyYg", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The 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.\n\nExperiments 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_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.", "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.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "FnWeeyLcfkv", "reviewer": "Reviewer_qdx1", "summary": "", "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": "Strengths:\n- 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.\n- 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.\n- The experiments are convincing are back up the author's analysis of the advantages of their methods\n- The paper is generally very well written and clear, I especially like the colour-coded variables like N,O,P etc\n\nWeaknesses (most of these are not huge or beyond the current scope of the paper):\n- 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.\n- 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.\n- 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.\n- 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.\n\nMinor points:\n- The legend in fig 1 left blocks the graph\n- 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.\n", "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.", "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.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "BYbuvXI-lDD", "reviewer": "Reviewer_tZAo", "summary": "", "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": "- 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.\n\n- 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?\n\n- Minor Issue:\n  - It would be great if more details of Automatic Differentiation (AD) operations to derive memory costs of JVP/VJP are provided.\n  - As assumed $O=\\mathcal{O}(LW)$, the time and memory costs in Section 3.2~3.4 can be reduced without $O$.\n  - The method names in Figure 1 are too small to identify. It would be great to make larger  font size.\n\n", "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.", "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.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

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