D EEP C URVATURE S UITE
Anonymous authors
Paper under double-blind review
A BSTRACT
The curvature of the loss, provides rich information on the geometry underlying neural networks, with applications in second order optimisation and Bayesian deep learning. However, accessing curvature information is still a daunting engineering challenge, inaccessible to most practitioners.
We hence provide a software package the Deep Curvature Suite , which allows easy curvature evaluation for large modern neural networks.
Beyond the calculation of a highly accurate moment matched approximation of the Hessian spectrum using Lanczos, our package provides: extensive loss surface visualisation , the calculation of the Hessian variance and stochastic second order optimisers .
We further address and disprove many common misconceptions in the literature about the Lanczos algorithm, namely that it learns eigenvalues from the top down.
We prove using high dimensional concentration inequalities that for speciﬁc matrices a single random vector is sufﬁcient for accurate spectral estimation, informing our spectral visualisation method.
We showcase our package practical utility on a series of examples based on realistic modern neural networks such as the VGG- 16 and Preactivated ResNets on the CIFAR- 10 / 100 datasets.
We further detail 3 speciﬁc potential use cases enabled by our software: research in stochastic second order optimisation for deep learning, learning rate scheduling using known optimality formulae for convex surfaces and empirical veriﬁcation of deep learning theory based on comparing empirical and theoretically implied spectra.
I NTRODUCTION
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.
; Dangel et al.
(2019), which we show in AppendixE can be very misleading.
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 ﬁrst 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. L OSS S URFACES
Recent neural network loss surfaces using full eigendecomposition (Sagun et al., 2016; 2017) have been limited to toy examples with less than ﬁve 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. 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 ﬂat and sharp minima have used the diagonal of the Fisher information matrix (Chaudhari et al., 2016), an assumption we will challenge in this paper.
Speciﬁcally we show in Appendix E, that diagonal approximations do not capture key properties in synthetic and real neural network examples.
From a practical perspective, speciﬁc properties of the loss surface are not captured by the aforementioned approaches. Examples include the ﬂatness as speciﬁed 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) ﬂatter 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 beneﬁts of large learning rate SGD Wu et al.
; Jastrzkebski et al.
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).
Second order optimisation methods solve the minimisation problem for the loss, L ( w ) 2 R associated with parameters w 2 R P ⇥ 1 and perturbation  w 2 R P ⇥ 1 to the second order in Taylor expansion,
[Equation]
Sometimes, such as in deep neural networks when the Hessian H
[Equation] L [Equation] w
[Equation] 2 R P ⇥ P is not positive deﬁnite, a positive deﬁnite surrogate is used. Note that Equation 1 is not lower bounded unless
H is positive semi-deﬁnite. 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 deﬁnite 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 ﬁnal layer f ( w ) . Hence the elements of the Hessian can be written as
[Equation]
The ﬁrst 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 difﬁcult 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 difﬁcult implementation.
As a small side note, Bayesian neural networks using the Laplace approximation feature the Hessian inverse multiplied by a vector (Bishop, 2006)
2. C ONTRIBUTIONS
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 Strakoˇs (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-deﬁnite surrogate to Hessian 1 .
• 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 shufﬂing.
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. R ELATED W ORK
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 [Table] batch normalization In Appendix E.5 we explicitly compare against their diagonal MC approximations on synthetic and real neural network examples and ﬁnd these approximations inadequate for spectral analysis.
Independent work has also uses the Lanczos algorithm for Hessian computation
Yao et al.
; Ghorbani et al.
; Papyan (2018); Izmailov et al. 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. ; 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.
and to predict the effect of minibatching on curvature using a random matrix theory model (Granziol, 2020).
We include an example ﬁgure where the theoretical prediction is validated using our software in Figure 1.
L ANCZOS , M ISCONCEPTIONS AND S PECTRAL D ENSITY A PPROXIMATIONS
The Lanczos Algorithm Meurant and Strakoˇs (2006), on which our package is based, (Algorithmin the Appendix) is an iterative algorithm for learning a subset of the eigenvalues/eigenvectors of any
Hermitian matrix, requiring only matrix vector products.
It is an adaptation of the power iteration method, where the Krylov subspace K ( H , v ) =
span { v , Hv , H 2 v , ... } , is orthogonalised using Gram-Schmidt, resulting in improved convergence. The Lanczos algorithm, is subject to many misconceptions in the literature, which limits its penetration in the research community. Two of which we state here and explicitly debunk in Appendix B.
• We can learn the negative and interior eigenvalues by shifting and inverting the matrix sign H !
H + µ I • Lanczos learns the largest m [  i , u i ] pairs of H 2 R P ⇥ P with high probability (Dauphin et al., 2014) The key properties of the Lanczos algorithm, which we summarise here (and expand upon in Appendix B) are: • When seeded with a vector v the Lanczos algorithm gives a discrete spectral density which matches the m moments vHv , vH 2 v , ..., vH m v . • If the vector v is zero mean unit variance (such as a Gaussian or Rademacher Hutchinson (1990)), this in expectation gives an m moment matched approximation to the spectral density p (  ) of the underlying matrix H .
• For a large dimensional matrix P ! , we expect a single random vector to give an m moment matched approximation of the underlying spectral density. We give the proof for this in Appendix C.This informs our visualisation technique of using a single random vector. • When seeded with a random vector v , Lanczos converges quickly to the end points of the spectral support of p (  ) , the spectral density of H .
We give an illustration on an example of using the Deep Curvature package.
We train a VGG16 network on CIFAR-100 for 100 epochs.
With the checkpoints generated, we may now compute analyse the eigenspectrum of the curvature matrix evaluated at the desired point in training.
To evaluate the
Hessian/GGN with 100 Lanczos iterations, we run the code given in Table 2.
train network (
[Equation]
data path
[Equation]
epochs
[Equation]
[Equation]
model
[Equation]
optimizer = ’SGD’ , optimizer kwargs= { ’ l r ’ : 0.03 , ’mom’ :
[Table]
.9 , ’ weight decay ’ : 5e −} ) compute eigenspectrum ( d a t a s e t = ’CIFAR100 ’ , data path = ’ data / ’ , model= ’VGG16’ , checkpoint path = ’ checkpoint .
pt ’ , save spectrum path = ’ spectrum − ggn ’ , save eigvec
= True , l a n c z o s i t
e r s =100 ,
[Table]
This function call saves the spectrum results (including eigenvalues, eigenvectors and other related statistics) in the save spectrum path path string deﬁned.
To visualise the spectrum as stem plot we simply run the code given in Table 3 with corresponding example Figure shown.
[Equation]
[Table]
[Equation]
[Table]
[Equation]
p l t . show ( ) Finally, with the eigenvalues and eigenvectors computed, we might be interested in knowing how sensitive the network is to perturbation along these directions.
To achieve this, we ﬁrst construct a loss landscape by setting the number of query points and maximum perturbation to apply.
To achieve that, we call the code given in Table 4.
In this example, we set the maximum perturbation to be [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
( dist argument) and number of query points along each direction to be 21 ( n points argument)
[Table]
The corresponding plots in Training and Testing loss are also shown. Although we choose to show loss, the package returns the results for accuracy also, which we show in Appendix D.
E XAMPLE R ESEARCH A PPLICATIONS
In this section we list 3 example scientiﬁc applications of our software package.
We investigate second order optimisation for deep neural networks using our inbuild stochastic Lanczos optimiser, online learning rate scheduling and verifying predictions made using random matrix theory in theoretical deep learning.
5. S IMPLICITY OF S ECOND O RDER O PTIMISATION WITH THE D EEP C URVATURE S UITE
To enable researchers to experiment with stochastic second order optimisation algorithms for deep neural networks, we implement a Lanczos based optimiser (which takes the absolute Hessian
Dauphin et al.
(2014), or Generalised Gauss Newton Martens (2010)
as input).
The code for running this optimiser is summarised in Table 5.
We plot the training error of the VGG- 16 network on train
[Equation]
[Equation]
[Table] [Equation]
CIFAR- 100 dataset, which has over 16 million parameters and hence is out of reach of full inversion methods, in Figure 2b.
We keep the ratio of damping constant to learning rate constant, where  = 10 ↵ , for a variety of learning rates in { 1 , 0 . 1 , 0 . 01 , 0 . 001 , 0 . 0001 } with a batch size of 128 for both the gradient and the curvature, all of which post almost identical performance.
We also compare against different learning rates of Adam and the best grid searched learning rate of SGD, both of which converge signiﬁcantly slower per iteration compared to our stochastic Newton methods.
We note that in terms of practical optimisation that the cost of running a stochastic second order optimiser is far greater than that of SGD.
For this example, where we use 20 Lanczos iterations, the per iteration cost scales as ⇡ 40 times that of SGD.
Hence we expect these algorithms simply to serve as useful baselines against more computationally efﬁcient approximations.
One example future direction could be to invert the matrices less frequently as done in Martens and Grosse (2015), reducing computational cost.
5. S PECTRAL S TOCHASTIC G RADIENT D ESCENT WITH THE D EEP C URVATURE S UITE
SGD is incredibly sensitive to the choice of initial learning rate and its scheduling.
Whilst adaptive methods are considered more robust to the initial learning rate choice and its schedule, they have been shown to perform poorly on validation or held out test data Wilson et al.
Hence practitioners typically experiment with a variety of schedules, such as the step, linear, exponential, polynomial and [-P mi =1  i < P Pi = m +
-]cosine annealing.
One potential application of our codebase is to use a stochastic estimate of the curvature and to use established optimal learning and momentum rates for the local quadratic approximation at the given point in weightspace Nesterov (2013), given in Equation 3.
[Equation]
Where  1 ,  P refer to the largest and smallest eigenvalues of the Hessian respectively.
The learning rates and momenta ↵ , ⇢ are given for both Polyak and Nesterov momentum respectively.
There are several complications to the above method.
Firstly the Hessian of neural networks is in general not positive deﬁnite and hence  P {+< 
+}
(these optimality formulae are strictly for convex methods).
Secondly the loss surface at each point in weight space is likely to be different and it is impractical
[Figure]
to recalculate the optimal step size using Equation 3 for every iteration, even if we sub-sample the data. We note from the work of Granziol (2020), that the loss surface when using a mini-batch varies to that of the full surface (it is sharper) and hence the learning rates when training using a mini-batch must be optimal for the minibatch surface and not that of the full dataset.
As an example use case for our software we consider whether the method of Izmailov et al. (2018), which combines SGD with Polyak averaging Polyak and Juditsky (1992) at the end of training (which they denote SWA ), to improve generalisation performance, could be improved.
Instead of ﬁne tuning the linear decay learning rate schedule by hand, we instead learn the learning rate using Equation 3 at regular intervals by estimating the curvature using a sub-sampled Lanczos spectral estimate. We run the preactivated ResNet- 110 on the CIFAR- 100 dataset for 180 epochs and learn the Polyak and Nesterov learning rates and momenta every 20 epochs using the same mini-batch size as the gradient batch size 128 for 20 Lanczos iterations. Since the smallest Ritz values are very very close to 0 , which would result in a momentum ⇢ = 1 , we use a heuristic to remove the smallest Ritz values, whereby if the Ritz value of largest spectral mass has more than 50% of the spectral mass, it is removed and the resulting density renormalised, forming the new spectral density of interest.
All methods employ Polyak Avearging at epoch 161 as in Izmailov et al. We plot the results in Figure 2a and show that the Nesterov variant of online learning rate and momentum learning performs well against the baseline (SWA). Interestingly the Polyak variant learns to drop the learning rate at regular intervals, similar to classical training methods, which does not perform as well as keeping the learning rate high and using Polyak averaging at the end of training. 5. V ERIFYING R ANDOM M ATRIX T HEORY IN D EEP L EARNING USING OUR P ACKAGE
The full spectrum of deep neural networks can be used to validate/invalidate novel theoretical contributions.
As an example.
analysis relating the loss surface to spin-glass models from condensed matter physics and random matrix theory (Choromanska et al., 2015b;a) rely on a number of unrealistic assumptions, such as input independence (i.e. for an input x that the feature x i is independent of x j , for images which have spatial correlations, such assumptions are clearly violated).
One potential veriﬁcation of the practical applicability of these results, that does not involve checking that the assumptions are satisﬁed, would be to visualise the spectra of large real networks and commonly used datasets.
This follows because spin glass models have Hessians given by random matrices (Arous and Guionnet, 1997) with known analytical forms (Tao, 2012; Akemann et al., 2011).
Hence statistical tests comparing the real world observed spectra and that underlying the speciﬁc spin glass model could be undertaken, showing that the results hold despite the assumptions not being met.
Two primary random matrix eigenvalue distributions to which the spectra of many classes of random matrices converge (and which feature in theoretical deep learning papers) are the semi-circle and the Marchenko-Pastur densities deﬁned below. [Figure] [Equation] [Figure]
[Equation] If there are N samples and P parameters and N < P the Marchenko-Pastur density has a degeneracy fraction 1  NP . In the example of Choromanska et al. (2015a), the corresponding Hessian spectral denisty is the semi-circle law, shown in Figure 3b.
This differs greatly from neural network spectra found us- [Figure] ing our package, such as that of the VGG- 16 with BN shown in Figure 3a and the Preactived ResNet- 110 , for the CIFAR- 100 datasets, shown in Figure 3c.
Major differences include the rankdegeneracies (we ﬁnd that over 90% of the spectral mass resides in the peak nearest the origin) and outliers.
The proof techniques used in Arous and Guionnet (1997), which underlie the claims in
Choromanska et al.
(2015a) speciﬁcally use the large deviation principle of the Gaussian Orthogonal Ensemble (which converges to the semi-circle law) and hence forbid outliers. The GGN has been modelled as the Marchenko-Pastur density in the literature (Pennington and Bahri, 2017; Sagun et al., 2017; 2016).
The Residual matrix R ( w ) , which is the difference of the Hessian and the GGN, has been modelled by the Wigner ensemble which gives the semi-circle law in Pennington and Bahri (2017). In Figure 4 we plot the GGN and Residual of a VGG- 16 network on the CIFAR- 100 dataset.
The GGN unlike the Marchenko-Pastur has no spectral gap (the density tails off exponentially) and many outliers. Similarly the shape and rank-degeneracy of the Residual matrix varies signiﬁcantly from the semi-circle and contains outliers. Whilst a single set of negative results does not invalidate an entire framework, it implies that further research is required to extend the claims of Choromanska et al. (2015a); Pennington and Bahri (2017) to practical deep learning.
[Figure]
These results using our package demonstrate that the underlying spectral densities of the matrices in question vary considerably from those given by the theoretical predictions.
This implies that the conclusions made may not necassarily hold in real neural networks and requires further investigation.
C ONCLUSION
We introduce the Deep Curvature suite in PyTorch framework, based on the Lanczos algorithm implemented in GPyTorch (Gardner et al., 2018), that allows deep learning practitioners to learn spectral density information as well as eigenvalue/eigenvector pairs of the curvature matrices at speciﬁc points in weight space. Our package also allows the user to evaluate the Generalised Gauss
Newton spectral density, the variance of both the Hessian and the Generalised Gauss Newton and includes two stochastic Lanczos based optimisers.
Together with the software, we also include a succinct summary of the linear algebra, iterative method theory including proofs of convergence and misconceptions and stochastic trace estimation that form the theoretical underpinnings of our work.
We also give 3 example scientiﬁc applications of our work, namely second order optimisation on large scale neural networks, learning rate scheduling using spectral information and evaluation of theoretical results which make predictions on spectral denisities of the Hessian of the neural networks which they are modelling.