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