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.

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

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

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

abstractText: 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.

authors:
Martı́n Abadi
Paul Barham
Jianmin Chen
Zhifeng Chen
Andy Davis
Jeffrey Dean
Matthieu Devin
Sanjay Ghemawat
Geoffrey Irving
Michael Isard

title: Tensorflow: A system for largescale machine learning

venue: In 12th USENIX Symposium on Operating Systems Design and Implementation

authors:
Ben Adlam
Jaehoon Lee
Lechao Xiao
Jeffrey Pennington
Jasper Snoek

title: Exploring the uncertainty properties of neural networks’ implicit priors in the infinite-width limit

venue: In International Conference on Learning Representations,

authors:
Sanjeev Arora
Simon S Du
Wei Hu
Zhiyuan Li
Russ R Salakhutdinov
Ruosong Wang

title: On exact computation with an infinitely wide neural net

venue: In Advances in Neural Information Processing Systems,

authors:
Sanjeev Arora
Simon S Du
Wei Hu
Zhiyuan Li
Ruosong Wang

title: Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks

venue: arXiv preprint arXiv:1901.08584,

authors:
Sanjeev Arora
Simon S. Du
Zhiyuan Li
Ruslan Salakhutdinov
Ruosong Wang
Dingli Yu

title: Harnessing the power of infinitely wide deep nets on small-data tasks

authors:
Yasaman Bahri
Ethan Dyer
Jared Kaplan
Jaehoon Lee
Utkarsh Sharma

title: Explaining neural scaling laws

venue: arXiv preprint arXiv:2102.06701,

authors:
Mikhail Belkin
Daniel Hsu
Siyuan Ma
Soumik Mandal

title: Reconciling modern machinelearning practice and the classical bias–variance trade-off

venue: Proceedings of the National Academy of Sciences,

authors:
James Bradbury
Roy Frostig
Peter Hawkins
Matthew James Johnson
Chris Leary
Dougal Maclaurin
Skye Wanderman-Milne

title: JAX: composable transformations of Python+NumPy programs, 2018

venue: URL http://github.com/google/jax

authors:
Andrew Brock
Soham De
Samuel L Smith

title: Characterizing signal propagation to close the performance gap in unnormalized resnets

venue: arXiv preprint arXiv:2101.08692,

authors:
Andrew Brock
Soham De
Samuel L Smith
Karen Simonyan

title: High-performance large-scale image recognition without normalization

venue: arXiv preprint arXiv:2102.06171,

authors:
Wuyang Chen
Xinyu Gong
Zhangyang Wang

title: Neural architecture search on imagenet in four gpu hours: A theoretically inspired perspective

authors:
Xiangning Chen
Cho-Jui Hsieh
Boqing Gong

title: When vision transformers outperform resnets without pretraining or strong data augmentations, 2021b

authors:
Yann Dauphin
Samuel S

title: Schoenholz. Metainit: Initializing learning by learning to initialize

authors:
Jia Deng
Wei Dong
Richard Socher
Li-Jia Li
Kai Li
Li Fei-Fei

title: Imagenet: A large-scale hierarchical image database

venue: IEEE conference on computer vision and pattern recognition,

authors:
Alexey Dosovitskiy
Lucas Beyer
Alexander Kolesnikov
Dirk Weissenborn
Xiaohua Zhai
Thomas Unterthiner
Mostafa Dehghani
Matthias Minderer
Georg Heigold
Sylvain Gelly
Jakob Uszkoreit
Neil Houlsby

title: An image is worth 16x16 words: Transformers for image recognition at scale

authors:
Simon S Du
Kangcheng Hou
Russ R Salakhutdinov
Barnabas Poczos
Ruosong Wang
Keyulu Xu

title: Graph neural tangent kernel: Fusing graph neural networks with graph kernels

venue: In Advances in Neural Information Processing Systems

authors:
Chelsea Finn
Pieter Abbeel
Sergey Levine

title: Model-agnostic meta-learning for fast adaptation of deep networks

venue: Proceedings of the 34th International Conference on Machine Learning,

authors:
Jean-Yves Franceschi
Emmanuel de Bézenac
Ibrahim Ayed
Mickaël Chen
Sylvain Lamprier
Patrick Gallinari

title: A neural tangent kernel perspective of gans

venue: arXiv preprint arXiv:2106.05566,

authors:
Roy Frostig
Matthew J Johnson
Dougal Maclaurin
Adam Paszke
Alexey Radul

title: Decomposing reverse-mode automatic differentiation

venue: arXiv preprint arXiv:2105.09469,

authors:
Adrià Garriga-Alonso
Laurence Aitchison
Carl Edward Rasmussen

title: Deep convolutional networks as shallow gaussian processes

authors:
Andreas Griewank
Andrea Walther

title: Evaluating Derivatives. Society for Industrial and Applied Mathematics, second edition, 2008. doi: 10.1137/1.9780898717761

title: Neural net training dynamics, January 2021. URL https://www.cs.toronto.edu/ ⇠rgrosse/courses/csc2541 2021/readings/L02 Taylor approximations.pdf

authors:
Boris Hanin
Mihai Nica

title: Finite depth and width corrections to the neural tangent kernel

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Jonathan Heek
Anselm Levskaya
Avital Oliver
Marvin Ritter
Bertrand Rondepierre
Andreas Steiner
Marc van Zee

title: Flax: A neural network library and ecosystem for JAX, 2020

venue: URL http://github.com/google/flax

authors:
Tom Hennigan
Trevor Cai
Tamara Norman
Igor

title: Babuschkin. Haiku: Sonnet for JAX, 2020

venue: URL http://github.com/deepmind/dm-haiku

authors:
Jiri Hron
Yasaman Bahri
Jascha Sohl-Dickstein
Roman Novak

title: Infinite attention: NNGP and NTK for deep attention networks

venue: In International Conference on Machine Learning,

authors:
Arthur Jacot
Franck Gabriel
Clement Hongler

title: Neural tangent kernel: Convergence and generalization in neural networks

authors:
Jaehoon Lee
Yasaman Bahri
Roman Novak
Sam Schoenholz
Jeffrey Pennington
Jascha Sohldickstein

title: Deep neural networks as gaussian processes

authors:
Jaehoon Lee
Lechao Xiao
Samuel S. Schoenholz
Yasaman Bahri
Roman Novak
Jascha SohlDickstein
Jeffrey Pennington

title: Wide neural networks of any depth evolve as linear models under gradient descent

authors:
Jaehoon Lee
Samuel S Schoenholz
Jeffrey Pennington
Ben Adlam
Lechao Xiao
Roman Novak
Jascha Sohl-Dickstein

title: Finite versus infinite neural networks: an empirical study

authors:
James Martens
Roger Grosse

title: Optimizing neural networks with kronecker-factored approximate curvature

venue: In International conference on machine learning,

authors:
Alexander G. de G. Matthews
Jiri Hron
Mark Rowland
Richard E. Turner
Zoubin Ghahramani

title: Gaussian process behaviour in wide deep neural networks

title: Solution directe de l’équation séculaire et de quelques problèmes analogues transcendants

venue: C. R. Acad. Sci. Paris,

title: Optimal accumulation of jacobian matrices by elimination methods on the dual computational graph

venue: Mathematical Programming,

title: Optimal jacobian accumulation is np-complete

title: Priors for infinite networks (tech

venue: rep. no. crg-tr-94-1). University of Toronto,

authors:
Timothy Nguyen
Zhourong Chen
Jaehoon Lee

title: Dataset meta-learning from kernel ridgeregression

venue: arXiv preprint arXiv:2011.00050,

authors:
Timothy Nguyen
Roman Novak
Lechao Xiao
Jaehoon Lee

title: Dataset distillation with infinitely wide convolutional networks

venue: arXiv preprint arXiv:2107.13034,

authors:
Roman Novak
Lechao Xiao
Jaehoon Lee
Yasaman Bahri
Greg Yang
Jiri Hron
Daniel A. Abolafia
Jeffrey Pennington
Jascha Sohl-Dickstein

title: Bayesian deep convolutional networks with many channels are gaussian processes

authors:
Roman Novak
Lechao Xiao
Jiri Hron
Jaehoon Lee
Alexander A. Alemi
Jascha Sohl-Dickstein
Samuel S. Schoenholz

title: Neural tangents: Fast and easy infinite neural networks in python

authors:
Boris N. Oreshkin
Pau Rodrı́guez López
Alexandre Lacoste

title: Tadam: Task dependent adaptive metric for improved few-shot learning

venue: In NeurIPS,

authors:
Daniel S Park
Jaehoon Lee
Daiyi Peng
Yuan Cao
Jascha Sohl-Dickstein

title: Towards nngpguided neural architecture search

venue: arXiv preprint arXiv:2011.06006,

authors:
Jeffrey Pennington
Yasaman Bahri

title: Geometry of neural network loss surfaces via random matrix theory

authors:
Samuel S Schoenholz
Justin Gilmer
Surya Ganguli
Jascha Sohl-Dickstein

title: Deep information propagation

venue: International Conference on Learning Representations,

authors:
Stefano Spigler
Mario Geiger
Stéphane d’Ascoli
Levent Sagun
Giulio Biroli
Matthieu Wyart

title: A jamming transition from under-to over-parametrization affects generalization in deep learning

venue: Journal of Physics A: Mathematical and Theoretical,

authors:
Andreas Steiner
Alexander Kolesnikov
Xiaohua Zhai
Ross Wightman
Jakob Uszkoreit
Lucas Beyer

title: How to train your vit? data, augmentation, and regularization in vision transformers

venue: arXiv preprint arXiv:2106.10270,

authors:
Matthew Tancik
Pratul P. Srinivasan
Ben Mildenhall
Sara Fridovich-Keil
Nithin Raghavan
Utkarsh Singhal
Ravi Ramamoorthi
Jonathan T. Barron
Ren Ng

title: Fourier features let networks learn high frequency functions in low dimensional domains

venue: NeurIPS,

authors:
Ilya Tolstikhin
Neil Houlsby
Alexander Kolesnikov
Lucas Beyer
Xiaohua Zhai
Thomas Unterthiner
Jessica Yung
Andreas Steiner
Daniel Keysers
Jakob Uszkoreit
Mario Lucic
Alexey Dosovitskiy

title: Mlp-mixer: An all-mlp architecture for vision, 2021

authors:
Lechao Xiao
Yasaman Bahri
Jascha Sohl-Dickstein
Samuel Schoenholz
Jeffrey Pennington

title: Dynamical isometry and a mean field theory of CNNs: How to train 10,000-layer vanilla convolutional neural networks

authors:
Lechao Xiao
Jeffrey Pennington
Samuel S Schoenholz

title: Disentangling trainability and generalization in deep learning

title: Non-Gaussian processes and neural networks at finite widths

venue: In Mathematical and Scientific Machine Learning Conference,

title: Tensor programs ii: Neural tangent kernel for any architecture

venue: arXiv preprint arXiv:2006.14548,

authors:
Greg Yang
Jeffrey Pennington
Vinay Rao
Jascha Sohl-Dickstein
Samuel S. Schoenholz

title: A mean field theory of batch normalization

authors:
Sergey Zagoruyko
Nikos Komodakis

title: Wide residual networks

venue: In British Machine Vision Conference,

authors:
Hongyi Zhang
Yann N Dauphin
Tengyu Ma

title: Fixup initialization: Residual learning without normalization

venue: arXiv preprint arXiv:1901.09321,

authors:
Yufan Zhou
Zhenyi Wang
Jiayi Xian
Changyou Chen
Jinhui Xu

title: Meta-learning with neural tangent kernels

authors:
Barret Zoph
Quoc V. Le

title: Neural architecture search with reinforcement learning. 2017

venue: URL https://arxiv.org/abs/1611.01578

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).

title: FAST FINITE WIDTH NEURAL TANGENT KERNEL