D EEP C URVATURE S UITE
Anonymous authors
Paper under double-blind review
A BSTRACT
Despite providing rich information into neural networks geometry aiding in their understanding and applications in second order optimisation, accessing curvature information is still a daunting engineering challenge and hence inaccessible to most practitioners.
In some cases, proxy sampled diagonal approximations, which we show on both real and synthetic examples can be arbitrarily bad, are employed instead.
We hence provide an open-source software package, the Deep Curvature
Suite to the community.
Our goal is to enable easy access to curvature information for real networks and datasets, not just toy examples.
Beyond the calculation of a highly accurate moment matched approximation of the Hessian spectrum using the Lanczos algorithm, our package provides extensive loss surface visualisation, the calculation of the gradient/Hessian variance and includes second order deep learning optimisers.
As a further contribution, we address and disprove many common misconceptions in the Machine Learning literature, namely that the Lanczos algorithm learns eigenvalues from the top down.
We also prove using high dimensional concentration inequalities that for speciﬁc classes of matrices a single random vector is sufﬁcient for accurate spectral estimation, which informs our Algorithm design choice and spectral visualisation method.
We showcase our package practical utility on a series of examples based on realistic modern neural networks tested on CIFAR- 10 / 100 datasets.
I NTRODUCTION
The success of Deep Neural Networks (DNNs) trained with stochastic gradient descent (SGD) based optimizers on a range of tasks, has led to an explosion in the availability of easy to use out of the box high performance software implementations.
Automatic differentiation packages such as TensorFlow (Abadi et al., 2016) and PyTorch (Paszke et al., 2017) have become widely adopted, with higher level packages allowing users to state their model, dataset and optimiser in a few lines of code (Chollet, 2015), effortlessly achieving state of the art performance.
However, the pace of development of software extracting second order information, representing the curvature at a point in weight space, has not kept abreast.
Researchers aspiring to evaluate or use curvature information need to implement their own libraries, which are rarely shared or kept up to date.
Naive implementations are computationally intractable for all but the smallest of models.
Hence, researchers typically either completely ignore curvature information or use highly optimistic approximations, such as the diagonal elements of the matrix or of a surrogate matrix, with limited justiﬁcation or empirical analysis of the harshness of the aforementioned approximations Chaudhari et al.
; Dangel et al.
The curvature at a point in weight-space 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).
The most common areas where curvature information is employed are analyses of the Loss Surface and Newton type methods in optimization.
Loss surface visualization 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 d -dimensional Gaussian distribution (Izmailov et al., 2018).
Whilst both of these approaches have been informative, it is not clear that these high dimensional loss surfaces, often containing millions or billions of dimensions, can be well captured in this manner.
Recent empirical analyses of the neural network loss surfaces invoking full eigen-decomposition (which captures the full dimensionality of the surface at a particular point in weight space)
(Sagun et al., 2016; 2017) have been limited to toy examples with < 5000 parameters.
Other works have used the diagonal of the Fisher information matrix (Chaudhari et al., 2016), an assumption we will challenge in this paper.
From a practical perspective, speciﬁc properties of the loss surface not captured by the aforementioned approaches, such as the ﬂatness such as the trace, Frobenius and spectral norm have been 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.
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.
From a theoretical stand point the full spectrum of Deep Neural Networks could in principle be used to validate/invalidate novel theoretical contributions.
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.
Since such an assumption is clearly not true, one potential veriﬁcation of the practical applicability of these results would be to visualise the spectra of large real networks and commonly used datasets to evaluate whether they match those of random matrices.
Such matrices have closed analytical formulas describing their spectra (Tao, 2012; Akemann et al., 2011) and statistical tests comparing the two could be undertaken .
Other important areas of loss surface investigation include understanding the effectiveness of batch normalization (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 associated with parameters p and perturbation d to the second order in Taylor expansion,
[Equation]
Where instead of the Hessian H =
rr L ( p ) 2 R n ⇥ n , a surrogate positive deﬁnite approximation to the Hessian ¯ H , such as the (Generalised) Gauss-Newton (Martens, 2010; Martens and Sutskever, 2012), is employed so to make sure the minimum is lower bounded; its solution is
[Equation]
where u i correspond to the generalised Hessian eigenvectors.
The parameters are updated with p = p  ↵ d , in which ↵ is the global learning rate.
For some common activations and loss functions typical in deep learning, such as the cross-entropy loss and sigmoid activation the generalised GaussNewton is equivalent to the Fisher information matrix (Pascanu and Bengio, 2013).
The Hessian may be expressed in terms of the activation  at the output of the ﬁnal layer f
( x ) using the chain rule as
H
[Equation]
The ﬁrst term on the LHS of Equation 3 is known as the Generalised Gauss-Newton matrix.
Despite the success of second order optimisation methods using the Generalised Gauss Newton 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).
1. C ONTRIBUTIONS
In this paper, we make calculating and visualizing curvature information as simple as calculating the gradient at a saved checkpoint in weight-space.
Speciﬁcally we combine fast Hessian vector products (Pearlmutter, 1994), advanced linear algebraic techniques (Golub and Meurant, 1994) and high dimensional geometry (Hutchinson, 1990), to derive highly accurate moment matched discrete approximations to the neural network spectrum in a fraction of the network training time.
The computational complexity of our approach is O ( mP ) , where m is the number of moments which we wish to match, P is the number of model parameters, as opposed to full exact eigendecomposition which has a numerical cost of O ( P 3 ) , which is infeasible for large neural networks.
We use the GPytorch implementation (Gardner et al., 2018) of the Lanczos algorithm (Meurant and Strakoˇs, 2006), which we introduce in Section 4 and discuss the most common misconceptions in the literature in
Section 4.
Our package supports the Hessian as well as the commonly used Generalsed Gauss Newton approximation (Martens, 2010; Martens and Sutskever, 2012; Henriques et al., 2019).
Our code allows for the use of batch normalization (Ioffe and Szegedy, 2015) in neural networks.
Whilst we keep batch norm training mode on as default, this default setting can be altered easily.
We provide additional novel theoretically justiﬁed visualisations for both the moment matched spectrum and the loss surface traversed in those directions.
Our package further allows for the calculation of the variance of the gradient and the variance of the Hessian, the latter has been related to the stability of SGD in Wu et al.
As a minor contribution, we also include two stochastic Lanczos based optimisers in our code.
Related Work:
Recent work making curvature information more available, again through diagonal approximations, explicitly disallows the use of batch normalization (Dangel et al., 2019).
Our software package extends seamlessly to batch normalization, allowing for evaluation in both train and eval mode.
Whilst the scope of Dangel et al.
(2019) is slightly different to our work, as their focus is primarily on the ability to compute these quantities at speed online for simple models, in
Section 5 we explicitly compare their diagonal MC approximations against our method for spectral visualisation on a range of synthetic and real neural networks and ﬁnd these approximations to be inadequate for spectral analysis.
Independent work has also considered using the Lanczos algorithm for Hessian computation Yao et al.
; Ghorbani et al.
; Papyan (2018); Izmailov et al.
However neither of these packages includes a similar level of focus on the visualisation aspect which we integrate.
Furthermore to the best of our knowledge, the computation of the Hessian variance Wu et al.
(2018) and in-built reference stochastic second order optimisers is not available in any other package.
D EEP C URVATURE
We introduce to 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:
• Network training and evaluation we provide a range of pre-built modern popular neural network structures, such as VGG and variants of ResNets, and various optimisation schemes in addition to the ones already present in the PyTorch frameworks, such as KFAC and SWATS.
These facilitates faster training and evaluation of the networks (although it is worth noting that any PyTorch-compatible optimisers or architectures can be easily integrated into our analysis framework).
• Eigenspectrum analysis of the curvature matrices Powered by the Lanczos 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 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 infer-
ence of the eigen-information of the Generalised Gauss-Newton 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.
While similar tools have been available, we would like to emphasise that one key difference is that, instead of the random directions as featured in some other packages, we explicitly perturb the weights in the eigenvector directions, which should yield more informative results.
Package Structure
The main interface functions are organsed as followed:
• ./core
The functions under core directories are the main analysis tools of the package.
train network allows network training and saving of the required statistics for subsequent spectrum learning.
Based on the output of it, we additionally include tools for spectrum analysis ( compute eigenspectrum ) and advanced loss statistics (such as covariance of gradients and second order information like Hessian variance) in compute loss stats and build loss landscape .
We provide some pre-built network architectures (such as VGG and ResNet architectures) and optimizers apart from PyTorch natives (such as K-FAC, SWATS optimizers).
We additionally support Stochastic Weight Averaging proposed in (Izmailov et al., 2018).
However, it is worth noting that any PyTorch compatible networks and optimizers can be easily integrated in our framework.
• ./visualise
This directory deﬁnes the various pre
-deﬁned visualisation functions for different purposes, including the visualisation of training, spectrum and the loss landscape.
To facilitate a quick start of our package, we have included an illustrated example of analysis on the
VGG-16 network on CIFAR-100 dataset.
A N I LLUSTRATED E XAMPLE
We give an illustration on an example of using the MLRG-DeepCurvature package and more details, including detailed documentation of each user function and the output of this particular example, can be found at our open-source repository.
We begin by importing the necessary functions and packages:
[Equation]
In this example, we train a VGG16 network on CIFAR-100 for 100 epochs.
In a test computer with
AMD Ryzen 3700X CPU and NVIDIA GeForce RTX 2080 Ti GPU, each epoch of training takes less than 10 seconds.
This step generates a number of training statistics ﬁles (starting with stats- ) and checkpoint ﬁles ( checkpoint-00XXX.pt , where XXX is the epoch number) that contain the state dict of the optimizer and the model.
We may additionally visualise the training processes by looking at the
[Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Table]
basic statistics by calling plot training function.
With the checkpoints generated, we may now compute analyse the eigenspectrum of the curvature matrix evaluated at the desired point of training.
For example, if we would like to evaluate the Hessian or Generalised Gauss Newton approximation to the Hessian matrix at the end of the training with 20 Lanczos iterations, we simply run the Eigenspetrum Computation code given in Table 1.
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 similar to Figure 2, we simply call:
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 3.
In this example, we set the maximum perturbation to be( dist argument) and number of query points along each direction to be 21 ( n points argument)
[Equation]
The corresponding plots in Training and Testing loss are also shown.
[Equation]
checkpoint
[Equation]
path
[Equation]
landscape − 100.
npz
[Equation] [Equation]
t .
show ( ) )
3. S IMPLICITY OF S ECOND O RDER O PTIMISATION
To enable researchers to experiment with second order optimisation algorithms for Deep Neural Networks
[Equation]
we implement a Lanczos based optimiser (which takes the absolute Hessian Dauphin et al.
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
(2014), or Generalised Gauss Newton as input).
The code for running this optimiser is summarised in Table 4.
We plot the training error of the VGG- 16 network on CIFAR- 100 dataset (Which has 16 m parameters) against epoch in Figure 1.
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.
L EARNING TO LOVE L ANCZOS
The Lanczos Algorithm, on which our package is based, (Algorithm 1) is an iterative algorithm for learning a subset of the eigenvalues/eigenvectors of any Hermitian matrix, requiring only matrix vector products.
It can be regarded as a far superior adaptation of the power iteration method, where the Krylov subspace K ( H , v ) =
span { v , H v , H 2 v ... } is orthogonalised using Gram-Schmidt.
Beyond having improved convergence to the power iteration method (Bai et al., 1996) by storing the intermediate orthogonal vectors in the corresponding Krylov subspace, Lanczos produces estimates of the eigenvectors and eigenvalues of smaller absolute magnitude, known as Ritz vectors/values.
Despite its known superiority to the power iteration method and relationship to orthogonal polynomials and hence when combined with random vectors the ability to estimate the entire spectrum of a matrix, these properties are often ignored or forgotten by practitioners, we hence include a full primer in Appendix E. We compare against diaognal approximations for random and synthetic matrices in Appendix C.
We also explicitly debunk some key persistent myths, given below.
• 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)
Since these two related beliefs are prevalent, we disprove them explicitly in this section, with Theorems 1 and 2.
Theorem 1.
The shift and invert procedure H !
H + µ
I , changes the Eigenvalues of the Tridiagonal matrix T (and hence the Ritz values) to  i =
i + µ
Proof.
Following the equations from Algorithm
[Equation]
Assuming this for m  1 , and repeating the above steps for m we prove by induction and ﬁnally arrive at the modiﬁed tridiagonal Lanczos matrix ˜ T
[Equation]
Remark.
No new Eigenvalues of the matrix H are learned.
Although it is clear that the addition of the identity does not change the Krylov subspace, such procedures are commonplace in code pertaining to papers attempting to ﬁnd the smallest eigenvalue .
This disproves the ﬁrst misconception.
Theorem 2.
For any matrix H 2 R P ⇥ P such that  1 >  2 > ..... >  P and
P mi =1  i <
P
Pi = m +
i in expectation over the set of random vectors v the m eigenvalues of the Lanczos Tridiagonal matrix T do not correspond to the top m eigenvalues of H
Proof.
Let us consider the matrix
˜ H
=
H   m
[Equation]
+
[Equation]
Under the assumptions of the theorem, Tr ( ˜ H ) < 0 and hence by Theorem 8 and Equation 24 there exist no w i > 0 such that
[Equation]
is satisﬁed for k = 1 , as the LHS is manifestly positive and the RHS is negative.
By Theorem 1 this holds for the original matrix H .
Remark.
Given that Theorem 2 is satisﬁed over the expectation of the set of random vectors, which by the CLT is realised by Monte Carlo draws of random vectors as d ! 1 the only way to really span the top m eigenvectors is to have selected a vector which lies in the m dimensional subspace of the P dimensional problem corresponding to those vectors, which would correspond to knowing those vectors a priori , defeating the point of using Lanczos at all.
Remark.
Given the relationship between the moments of the spectral density and the expectation over the set of random vectors, one may be curious to ask how we expect the deviation to be depending on the number of random vectors actually used in practice.
Whilst usually packages using
Lanczos use a number of random vectors (increasing the corresponding computational cost), in
Appendix A we demonstrate that we expect the difference to reduce as we increase the problem dimension P .
This informs our choice of only using a single random vector in our package, we compare results for different random vectors in Appendix A and ﬁnd minimal differences.
N EURAL N ETWORK E XAMPLES
We showcase our spectral learning algorithm and visualization tool on real networks trained on real data-sets and we test on VGG networks (Simonyan and Zisserman, 2014).
We train our neural networks using stochastic gradient descent with momentum ⇢ = 0 .
9 , using a linearly decaying learning rate schedule.
The learning rate at the t -th epoch is given by:
[Equation]
where ↵ 0 is the initial learning rate.
T = 300 is the total number of epochs budgeted for all experiments.
We set r = 0 .
We explicitly give an example code run in F
We compare our method against recently developed open-source tools which calculate on the ﬂy diagonal Hessian and Generalised Gauss-Newton diagonal approximations (Dangel et al., 2019).
5. VGG- 16 CIFAR-100 D ATASET
We train a 16-layer VGG network, comprising of P = 15 , 291 , 300 parameters on the CIFARdataset, using ↵ 0 = 0 .
Even for this relatively small model, the open-source Hessian and GGN exact diagonal computations require over 125 GB of GPU memory and so to avoid re-implementing the library to support multiple GPUs and node communication we use the Monte Carlo approximation to the GGN diagonal against both our GGN-Lanczos and Hessian-Lanczos spectral visualizations.
We plot a histogram of the Monte Carlo approximation of the diagonal GGN (Diag-GGN) against both the Lanczos GGN (Lanc-GGN) and Lanczos Hessian (Lanc-Hess) in Figure 2.
Note that as the Lanc-GGN and Lanc-Hess are displayed as stem plots (with the discrete spectral density summing to 1 as opposed to the histogram area summing to 1 ).
[Figure]
We note that the Gauss-Newton approximation quite closely resembles its Hessian counterpart, capturing the majority of the bulk and the outlier eigenvectors at  1 ⇡ 6 . 88 and the triad near  i ⇡ 2 . 29 .
The Hessian does still have signiﬁcant spectral mass on the negative axis, around 37% .
However most of this is captured by a Ritz value at  0 .
0003 , with this removed, the negative spectral mass is only 0 .
However the Diag-GGN gives a very poor spectral approximation.
It vastly overestimates the bulk region, which extends well beyond  ⇡ 1 implied by Lanczos and adds many spurious outliers between 3 and the misses the largest outlier of 6 . 88 .
Computational Cost Using a single NVIDIA GeForce GTX 1080 Ti GPU, the Gauss-Newton takes an average 26 . 5 seconds for each Lanczos iteration with the memory useage 2850 Mb.
Using the Hessian takes an average of 27 .
9 seconds for each Lanczos iteration with 2450 Mb memory usage.
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.
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.
Finally, we also included various examples of our package of analysis of both synthetic data and real data with modern neural network architectures.