NNGeometry: Easy and Fast Fisher Information Matrices and Neural Tangent Kernels in PyTorch

Practical and theoretical advances in deep learning have been accelerated by the development of an ecosystem of libraries allowing practitioners to focus on developing new techniques instead of spending weeks or months re-implementing the wheel . In particular , automatic differentiation frameworks such as Theano ( Bergstra et al. , 2011 ) , Tensorflow ( Abadi et al. , 2016 ) or PyTorch ( Paszke et al. , 2019 ) have been the backbone for the leap in performance of last decade ’ s increasingly deeper neural networks as they allow to compute average gradients efficiently , used in the stochastic gradient algorithm or variants thereof . While being versatile in neural networks that can be designed by varying the type and number of their layers , they are however specialized to the very task of computing these average gradients , so more advanced techniques can be burdensome to implement . While the popularity of neural networks has grown thanks to their always improving performance , other techniques have emerged , amongst them we highlight some involving Fisher Information Matrices ( FIM ) and Neural Tangent Kernels ( NTK ) . Approximate 2nd order ( Schraudolph , 2002 ) or natural gradient techniques ( Amari , 1998 ) aim at accelerating training , elastic weight consolidation ( Kirkpatrick et al. , 2017 ) proposes to fight catastrophic forgetting in continual learning and WoodFisher ( Singh & Alistarh , 2020 ) tackles the problem of network pruning so as to minimize its computational footprint while retaining prediction capability . These 3 methods all use the Fisher Information Matrix while formalizing the problem they aim at solving , but resort to using different approximations when going to implementation . Similarly , following the work of Jacot et al . ( 2018 ) , a line of work study the NTK in either its limiting infinite-width regime , or during training of actual finite-size networks . All of these papers start by formalizing the problem at hand in a very concise math formula , then face the experimental challenge that computing the FIM or NTK involves performing operations for which off-the-shelf automatic differentiation libraries are not well adapted . An even greater turnoff comes from the fact that these matrices scale with the number of parameters ( for the FIM ) or the number of examples in the training set ( for the empirical NTK ) . This is prohibitively large for modern neural networks involving millions of parameters or large datasets , a problem circumvented by a series of techniques to approximate the FIM ( Ollivier , 2015 ; Martens & Grosse , 2015 ; George et al. , 2018 ) . NNGeometry aims at making use of these approximations effortless , so as to accelerate development or analysis of new techniques , allowing to spend more time on the theory and less time in fighting development bugs . NNGeometry ’ s interface is designed to be as close as possible to maths formulas . In summary , this paper and library contribute : • We introduce NNGeometry by describing and motivating design choices . – A unified interface for all FIM and NTK operations , regardless of how these are ap- proximated . – Implicit operations for ability to scale to large networks .. • Using NNGeometry , we get new empirical insights on FIMs and NTKs : – We compare different approximations in different scenarios . – We scale some NTK evolution experiments to TinyImagenet . 1 PRELIMINARIES . 1.1 NETWORK LINEARIZATION . Neural networks are parametric functions f ( x , w ) : X × Rd → Rc where x ∈ X are covariates from an input space , and w ∈ Rd are the network ’ s parameters , arranged in layers composed of weight matrices and biases . The function returns a value in Rc , such as the c scores in softmax classification , or c real values in c-dimensional regression . Neural networks are trained by iteratively adjusting their parameters w ( t+1 ) ← w ( t ) + δw ( t ) using steps δw ( t ) typically computed using the stochastic gradient algorithm or variants thereof , in order to minimize the empirical risk of a loss function . In machine learning , understanding and being able to control the properties of the solution obtained by an algorithm is of crucial interest , as it can provide generalization guarantees , or help design more efficient or accurate algorithms . Contrary to ( kernelized ) linear models , where closed-form expressions of the empirical risk minimizer exist , deep networks are non-linear functions , whose generalization properties and learning dynamics is not yet fully understood . Amongst the recent advances toward improving theory , is the study of the linearization ( in w ) of the deep network function f ( x , w ) : f ( x , w + δw ) = f ( x , w ) + J ( x , w ) δw + o ( ‖δw‖ ) 1 ( 1 ) where J ( x , w ) = ∂f ( x , w ) ∂w is the Jacobian with respect to parameters w , computed in ( w , x ) , mapping changes in parameter space δw to corresponding changes in output space using the identity δf ( x , w , δw ) = J ( x , w ) δw . For tiny steps δw , we neglect the term o ( ‖δw‖ ) thus f is close to its linearization . It happens for instance at small step sizes , or in the large-width limit with the specific parameter initialization scheme proposed by Jacot et al . ( 2018 ) . 1The Landau notation o ( pronounced ” little-o ” ) means a function whose exact value is irrelevant , with the property that limx→0 o ( x ) x = 0 , or in other words that is negligible compared to x for small x . 1.2 PARAMETER SPACE METRICS AND FISHER INFORMATION MATRIX . While neural networks are trained by tuning their parameters w , the end goal of machine learning is not to find the best parameter values , but rather to find good functions , in a sense that is dependent of the task at hand . For instance different parameter values can represent the same function ( Dinh et al. , 2017 ) . On the contrary 2 parameter space steps δw1 and δw2 with same euclidean norm can provide very different changes in a function ( δf ( x , w , δw1 ) 6= δf ( x , w , δw2 ) ) . In order to quantify changes of a function , one generally defines a distance2 on the function space . Examples of such distances are the Lk-norms , Wasserstein distances , or the KL divergence used in information geometry . To each of these function space distances correspond a parameter space metric . We continue our exposition by focusing on the KL divergence , which is closely related to the Fisher Information Matrix , but our library can be used for other function space distances . Suppose f is interpreted as log-probability of a density p : log p ( x , w ) = f ( x , w ) , the KL divergence gives a sense of how much the probability distribution changes when adding a small increment δw to the parameters of f ( x , w ) . We can approximate it as : KL ( p ( x , w ) ‖p ( x , w + δw ) ) = ∫ x∈X log ( p ( x , w ) p ( x , w + δw ) ) dp ( x , w ) ( 2 ) = 1 2 ∫ x∈X ( 1 p ( x , w ) J ( x , w ) δw ) 2 dp ( x , w ) + o ( ‖δw‖2 ) ( 3 ) where we used this form ( derived in appendix ) in order to emphasize how steps in parameter space δw affect distances measured on the function space : equation 3 is the result of i ) taking a step δw in parameter space ; ii ) multiplying with J ( x , w ) to push the change to the function space ; iii ) weight this function space change using p ( x , w ) −1 ; iv ) square and sum . In particular , because of the properties of the KL divergence , there is no second derivative of f involved , even if equation 3 is equivalent to taking the 2nd order Taylor series expansion of the KL divergence . We can rewrite in a more concise way : KL ( f ( x , w ) ‖f ( x , w + δw ) ) = δw > Fwδw + o ( ‖δw‖2 ) ( 4 ) which uses the d × d FIM Fw = ∫ x∈X 1 p ( x , w ) 2 J ( x , w ) > J ( x , w ) dp ( x , w ) . In particular , we can now define the norm ‖δw‖Fw = δw > Fwδw used in the natural gradient algorithm ( Amari ( 1998 ) , also see Martens ( 2020 ) for a more thorough discussion of the FIM ) , in elastic weight consolidation ( Kirkpatrick et al. , 2017 ) , or in pruning ( Singh & Alistarh , 2020 ) . Other quantities also share the same structure of a covariance of parameter space vectors , such as the covariance of loss gradients in TONGA ( Roux et al. , 2008 ) , the second moment of loss gradients3 ( Kunstner et al. , 2019 ; Thomas et al. , 2020 ) , or posterior covariances in bayesian deep learning ( e.g . in Maddox et al . ( 2019 ) ) . 1.3 NEURAL TANGENT KERNEL . Another very active line of research around the linearization of equation 1 is to take inspiration from the rich literature on kernel methods by defining the neural tangent kernel ( NTK ) : kw ( x , y ) = J ( x , w ) J ( y , w ) > ( 5 ) In the limit of networks infinite width , Jacot et al . ( 2018 ) have shown that the tangent kernel remains constant through training using gradient descent , which allows to directly apply kernel learning theory to deep learning . While this regime is of theoretical interest , it arguably does not explain what happens at finite width , where the NTK evolves during training . While kernels are functions of the whole input space X ×X , we often only have access to a limited number of samples in a datasets . We thus resort to using the kernel evaluated at points xi of a 2We here use the notion of distance informally . 3The second moment of loss gradients is sometimes called empirical Fisher . training or a test set , called the Gram Matrix ( Kw ) ij = kw ( xi , xj ) . Note that in the case where the output space is multidimensional with dimension c , then Kw is in fact a 4d tensor . 2 DESIGN AND IMPLEMENTATION . 2.1 DIFFICULTIES . Current deep learning frameworks such as PyTorch and Tensorflow are well adapted to neural network training , i.e . computing average gradients over parameters , used in optimizers such as Adam and others . However , when going to more advanced algorithms or analysis techniques involving FIMs and NTKs , practitioners typically have to hack the framework ’ s internal mechanisms , which is time consuming , error prone , and results in each project having its own slightly different implementation of the very same technique . We here list the difficulties in computing FIMs and NTKs using current frameworks : Per-example gradient FIMs and NTKs require per-example Jacobians J ( xi , w ) of a dataset ( xi ) i . This can be obtained by looping through examples x , but at the cost of not using mini-batched operations , thus missing the benefit of using GPUs . NNGeometry ’ s Jacobian generator extensively use efficient techniques such as Goodfellow ( 2015 ) Memory usage and computational cost A FIM matrix is d × d where d is the total number of parameters . With a memory cost in O ( d2 ) , this is prohibitively costly even for moderate size networks . Typical linear algebra operations have a computational cost in eitherO ( d2 ) ( e.g . matrixvector product ) or evenO ( d3 ) ( e.g . matrix inverse ) . NNGeometry instead comes with recent lower memory intensive approximations .

Summary: This paper describes a new PyTorch package, NNGeometry, for computing complicated neural network objects, such as the Fisher Information Matrix (FIM) and the Neural Tangent Kernel (NTK). The package uses an abstract representation to allow the user to implicitly choose between different approximations to these objects and automatically makes a bunch of efficient choices "under the hood" for the user. The paper concludes with a selection of experiments that compare these approximations as well as verify their validity against Monte-Carlo estimates of these matrices.