F AST F INITE W IDTH N EURAL T ANGENT K ERNEL
Anonymous authors
Paper under double-blind review
A BSTRACT
The Neural Tangent Kernel (NTK), deﬁned as the outer product of the neural network (NN) Jacobians, ⇥ ✓ ( x 1 , x 2 ) = ⇥ @
f ( ✓ , x 1 )  @✓ ⇤ ⇥ @
f ( ✓ , x 2 )  @✓ ⇤
T , has emerged as a central object of study in deep learning.
In the inﬁnite width limit, the NTK can sometimes be computed analytically and is useful for understanding training and generalization of NN architectures.
At ﬁnite widths, the NTK is also used to better initialize NNs, compare the conditioning across models, perform architecture search, and do meta-learning.
Unfortunately, the ﬁnite width NTK is notoriously expensive to compute, which severely limits its practical utility.
We perform the ﬁrst in-depth analysis of the compute and memory requirements for NTK computation in ﬁnite 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 ﬁnite width NTK, dramatically improving efﬁciency.
We open-source [github.com/iclr2022anon/fast ﬁnite 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.
I NTRODUCTION
The past few years have seen signiﬁcant 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 signiﬁcant 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 inﬁ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 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 inﬁnite width limit and help explain the impressive generalization capabilities of NNs.
Insights from the wide network limit have had signiﬁcant practical impact.
The conditioning of the
NTK has been shown to signiﬁcantly 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 efﬁcient 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 ﬁelds (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 signiﬁcant promise of theory based on the NTK, computing the NTK in practice is challenging.
In the inﬁnite width limit, the NTK can sometimes be computed analytically.
However, it remains intractable for many architectures, and ﬁnite width corrections can be important to describe actual NNs used in practice.
The NTK matrix can be computed for ﬁnite width networks as the outer product of Jacobians using forward or reverse mode automatic differentiation (AD), ⇥ ✓ ( x 1 , x 2 ) | {z } O ⇥ O := ⇥ @ f ( ✓ , x 1 )  @✓ ⇤ | {z } O ⇥ P ⇥ @
f ( ✓ , x 2 )  @✓ ⇤ T | {z } P ⇥ O , (1) where f is the forward pass NN function producing outputs in R O , ✓ 2 R P are all trainable parameters, and x 1 and x 2 are two inputs to the network. If inputs are batches of sizes N 1 and N 2 , the NTK is an N 1 O ⇥ N 2 O matrix.
Unfortunately, evaluating Eq. (1) is often infeasible due to time and memory requirements.
In this paper, we perform the ﬁrst in-depth analysis of the compute and memory requirements for the NTK as in Eq.
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 efﬁcient 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.
R ELATED W ORK
The ﬁnite 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 efﬁcient NTK computation as their central goal).
In order to compare performance of models based on the NTK and the inﬁnite width NTK, Arora et al. (2019a, Table 2) compute the NTK of up to 20-layer, 128-channel CNN in a binary CIFARclassiﬁcation 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 2 12 -wide or 1-layer 2 14 -wide FCNs.
In the same setting with the same hardware (NVIDIA V100), we can reach widths of at least 2 14 and 2 18 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 inﬁnite 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/
(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,which impacts the quality of the respective trainability/generalization proxy.
[Table]
N 2 LO 2 W
{+NOW 2 + N 2 O 
 ++} NLW + LW
Don’t
[Table]
N {+2 O 2 W + N 2 LOW NOW 2 + N 2 O 
 + NLW + LW O > W or N =+}
[Table]
N 2 LO 2 W + N LOW
{+NOW + N 2 O 
+}
+ NLW + LW
O < W or L =
[Table]
have a cheap FP , as it consists of a single matrix multiply with the parameter matrix.
CNNs, notably when the number of output pixels D is large, have a costly FP , since it amounts to D matrix multiplies with the parameters. For precise analysis, see Table 1 for FCN and
Table 7 for CNN.
In contrast, in this work we can compute the full 1000 ⇥ 1000 NTK on the same models (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 inﬁnite width/linearization approximation or higher-quality trainability/generalization proxy respectively, which makes our work especially relevant to modern research.
E FFICIENT FINITE WIDTH NTK S 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 time and memory complexities for FCN in Table 1, CNN in Table 7, and a more general takeaway in Table 2.
Setting.
Consider an L -layer FCN f
( ✓ , x ) = ✓ L  [Equation]
R O , 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 R LW 2 + OW .
We further deﬁne x l :=   y l  as post-activations (with x 0 := x ), and y l := ✓ l x l as pre-activations with y L = 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 x 1 and x 2 is an O ⇥ O matrix deﬁned as
[Equation]
where we have deﬁned ⇥ l ✓ to be the summands.
We omit dependence on x 1 , x 2 , 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 ﬁrst contribution, that consists in observing that many intermediate operations on weights performed by NNs possess a certain structure, that can allow linear algebra simpliﬁcations 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 beneﬁcial in many (but not all) settings.
3. J ACOBIAN - VECTOR PRODUCTS AND VECTOR -J ACOBIAN PRODUCTS
We begin by deﬁning Jacobian-vector products and vector-Jacobian products:
[Equation]
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.
Below we present a brief intuition for why this is the case, and refer the reader to the JAX Autodiff Cookbook for more details, and (Griewank & Walther, 2008, Section 3) for a rigorous treatment of the subject. JVP can be computed by traversing a computational graph of the same topology as FP , except for primitive nodes in the graph need to be augmented to compute not only the forward pass of the node, but also the JVP of the node (see Fig. 0). Due to identical topology and order of evaluation, asymptotically time and memory costs remain unchanged. However, constructing the augmented nodes in the JVP graph, and their consequent evaluation results in extra time cost proportional to the size of the graph.
Therefore in practice JVP costs about 3 ⇥ FP time and 2 ⇥ FP memory. VJP , as a linear function of cotangents f c , is precisely the transpose of the linear function JVP. As such, it can be computed by traversing the transpose of the JVP graph (Fig. 0, right), with each JVP node replaced by its transposition as well.
This results in identical time and memory costs, as long as node transpositions are implemented efﬁciently.
However, their evaluation requires primal outputs y k (now inputs to the transpose nodes), which is why VJP necessitates an extra FP time cost to compute them (hence costlier than JVP, but still inconsequential asymptotically) and extra memory to store them, which can generally increase asymptotic memory requirements.
For the case of FCNs, the time cost 2 of both operations is therefore [
[Equation] [ cost of all intermediate layers [Equation] of the top
[Equation]
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 ] =
[Equation]
[ size of all weights ] + [ size of activations in all layers ] =
[Equation]
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. [Table] JVP ( f
[Table]
✓ [Table] x ) : ✓ t 2 R P [Table]
f ( ✓
[Table]
x ) @✓ ✓ t 2 R O
[Table]
VJP ( f
[Table]
✓
[Table]
x
[Table]
f c 2 R
O
[Table]
/WaM02zoa0hcTJ1ReAFesOjktZOinsxp1mAjDGcilF4QZ4ZsP2ZwYwsA/yH8usahu0W5jxS+YzjKiEodVOY4mzuF10eWFyf0YHKZp+L4syzqbJDUYOJEV+raB5jXS8KTEwBdAUzdosP4hr8JUFpsOPjTQfr+GakPUbAP3G/BFjWVLoipytCL96sTbi9IUpyfd+Fn36Ue/Q69RFXvoITpARyhGz1EPvUMDNEQMfUc/0E/0KwiC4yAKTiq0tbPxPEC1CF79Bm8HMCc=</latexit> f (
[Figure]
[Figure] , in order to create and grow Batched inputs. If x is a batch of inputs of size N , the time cost of JVP and VJP increases linearly to NLW 2 .
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 ]
[Equation]
[ size of all weights ] + N [ size of activations in all layers ] + N [ size of all weight matrices ]
[Equation]
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 NLW 2 time and LW 2 + NW + NO memory.
• VJP costs NLW 2 time and LW 2 + NLW
+ NW 2 + NOW memory.
3. J ACOBIAN COMPUTATION
For neural networks, the Jacobian is most often computed by evaluating the VJP on rows of the identity matrix
I O , i.e.
[Equation]
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, the time and memory costs to compute the Jacobian are respectively,
ON ([ cost of all intermediate layers ] + [ cost of the top layer ])
[Equation]
[ size of all weights ] + N [ size of activations in all layers ] + ON [ size of a single weight matrix ]
[Equation]
Therefore, asymptotically,
[Equation]
3. J ACOBIAN 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.
[Equation]
The time cost of this contraction is O 2 W 2 , and the memory necessary to instantiate each factor and the result is OW 2 + O 2 .
Repeating the above operation for each ✓ l , we arrive at LO 2 W 2 time cost and unchanged memory, due to being able to process summands sequentially.
Batched inputs.
If we consider x 1 and x 2 to be input batches of size N , then the resulting NTK is a matrix of shape
NO ⇥ NO , and the time cost becomes N 2 LO 2 W 2 , while memory grows to [ NTK matrix size ] + [ factors size ]
= N 2 O 2 + NOW 2 .
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 N 2 LO 2 W 2 time and N 2 O 2 + NOW [-
-]2 + NO 2 W + LW 2 + NLW memory.
3. L EVERAGING STRUCTURED DERIVATIVES FOR COMPUTING THE NTK
We can rewrite ⇥ l ✓ in Eq. (35) using the chain rule and our pre- and post-activation notation as: ⇥ l ✓
[Equation] [Equation]
At face value, rewriting Eq. (35) this way is unhelpful as it appears to have introduced additional costly contractions.
However, recall that y l = ✓ l x l , and therefore
[Equation]
[Equation] where ⌦ is the Kronecker product.
Plugging Eq. (8) into Eq. (7) we obtain (see §G)
[Equation]
|{z
[Equation]
|{z}
W ⇥
[Equation]
CA
[Equation] [Equation]
| {z }
[Equation]
y lx
[Equation]
T | {z }
W ⇥ O  , (9) and observe that it takes only O 2 W time and OW + O 2 memory.
Accounting for depth, time cost increases by a factor of depth L and becomes LO 2 W , while memory does not change since the summands can be processed sequentially.
Batched inputs.
In the batched setting, the time cost grows quadratically with the size of the NTK to N 2 LO 2 W , while the memory cost increases to N 2 O 2 + NOW to store the result, ⇥ l ✓ ( x 1 , x 2 ) , and factors, @ f ( ✓ , x )  @ y lx , respectively.
Finally, we need to account for the cost of computing the derivatives, @
f  @ y l , and post-activations, x l . Notice that both x l and @ f  @ y l 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 ]
[Equation]
The extra time cost is asymptotically the cost of O forward passes, NLOW 2 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 .
Putting everything together we ﬁnd the following costs,
By leveraging Structured derivatives in NN computations, we have reduced the cost of NTK to N 2 LO 2 W + NLOW 2 time and N 2 O 2 + NOW + LW 2 + NLW memory.
The key insight was to leverage the constant block-diagonal structure of the pre-activation derivatives @
y l  @✓ 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 efﬁciently compute single evaluations of the VJP and JVP, this knowledge is lost once the (structureless) Jacobian @ f ( ✓ , x 1 ) / @ ✓ l is instantiated, and cannot be taken advantage of in the following contraction with @ f ( ✓ , x 2 ) / @✓ l .
We discuss how we impose this structure to compute the NTK for general neural networks in §4.
3. NTK VIA NTK- VECTOR PRODUCTS
Computing the Jacobian contraction using Jacobian ﬁrst 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.
In this section we avoid performing a contraction altogether at the cost of extra VJP/JVP calls; this ends up being beneﬁcial for FCNs.
We introduce the linear function performing the NTK-vector product:
[Equation]
Applying this function to O columns of the identity matrix I O allows us to compute the NTK, i.e.
⇥ ✓
I O = ⇥ ✓ .
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
[Equation]
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 LOW 2 + O 2 W time and LW 2 + OW 2 + O 2 W memory.
Batched inputs.
In the batched setting Eq. (10) is repeated for each pair of inputs, and therefore time increases by a factor of N 2 to become N 2 LOW 2 + N 2 O 2 W .
However, the memory cost grows only linearly in N (except for the cost of storing the NTK of size N 2 O 2 ), since intermediate activations and derivatives necessary to compute the JVP and VJP can be computed for each batch x 1 and x 2 separately; these quantities are then reused for every pairwise combination resulting in a memory cost equal to the cost of computing the Jacobian over a batch, i.e.
N 2 O 2 +
[Equation]
NTK computation as a sequence of NTK-vector products costs N 2 LOW 2 + N 2 O 2 W time
and N 2 O 2 + NOW 2 + LW 2 + NLW memory.
3. S UMMARY
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 beneﬁcial for wide networks, with large W , and NTK-vector products are beneﬁcial for networks with large outputs O .
We conﬁrm
our predictions with FLOPs measurements in Fig.
We further conﬁrm our methods can provide orders of magnitude speed-ups and memory savings on all major hardware platforms in Fig. 1 (right) and Fig.
However, we notice that our wall-clock time measurements often deviate from predictions due to unaccounted constant overheads of various methods, hardware speciﬁcs, padding, and the (largely black-box) behavior of the XLA compiler.
Notably, in practice, we ﬁnd Structured derivatives almost always outperform NTK-vector products.
Finally, we evaluate our methods in the wild, and conﬁrm computational beneﬁts 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.
; Steiner et al.
(2021); and
MLP-Mixers Tolstikhin et al. Computing the full O ⇥ O = 1000 ⇥ NTK of many of these models on modern accelerators is only possible with Structured derivatives.  S TRUCTURED DERIVATIVES FOR GENERIC FUNCTIONS
Here we generalize the idea of leveraging structure in subexpressions presented in §3.4.
This section (and our implementation) is not speciﬁc to NNs and applies to any differentiable function.
Consider two differentiable functions deﬁned on a common input domain: [Equation]
[Equation]
For NNs, typically  ✓ 0 , . .
. , ✓ L  correspond to trainable parameters in layers 0 , . . .
, L , and f i  ✓ 0 , . . .
, ✓ L  = f  ✓ 0 , . .
. , ✓ L , x i  , x i being network inputs, O i
= O being the number of outputs (logits, classes).
The NTK is deﬁned as
[Equation]
[Equation] Assume the following decompositions of f i into computational graphs made of primitives y i : [Equation]
[Equation]
12) with y ki  ✓ 0 , . .
, ✓ L  2 R Y ki .
In common NNs, y k
i i would correspond to pre-activations evaluated on inputs x i in layer
k i , and, without weight sharing, typically K 1 = K 2 =
L .
However, we do not impose any relationship between the number of parameter variables L and number of primitives Kand K 2 , allowing arbitrary weight sharing.
We can then use the chain rule in Eq.
(2) to obtain: ⇥ ✓ ( f 1 , f 2 ) = L , K 1 , K [Equation] l,k 1 ,k
[Equation]
˜ f
[Equation]
y k
[Equation]
y k
[Equation]
@✓ l ! @ ˜ f
[Equation]
y k
[Equation]
y k
[Equation]
@✓ l !
T =
L
[Equation]
K
[Equation]
K
[Equation]
k
[Equation]
k
[Equation]
˜
f@ y k 1@
y k [Equation] 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 ﬁrst computes @
f i  @✓ l terms with VJPs followed by their contraction.
As discussed in §3.3, this costs (in a batched setting)
NO [ FP ] + N 2 O 2 P , where FP is the cost of the forward pass.
NTK-vector products use both JVPs and VJPs to contract “Right-to-left”, i.e. ﬁrst compute @
f 2 / @✓ l as an implicit contraction of @ f @
y 2 with @ y @✓ l via VJP, followed by an implicit contraction
(6) All methods, as expected, scale quadratically with width W
Wall clock
: In real appli , given constraints and delicate with the XLA compiler we observe that products improve upon contraction for O , but the effect is not (see for small and Fig. 3, notably GPU platforms)
Structured derivatives robustly outperform all methods, including simply computing the
, as discussed in Structured derivatives have lower and to up (see platforms §H for details.
of the result with @ y @✓ l via a JVP, followed by another implicit contraction with @ f @
y 1 with another JVP.
Per §3.5 this costs N 2 O
[ FP ] .
The idea of Structured derivatives is to design rules for efﬁcient computation of such contractions, similarly to how JAX and other packages have rules for efﬁcient JVPs and VJPs.
From Eq. (13), in the general case this requires hand-made rules for all pairwise combinations of primitives y 1 and y 2 .
Due to quadratic scaling in the number of primitives, we restrict the current implementation to rules that operate on individual primitives y .
This still provide substantial computational beneﬁt.
Speciﬁcally, 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.
In practice this amounts to replacing the inner terms @
y k 1 @✓ l and @ y k 2 @✓ l with their (much) smaller subarrays, and modifying the instructions passed to np.einsum that contracts all 4 terms.
In §C we provide speciﬁc descriptions of our rules and their impact on the computational complexity of Eq.
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 beneﬁcial in a much wider set of operations used by contemporary ImageNet models in Fig. 2 and Fig.
For both Structured derivatives and NTK-vector products a fully general and rigorous comparison of complexities is not feasible since it will rely upon speciﬁcs of the model architecture and the pairs of primitives, y 1 and y 2 , present in the network.
Nonetheless, we can offer heuristics that suggest when each method will be beneﬁcial.
The time-complexity of Structured derivatives is at most NO [ g ( Y , P )]
+ N 2 O 2 min
( Y , P ) ,
where max( Y , P )  g ( Y , P )  YP (Table 3), and the exact complexity depends on the structure present in y 1 and y 2 .
This is guaranteed to be no worse than
Jacobian contraction if g ( Y , P )  [ FP ] , which is the case 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 efﬁcient one.
I MPLEMENTATION
Both algorithms are implemented in JAX (Bradbury et al., 2018) as the following function transformation ntk_fn : [ f : ( ✓ , x ) 7! f
( ✓ , x )] 7! [ ⇥ : ( x 1 , x 2 , ✓ ) 7! ⇥
✓ ( x 1 , x 2 )] , i.e. our function accepts any function f with the above signature and returns the efﬁcient NTK kernel function operating on inputs x 1 and x 2 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 I O matrix and to parallelize the computation over pairwise combinations of N inputs in each batch x 1 and x 2 .
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.
Each summand is a contraction of 4 factors: @
˜ f @
y 1 , @ y 1 / @✓ , @ y 2 / @
✓ , @ ˜ f @
y 2 .
First, we linearize f to obtain a computational graph constructed out of a limited set (54, 3 see
Table 5) of linear primitives y 1 , . .
, y K .
Then, we can obtain two factors @
˜ f @
y 1 , @ ˜ f
y 2 as part of a backward pass almost identical to calling jax.jacobian ( f )( ✓ , x ) .
To contract these terms with @ y 1 / @
✓ and @ y 2 / @✓ , 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.