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 work proposes a curvature model VIVIT based on generalized Gauss-Newton (GGN) approximation for the training of neural networks with a convex loss function. The low-rank structure of VIVIT allows for efficient eigen-value decomposition, which also gives per-sample directional derivatives and curvatures. To further improve the efficiency, sampling within the mini-batch and among the coordinates of the prediction can be applied to trade off computational cost with accuracy. As an application example, VIVIT is used to provide noise-aware directional damping which improves the stability of second-order methods.