B ACK PACK: P ACKING MORE INTO B ACKPROP
kunstner@
A BSTRACT
Automatic differentiation frameworks are optimized for exactly one thing: computing the average mini-batch gradient.
Yet, other quantities such as the variance of the mini-batch gradients or many approximations to the Hessian can, in theory , be computed efﬁciently, and at the same time as the gradient.
While these quantities are of great interest to researchers and practitioners, current deep-learning software does not support their automatic calculation.
Manually implementing them is burdensome, inefﬁcient if done na ¨ ıvely, and the resulting code is rarely shared.
This hampers progress in deep learning, and unnecessarily narrows research to focus on gradient descent and its variants; it also complicates replication studies and comparisons between newly developed methods that require those quantities, to the point of impossibility.
To address this problem, we introduce B ACK PACK 1 , an efﬁcient framework built on top of P Y T ORCH , that extends the backpropagation algorithm to extract additional information from ﬁrst- and second-order derivatives.
Its capabilities are illustrated by benchmark reports for computing additional quantities on deep neural networks, and an example application by testing several recent curvature approximations for optimization.
I NTRODUCTION
The success of deep learning and the applications it fuels can be traced to the popularization of automatic differentiation frameworks.
Packages like T ENSOR F LOW (Abadi et al., 2016),
C HAINER (Tokui et al., 2015), MXN ET (Chen et al., 2015), and P Y T ORCH (Paszke et al., 2019) provide efﬁcient implementations of parallel, GPU-based gradient computations to a wide range of users, with elegant syntactic sugar.
However, this specialization also has its shortcomings: it assumes the user only wants to compute gradients or, more precisely, the average of gradients across a mini-batch of examples.
Other quantities can also be computed with automatic differentiation at a comparable cost or minimal overhead to the gradient backpropagation pass; for example, approximate second-order information or the variance of gradients within the batch.
These quantities are valuable to understand the geometry of deep neural networks, for the identiﬁcation of free parameters, and to push the development of more efﬁcient optimization algorithms.
But researchers who want to investigate their use face a chickenand-egg problem: automatic differentiation tools required to go beyond standard gradient methods are not available, but there is no incentive for their implementation in existing deep-learning software as long as no large portion of the users need it.
Second-order methods for deep learning have been continuously investigated for decades (e.g.,
Becker & Le Cun, 1989; Amari, 1998; Bordes et al., 2009; Martens & Grosse, 2015).
But still, the standard optimizers used in deep learning remain some variant of stochastic gradient descent ( SGD ); more complex methods have not found wide-spread, practical use.
This is in stark contrast to domains like convex optimization and generalized linear models, where second-order methods are the default.
There may of course be good scientiﬁc reasons for this difference; maybe second-order methods do not work well in the (non-convex, stochastic) setting of deep learning.
And the computational cost associated with the high dimensionality of deep models may offset their beneﬁts.
Whether these are the case remains somewhat unclear though, because a much more direct road-block is that these methods are so complex to implement that few practitioners ever try them out.
Recent approximate second-order methods such as KFAC (Martens & Grosse, 2015) show promising results, even on hard deep learning problems (Tsuji et al., 2019).
Their approach, based on the earlier work of Schraudolph (2002), uses the structure of the network to compute approximate secondorder information in a way that is similar to gradient backpropagation.
This work sparked a new line of research to improve the second-order approximation (Grosse & Martens, 2016; Botev et al., 2017; Martens et al., 2018;
George et al., 2018).
However, all of these methods require low-level applications of automatic differentiation to compute quantities other than the averaged gradient.
It is a daunting task to implement them from scratch.
Unless users spend signiﬁcant time familiarizing themselves with the internals of their software tools, the resulting implementation is often inefﬁcient, which also puts the original usability advantage of those packages into question.
Even motivated researchers trying to develop new methods, who need not be expert software developers, face this problem.
They often end up with methods that cannot compete in runtime, not necessarily because the method is inherently bad, but because the implementation is not efﬁcient.
New methods are also frequently not compared to their predecessors and competitors because they are so hard to reproduce.
Authors do not want to represent the competition in an unfair light caused by a bad implementation.
Another example is offered by a recent string of research to adapt to the stochasticity induced by mini-batch sampling.
An empirical estimate of the (marginal) variance of the gradients within the batch has been found to be theoretically and practically useful for adapting hyperparameters like learning rates (Mahsereci & Hennig, 2017) and batch sizes (Balles et al., 2017), or regularize ﬁrstorder optimization (Le Roux et al., 2007; Balles & Hennig, 2018; Katharopoulos & Fleuret, 2018).
To get such a variance estimate, one simply has to square, then sum, the individual gradients after the backpropagation, but before they are aggregated to form the average gradient.
Doing so should have negligible cost in principle , but is programmatically challenging in the standard packages.
Members of the community have repeatedly asked for such features
but the established automatic differentiation frameworks have yet to address such requests, as their focus has been—rightly—on improving their technical backbone.
Features like those outlined above are not generally deﬁned for arbitrary functions, but rather emerge from the speciﬁc structure of machine learning applications.
General automatic differentiation frameworks can not be expected to serve such specialist needs.
This does not mean, however, that it is impossible to efﬁciently realize such features within these frameworks: In essence, backpropagation is a technique to compute multiplications with Jacobians.
Methods to extract second-order information (Mizutani & Dreyfus, 2008) or individual gradients from a mini-batch (Goodfellow, 2015) have been known to a small group of specialists; they are just rarely discussed or implemented.
1. O UR CONTRIBUTION
To address this need for a specialized framework focused on machine learning, we propose a framework for the implementation of generalized backpropagation to compute additional quantities.
The structure is based on the conceptual work of Dangel et al. for modular backpropagation.
This framework can be built on top of existing graph-based backpropagation modules; we provide an implementation on top of P Y T ORCH , coined B ACK PACK, available at https://f-dangel.github.io/backpack/ .
The initial release supports efﬁcient computation of individual gradients from a mini-batch, their (cid:96)norm, an estimate of the variance, as well as diagonal and Kronecker factorizations of the generalized Gauss-Newton ( GGN ) matrix (see Tab. 1 for a feature overview).
The library was designed to be minimally verbose to the user, easy to use (see Fig. 1), and to have low overhead (see § 3).
While other researchers are aiming to improve the ﬂexibility of automatic differentiation systems (Innes, 2018a;b; Bradbury et al., 2018), our goal with this package is to provide access to quantities that are only byproducts of the backpropagation pass, rather than gradients themselves. Computing the gradient with P Y T ORCH ...
[Figure] ...and the variance with B ACK PACK [Figure] To illustrate the capabilities of B ACK PACK, we use it to implement preconditioned gradient descent optimizers with diagonal approximations of the GGN and recent Kronecker factorizations KFAC (Martens & Grosse, 2015), KFLR , and KFRA (Botev et al., 2017).
Our results show that the curvature approximations based on Monte-Carlo ( MC ) estimates of the GGN , the approach used by KFAC , give similar progress per iteration to their more accurate counterparts, but being much cheaper to compute.
While the na ¨ ıve update rule we implement does not surpass ﬁrst-order baselines such as
SGD with momentum and Adam (Kingma & Ba, 2015), its implementation with various curvature approximations is made straightforward.
T HEORY AND I MPLEMENTATION
We will distiguish between quantities that can be computed from information already present during a traditional backward pass (which we suggestively call ﬁrst-order extensions ), and quantities that need additional information (termed second-order extensions ).
The former group contains additional statistics such as the variance of the gradients within the mini-batch or the (cid:96) 2 norm of the gradient for each sample.
Those can be computed with minimal overhead during the backprop pass.
The latter class contains approximations of second-order information, like the diagonal or Kronecker factorization of the generalized Gauss-Newton ( GGN ) matrix, which require the propagation of additional information through the graph.
We will present those two classes separately:
First-order extensions Extract more from the standard backward pass.
– Individual gradients from a mini-batch – (cid:96) 2 norm of the individual gradients – Diagonal covariance and 2 nd moment Second-order extensions Propagate new information along the graph.
– Diagonal of the GGN and the Hessian – KFAC (Martens & Grosse, 2015) – KFRA and KFLR (Botev et al., 2017)
These quantities are only deﬁned, or reasonable to compute, for a subset of models: The concept of individual gradients for each sample in a mini-batch or the estimate of the variance requires the loss for each sample to be independent.
While such functions are common in machine learning, not all neural networks ﬁt into this category.
For example, if the network uses Batch Normalization (Ioffe & Szegedy, 2015), the individual gradients in a mini-batch are correlated.
Then, the variance is not meaningful anymore, and computing the individual contribution of a sample to the mini-batch gradient or the GGN becomes prohibitive.
For those reasons, and to limit the scope of the project for version 1.0, B ACK PACK currently restricts the type of models it accepts.
The supported models are traditional feed-forward networks that can be expressed as a sequence of modules , for example a sequence of convolutional, pooling, linear and activation layers.
Recurrent networks like LSTM s (Hochreiter & Schmidhuber, 1997) or residual networks (He et al., 2016) are not yet supported, but the framework can be extended to cover them.
We assume a sequential model f : Θ × X → Y and a dataset of N samples ( x n , y n ) ∈
X × Y with n = 1 , . .
N .
The model maps each sample x n to a prediction ˆ y n using some parameters θ ∈ Θ .
The predictions are evaluated with a loss function (cid:96) :
Y × Y → R , for example the cross-entropy, [Equation] which compares them to the ground truth y n .
This leads to the objective function L : Θ → R ,
[Equation]
As a shorthand, we will use (cid:96) n ( θ ) =
(cid:96) ( f ( θ , x n ) , y n ) for the loss and f n ( θ ) =
f
( θ , x n ) for the model output of individual samples.
Our goal is to provide more information about the derivatives of { (cid:96) n } Nn =1 with respect to the parameters θ of the model f .
2. P RIMER ON BACKPROPAGATION
Machine learning libraries with integrated automatic differentiation use the modular structure of f n ( θ ) to compute derivatives (see Baydin et al.
(2018) for an overview).
If f n is a sequence of L transformations, it can be expressed as
[Equation]
[Equation]
) where T ( i ) θ ( i ) is the i th transformation with parameters θ ( i ) , such that θ =
[ θ (1) , . .
, θ ( L ) ] .
The loss function can also be seen as another transformation, appended to the network.
Let z ( i − 1) n , z
( i ) n denote the input and output of the operation T ( i )
θ ( i ) for sample n , such that z (0) n is the original data and z (1) n , · · · , z ( L ) n represent the transformed output of each layer, leading to the computation graph
[Equation]
To compute the gradient of (cid:96) n with respect to the θ ( i ) , one can repeatedly apply the chain rule, [Equation]
[Equation]
) where J a b is the Jacobian of b with respect to a , [J a b ] ij = ∂
[ b ] i /∂
[ a ] j .
A similar expression exists for the module inputs z
( i − 1) n :
[Equation]
z
[Equation]
θ )) .
This recursive structure makes it possible to extract the gradient by propagating the gradient of the loss.
In the backpropagation algorithm, a module i receives the loss gradient with respect to its output, ∇ z ( i ) n (cid:96) n ( θ ) .
It then extracts the gradient with respect to its parameters and inputs, ∇ θ ( i ) (cid:96) n ( θ ) and ∇ z ( i − 1) n (cid:96) n ( θ ) , according to Eq. 3.
The gradient with respect to its input is sent further down the graph.
This process, illustrated in Fig.
2, is repeated for each transformation until all gradients are computed.
To implement backpropagation, each module only needs to know how to multiply with its Jacobians.
For second-order quantities, we rely on the work of Mizutani & Dreyfus (2008) and Dangel et al.
(2019), who showed that a scheme similar to Eq.
3 exists for the block-diagonal of the Hessian.
A block with respect to the parameters of a module, ∇ 2 θ ( i ) (cid:96) n ( θ ) , can be obtained by the recursion ∇
[Equation]
(4) and a similar relation holds for the Hessian with respect to each module’s output, ∇ 2 z ( i ) n (cid:96) n ( θ ) .
Both backpropagation schemes of Eq.
3 and Eq.
4 hinge on the multiplication by Jacobians to both vectors and matrices.
However, the design of automatic differentiation limits the application of Jacobians to vectors only.
This prohibits the exploitation of vectorization in the matrix case, which is needed for second-order information.
The lacking ﬂexibility of Jacobians is one motivation for our work.
Since all quantities needed to compute statistics of the derivatives are already computed during the backward pass, another motivation is to provide access to them at minor overhead. [Figure] [Figure] [Figure]
[Figure]
2. F IRST ORDER EXTENSIONS
As the principal ﬁrst-order extension, consider the computation of the individual gradients in a batch of size N .
These individual gradients are implicitly computed during a traditional backward pass because the batch gradient is their sum, but they are not directly accessible.
The na¨ıve way to compute N individual gradients is to do N separate forward and backward passes, This (inefﬁciently) replaces every matrix-matrix multiplications by N matrix-vector multiplications.
B ACK PACK’s approach batches computations to obtain large efﬁciency gains, as illustrated by Fig.
[Figure]
As the quantities necessary to compute the individual gradients are already propagated through the computation graph
[Figure]
we can reuse them by inserting code in the standard backward pass.
With access to this information, before it is cleared for memory efﬁciency, B ACK PACK computes the Jacobianmultiplications for each sample
[Equation]
without summing the result—see Fig.
4 for a schematic representation.
This duplicates some of the computation performed by the backpropagation, as the Jacobian is applied twice (once by P YT ORCH and B ACK PACK with and without summation over the samples, respectively).
However, the associated overhead is small compared to the for-loop approach: The major computational cost arises from the propagation of information required for each layer, rather than the formation of the gradient within each layer.
This scheme for individual gradient computation is the basis for all ﬁrst-order extensions.
In this direct form, however, it is expensive in memory: if the model is D -dimensional, storing O ( ND ) elements is prohibitive for large batches.
For the variance, 2 nd moment and (cid:96) 2 norm, B ACK PACK takes advantage of the Jacobian’s structure to directly compute them without forming the individual gradient, reducing memory overhead.
See Appendix A.1 for details.
2. S ECOND - ORDER EXTENSIONS
Second-order extensions require propagation of more information through the graph.
As an example, we will focus on the generalized Gauss-Newton ( GGN ) matrix (Schraudolph, 2002).
It is guaranteed to be positive semi-deﬁnite and is a reasonable approximation of the Hessian near the minimum, which motivates its use in approximate second-order methods. For popular loss functions, it coincides with the Fisher information matrix used in natural gradient methods (Amari, 1998); for a more in depth discussion of the equivalence, see the reviews of Martens (2014) and Kunstner et al.
For an objective function that can be written as the composition of a loss function (cid:96) and a model f , such as Eq. 1, the GGN of 1 N (cid:80) n (cid:96) ( f ( θ , x n ) , y n ) is [Equation]
[Equation]
The full matrix is too large to compute and store.
Current approaches focus on its diagonal blocks, where each block corresponds to a layer in the network.
Every block itself is further approximated, { z ( i − 1) n } Nn = { z ( i ) n } Nn = {∇ z ( i − 1) n ‘ n } Nn = {∇ z ( i ) n ‘ n } Nn = { [J z ( i ) n f n ] > S n } Nn =1 { [J z ( i − 1) n f n ] > S n } Nn = { [J θ ( i ) f n ] > S n } Nn = ∇ θ ( i )
L θ ( i )
T ( i ) θ ( i ) (cid:0) z
( i − 1) n (cid:1) =
z ( i ) n
G ( θ ) =
(cid:80) [Figure] for example using a Kronecker factorization. The approach used by B ACK PACK for their computation is a reﬁnement of the Hessian Backpropagation equations of Dangel et al.
It relies on two insights: Firstly, the computational bottleneck in the computation of the GGN is the multiplication with the Jacobian of the network, J θ f n , while the Hessian of the loss with respect to the output of the network is easy to compute for most popular loss functions.
Secondly, it is not necessary to compute and store each of the N [ D × D ] matrices for a network with D parameters, as Eq.
6 is a quadratic expression.
Given a symmetric factorization S n of the Hessian, S n S (cid:62) n = ∇ 2 f (cid:96) ( f ( θ , x n ) , y n ) , it is sufﬁcient to compute [J θ f n ] (cid:62) S n and square the result.
A network output is typically small compared to its inner layers; networks on CIFAR -100 need C =class outputs but could use convolutional layers with more than 100,000 parameters.
The factorization leads to a [ D × C ] matrix, which makes it possible to efﬁciently compute GGN block diagonals.
Also, the computation is very similar to that of a gradient, which computes [J θ f n ] (cid:62) ∇ f n (cid:96) n .
A module T ( i ) θ ( i ) receives the symmetric factorization of the GGN with respect to its output, z ( i ) n , and multiplies it with the Jacobians with respect to the parameters θ ( i ) and inputs z ( i − 1) n to produce a symmetric factorization of the GGN with respect to the parameters and inputs, as shown in Fig.
This propagation serves as the basis of the second-order extensions.
If the full symmetric factorization is not wanted, for memory reasons, it is possible to extract more speciﬁc information such as the diagonal.
If B is the symmetric factorization for a GGN block, the diagonal can be computed as [ BB (cid:62) ] ii = (cid:80)
j [ B ] 2 ij , where [ · ] ij denotes the element in the i th row and j th column.
This framework can be used to extract the main Kronecker factorizations of the GGN , KFAC and
KFLR , which we extend to convolution using the approach of Grosse & Martens (2016).
The important difference between the two methods is the initial matrix factorization S n .
Using a full symmetric factorization of the initial Hessian, S n S (cid:62) n = ∇ 2 f n (cid:96) n , yields the KFLR approximation.
KFAC uses an MC -approximation by sampling a vector s n such that E s n [ s n s (cid:62) n ] = ∇ 2 f n (cid:96) n .
KFLR is therefore more precise but more expensive than KFAC , especially for networks with high-dimensional outputs, which is reﬂected in our benchmark on CIFAR -100 in Section 3.
The technical details on how Kronecker factors are extracted and information is propagated for second-order B ACK PACK extensions are documented in Appendix A.2.
E VALUATION AND BENCHMARKS
We benchmark the overhead of B ACK PACK on the CIFAR -10 and CIFAR -100 datasets, using the 3 C 3 D network 3 provided by D EEP OBS (Schneider et al., 2019) and the A LL -CNN-C 4 network of
Springenberg et al. (2015).
The results are shown in Fig.
For ﬁrst-order extensions, the computation of individual gradients from a mini-batch adds noticeable overhead due to the additional memory requirements of storing them.
But more speciﬁc quantities such as the (cid:96) 2 norm, 2 nd moment and variance can be extracted efﬁciently.
Regarding secondorder extensions, the computation of the GGN can be expensive for networks with large outputs like
CIFAR -100, regardless of the approximation being diagonal of Kronecker-factored.
Thankfully, the
MC approximation used by KFAC , which we also implement for a diagonal approximation, can be computed at minimal overhead—much less than 2 backward passes.
This last point is encouraging, as our optimization experiment in Section 4 suggest that this approximation is reasonably accurate. [Figure]
[Figure]
Epoch
T r a i n acc u r ac y
E XPERIMENTS
To illustrate the utility of B ACK PACK, we implement preconditioned gradient descent optimizers using diagonal and Kronecker approximations of the GGN .
To our knowledge, and despite their apparent simplicity, results using diagonal approximations or the na¨ıve damping update rule we chose have not been reported in publications so far.
However, this section is not meant to introduce a bona-ﬁde new optimizer.
Our goal is to show that B ACK PACK can enable research of this kind.
The update rule we implement uses a curvature matrix G ( θ t ( i ) ) , which could be a diagonal or Kronecker factorization of the GGN blocks, and a damping parameter λ to precondition the gradient:
[Equation]
[Table] We run the update rule with the following approximations of the generalized Gauss-Newton: the exact diagonal (Diag GGN ) and an MC estimate (Diag GGN - MC ) [Table] and the Kronecker factorizations
KFAC (Martens & Grosse, 2015), KFLR and KFRA 5 (Botev et al., 2017).
The inversion required by the update rule is straightforward for the diagonal curvature.
For the
Kronecker-factored quantities, we use the approximation introduced by Martens & Grosse (2015)
(see Appendix C.3).
These curvature estimates are tested for the training of deep neural networks by running the corresponding optimizers on the main test problems of the benchmarking suite D EEP OBS (Schneider et al., 2019).
6 We use the setup (batch size, number of training epochs) of D EEP OBS’ baselines, and tune the learning rate α and damping parameter λ with a grid search for each optimizer (details in Appendix C.2).
The best hyperparameter settings is chosen according to the ﬁnal accuracy on a validation set.
We report the median and quartiles of the performance for ten random seeds.
Fig.
7a shows the results for the 3 C 3 D network trained on CIFAR -10.
The optimizers that leverage
Kronecker-factored curvature approximations beat the baseline performance in terms of per-iteration progress on the training loss, training and test accuracy.
Using the same hyperparameters, there is little difference between KFAC and KFLR , or Diag GGN and Diag GGN - MC .
Given that the quantities based on MC -sampling are considerably cheaper, this experiment suggests it being an important technique for reducing the computational burden of curvature approximations.
Fig. 7b shows benchmarks for the A LL -CNN-C network trained on CIFAR -100.
Due to the highdimensional output, the curvatures using a full matrix propagation rather than an MC sample cannot be run on this problem due to memory issues.
Both Diag GGN - MC and KFAC can compete with the baselines in terms of progress per iteration.
As the update rule we implemented is simplistic on purpose, this is promising for future applications of second-order methods that can more efﬁciently use the additional information given by curvature approximations.
C ONCLUSION
Machine learning’s coming-of-age has been accompanied, and in part driven, by a maturing of the software ecosystem.
This has drastically simpliﬁed the lives of developers and researchers alike, but has also crystallized parts of the algorithmic landscape.
This has dampened research in cutting-edge areas that are far from mature, like second-order optimization for deep neural networks.
To ensure that good ideas can bear fruit, researchers must be able to compute new quantities without an overwhelming software development burden.
To support research and development in optimization for deep learning, we have introduced B ACK PACK, an efﬁcient implementation in P Y T ORCH of recent conceptual advances and extensions to backpropagation (Tab. 1 lists all features).
B ACK PACK enriches the syntax of automatic differentiation packages to offer additional observables to optimizers beyond the batch-averaged gradient.
Our experiments demonstrate that B ACK PACK’s implementation offers drastic efﬁciency gains over the kind of na¨ıve implementation within reach of the typical researcher.
As a demonstrative example, we “invented” a few optimization routines that, without
B ACK PACK, would require demanding implementation work and can now be tested with ease.
We hope that studies like this allow B ACK PACK to help mature the ML software ecosystem further. A CKNOWLEDGMENTS
The authors would like to thank Aaron Bahde, Ludwig Bald, and Frank Schneider for their help with D EEP OBS and Lukas Balles, Simon Bartels, Filip de Roos, Tim Fischer, Nicolas Kr¨amer, Agustinus Kristiadi, Frank Schneider, Jonathan Wenger, and Matthias Werner for constructive feedback. Felix Dangel and Philipp Hennig gratefully acknowledge ﬁnancial support by the European Research Council through ERC Action StG 757275 / PANAMA.
Felix Dangel kindly acknowledges the support of the International Max Planck Research School for Intelligent Systems (IMPRS-IS).
– § A: B ACK PACK extensions – § A.1: First-order quantities – § A.2: Second-order quantities based on the generalized Gauss-Newton – § A.3: The exact Hessian diagonal – § B: Additional details on benchmarks – § C: Additional details on experiments – § D: B ACK PACK cheat sheet
[Table]
Quadratic model of the objective at θ , constructed at θ t using G as the curvature
[Table]
of layer
[Table]
sample
[Table]
layer
[Table]
layer i
[Table]
( z (0) n = x n , z ( L ) n = f ( x n , θ ) )
[Table]
A B ACK PACK EXTENSIONS
This section provides more technical details on the additional quantities extracted by B ACK PACK.
Notation: Consider an arbitrary module T ( i ) θ ( i ) of a network i = 1 , . . .
, L , parameterized by θ ( i ) .
It transforms the output of its parent layer for sample n , z
( i − 1) n , to its output z ( i ) n , i.e.
[Equation]
(8) where N is the number of samples.
In particular, z (0) n = x n and z ( L ) n ( θ ) = f
( x n , θ ) , where f is the transformation of the whole network.
The dimension of the hidden layer i ’s output z ( i ) n is written h
( i ) and θ ( i ) is of dimension d ( i ) .
The dimension of the network output, the prediction z ( L ) , is h ( L ) =
C .
For an image classiﬁcation task, C corresponds to the number of classes.
All quantities are assumed to be vector-shaped.
For image-processing transformations that usually act on tensor-shaped inputs, we can reduce to the vector scenario by vectorizing all quantities; this discussion does not rely on a speciﬁc ﬂattening scheme.
However, for an efﬁcient implementation, vectorization should match the layout of the memory of the underlying arrays.
Jacobian: The Jacobian matrix J a b of an arbitrary vector b ∈ R B with respect to another vector a ∈ R A is an [ A × B ] matrix of partial derivatives, [J a b ] ij = ∂
[ b ] i /∂
[ a ] j .
A. F IRST - ORDER QUANTITIES
The basis for the extraction of additional information about ﬁrst-order derivatives is given by Eq. 3, which we state again for multiple samples,
[Equation]
During the backpropagation step of module i , we have access to ∇ z ( i ) n (cid:96) ( θ ) , i = 1 , . . .
, N .
To extract more quantities involving the gradient, we use additional information about the transformation T ( i ) θ ( i ) within our custom implementation of the Jacobian J θ ( i )
z ( i ) n and transposed Jacobian (J θ ( i )
z
( i ) n ) (cid:62) .
Individual gradients: The contribution of each sample to the overall gradient, 1 N ∇ θ
( i ) (cid:96) n ( θ ) , is computed by application of the transposed Jacobian,
[Equation]
For each parameter θ ( i ) the individual gradients are of size ( N, d ( i ) ) .
Individual gradient (cid:96) 2 norm: The quantity (cid:13)(cid:13) 1 N ∇ θ ( i ) (cid:96) n ( θ ) (cid:13)(cid:13) 2 2 , for n = 1 , ..., N , could be extracted from the individual gradients (Eq. 9) as
[Equation]
which is an N -dimensional object for each parameter θ ( i ) .
However, this is not memory efﬁcient as the individual gradients are an [ N × d ( i ) ] tensor.
To circumvent this problem, B ACK PACK uses the structure of the Jacobian whenever possible.
For a speciﬁc example, take a linear layer with parameters θ as an [ A × B ] matrix.
The layer transforms the inputs z ( i − 1) n , an [ N × A ] matrix which we will now refer to as A .
During the backward pass, it receives the gradient of the individual losses with respect to its output, { 1 N ∇ z ( i ) n (cid:96) n } Nn =1 , as an [ N × B ] matrix which we will refer to as B .
The overall gradient, an [ A × B ] matrix, can be computed as A (cid:62) B , and the individual gradients are a set of N [ A × B ] matrices, { A [ n, :] B [ n, :] (cid:62) } Nn =1 .
We want to avoid storing that information.
To reduce the memory requirement, note that the individual gradient norm can be written as
[Equation]
=
[Equation]
i (cid:88) j
( A [ n, i ] B [ n, j ]) 2 , and that the summation can be done independently for each matrix, as (cid:80) i (cid:80)
j ( A [ n, i ] B [ n, j ]) 2 =
( (cid:80) i A [ n, i ]) 2 ( (cid:80) j B [ n, j ] 2 ) .
Therefore, we can square each matrix (element-wise) and sum over non-batch dimensions.
This yields vectors a , b of N elements, where a [ n ] = (cid:80) i
A [ n, i ] 2 .
The individual gradients’ (cid:96) 2 norm is then given by a ◦ b where ◦ is element-wise multiplication.
Second moment: The gradient second moment (or more speciﬁcally, the diagonal of the second moment) is the sum of the squared elements of the individual gradients in a mini-batch, i.e.
[Equation]
It can be used to evaluate the variance of individual elements of the gradient (see below).
The second moment is of dimension d ( i ) , the same dimension as the layer parameter θ
( i ) .
Similarly to the (cid:96)norm, it can be computed from individual gradients, but is more efﬁciently computed implicitly.
Revisiting the example of the linear layer from the individual (cid:96) 2 norm computation, the second moment of the parameters θ [ i, j ] is given by (cid:80) n ( A [ n, i ] B [ n, j ]) 2 , which can be directly computed by taking the element-wise square of A and B element-wise, A 2 , B 2 , and computing A 2 (cid:62) B 2 .
Variance: Gradient variances over a mini-batch (or more precisely, the diagonal of the covariance) can be computed using the second moment and the gradient itself,
[Equation]
The elementwise gradient variance of same dimension as the layer parameter θ
( i ) , i.e. d ( i ) .
A. S ECOND - ORDER QUANTITIES BASED ON THE GENERALIZED G AUSS -N EWTON
The computation of quantities that originate from the approximations of the Hessian require an additional backward pass (see Dangel et al.
Most curvature approximations supported by BackPACK rely on the generalized Gauss-Newton ( GGN ) matrix (Schraudolph, 2002)
[Equation]
One interpretation of the GGN is that it is the Hessian of the empirical risk when the model f is approximated with its ﬁrst-order Taylor expansion, i.e., linearizing the network and ignoring second order effects.
The effect of the curvature of each module in the recursive scheme of equation 4 can be ignored to obtain the simpler expression
[Equation]
for the exact block-diagonal of the full generalized Gauss-Newton matrix.
In analogy to G ( θ ( i ) )
we introduce the ( d ( i ) , d ( i ) )
-dimensional quantity
[Equation]
that needs to be backpropagated.
The curvature backpropagation also follows from equation 4 as
[Equation]
and is initialized with the Hessian of the loss function with respect to the network prediction, i.e.
[Equation]
Although this scheme is exact, it is computationally infeasible as it requires the backpropagation of
N [ h ( i ) , h ( i ) ]
matrices between module i +
and i .
Even for small N , this is not possible for networks containing large convolutions.
As an example, the ﬁrst layer of the All-CNN-C network outputs 29 × 29 images with 96 channels, which already gives h ( i ) = 80,736, which leads to half a Gigabyte per sample.
Moreover, storing all the ( d ( i ) , d ( i ) )
-dimensional blocks G ( θ ( i ) ) is not possible.
BackPACK implements different approximation strategies, developed by Martens & Grosse (2015) and Botev et al.
(2017) that address both of these complexity issues from different perspectives.
Symmetric factorization scheme: One way to improve the memory footprint of the backpropagated matrices in the case where the model prediction’s dimension C (the number of classes in an image classiﬁcation task) is small compared to all hidden feature is to propagate a symmetric factorization of the generalized Gauss-Newton instead.
It relies on the observation that the loss function itself is convex, even though its composition with the network might not be, and its Hessian with respect to the network output can be decomposed as
[Equation]
with the ( C, C ) -dimensional matrix factorization of the loss Hessian, S ( z ( L ) n ) , for sample n .
Consequently, the generalized Gauss-Newton in equation 12 reduces to an outer product,
[Equation]
The analogue for the diagonal blocks follows from equation 13 and reads
[Equation]
where we deﬁne the ( h ( i ) , C ) -dimensional matrix square root S ( z ( i ) n ) =
(J z ( i ) n f ) (cid:62) S ( z ( L ) n ) .
Instead of having layer i backpropagate N objects of shape ( h ( i ) , h ( i ) )
according to equation 14, one can instead backpropagate the matrix square root as
[Equation]
initialized with the condition speciﬁed by equation 15.
This reduces the backpropagated matrix’s size of layer i to ( h ( i ) , C ) for each sample.
A.2. D IAGONAL CURVATURE APPROXIMATIONS
Diagonal of the generalized Gauss-Newton:
The factorization trick for the loss Hessian reduces the size of the backpropagated quantities, but does not address the intractable size of the generalized
Gauss-Newton blocks G ( θ ( i ) ) .
In BackPACK , we can extract diag (cid:0) G ( θ ( i ) ) (cid:1) given the backpropagated quantities S ( z ( i ) n ) , i = 1 , . . .
, N , without building up the matrix representation of equation 17.
For completeness, we compute
[Equation]
Diagonal of the generalized Gauss-Newton with Monte-Carlo sampled loss Hessian: We can use the same backpropagation strategy of equation 18, replacing the symmetric factorization of equation 15 with an approximation using a smaller matrix ˜ S
( z ( L ) n ) of size ( C, ˜ C ) and ˜ C < C ,
[Equation]
This further reduces the size of backpropagated curvature quantities.
Martens & Grosse (2015) introduced such a sampling scheme with KFAC based on the connection between the generalized
Gauss-Newton and the Fisher.
Most loss functions used in machine learning have a probabilistic interpretation, as the log-likelihood of a probabilistic model.
The squarred error of regression is equivalent to a Gaussian noise assumption and the cross-entropy is linked to the categorical distribution.
The Hessian of the loss function with respect to the output of the network is equal, in expectation, to the outer products of gradients if the output of the network is sampled according to a particular distribution , p f ( x ) , deﬁned by the network output f ( x ) .
Sampling outputs ˆ y ∼ p , we have that
[Equation]
Sampling one such gradient leads to a rank-1 Monte-Carlo approximation of the initial Hessian.
With the substitution S ↔ ˜ S , we compute a MC-sampled approximation of the generalized GaussNewton matrix’s diagonal in BackPACK as
[Equation]
A.2. K RONECKER - FACTORED CURVATURE APPROXIMATIONS
A different approach to reduce memory complexity of the generalized Gauss-Newton blocks
G ( θ ( i ) ) , apart from diagonal curvature approximations, is representing them as Kronecker products (KFAC for linear and convolution layer by Martens & Grosse (2015); Grosse & Martens (2016)
KFLR and KFRA for linear layers by Botev et al.
[Equation]
For both linear and convolution layers, the ﬁrst factor A ( i ) is straightforwardly obtained from the inputs z ( i − 1) n to layer i .
Instead of repeting the technical details of the aforementioned references, we would like to focus on how they differ in (i) the backpropagated quantities and (ii) the backpropagation strategy.
As a result, we are able to extend the KFLR and KFRA formulations to layers used in convolutional neural networks 7 .
KFAC and KFLR: KFAC uses an MC-sampled estimate of the loss Hessian with a square root factorization ˜ S ( z ( L ) n ) like in equation 20.
The backpropagation procedure for this quantity is equivalent to the computation of the GGN diagonal.
For the GGN of the weights of a linear layer i , the second Kronecker term is given by
[Equation]
which at the same time corresponds to the GGN of the layer’s bias 8 .
In contrast to KFAC, the KFLR approximation backpropagates the exact square root factorization
S ( z ( L ) n ) , i.e. for the weights of a linear layer one ﬁnds 8 (see Botev et al.
(2017) for more details)
[Equation]
KFRA:
The backpropagation strategy for KFRA eliminates the scaling of the backpropagated curvature quantities with the batch size N in equation 14.
Instead of having layer i receive the
N exact ( h ( i ) , h ( i ) )
-dimensional matrices G
( z ( i ) n ) , n = 1 , . . .
, N , only a single averaged object, which is denoted G ( i ) , is used as an approximation.
In particular, the recursion changes to
[Equation]
and is initialized with the batch-averaged loss Hessian
[Equation]
For a linear layer, KFRA uses see Botev et al.
(2017) for more details)
[Equation]
A. T HE EXACT H ESSIAN DIAGONAL
For neural networks consisting only of piecewise linear activation functions, computing the diagonal of the Hessian is equivalent to compute the diagonal of the GGN.
This is because these activations make the second term in the Hessian backpropagation recursion (equation 4) vanish.
However, for activation functions with non-vanishing second derivative, these residual terms have to be accounted for in the backpropagation scheme.
For completeness, the Hessian backpropagation for module i read
[Equation]
for n = 1 , . . .
, N .
Those ( h ( i ) , h ( i ) )
-dimensional residual terms are deﬁned as
[Equation]
For the common parameterized layers, such as linear and convolution transformations, the residual vanishes, R ( i ) n ( θ ( i ) ) = 0 .
Regarding R ( i ) n ( z ( i
[Equation]
[Equation]
) , if the activation function is applied elementwise, the residual terms are diagonal matrices.
Storing these quantities becomes very memory-intensive for high-dimensional nonlinear activation layers.
In BackPACK , this complexity is reduced by application of the aforementioned matrix square root factorization trick.
To do so, we express the symmetric factorization of R ( i ) n ( z ( i − 1) n ) as
[Equation]
where P ( i ) n ( z ( i − 1) n ) , N ( i ) n ( z ( i − 1) n ) represent the matrix square root of R ( i ) n ( z ( i − 1) n ) projected on its positive and negative eigenspace, respectively.
This composition allows for the extension of the backpropagation for the GGN:
In addition to
S ( z ( i ) n ) , the decompositions P ( i ) n ( z ( i − 1) n ) , N ( i ) n ( z ( i − 1) n ) for the residual parts also have to be backpropagated according to equation 18.
All diagonals are extracted from the backpropagated matrix square roots (see equation 19).
To obtain the correct diagonal, all diagonals stemming from decompositions of the negative residual eigenspace have to be weighted by a factor of − 1 before summation.
In terms of complexity, one backpropagation step for R ( i ) n ( z ( i − 1) ) changes the dimensionality in the following scheme
[Equation]
− 1) )
→ ( h ( i − 2) , h ( i − 2) )
→ . . .
With the square root factorization, one instead obtains
P
[Equation]
[Equation]
h ( i ) , h ( i ) )
→ ( h ( i − 1) , h
[Equation]
i ) )
→ ( h ( i − 2) , h ( i ) )
→ . . . ,
N ( i ) n ( z ( i − 1) n ) : ( h ( i ) , h ( i ) )
→ ( h ( i − 1) , h ( i ) ) → ( h ( i − 2) , h ( i ) )
→ . . . .
Roughly speaking, the latter scheme is more efﬁcient whenever the hidden dimension of a nonlinear activation layer deceeds the largest hidden dimension of the network.
Example: Consider one backpropagation step of module i .
Assume R ( i ) n ( θ ( i ) )
= 0 , i.e. a linear, convolution, or non-parameterized layer.
Then the following computations are performed in the protocol for the diagonal Hessian:
• Receive the following quantities from the child module i + 1 (for n = 1 , . . .
, N )
[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]
• Extract the module parameter Hessian diagonal, diag (cid:0) ∇ 2 θ ( i ) L ( θ
[Equation]
(cid:1) –
For each quantity A ∈ Φ extract the diagonal from the square root factorization and sum over the samples, i.e. compute
[Equation]
Multiply the expression by − 1 if A stems from backpropagation of a residual’s negative eigenspace’s factorization.
– Sum all expressions to obtain the block Hessian’s diagonal, i.e. diag (cid:0) ∇ 2 θ ( i ) L ( θ ) (cid:1) • Backpropagate the received quantities to the parent module i −– For each quantity A n ∈ Φ , apply (J z ( i − 1) n z ( i ) n ) (cid:62) A n – Append P ( i +1) n ( z ( i ) n ) and N ( i +1) n ( z ( i ) n ) to Φ
B A DDITIONAL DETAILS ON BENCHMARKS
KFAC vs. KFLR As the KFLR method of Botev et al.
(2017) is orders of magnitude more expensive to compute than the KFAC method of Martens & Grosse (2015) on CIFAR -100, it was not included in the main plot.
This is not an implementation error; it follows from the deﬁnition of those methods.
To approximate the Generalized Gauss-Newton, G = (cid:80) n
[J
[Equation]
f n ] (cid:62) ∇ [Equation] n
[Equation]
n [J
[Equation]
f n ] , KFAC uses a rank-one approximation for each of the inner Hessian G n = ∇ 2 f n (cid:96) n
= s n s (cid:62) n , and needs to propagate a vector through the computation graph for each sample.
KFLR uses the complete inner
Hessian instead.
For CIFAR100, the network has 100 output nodes—one for each class—and the inner Hessians are [100 × 100] matrices.
KFLR needs to propagate a matrix through the computation graph for each sample, which is 100 × more expensive as shown in Fig. 8.
[Figure]
Diagonal of the GGN vs. Diagonal of the Hessian Most networks used in deep learning use
ReLU activation functions.
ReLU functions have no curvature as they are piece-wise linear functions.
Because of this, the diagonal of the GGN is equivalent to the diagonal of the Hessian (Martens, 2014).
However, for networks that use non piece-wise linear activation functions like sigmoids or tanh functions, computing the diagonal of the Hessian can be much more expensive than the diagonal of the GGN.
To illustrate this point, we modify the smaller network used in our benchmarks to include a single sigmoid activation function before the last classiﬁcation layer.
The results in Fig.show that the computation of the diagonal of the Hessian is already an order of magnitude more expensive than the diagonal of the GGN.
[Figure]
more
C A DDITIONAL DETAILS ON EXPERIMENTS C. P ROTOCOL
The experiments on optimizer performance are performed according to the protocol suggested by the maintainers of DeepOBS:
• Train the neural network with the investigated optimizer and vary its hyperparameters on a speciﬁed grid.
This training is performed for a single random seed only.
• DeepOBS evaluates metrics during the training procedure (estimates of the loss function and classiﬁcation accuracy on the training, validation, and test data set).
From all runs of the grid search, it selects the best run automatically.
The results shown in this work were obtained with the default strategy, favoring the run with highest ﬁnal accuracy on the validation set.
• For a better understanding of the optimizer performance with respect to randomized routines in the training process, DeepOBS reruns the best hyperparameter setting for ten different random seeds.
The results show mean values over these repeated runs, with standard deviations as uncertainty indicators.
• Along with the benchmarked optimizers, we show the DeepOBS base line performances for Adam and momentum SGD (Momentum).
They are provided by the DeepOBS maintainers.
The optimizer’s built upon B ACK PACK’s curvature estimates were benchmarked on the DeepOBS image classiﬁcation problems summarized in table 3. Schneider
[Table]
C. G RID SEARCH AND BEST HYPERPARAMETER SETTING
Both the learning rate α and damping λ are tuned over the following grid:
• α ∈ (cid:8) 10 − 4 , 10 − 3 , 10 − 2 , 10 − 1 , 1 (cid:9) • λ ∈ (cid:8) 10 − 4 , 10 − 3 , 10 − 2 , 10 − 1 , 1 , 10 (cid:9) .
[Table]
For all test problems, the same batch size ( B = 128 for all problems, except B = 256 for AllCNNC on CIFAR -100) as used for the base lines is used and the optimizers run for the identical number of epochs.
The best hyperparameter settings are summarized in table 4.
C. U PDATE RULE
For the experiments in the main text, we use a simple update rule with a constant damping parameter λ .
Consider the parameters θ of a single module in a neural network, with L 2 -regularization of strength η .
Let G ( θ k ) denote the curvature matrix and ∇ θ L ( θ k ) the gradient at step k .
One iteration of the optimizer applies the update
[Equation]
The inverse cannot be computed exactly (in a reasonable time) for the Kronecker-factored curvatures
KFAC , KFLR , and KFRA .
We use the scheme introduced by Martens & Grosse (2015) to approximately invert G ( θ k ) + ( λ + η ) I if G ( θ k ) is Kronecker-factored; G ( θ k )
=
A ( θ k ) ⊗ B ( θ k ) .
It replaces the expression ( λ + θ ) I by diagonal terms added to each Kronecker factor.
In summary, this replaces
[Equation]
A principled choice for the parameter π k is given by the expression π k = (cid:113)
[Equation] [Equation]
B
[Equation]
k ) (cid:107) for an arbitrary matrix norm (cid:107)·(cid:107) .
We follow Martens & Grosse (2015) and choose the trace norm,
[Equation] [Equation]
C. A DDITIONAL RESULTS
This section presents the results for MNIST using a Logistic Regression in Fig.
10 and F ASHIONMNIST using the 2C2D network, composed of two convolution and two linear layers, in Fig.
[Figure] [Figure]
D B ACK PACK CHEAT SHEET
• Assumptions – Feedforward network
[Equation]
– Dimension of parameter θ ( i ) : dim( θ
[Equation]
= d
[Equation]
– Empirical risk
[Equation]
• Shorthands
[Equation]
• Generalized Gauss-Newton matrix
[Equation]
• Approximative GGN via MC sampling
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Table]