FAST FINITE WIDTH NEURAL TANGENT KERNEL
Year: 2021

Abstract:
The Neural Tangent Kernel (NTK), defined as the outer product of the neural network (NN) Jacobians, ⇥✓(x1, x2) = ⇥ @f(✓, x1) @✓ ⇤ ⇥ @f(✓, x2) @✓ ⇤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 [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.

text: ⇥ @f(✓, x1) @✓ ⇤ ⇥ @f(✓, x2) @✓ ⇤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 [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.

heading: 1 INTRODUCTION
text: The past few years have seen significant progress towards a theoretical foundation for deep learning. Much of this work has focused on understanding the properties of random functions in high dimensions. One significant line of work (Neal, 1994; Lee et al., 2018; Matthews et al., 2018; Novak et al., 2019; Garriga-Alonso et al., 2019; Hron et al., 2020; Yang, 2019) established that in the limit of infinite 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 et al., 2019) showed that the ubiquitous bias-variance decomposition breaks down as high-dimensional models enter the socalled 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 et al., 2017; Xiao et al., 2018; 2020). This notion inspired initialization schemes like Fixup (Zhang et al., 2019), MetaInit (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., 2021b). The NTK has additionally given insight into a wide range of phenomena such as: behavior of Generative 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 distillation (Nguyen et al., 2020; 2021) and uncertainty prediction (He et al., 2020; Adlam et al., 2020).
Despite the significant promise of theory based on the NTK, computing the NTK in practice is challenging. 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
product of Jacobians using forward or reverse mode automatic differentiation (AD),
⇥✓(x1, x2)| {z } O⇥O
:= ⇥ @f(✓, x1) @✓ ⇤
| {z } O⇥P
⇥ @f(✓, x2) @✓ ⇤T | {z } P⇥O , (1)
where f is the forward pass NN function producing outputs in RO, ✓ 2 RP are all trainable parameters, and x1 and x2 are two inputs to the network. If inputs are batches of sizes N1 and N2, the NTK is an N1O ⇥ N2O matrix. Unfortunately, evaluating Eq. (1) is often infeasible due to time and memory requirements.
In this paper, we 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 is asymptotically more efficient in both settings. We compute the NTK over a 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.

heading: 2 RELATED WORK
text: 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 a 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 an equivalent setting with the same hardware (NVIDIA V100), we are able to compute the NTK of a 2048-channel CNN, i.e. 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 212-wide or 1-layer 214-wide FCNs. In the same setting with the same hardware (NVIDIA V100), we can reach widths of at least 214 and 218 respectively, i.e. handle networks with at least 16 times more parameters.
To investigate convergence of a 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 to reach the widening factor of at least 64, i.e. work with models at least 4 times larger.
To meta-learn NN parameters for transfer learning in a MAML-like (Finn et 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 MiniImageNet (Oreshkin et 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 et 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 a different proxy. In another NAS setting, Chen et al. (2021a, §3.1.1) use the condition number of NTK to predict a model’s trainability. Remarking 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 a pseudo-NTK, i.e. an NTK of a scalar-valued function,1 which impacts the quality of the respective trainability/generalization proxy.
1Precisely, 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 a single diagonal block or the sum of its 1000 diagonal blocks (finite width NTK is a dense matrix, not block-diagonal).
In contrast, in this work we can compute the full 1000 ⇥ 1000 NTK on the same models (1000 classes), i.e. 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 trainability/generalization proxy respectively, which makes our work especially relevant to modern research.

heading: 3 EFFICIENT FINITE WIDTH NTKS IN A SIMPLIFIED SETTING
text: 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.
Setting. Consider an L-layer FCN f (✓, x) = ✓L ✓L 1 . . . ✓1 ✓0x . . . 2 RO, where O is the number of logits. We denote individual weight matrices as ✓l with shapes W ⇥ W (except for toplayer ✓L of shape O ⇥ W), where W is the width of the network, and write the set of all parameters as ✓ = vec ⇥ ✓0, . . . , ✓L ⇤ 2 RLW2+OW. We further define xl := yl 1 as post-activations (with x0 := x), and yl := ✓lxl as pre-activations with yL = f (✓, x). See Fig. 5 for a visual schematic of these quantities. For simplicity, we assume that inputs x 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 at qualitatively identical conclusions.
The NTK of f evaluated at two inputs x1 and x2 is an O ⇥ O matrix defined as
⇥✓ := @f(✓, x1)
@✓
@f(✓, x2)
@✓
T = LX
l=0
@f (✓, x1) @✓l @f (✓, x2) @✓l
T =: LX
l=0
⇥l✓ 2 RO⇥O, (2)
where we have defined ⇥l✓ to be the summands. We omit dependence on x1, x2, 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-
rin et al.) and describe Jacobian-vector products (JVP), vector-Jacobian products (VJP), as well as the cost of computing the Jacobian @f(✓, x) @✓.
In §3.3, we describe the baseline complexity of evaluating the NTK, by computing two Jacobians and contracting them. This approach is used in 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 a certain structure, that can allow linear algebra simplifications of the NTK expression, leading to a cheaper contraction and smaller memory footprint.
In §3.5 we present our second contribution, where we rephrase the NTK computation as instantiating itself row-by-row by applying the NTK-vector product function to columns of an identity matrix. As we will show, this trades off Jacobian contraction for more forward passes, which proves beneficial in many (but not all) settings.

heading: 3.1 JACOBIAN-VECTOR PRODUCTS AND VECTOR-JACOBIAN PRODUCTS
text: We begin by defining Jacobian-vector products and vector-Jacobian products:
JVP(f,✓,x) : ✓t 2 RLW 2+OW 7! @f (✓, x)
@✓ ✓t 2 RO, (3)
VJP(f,✓,x) : fc 2 RO 7! @f (✓, x)
@✓
T
fc 2 RLW 2+OW. (4)
The JVP can be understood as pushing forward a tangent vector in weight space to a tangent vector in the space of outputs; by contrast the VJP pulls back a 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 as a basis for typical AD computations such as 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 cost2 of both operations is therefore
[FP] = [cost of all intermediate layers] + [cost of the top layer] = ⇥ LW2 ⇤ + [OW] ⇠ LW2.
For a single input, the memory cost of computing both the JVP and the VJP are respectively, [size of all weights] + [size of activations at a single layer] = ⇥ LW2 + OW ⇤ + [W + O] ⇠ LW2,
[size of all weights] + [size of activations in all layers] = ⇥ LW2 + OW ⇤ + [LW + O] ⇠ LW2.
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 to 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 a single layer] = ⇥ LW2 + OW ⇤ + N [W + O] ⇠ LW2 + NW + NO,
[size of all weights] + N [size of activations in all layers] + N [size of all weight matrices] = ⇥ LW2 + OW ⇤ + N [LW + O] + 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 of computing the NTK, we will be contracting Jacobians layerwise and so we will only need to store one cotangent weight matrix, @f @✓l, at a time. Thus, for the purposes of this work we end up with the following costs:
• JVP costs NLW2 time and LW2 + NW + NO memory. • VJP costs NLW2 time and LW2 + NLW + NW2 + NOW memory.
2To declutter notation, we omit the O symbol to indicate asymptotic complexity in this work.

heading: 3.2 JACOBIAN COMPUTATION
text: For neural networks, the Jacobian is most often computed by evaluating the VJP on rows of the identity matrix IO, i.e.
⇥ @f (✓, x) @✓ ⇤T = ⇥ @f (✓, x) @✓ ⇤T IO 2 R(LW 2+OW)⇥O. (5)
It follows that computing the Jacobian takes O evaluations of the VJP. However, as mentioned in §3.1, we only need to store one @f @✓l at a 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 ⇥ LW2 ⇤ + [OW] ⇠ NLOW2 + NO2W,
[size of all weights] + N [size of activations in all layers] + ON [size of a single weight matrix] = ⇥ LW2 + OW ⇤ + N [LW + O] + ON ⇥ W2 + OW ⇤ ⇠ LW2 + NLW + NOW2 + NO2W.
Therefore, asymptotically,
Jacobian costs NLOW2 + NO2W time and LW2 + NLW + NOW2 + NO2W memory.

heading: 3.3 JACOBIAN CONTRACTION
text: 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):
⇥l✓|{z} O⇥O
= @f (✓, x1)
@✓l| {z } O⇥(W⇥W)
@f (✓, x2)
@✓l
T
| {z } (W⇥W)⇥O
. (6)
The time cost of this contraction is O2W2, and the memory necessary to instantiate each factor and the result is OW2 + O2. Repeating the above operation for each ✓l, we arrive at LO2W2 time cost and unchanged memory, due to being able to process summands sequentially.
Batched inputs. If we consider x1 and x2 to be input batches of size N, then the resulting NTK is a matrix of shape NO ⇥ NO, and the time cost becomes N2LO2W2, while memory grows to [NTK matrix size] + [factors size] = N2O2 + NOW2.
What remains is to account for the cost of computing and storing individual cotangents @f @✓l, which is exactly the cost of computing the Jacobian (§3.2). Adding the costs up we obtain
Jacobian contraction costs N2LO2W2 time and N2O2 + NOW2 + NO2W + LW2 + NLW memory.

heading: 3.4 LEVERAGING STRUCTURED DERIVATIVES FOR COMPUTING THE NTK
text: We can rewrite ⇥l✓ in Eq. (6) using the chain rule and our pre- and post-activation notation as:
⇥l✓ =
" @f (✓, x1)
@ylx1
@ylx1 @✓l
#" @f (✓, x2)
@ylx2
@ylx2 @✓l
#T = @f (✓, x1)
@ylx1| {z } O⇥W
@ylx1 @✓l| {z }
W⇥(W⇥W)
@ylx2 @✓l
T
| {z } (W⇥W)⇥W
@f (✓, x2)
@ylx2
T
| {z } W⇥O
.
(7) At face value, rewriting Eq. (6) this way is unhelpful as it appears to have introduced additional costly contractions. However, recall that yl = ✓lxl, and therefore
@ylx1 @✓l = IW ⌦ x l 1 T , @ylx2 @✓l = IW ⌦ x l 2 T , (8)
where ⌦ is the Kronecker product. Plugging Eq. (8) into Eq. (7) we obtain (see §G)
⇥l✓ =
0
B@xl1 T
|{z} 1⇥W xl2|{z} W⇥1
1
CA
2
6664 @f (✓, x1)
@ylx1| {z } O⇥W
@f (✓, x2)
@ylx2
T
| {z } W⇥O
3
7775 , (9)
and observe that it takes only O2W time and OW + O2 memory. Accounting for depth, time cost becomes LO2W, 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 N2LO2W, while the memory cost increases to N2O2 + NOW to store the resulting NTK and @f (✓, x)
@ylx factors.
Finally, we need to account for the cost of computing the derivatives, @f @yl, and post-activations,
xl. Notice that both xl and @f @yl arises naturally when computing the Jacobian as the primals and cotangents in layer l 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 of all weights] + N [size of activations in all layers] = ⇥ LW2 + OW ⇤ + N [LW + O] ⇠ LW2 + NLW.
The extra time cost is asymptotically the cost of O forward passes, NLOW2 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 to not computing the weight space cotangents @f/@✓l. Altogether,
By leveraging Structured derivatives in NN computations, we have reduced the cost of NTK to N2LO2W + NLOW2 time and N2O2 + NOW + LW2 + NLW memory.
The key insight was to leverage the constant block-diagonal structure of the pre-activation derivatives @yl @✓l. This idea is quite general; as we 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 evaluations of the VJP and JVP, this knowledge is lost once the (structureless) Jacobian @f(✓, x1)/@✓l is instantiated, and cannot be taken advantage of in the following contraction with @f(✓, x2)/@✓l. We discuss how we impose this structure to compute the NTK for general neural networks in §4.

heading: 3.5 NTK VIA NTK-VECTOR PRODUCTS
text: 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 we avoid performing a contraction altogether at the cost of extra VJP/JVP calls; this ends up being beneficial for FCNs.
We introduce the linear function performing the NTK-vector product: ⇥VP : v 2 RO 7! ⇥✓v 2 RO. Applying this function to O columns of the identity matrix IO allows us to compute the NTK, i.e. ⇥✓IO = ⇥✓. The cost of evaluating the NTK in this fashion is equal to O times the cost of a single NTK-vector product evaluation ⇥VP(v). We now expand ⇥VP(v) = ⇥✓v as
@f (✓, x1)
@✓
@f (✓, x2)
@✓
T
v = @f (✓, x1)
@✓ VJP(f,✓,x2) (v) = JVP(f,✓,x1)
⇥ VJP(f,✓,x2) (v) ⇤ , (10)
where we have observed that, if contracted from right to left, the NTK-vector product can be expressed as a composition of a 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 LOW2 + O2W time and LW2 + OW2 + O2W memory. Batched inputs. In the batched setting Eq. (10) is repeated for each pair of inputs, and therefore time increases by a factor of N2 to become N2LOW2 + N2O2W. However, the memory cost grows only linearly in N (except for the cost of storing the NTK of size N2O2), since intermediate activations and derivatives necessary to compute the JVP and VJP can be computed for each batch x1 and x2
separately; these quantities are then reused for every pairwise combination resulting in a memory equivalent to the Jacobian, i.e. N2O2 + LW2 + NOW2 + NO2W + NLW , resulting in
NTK computation as a sequence of NTK-vector products costs N2LOW2 + N2O2W time and N2O2 + NOW2 + LW2 + NLW memory.

heading: 3.6 SUMMARY
text: Structured derivatives and NTK-vector products allow a reduction in the time cost of NTK computation 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. 1. We further confirm our methods can provide orders of magnitude speed-ups and memory savings on all major hardware platforms in Fig. 1 (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-ResNet hybrids Dosovitskiy et al. (2021); Steiner et al. (2021); and MLP-Mixers Tolstikhin et al. (2021)). Computing the full O ⇥ O = 1000 ⇥ 1000 NTK for many of these models on modern accelerators is only possible with Structured derivatives.

heading: 4 STRUCTURED DERIVATIVES FOR GENERIC FUNCTIONS
text: Here we 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 a common input domain: fi : ✓0, . . . , ✓L 2 RP0⇥···⇥PL 7! fi ✓0, . . . , ✓L 2 ROi (i 2 {1, 2}).
For NNs, typically ✓0, . . . , ✓L correspond to trainable parameters in layers 0, . . . , L, and fi ✓0, . . . , ✓L = f ✓0, . . . , ✓L, xi , xi being network inputs, Oi = O being the number of outputs (logits, classes). The NTK is defined as
⇥✓ (f1, f2) := LX
l=0
@f1 @✓l @f2 @✓l
T
2 RO1⇥O2 . (11)
Assume the following decompositions of fi into computational graphs made of primitives yi: fi ✓0, . . . , ✓L = f̃i y1i (✓ 0, . . . , ✓L), . . . , yKii (✓ 0, . . . , ✓L) (i 2 {1, 2}). (12)
with yki ✓0, . . . , ✓L 2 RYki . In common NNs, ykii would correspond to pre-activations evaluated on inputs xi in layer ki, and, without weight sharing, typically K1 = K2 = L. However, we do not impose any relationship between the number of parameter variables L and number of primitives K1 and K2, allowing arbitrary weight sharing. We can then use the chain rule in Eq. (2) to obtain:
⇥✓ (f1, f2) = L,K1,K2X
l,k1,k2
@f̃1
@yk11
@yk11 @✓l
! @f̃2
@yk22
@yk22 @✓l
!T = L,K1,K2X
l,k1,k2
@f̃1
@yk11
@yk11 @✓l @yk22 @✓l T @f̃2
@yk22
T
. (13)
All methods from §3 perform the sum of contractions in Eq. (13) one way or another. Jacobian contraction uses VJPs to implicitly contract each summand “outside-in”, i.e. it first computes @fi @✓l terms with VJPs followed by their contraction. As discussed in §3.3, this costs NO [FP] + N2O2P. NTK-vector products use both JVPs and VJPs to contract “Right-to-left”, i.e. first compute @f2/@✓l as an implicit contraction of @f2 @y2 with @y2 @✓l via VJP, followed by an implicit contraction of the result with @y1 @✓l via a JVP, followed by another implicit contraction with @f1 @y1 with another JVP. Per §3.5 this costs N2O [FP].
However, recall from §3.4, while JVPs and VJPs themselves are computationally optimal, higherorder computations like their contraction (Jacobian contraction) or composition (NTK-vector products) 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 y1 and y2. 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 @y @✓l, that allow us to simplify the contraction in Eq. (13). In practice this
amounts to replacing the inner terms @yk11 @✓l and @yk22 @✓l with their (much) smaller subarrays, and modifying the instructions passed to np.einsum that contracts all 4 terms. In §C we provide specific descriptions of our rules and their impact on the computational complexity of Eq. (13).
In Table 1 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 a 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, y1 and y2, present in the network. Nonetheless, we can offer heuristics that suggest when each method will be beneficial. The time complexity of Structured derivatives has the form of NO [FP] + NOG + N [J OP], where G is related to the cost of contraction, and J to the cost of computing @y @✓l (exact values depend on the structure present in y1 and y2). 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-offs 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, we include an automated setting which will perform FLOPs analysis for each method at compilation time and automatically select the most efficient one.

heading: 5 IMPLEMENTATION
text: Both algorithms are implemented in JAX (Bradbury et al., 2018) as the following function transformation ntk_fn : [f : (✓, x) 7! f(✓, x)] 7! [⇥ : (x1, x2, ✓) 7! ⇥✓(x1, x2)] , i.e. our function accepts any function f with the above signature and returns the efficient NTK kernel function operating on inputs x1 and x2 and parameterized by ✓. Inputs x, parameters ✓, and outputs f(✓, x) can be arbitrary PyTrees. We rely on many utilities from JAX and Neural Tangents (Novak et 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 IO matrix and to parallelize the computation over pairwise combinations of N inputs in each batch x1 and x2. Structured derivatives algorithm is implemented as a Jaxpr interpreter, built on top of the default JAX reverse mode AD interpreter. On a high level, the algorithm performs the sum in Eq. (13). Each summand is a contraction of 4 factors: @f̃1 @y1, @y1/@✓, @y2/@✓, @f̃2 @y2.
First, we linearize f to obtain a computational graph constructed out of a limited set (54,3 see
Table 5) of linear primitives y1, . . . , yK. Then, we can obtain two factors @f̃1 @y1, @f̃2 @y2 as part
of a backward pass almost identical to calling jax.jacobian (f)(✓, x). To contract these terms with @y1/@✓ and @y2/@✓, as described above, we query a dictionary of rules which map primitives to a structural description (§C.8); for a given pair of primitives, these rules allow us to analytically simplify the contraction and avoid explicitly instantiating the derivatives.
3JAX leverages a similar approach to implement only 54 transpose rules for linear primitives for reverse mode differentiation instead of 131 VJP rules (Frostig et al., 2021).

Review summaries:
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.

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.

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.

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.