Fast Finite Width Neural Tangent Kernel

⇥ @ 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 . 1 INTRODUCTION . 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 . 2 RELATED WORK . The finite width NTK ( denoted as simply NTK throughout this work ) has been used extensively in many recent works , but to our knowledge implementation details and compute costs were rarely made public . Below we draw comparison to some of these works , but we stress that it only serves as 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 . 3 EFFICIENT FINITE WIDTH NTKS IN A SIMPLIFIED SETTING . To gain intuition for the problem , we start by analyzing and improving the cost of computing the NTK for a simple FCN . See §F for an equivalent analysis of CNNs . We summarize the resulting complexities for FCN in Table 1 , CNN in Table 7 , and a general takeaway in Table 2 . 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 .

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