ViViT: Curvature access through the generalized Gauss-Newton's low-rank structure

1 INTRODUCTION & MOTIVATION . The large number of trainable parameters in deep neural networks imposes computational constraints on the information that can be made available to optimization algorithms . Standard machine learning libraries ( Abadi et al. , 2015 ; Paszke et al. , 2019 ) mainly provide access to first-order information in the form of average mini-batch gradients . This is a limitation that complicates the development of novel methods that may outperform the state-of-the-art : They must use the same objects to remain easy to implement and use , and to rely on the highly optimized code of those libraries . There is evidence that this has led to stagnation in the performance of first-order optimizers ( Schmidt et al. , 2021 ) . Here , we thus study how to provide efficient access to richer information , namely higher-order derivatives and full statistics of the mini-batch loss . Recent advances in automatic differentiation ( Bradbury et al. , 2020 ; Dangel et al. , 2020 ) have made such information more readily accessible through vectorization of algebraic structure in the differentiated loss . We leverage and extend this functionality to efficiently access curvature in form of the Hessian ’ s generalized Gauss-Newton ( GGN ) approximation . It offers practical advantages over the Hessian and is established for training ( Martens , 2010 ; Martens & Grosse , 2015 ) , compressing ( Singh & Alistarh , 2020 ) , or adding uncertainty to ( Ritter et al. , 2018b ; a ; Kristiadi et al. , 2020 ) neural nets . It is also linked theoretically to the natural gradient method ( Amari , 2000 ) via the Fisher information matrix ( Martens , 2020 , Section 9.2 ) , and has been used to investigate the generalization of neural networks ( Jastrzebski et al. , 2020 ; Thomas et al. , 2020 ) . Traditional ways to access curvature fall into two categories . Firstly , repeated automatic differentiation allows for matrix-free exact multiplication with the Hessian ( Pearlmutter , 1994 ) and GGN ( Schraudolph , 2002 ) . Iterative linear and eigensolvers can leverage such functionality to compute Newton steps ( Martens , 2010 ; Zhang et al. , 2017 ; Gargiani et al. , 2020 ) and spectral properties ( Sagun et al. , 2017 ; 2018 ; Adams et al. , 2018 ; Ghorbani et al. , 2019 ; Papyan , 2019b ; Yao et al. , 2019 ; Granziol et al. , 2021 ) on arbitrary architectures thanks to the generality of automatic differentiation . However , repeated matrix-vector products represent a critical factor for performance . Secondly , Kronecker-factored approximate curvature ( K-FAC ) ( Martens & Grosse , 2015 ; Grosse & Martens , 2016 ; Botev et al. , 2017 ; Martens et al. , 2018 ) constructs an explicit light-weight representation of the GGN based on its algebraic Kronecker structure . The computations are streamlined via gradient backpropagation and the resulting matrices are cheap to store and invert . This allows K-FAC to scale : It has been used successfully with very large mini-batches ( Osawa et al. , 2019 ) . One reason for this efficiency is that K-FAC only approximates the GGN ’ s block diagonal , neglecting interactions across layers . Such terms could be useful , however , for other applications , like uncertainty quantification with Laplace approximations ( Ritter et al. , 2018b ; a ; Kristiadi et al. , 2020 ) that currently rely on K-FAC . Moreover , due to its specific design for optimization , the Kronecker representation does not become more accurate with more data . It remains a simplification , exact only under assumptions unlikely to be met in practice ( Martens & Grosse , 2015 ) . This might be a downside for applications that depend on a precise curvature proxy . Here , we propose VIVIT ( inspired by V V > in Equation ( 3 ) ) , a vivid curvature model that leverages the GGN ’ s low-rank structure . Like K-FAC , its representation is computed in parallel with gradients . But it allows a cost-accuracy trade-off , ranging from the exact GGN to an approximation that has the cost of a single gradient computation . Our contributions are as follows : • We highlight the GGN ’ s low-rank structure , and with it a structural limit for the inherent curvature information contained in a mini-batch . • This low-rank structure allows for efficient computation of various GGN properties : The exact eigenvalue spectrum , including eigenvectors , and per-sample directional derivatives . They enable VIVIT to model curvature noise in a mini-batch , in contrast to existing methods . • Approximations allow VIVIT to flexibly trade off computational cost and accuracy . We empirically demonstrate scalability on deep neural networks and provide a fully-featured efficient implementation in PYTORCH ( Paszke et al. , 2019 ) on top of the BACKPACK ( Dangel et al. , 2020 ) package.1 Using VIVIT , we illustrate that noise in deep learning poses a challenge for the stability of secondorder methods and give a simple example how its quantities can be used to address this problem . 2 NOTATION & METHOD . Consider a model f : Θ × X → Y and a dataset { ( xn , yn ) ∈ X × Y } Nn=1 . For simplicity we use N for both the mini-batch and training set size . The network , parameterized by θ ∈ Θ , maps a sample xn to a prediction ŷn . Predictions are scored by a convex loss function ` : Y × Y → R ( e.g . cross-entropy or square loss ) , which compares to the ground truth yn . The training objective L : Θ→ R is the empirical risk L ( θ ) = 1N ∑N n=1 ` ( f ( θ , xn ) , yn ) . ( 1 ) We use ` n ( θ ) = ` ( f ( θ , xn ) , yn ) and fn ( θ ) = f ( θ , xn ) for per-sample losses and predictions . For gradients , we write gn ( θ ) = ∇θ ` n ( θ ) and g ( θ ) = ∇θL ( θ ) , suppressing θ if unambiguous . We also set Θ = RD and Y = RC with D , C the model parameter and prediction space dimensions , respectively . For classification , C is the number of classes . Hessian & GGN : Two-fold chain rule application to the split ` ◦ f decomposes the Hessian of Equation ( 1 ) into two parts∇2θL ( θ ) = G ( θ ) +R ( θ ) ∈ RD×D ; the positive semi-definite GGN G = 1N ∑N n=1 ( Jθfn ) > ( ∇2fn ` n ) ( Jθfn ) = 1 N ∑N n=1Gn ( 2 ) and a residual R = 1/N ∑N n=1 ∑C c=1 ( ∇2θ [ fn ] c ) [ ∇fn ` n ] c. Here , we use the Jacobian Jab that contains partial derivatives of b with respect to a , [ Jab ] ij = ∂ [ b ] i/∂ [ a ] j . As the residual may alter the Hessian ’ s definiteness – an undesirable property for many applications – we focus on the GGN . Low-rank structure : By basic inequalities , Equation ( 2 ) has rank ( G ) ≤ NC.2 To make this explicit , we factorize the positive semi-definite Hessian ∇2fn ` n = ∑C c=1 sncs > nc , where snc ∈ RC and denote its backpropagated version by vnc = [ Jθfn ] > snc ∈ RD . Absorbing sums into matrix multiplications , we arrive at the GGN ’ s outer product representation that lies at the heart of VIVIT , G = 1N ∑N n=1 ∑C c=1 vncv > nc = V V > with V = 1√ N ( v11 v12 . . . vNC ) ∈ RD×NC . ( 3 ) 1Code available at https : //github.com/PwLo3K46/vivit . 2We assume the overparameterized deep learning setting ( NC < D ) and suppress the trivial rank bound D. V allows for exact computations with the explicit GGN matrix , at linear rather than quadratic memory cost in D. We first formulate the extraction of relevant GGN properties from this factorization , before addressing how to further approximate V to reduce memory and computation costs . 2.1 COMPUTING THE FULL GGN EIGENSPECTRUM . Each GGN eigenvalue λ ∈ R satisfies the characteristic polynomial det ( G − λID ) = 0 with identity matrix ID ∈ RD×D . Leveraging the VIVIT factorization of Equation ( 3 ) and the matrix determinant lemma , the D-dimensional eigenproblem reduces to that of the much smaller Gram matrix G̃ = V > V ∈ RNC×NC which contains pairwise scalar products of vnc ( see Appendix A.1 ) , det ( G− λID ) = 0 ⇔ ( −λ ) D−NC det ( G̃− λINC ) = 0 . ( 4 ) With at least D −NC trivial solutions that represent vanishing eigenvalues , the GGN curvature is flat along most directions in parameter space . Nontrivial solutions that give rise to curved directions are fully-contained in the Gram matrix , and hence much cheaper to compute . For example , the left panel of Figure 1a visualizes the full , exact GGN ’ s empirical spectral density on a mini-batch for a deep convolutional neural net on CIFAR-10 . It reproduces the characteristics that have been reported by numerous works , e.g . Sagun et al . ( 2018 ) : An extensive amount of vanishing or small eigenvalues and a small number of larger outliers . Despite these various Hessian spectral studies which rely on iterative eigensolvers and implicit matrix multiplication ( Sagun et al. , 2017 ; 2018 ; Adams et al. , 2018 ; Ghorbani et al. , 2019 ; Papyan , 2019b ; Yao et al. , 2019 ; Granziol et al. , 2021 ) , we are not aware of works that extract the exact GGN spectrum from its Gram matrix . In contrast to those techniques , this matrix can be computed in parallel with gradients in a single backward pass , which results in less sequential overhead . In fact , our approach allows for plots like Figure 1 to be efficiently live-monitored during training , which may be interesting for practitioners that seek to better understand their model ( Schneider et al. , 2021 ) . Eigenvalues themselves can help identify reasonable hyperparameters , like learning rates . But we can also reconstruct the associated eigenvectors in parameter space . These are directions along which curvature information is contained in the mini-batch . Let S̃+ = { ( λk , ẽk ) |λk 6= 0 , G̃ẽk = λkẽk } Kk=1 denote the nontrivial Gram spectrum with orthonormal eigenvectors ẽ > j ẽk = δjk ( δ represents the Kronecker delta and K = rank ( G ) ) . Then , the transformed set of vectors ek = 1/ √ λkV ẽk are orthonormal eigenvectors ofG associated to eigenvalues λk ( see Appendix A.2 ) , ∀ ( λk , ẽk ) ∈ S̃+ : G̃ẽk = λkẽk =⇒ GV ẽk = λkV ẽk . ( 5 ) The eigenspectrum provides access to the GGN ’ s pseudo-inverse based on V and S̃+ , required by e.g . second-order methods ( see Section 4.2 ) . As we show next , quadratic models defined through the GGN naturally decompose along the eigenvectors from S+ = { ( λk , ek ) |λk 6= 0 , Gek = λkek } Kk=1 .

Summary: This paper highlights how the low-rank structure of the generalized Gauss-Newton (GGN) approximation of the Hessian can be used as a computationally efficient tool to study the loss landscape of deep neural networks. In particular, authors discuss methods to compute the full spectrum of the GGN and thus providing access to per-sample directional gradients and curvature approximation. Through the lens of the GGN spectrum, authors make observations on the geometry of the loss landscape and its evolution during training, and propose an adaptive damping technique for second-order optimizers that utilizes the GGN curvature information.