Deep Curvature Suite

1 INTRODUCTION . The success of deep neural networks trained with gradient based optimisers in speech and object recognition ( LeCun et al. , 2015 ) , has led to an explosion in easy to use high performance software implementations . Automatic differentiation packages such as TensorFlow ( Abadi et al. , 2016 ) and PyTorch ( Paszke et al. , 2017 ) have become widely adopted . Higher level packages , such as Keras ( Chollet , 2015 ) allow practitioners users to state their model , dataset and optimiser in a few lines of code , effortlessly achieving state of the art performance . However , software for extracting second order information , representing the curvature of the loss at a point in weight space , has not kept abreast . Researchers aspiring to evaluate curvature information need to implement their own libraries , which are rarely shared or kept up to date . Naive implementations , which rely on full eigendecomposition ( cubic cost in the parameter count ) are computationally intractable for all but the smallest of models . Hence , researchers typically ignore curvature information or use highly optimistic approximations . Examples in the literature include the diagonal elements of the matrix or of a surrogate matrix , Chaudhari et al . ( 2016 ) ; Dangel et al . ( 2019 ) , which we show in AppendixE can be very misleading . 2 MOTIVATION . The curvature of the loss informs us about the local conditioning of the problem ( i.e the ratio of the largest to smallest eigenvalues λ1λP ) . This determines the rate of convergence for first order methods and informs us about the optimal learning and momentum rates ( Nesterov , 2013 ) . Hence easily accessible curvature information could allow practitioners to scale their learning rates in an optimal way throughout training , instead of relying on expert scheduling , we investigate this using our software in Section 5.2 . Research areas where curvature information features most prominently are analyses of the Loss Surface and Newton type optimization methods . 2.1 LOSS SURFACES . Recent neural network loss surfaces using full eigendecomposition ( Sagun et al. , 2016 ; 2017 ) have been limited to toy examples with less than five thousand parameters . Hence , loss surface visualisation of deep neural networks have often focused on two dimensional slices of random vectors ( Li et al. , 2017 ) or the changes in the loss traversing a set of random vectors drawn from the ddimensional Gaussian distribution ( Izmailov et al. , 2018 ) . It is not clear that the loss surfaces of modern expressive neural networks , containing millions or billions of dimensions , can be well captured in this manner . Small experiments have shown neural networks have a large rank degeneracy Sagun et al . ( 2016 ) with a small number of large outliers . However high dimensional concentration theorems Vershynin ( 2018 ) guarantee that even a large number of randomly sampled vectors are unlikely to encounter such outliers and hence have limited ability to discern the geometry between various solutions . Other works , that try to distinguish between flat and sharp minima have used the diagonal of the Fisher information matrix ( Chaudhari et al. , 2016 ) , an assumption we will challenge in this paper . Specifically we show in Appendix E , that diagonal approximations do not capture key properties in synthetic and real neural network examples . From a practical perspective , specific properties of the loss surface are not captured by the aforementioned approaches . Examples include the flatness as specified by the trace , Frobenius and spectral norm . These measures have been extensively used to characterise the generalisation of a solution found by SGD ( Wu et al. , 2018 ; Izmailov et al. , 2018 ; He et al. , 2019 ; Jastrzkebski et al. , 2017 ; 2018 ; Keskar et al. , 2016 ) . Under a Bayesian and minimum description length argument ( Hochreiter and Schmidhuber , 1997 ) flatter minima should generalise better than sharper minima . The magnitude of these outliers have been linked to poor generalisation performance Keskar et al . ( 2016 ) and as a consequences the generalisation benefits of large learning rate SGD Wu et al . ( 2018 ) ; Jastrzkebski et al . ( 2017 ) . These properties are extremely easy in principle to estimate , at a computational cost of a small multiple of gradient evaluations . However the calculation of these properties are not typically included in standard deep learning frameworks , which limits the ability of researchers to undertake such analysis . Other important areas of loss surface investigation include understanding the effectiveness of batch normalisation ( Ioffe and Szegedy , 2015 ) . Recent convergence proofs ( Santurkar et al. , 2018 ) bound the maximal eigenvalue of the Hessian with respect to the activations and bounds with respect to the weights on a per layer basis . Bounds on a per layer basis do not imply anything about the bounds of the entire Hessian and furthermore it has been argued that the full spectrum must be calculated to give insights on the alteration of the landscape ( Kohler et al. , 2018 ) . 2.2 SECOND ORDER OPTIMISATION METHODS . Second order optimisation methods solve the minimisation problem for the loss , L ( w ) ∈ R associated with parameters w ∈ RP×1 and perturbation δw ∈ RP×1 to the second order in Taylor expansion , δw∗ = argminδwL ( w + δw ) = argminδwL ( w + δw ) = −H̄−1∇L ( w ) . ( 1 ) Sometimes , such as in deep neural networks when the HessianH = ∇2L ( w ) ∈ RP×P is not positive definite , a positive definite surrogate is used . Note that Equation 1 is not lower bounded unless H is positive semi-definite . Often either a multiple of the identity is added ( known as damping ) to the Hessian ( H → H + γI ) Dauphin et al . ( 2014 ) or a surrogate positive definite approximation to the Hessian H̄ , such as the Generalised Gauss-Newton ( GGN ) ( Martens , 2010 ; Martens and Sutskever , 2012 ) is employed . To derive the GGN , we express the loss in terms of the activation L ( w ) = σ ( f ( w ) ) of the output of the final layer f ( w ) . Hence the elements of the Hessian can be written as H ( w ) ij = dy∑ k=0 dy∑ l=0 ∂2σ ( f ( w ) ) ∂fl ( w ) ∂fk ( w ) ∂fl ( w ) ∂wj ∂fk ( w ) ∂wi + dy∑ k=0 ∂σ ( f ( w ) ) ∂wj ∂2fk ( w ) ∂wj∂wi . ( 2 ) The first term on the LHS of Equation 2 is known as the Generalised Gauss-Newton matrix . Despite the success of second order optimisation methods using the GGN for difficult problems on which SGD is known to stall , such as recurrent neural networks ( Martens and Sutskever , 2012 ) , or auto-encoders ( Martens , 2016 ) . Researchers wanting to implement second order methods such as ( Vinyals and Povey , 2012 ; Martens and Sutskever , 2012 ; Dauphin et al. , 2014 ) face the aforementioned problems of difficult implementation . As a small side note , Bayesian neural networks using the Laplace approximation feature the Hessian inverse multiplied by a vector ( Bishop , 2006 ) . 2.3 CONTRIBUTIONS . We introduce our package , the Deep Curvature suite , a software package that allows analysis and visualisation of deep neural network curvature . The main features and functionalities of our package are : • Eigenspectrum analysis of the curvature matrices Powered by the Lanczos Meurant and Strakoš ( 2006 ) techniques implemented in GPyTorch ( Gardner et al. , 2018 ) and outlined in Section 3 , with a single random vector we use the Pearlmutter matrix-vector product trick Pearlmutter ( 1994 ) for fast inference of the eigenvalues and eigenvectors of the common curvature matrices of the deep neural networks . In addition to the standard Hessian matrix , we also include the feature for inference of the eigen-information of the Generalised GaussNewton matrix , a commonly used positive-definite surrogate to Hessian1 . • Advanced Statistics of Networks In addition to the commonly used statistics to evaluate network training and performance such as the training and testing losses and accuracy , we support computations of more advanced statistics : For example , we support squared mean and variance of gradients and Hessians ( and GGN ) , squared norms of Hessian and GGN , L2 and L-inf norms of the network weights and etc . These statistics are useful and relevant for a wide range of purposes such as the designs of second-order optimisers and network architecture . • Visualisations For all main features above , we include accompanying visualisation tools . In addition , with the eigen-information obtained , we also feature visualisations of the loss landscape by studying the sensitivity of the neural network to perturbations of weights . • Second Order Optimisation We include out of the box stochastic lanczos second order optimisers which use the absolute Hessian Dauphin et al . ( 2014 ) or the Generalised Gauss Newton Martens ( 2010 ) as curvature matrices • Batch Normalisation Support Our code allows for the use of batch normalisation ( Ioffe and Szegedy , 2015 ) in neural networks , which is integrated by default on most modern convolutional neural networks . Since the way in watch batch normalisation is calculated has a big effect on the resulting curvature calculations and can be dependent on the order in which the samples are processed , we offer two models of evaluation . Train mode which depends on the order in which the samples are processed and is similar to what the optimiser sees during training and evaluation mode , in which the statistics are pre-computed and hence invariant to sample shuffling . We expand on the differences between these methods in Appendix A . Our software makes calculating and visualising curvature information as simple as calculating the gradient at a saved checkpoint in weight-space . The computational complexity of our approach is O ( mP ) , where m is the number of Lanczos steps 2 and P is the number of model parameters . This in stark contrast full exact eigendecomposition which has a numerical cost of O ( P 3 ) , which is infeasible for large neural networks . 2.4 RELATED WORK . Recent work making curvature information more available using diagonal approximations , disallows the use of batch normalisation ( Dangel et al. , 2019 ) . Our software extends seamlessly to support 1The computation of the GGN-vector product is similar with the computational cost of two backward passes in the network . Also , GGN uses forward-mode automatic differentiation ( FMAD ) in addition to the commonly employed backward-mode automatic differentiation ( RMAD ) . In the current PyTorch framework , the FMAD operation can be achieved using two equivalent RMAD operations . 2which corresponds to the number of moments of the underlying Hessian/GGN density which we wish to match Functionality Deep Curvature Suite PyHessian Hessian Density Estimation 3 3 GGN Density Estimation 3 7 Loss Surface Visualisation 3 3 Gradient Variance 3 7 Hessian/GGN Variance 3 7 Second Order Optimisers 3 7 Table 1 : Table of Contributions : comparison between our released software and that of another PyTorch based Hessian tool PyHessian batch normalization In Appendix E.5 we explicitly compare against their diagonal MC approximations on synthetic and real neural network examples and find these approximations inadequate for spectral analysis . Independent work has also uses the Lanczos algorithm for Hessian computation Yao et al . ( 2019 ) ; Ghorbani et al . ( 2019 ) ; Papyan ( 2018 ) ; Izmailov et al . ( 2019 ) . We compare the functionality of our package to that of PyHessian Yao et al . ( 2018 ) in Table 1 . One major extra functionality we offer is the calculation of the GGN . For common activatios and loss functions , such as the cross-entropy loss and sigmoid activation , the generalised Gauss-Newton is equivalent to the Fisher information matrix ( Pascanu and Bengio , 2013 ) . Hence this approximation to the Hessian is interesting in its own right and many theoretical analyses , consider the generalised Gauss-Newton or Fisher information matrices as opposed to the Hessian Pennington and Worah ( 2018 ) ; Karakida et al . ( 2019 ) ; Papyan ( 2019 ) . Hence accessing the GGN spectrum could be a major asset to researchers . Another extra functionality we offer is the calculation of the Hessian/GGN variance . This quantity has been used to derive optimal learning rates in Wu et al . ( 2018 ) and to predict the effect of minibatching on curvature using a random matrix theory model ( Granziol , 2020 ) . We include an example figure where the theoretical prediction is validated using our software in Figure 1 .

Summary: This paper proposes a software package to ease and provide a standard way for Hessian-related computation, both for loss analysis and second order optimization. It also provides analysis on why Lanczos algorithm is a better choice to estimate Hessian eigenvalue compared to Power Iterations. Finally, it empirically shows the advantage of using Hessian approximation that goes beyond diagonal approximation for spectrum computation.