BackPACK: Packing more into Backprop
Authors: Felix Dangel
Frederik Kunstner
Philipp Hennig
Venue:  (2020)

Abstract: 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 efficiently, 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, inefficient if done naively, 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  BackPACK, an efficient framework built on top of  PyTorch, that extends the backpropagation algorithm to extract additional information from first-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.

Published as a conference paper at ICLR 2020
BACKPACK: PACKING MORE INTO BACKPROP
Felix Dangel
Frederik Kunstner" University of Tuebingen kunstner@cs ubc ca
University of Tuebingen
fdangel (uue. mpg. de
Philipp Hennig
University of Tuebingen and MPI for Intelligent Systems, Tuebingen ph@tue. mpg. de
ABSTRACT
Automatic differentiation frameworks are optimized for exactly one thing: com- puting 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, computed efficiently, and at the same time as the gradient. While these quan- tities are of great interest to researchers and practitioners, current deep-learning software does not support their automatic calculation. Manually implementing them is burdensome, inefficient if done naïvely, and the resulting code is rarely shared, This hampers progress in deep learning, and unnecessarily narrows re- search to focus gradient descent and its variants; it also complicates repli- cation studies and comparisons between newly developed methods that require those quantities, the point of impossibility. To address this problem, we intro- duce BACKPACK! an efficient framework built on top of PYTORCH, that extends the backpropagation algorithm
extract additional information from first- and second-order derivatives. Its capabilities are illustrated by benchmark reports for computing additional quantities on deep neural networks, and an example appli- cation by testing several recent curvature approximations for optimization.
INTRODUCTION
The success of deep learning and the applications it fuels be traced to the populariza- tion of automatic differentiation frameworks. Packages like TENSORFLow (Abadi et al., 2016), CHAINER (Tokui al., 2015), MXNET (Chen al., 2015), and PYTORCH (Paszke et al., 2019) provide efficient 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 mini-batch of examples. Other quan- tities 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 identification of free parameters, and to push the development of more efficient optimization algorithms. But researchers who want to investigate their use face a chicken- and-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
Equul contributions httpps:// htttps:/ffddange1 gi thub. io/backpack

---

Published as a conference paper at ICLR 2020
the default. There may of course be good scientific reasons for this difference; maybe second-ordes methods do not work well in the (non-convex, stochastic) setting of deep learning. And the computa- tional cost associated with the high dimensionality of deep models may offset their benefits. Whether these are the case remains somewhat unclear though, because a much more direct road-block is that these methods 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 al., 2019). Their approach, based on the earlier work of Schraudolph (2002), uses the structure of the network to compute approximate second- order information in a way that is similar to gradient backpropagation. This work sparked 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 a daunting task to implement them from scratch. Unless users spend significant time familiarizing themselves with the internals of their software tools, the resulting implementation is often inefficient, 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 efficient. New methods 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 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 first- order 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-righhly-o improving their technical backbone. Features like those outlined above are not generally defined for arbitrary functions, but rather emerge from the specific 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 efficiently realize such features within these frameworks: In essence, backpropagation is technique 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 small group of specialists; they are just rarely discussed or implemented.
1.1 OUR 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. (2019) for modular backpropagation. This framework can be built on top of existing graph-based based backpropagation modules; we provide an implementation on top of PYTORCH, coined BACKPACK, available at
https: / ha////ddaagge-eee/// github 10/ http///ack/////////////////////.
The initial release supports efficient computation of individual gradients from a mini-batch, their & norm, an estimate of the variance, as well as diagonal and Kronecker factorizations of the generalized Gauss-Newton (GGN) matrix (see Tab. 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 flexibility of automatic differentiation systems (Innes, 201 8a;;; Bradbury 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.
2 See, e.g., the Github issues gi thub. wttpbbccom//yyeeeee www.:////pttroc/yttttooooo and forum discussions d scuss. wyt.sss.tttrcc..ooooooo
wsssssssssss... 1077, 7786, 8897 8405, 1 52700, 17204, 19350, 24955.
2

---

Published as a conference paper at ICLR 2020
Computing the gradient with PYTORCH.
and the variance with BACKPACK
X, y
gead.mnist.date Linear (784, 10)
X, y model extend (Linear (784, 101
model
lossfure
croasintropyLoos ()
lossfun: extend (Cronssntrogylooss] (1)
loss
losstunc (model (X),
loss
loas Func (model (X), y)
backpack (ariance()):
loss. (ackward
loas backward
for pazam modell.prameee ()::
for param in peddel.aasmmterss
print (param.gradi
print (param. print (param. var)
Figure I: BACKPACK integrates with PYTORCII to seamlessly extract more information from the backward pass. Instead of the variance (or alongside it, in the same pass), BACKPACK can compute individual gradients in the
their &g norm and 2 moment. It can also compute curvature approximations like diagonal or Kronecker factorizations of the GGN such as KFAC, KFLR R & KFRA.
To illustrate the capabilities of BACKPACK, we it to implement preconditioned gradient descent optimizers with diagonal approximations of the GGN and recent Kronecker factorizations KFAC (Martens& & Grosse, 2015), KFI.R, 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 first-order baselines such as SGD with momentum and Adam (Kingma & Ba, 2015), its implementation with various curvature approximations is made straightforward.
2 THEORY AND IMPLEMENTATION
We will distiguish between quantities that can be computed from information already present during traditional backward pass (which we suggestively call first-order extensions), and quantities that need additional information (termed second-order extensions). The former group contains additional statistics such the variance of the gradients within the mini-batch or the l 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
Second-order extensions Propagate new information along the graph.
Extract more from the standard backward pass.
Individual gradients from mini-batch
Diagonal of the GGN and the Hessian
l2 norm of the individual gradients
KFAC (Martens & Grosse, 2015) KFRA and KFLR (Botev et al., 2017)
Diagonal covariance and 2nd moment
These quantities are only defined, or reasonable to compute, for subset of models: The concept of individual gradients for each sample 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 fit into this category. For example, if the network uses Batch Normalization (loffe & Szegedy, 2015), the individual gradients in a mini-batch are correlated. Then, the variance not meaningful anymore, and computing the individual contribution of sample to the mini-batel gradient or the GGN becomes prohibitive. For those reasons, and to limit the scope of the project version 1.0, BACKPACK currently restricts the type of models it accepts. The supported models traditional feed-forward networks that can be expressed a sequence of convolutional, pooling, linear and activation layers. Recurrent networks like LSTMS (Hochreiter & Schmidhuber, 1997) or residual networks (He al., 2016) the framework can be extended to cover them.
sequence of modules, for example
yet supported, but
We assume sequential model θ X Y and dataset of N samples (æ,, 3n) X Y with n = 1,..., N. ... The model maps each sample H,, to a prediction 9n. using some parameters 0 E e. The predictions evaluated with loss function l V R. for example the cross-entropy,
3

---

Published as a conference paper at ICLR 2020
Verc C
Figure 2: Schematic rep- resentation of the standard backpropagation pass for module with N samples.
(zh (=!--1))a-1) (1 fn mam1
(= (=
1000 1) =!!")
(Vigo fn
which compares them to the ground truth yn- This leads to the objective function L : O - I.
( N E ((f 2n). yn)
As a shorthand, we will use e, (0) ((, x. yn) for the loss and (0) f(0, a.,)) for the model output of individual samples. Our goal is to provide more information about the derivatives of {4 fn with respect to the parameters of the model f.
2.1 PRIMER ON BACKPROPAGATION
Machine learning libraries with integrated automatic differentiation use the modular structure of (0) to compute derivatives (see Baydin et al. (2018) for an overview). If fn is a sequence of L transformations, it can be expressed as
(0) Tel) o Tull (mn),
(2)
where T:(1] g(o) is the ith transformation with parameters ( such that A = loss function can also be seen as another transformation, appended to the network. Let 2 (i-1) (8)
0(0) The
2n
denote the input and output of the operation T/0 for sample n, such that 2 is the original data and 2n (L) represent the transformed output of each layer, leading the computation graph
e(l) (=1)
T
(z()) 0(2)
T(L) etc)
(=("
(1) an)
(0)
(1)
(L)
e (0).
To compute the gradient of e, with respect to the O(0), one can repeatedly apply the chain rule,
V a (()) = (Jan (1) TJ 20 2 ((tt))
(J (1-- =1) (1)) (V e En (0))
(3)
= 2J (6) 2 (0))
where Jab is the Jacobian of b with respect to a, [Jabls = 0bb):/0aa).. A similar expression exists for the module inputs 2(:-1), V # 1, (0) (J # () ) (V (V This recursive structure makes it possible to extract the gradient by propagating the gradient of the loss. In the backprop- V_ Ln6n (0). It then extracts the gradient with respect to its parameters and inputs, (z)) and V 2! (z))), according to Eq. 3. The gradient with respect to its input is sent further down the graph. This pro- cess, 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.
agation algorithm, module i receives the loss gradient with respect to its output, V
For second-order quantities, rely the work of Mizutani & Dreyfus (2008) and Dangel et al. (2019), who showed that scheme similar to Eq. exists for the block-diagonal of the Hessian. A block with respect to the parameters of module, V2 (0), can be obtained by the recursion
v 22. (0) = (Jore) 2m
(1) (J) 21 + E, ( OCo) 11)
|
Cn (0))
(4)
and similar relation holds for the Hessian with respect to each module's output, V (0).
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 Ja- cobians to vectors only. This prohibits the exploitation of vectorization in the matrix case, which is needed for second-order information. The lacking flexibility of Jacobians is motivation for our work. Since all quantities needed to compute statistics the derivatives are already computed during the backward pass, another motivation is to provide access to them minor overhead.
4

---

Published as a conference paper at ICLR 2020
Batch gradients
Figure 3: Computing individual gradients in batch using a for-loop (i.c. one individual forward and backward pass per sample) using vectorized operations with BACK- PACK. The plot shows computation time, comparing to a traditional gradient computation, on the 3C3D network (See $4) for the CIFAR-10 10 dataset (Schneider et al., 2019).
Batch size 64
Time [mus]
400
200
0
For-Joop
BackPACK
Gradient (Ref)
Figure 4: Schematic representa- tion of the individual gradients' extraction in addition to the stan- dard backward pass at the ith module for N samples.
(2b
Telit (all- 1))
2.2 FIRST ORDER EXTENSIONS
As the principal first-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 not directly accessible. The na've way com- pute individual gradients is to do N separate forward and backward passes, This (inefficiently) replaces every matrix-matrix multiplications by N matrix-vector multiplications. BACKPACK's ap- proach batches computations to obtain large efficiency gains, as illustrated by Fig. 3.
As the quantities necessary compute the individual gradients are already propagated through the computation graph, we can reuse them by inserting code in the standard backward pass. With access this information, before it is cleared for memory efficiency, BACKPACK computes the Jacobian- multiplications for each sample
(V C. (0))
{Jor 2m
( (0))
(5)
without summing the result- see Fig. 4 for a schematic representation. This duplicates some of the computation performed by the backpropagation, ToRCH and BACKPACK 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.
the Jacobian is applied twice (once by Py-
This scheme for individual gradient computation is the basis for all first-order extensions. In this direct form, however,
expensive in memory: if the model is D-dimensional, storing O(N D) elements is prohibitive for large batches. For the variance, 2nd moment and l norm, BACKPACK takes advantage of the Jacobian's structure to directly compute them without forming the individual gradient, reducing memory overhead. See Appendix A. for details.
2.3 SECOND-ORDER EXTENSIONS
Second-order extensions require propagation of more information through the graph. As example, we will focus on the generalized Gauss-Newton (GGN) matrix (Schraudolph, 2002). It is guaranteed be positive semi-definite 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 coin- cides with the Fisher information matrix used in natural gradient methods (Amari, 1998); for more in depth discussion of the equivalence, see the reviews of Martens (2014) and Kunstner et al. (2019). For an objective function that can be written as the composition of a loss function and model f. such as Eq. 1, the GGN of VV Σn {(((0, xn), yn) is
(( ト Σ. [o ar,) VS ((((0, x,), yn) (0,
The full matrix is too large to compute and store. Current approaches focus on its diagonal blocks, where each block corresponds layer in the network. Every block itself is further approximated,
5

---

Published as a conference paper at ICLR 2020
e
Egro
Figure 5: Schematic of the addi- tional backward pass to compute symmetric factorization of the GGN, C( [. S, fJ) alongside the gradient at the ith module, for N samples.
win (=l
(g Co famm
[111-11 [-]" S.
(1a22) ful" S.).
for example using Kronecker factorization. The approach used by BACKPACK for their com- putation is a refinement of the Hessian Backpropagation equations of Dangel et al. (2019). It re- lies on two insights: Firstly, the computational bottleneck in the computation of the GGN is the multiplication with the Jacobian of the network, Jn 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 x D] matrices for a network with D pa- rameters, Eq. 6 is quadratic expression. Given symmetric factorization S, of the Hessian, S, = v (ff In), yn). it is sufficient to compute [Jl I S, and square the result. A net- work output is typically small compared to its inner layers; networks on CIFAR 100 need C 100 class outputs but could use convolutional layers with more than 100,000 parameters.
The factorization leads to a [DxC [X] matrix, which makes possible to efficiently compute GGN block diagonals. Also, the computation is very similar to that of a gradient, which computes [J Vy,ln- module ( g(t) receives the symmetric factorization of the GGN with respect to its output, = and
multiplies it with the Jacobians with respect to the parameters (()) and inputs zn (t-1) to produce a
symmetric factorization of the GGN with respect to the parameters and inputs, shown in Fig. 5.
This propagation serves as the basis of the second-order extensions. If the full symmetric factoriza- tion is not wanted, for memory reasons, it is possible to extract more specific information such as the diagonal. If is the symmetric factorization for a GGN block, the diagonal can be computed as [BB]] A Σ, BB3 where Hly denotes the element in the ith row and jth column.
This framework can be used extract the main Kronecker factorizations of the GGN, KFAC and KFLR, which we extend to convolution using the approach of Grosse Martens (2016). The im- portant difference between the two methods is the initial matrix factorization S,, Using full sym- metric factorization of the initial Hessian, S, MC-approximation by sampling vector s, such that E.,, [ fore more precise but more expensive than KFAC, especially for networks with high-dimensional outputs, which is reflected in benchmark CIFAR-100 in Section 3. The technical details how Kronecker factors are extracted and information is propagated for second-order BACKPACK extensions are documented in Appendix A.2.
Ví. fn. yields the KFLR approximation. KFAC
Vi, ln KFLR is there-
3 EVALUATION AND BENCHMARKS
We benchmark the overhead of BACKPACK the CIFAR-10 and CIFAR-100 datasets, using the 3C3D network3 provided by DEEPOBS (Schneider et al., 2019) and the ALL.-CN--C network of Springenberg et al. (2015). The results shown in Fig. 6.
For first-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 specific quantities such as the l2 norm, 2nd moment and variance can be extracted efficiently. Regarding second-order extensions, the computation of the GGN be expensive for networks with large outputs like CIFAR- 100, regardless of the approximation being diagonal of poonecker-factoree Thankfully, the MC approximation used by KFAC, which we also implement for a diagonal approximation, can be com- puted at minimal overhead much less than two backward passes. This last point is encouraging, as our optimization experiment in Section suggest that this approximation is reasonably accurate.
$3C3D is sequence of 3 convolutions and dense linear layers with 895,210 parameters. 'ALL-CNN-C sequence of convolutions with 1,387,108 parameters.

---

Published as a conference paper at ICLR 2020
3C3D on CIFARIO
AlI-CNN-C on CIFAR100
Batch 32
488
100
Batch 32
T1me (ms)
50
2nd Moment DiagGGN MC
Indiv. Grad
DiagGGN
2nd Moment
DiagGGN
Indiv. Grad
12
KFLR
Figure 6: Overhead benchmark for computing the gradient and first- second-order extensions real networks, compared to just the gradient. Most quantities add little overhead. KFLR and DiagGGN propagate x00x more information than KFAC and DiagGGN-MC on CIFAR- 100 and are two orders of magnitude slower. We report benchmarks on those, and the Hessian's 's diagonal, in Appendix B.
4 EXPERIMENTS
To illustrate the utility of BACKPACK, we implement preconditioned gradient descent optimizers using diagonal and Kronecker approximations the GGN. To our knowledge, and despite their apparent simplicity, results using diagonal approximations chose have not been reported in publications so far. However, this section is not meant to introduce a bona-fide new optimizer. Our goal is to show that BACKPACK can enable research of this kind. The update rule we implement uses a curvature matrix G( which could be a diagonal or Kronecker factorization of the GGN blocks, and damping parameter À to precondition the gradient:
the naive damping update rule
AX) 1+1 = - -
We run the update rule with the following approximations of the generalized Gauss-Newton: the exact diagonal (DiagGGN) and an MC estimate (DiagGGN-MC), and the Kronecker factorizations KFAC (Martens Grosse, 2015), KFLR and KFRA KFRte et al., 2017). The inversion required by the update rule is straightforward for the diagonal curvature. For the cronecke-ffactored 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 cor- responding optimizers on the main test problems of the benchmarking suite DEEPOBS (Schneider et al., 2019). We use the setup (batch size, number of training epochs) of DEEPOBS' baselines, and tune the learning rate a and damping parameter À with a grid search for each optimizer (details in Appendix C.2). The best hyperparameter settings is chosen according to the final accuracy validation set. We report the median and quartiles the performance for ten random seeds.
Fig. 7a shows the results for the 3C3D network trained on CIFAR- 10. The optimizers that leverage tronecker-factored curvature approximations beat the baseline performance terms of per-iteration progress the training loss, training and test accuracy. Using the same hyperparameters, there is little difference between KFAC and KFLR,, or DiagGGN and DiagGGN-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 ALL-CNN-C network trained CIPAR-100. Due the high- dimensional output, the curvatures using a full matrix propagation rather than an MC sample cannot be run on this problem due to memory issues. Both DiagGGN-MC 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 efficiently use the additional information given by curvature approximations.
5 KFRA was not originally designed for convolutions; we extend it using the Kronecker factorization of Grosse & Martens (2016). While it can be computed for small networks MNIST, which report in Ap- pendix C.4, the approximate backward pass of KFRA does
scale large convolution layers. 'htpps: / /deepbbs. github 1o/. / We cannot run BACKPACK on all test problems in this bench- mark due to the limitations outlined in Section 2. Despite this limitation, we still run on models spanning representative range of image classification problems.

---

Published as a conference paper at ICLR 2020
(a)
CIFAR-10: 3C3D
DaggiGN
DiggGGN-MMC
1.5
KRAC
KFIR
1.5
Train loss
Test
Adam
Momentum
0.5
0.5
Test accurecy
0.9
Train accuracy
0.8
0.7
0.7
20
40
60
80
100
20
10
60
80
100
Epoch
Epoch
(b)
CIFAR-100: ALL.-CNN-C
Text
Train
DaggGNN-MC MC
KFAC
Adam
Momentum
Test accuricy
0.8
0.8
Train accurcc
0.4
0.4
100
200
300
100
200
300
Epoch
Epoch
Figure 7: Median performance with shaded quartiles of the DEEPOBS benchmark for (a) 3C3D network (895,210) parameters) on CIFAR-10 and (b) AL..- network (1,387,10 parameters) on 11FAR-100. Solid lines show baselines of momentum SGD and Adam provided by DEEPOBS.
CONCLUSION
Machine learning's coming-of-age has been accompanied, and in part driven, by a maturing of the software ecosystem. This has drastically simplified the lives of developers and researchers alike, but has also crystallized parts the algorithmic landscape. This has dampened research in cutting- edge areas that 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 over- whelming software development burden. To support research and development in optimization for deep learning, we have introduced BACKPACK, an efficient implementation PYTORCH recent conceptual advances and extensions backpropagation (Tab. lists all features). BACKPACK en- riches the syntax of automatic differentiation packages to offer additional observables to optimizers beyond the batch-averaged gradient. Our experiments demonstrate that BACKPACK's implementa- tion offers drastic efficiency 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 BACKPACK, would require demanding implementation work and can now be tested with ease. We hope that studies like this allow BACKPACK to help mature the ML software ecosystem further.
ACKNOWLEDGMENTS
The authors would like to thank Aaron Bahde, Ludwig Bald, and Frank Schneider for their help with DEEPOBS and Lukas Balles, Simon Bartels, Filip de Roos, Tim Fischer, Nicolas Krämer, Agustinus Kristiadi, Frank Schneider, Jonathan Wenger, and Matthias Werner for constructive feedback.
The authors gratefully acknowledge financial support by the European Research Council through ERC StG Action 757275 / PANAMA; the DFG Cluster of Excellence "Machine Leaming New
8

---

Published as a conference paper at ICLR 2020
Table 1: Overview of the features supported in the first release of BACKPACK.
Feature
Details
Individual gradients
N V (()))
7 I.
Batch variance
ΣN n=
(0) - Vo [(0)] (
N
2nd moment
E
j= 1 dlo)
N
Indiv. gradient l norm
V (0)
n= ===1... 1....,N N
DiagGGN
diag
(G()()))
DiagGGN-MC
diag
G((())
Hessian diagonal
diag (V2 (())
((0()) R ( 69 B (8)
KFAC
KEAC
KFLR
G(O ~ A
(8)
KFRA
G( ~ (
Perspectives for Science", EXC 2064/1, project number 390727645; the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: (IIS1803A) and funds from the Ministry of Science, Research and Arts of the State of Baden- Württembeer,, F. D. is grateful to the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for support.
REFERENCES
Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthicu Devin, Sanjay Ghemawat, Geoffrey Irving, Michacl Isard, Manjunath Kudlur, Josh Levenberg, Rajat Monga, Sherry Moore, Derek Gordon Murray, Benoit Steiner, Paul A. Tucker, Vijay Vasudevan, Pete Warden, Martin Wieke, Yuan Yu, and Xiaoqiang Zheng. Tensorflow: system for large-scale machine learning. In 12th USENIX Sympo- sium on Operating Systems Design and Implementation, 2016.
Shun-ichi Amari Natural gradient works efficiently in learning. Neural Computation, 10(2), 1998.
Lukas Balles and Philipp Hennig. Dissecting Adam:
The sign, magnitude and variance of stochastic gradients. In Proceedings of the 35th International Conference Machine Learning, 2018.
Lukas Balles, Javier Romero, and Philipp Hennig. Coupling adaptive batch sizes with learning rates. In Proceedingy of the 33rd Conference Uncertainty Artificial Intelligence, 2017.
Atilim Gunes Baydin, Barak Pearlmutter, Alexey Andreyevich Radul, and Jeffrey Mark Siskind. Automatic diffcrentiation in machine learning: A survey, Journal of Machine Learning Research, 18(153), 2018.
Sue Becker and Yann Le Cun. Improving the convergence of back-propagation learning with second order methods. In Proceedings of the 1988 Connectionist Models Summer School, 1989.
Antoine Bordes, Léon Bottou, and Patrick Gallinari. SGD-QN: careful quasi-Newton stochastic gradient de- scent. Mach. Learn. Res., 10, 2009.
Aleksandar Botev, Hippolyt Ritter, David Barber. Practical Newss- optimisation for deep learning. In Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, 2017.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson. Chris Leary. Dougal Maclaurin, and Skye Wanderman-Miile JAX: Composable transformations of Python+NumPy programs, 2018.
Tianqi Chen, Mu Yutian Li, Min Lin, Naiyan Wang, Minjic Wang. Tianjun Xiao, Bing Xu, Chiyuan Zhang, and Zheng Zhang. MXNet: A flexible and efficient machine learning library for heterogeneous distributed systems. In 31st Conference on Neural Information Processing Systems, Workshop on Machine Learning Systems, 2015.
Felix Dangel, Stefan Harmeling, and Philipp Hennig. A modular approach to block-diagonal Hessian approxi- mations for second-order optimization methods. CoRR, abs/1 900.01111.. 2019.

---

Published as a conference paper at ICLR 2020
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in Kronecker-factoree eigenbasis. 2018.
Ian J. Goodfellow. Efficient per-example gradient computations. CoRR, abs/ 2015.
Roger B. Grosse James Martens. Kronecker-factored approximate Fisher matrix for convolution layers.
In Proceedings of the 33rd International Conference on Machine Learning, volume 48 of JMLR Workshop and Conference Proceedings. 2016.
Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Sun. Deep residual learning for image recognition. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9(8), 1997.
Michael Innes. Flux: Elegant machine learning with Julia. Journal of Open Source Software, 3(25), 2018a.
Michael Innes. Don't unroll adjoint: Differentiating SSA- form programs. CoRR, abs/1810.07951, 2018b.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing in- ternal covariate shift. In Proceedings of the 32nd International Conference Machine Leurning, volume 37 of JMLR Workshop and Conference Proceedings, 2015.
Angelos Katharopoulos and François Fleuret. Not all samples are created equal: Deep learning with importance sampling. In Proceedings the 35th International Conference Machine Learning, 2018
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd International Confer- ence on Learning Representations, 2015.
Frederik Kunstner, Lukas Balles, and Philipp Hennig. Limitations of the empirical Fisher approximation. In Advances in Neural Information Processing Systems 32, 2019.
Nicolas Le Roux, Pierre-Antoine Manzagol, and Yoshua Bengio. Topmoumoute online natural gradient algo- rithm. In Advances in Neural Information Processing Systems 20, 2007.
Maren Mahsereci and Philipp Hennig. Probabilistic line searches for stochastic optimization. Journal of Machine Learning Research, 18, 2017.
James Martens, New perspectives on the natural gradient method. CoRR, abss/141..111 193, 2014.
James Martens and Roger B. Grosse. Optimizing neural networks with Kronecker-factored approximate cur- vature. In Proceedings of the 32nd International Conference on Machine Learning, volume 37 of JMLR Workshop and Conference Proceedings, 2015.
James Martens, Jimmy Ba, and Matt Johnson. Kronecker-factored curvature approximations for recurrent neural networks. In 6th International Conference on Learning Representations, 2018.
Fiji Mizutani and Stuart F. Dreyfus. Second-order stagewise hackpropagation for Hessian-mant matrix analyses and investigation negative curvature. Neurul Networks, 21(2-3), 2008
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high perfforrre deep learning library. In Advances in Neural Information Processing Systems 32. 2019.
Frank Schneider, Lukas Balles, and Philipp Hennig. DeepOBS: deep learning optimizer benchmark suite. In 7th International Conference on Learning Representations, 2019.
Nicol N. Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural Com- putation, 14(7), 2002.
Jost Tobias Springenberg. Alexey Dosovitskiy, Thomas Brox, and Martin A. Riedmiller. Striving for simplicity: The all convolutional net. In 3rd International Conference Learning Representations, 2015.
Sciya Tokui, Kenta Oono, Shohei Hido, and Justin Clayton. Chainer: A next-generation open source framework for deep learning. In 29th Conference Neural Information Processing Systems, Workshop on Machine Learnnng Systems, 2015.
Yohei Tsuji. Kazuki Osawa. Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. Performance op- timizations and analysis of distributed deep learning with approximated second-order optimization method. In 48th International Conference on Parallel Processing, Workshop Proceedings, 2019.
10

---

Published as a conference paper at ICLR 2020
BACKPACK: PACKING MORE INTO BACKPROP
SUPPLEMENTARY MATERIAL
Table of Content
SA: BACKPACK extensions
$:.1: First-order quantities
$A.2: Second-order quantities based the generalized Gauss-Newton
$A.3: The exact Hessian diagonal
JB: Additional details benchmarks
SC: Additional details on experiments
3D: BACKPACK cheat sheet
A BACKPACK EXTENSIONS
This section provides more technical details on the additional quantities extracted by BACKPACK.
Notation: Consider an arbitrary module T(l) of a network = .....,, I., parameterized by (()), gis)
(i-1) to its output zn i.e.
transforms the output of its parent layer for sample n, zn
( T (=) n 1. I,
(8)
T2
where N the number of samples. In particular, 2((0) = n. and 2n 21 swx, 0), where f is the transformation of the whole network. The dimension of the hidden layer i's output 2n, is written (8) and (*) is of dimension d( The dimension of the network output, the prediction z (1.) (LL = C. For an image classification task, C' corresponds to the number of classes.
All quantities assumed to be vector-shaped. For image-processing transformations that usually act on lensor-shaped inputs, we can reduce to the vector scenario by vectorizing all quantities; this discussion does not rely on specific flattening scheme. However, for efficient implementation, vectorization should match the layout of the memory of the underlying arrays.
Jacobian:
The Jacobian matrix J.b of an arbitrary vector 18. with respect to another vector E RA is an [A x B] matrix of partial derivatives, [Jabliy è 10)
A. FIRST-ORDER QUANTITIES
The basis for the extraction of additional information about first-order derivatives is given by Eq. 3, which we state again for multiple samples,
oco L(0)
on (0)
N Σ (J ( (V_())
N
n=1
During the backpropagation step of module i, we have access to (()) = 1,...., N. To extract
more quantities involving the gradient, we use additional information about the transformation (5)
within our custom implementation of the Jacobian Joco z and transposed Jacobian (Joro 2n (6)
Individual gradients: The contribution of each sample to the overall gradient, N V gte) e, (0), is computed by application of the transposed Jacobian,
Vomen en N (Jarn (0)
(0))
n= 1, N.
(9)
N
For each parameter 0 the individual gradients of size [N d(0).

---

Published as a conference paper at ICLR 2020
Individual gradient &2 norm: The quantity > V oco (n for n 1,,,, N, could be ex- tracted from the individual gradients (Eq. 9) as
1 v e Cn (0)
|
I (J (V 1 (o) H (Jocn (V (o)
which is an N-dimensional object for each parameter 0()) However, this is not memory efficient as the individual gradients an [N d( tensor. To circumvent this problem, BACKPACK uses the structure of the Jacobian whenever possible.
For specific example, take a linear layer with parameters 0 an [A B] matrix. The layer trans- forms the inputs 2/ (((11) an [N x A] matrix which we will now refer to as A. During the backward pass, receives the gradient of the individual losses with respect its output, { V x Jn as an [N x B) matrix which we will refer to as B. The overall gradient, an [A x B] matrix, can be com- puted A B, and the individual gradients We want to avoid storing that information. To reduce the memory requirement, note that the indi- vidual gradient norm can be written
set of N [Ax B] matrices, {A
[(n,:) ) !==1-
Voln
(m,
and that the summation can be done independently for each matrix, as E, EA(n,, Bm, j1)
A(n, Alm,d))²(E) EE, J13. Therefore, can square each matrix (element-wise) and sum non-bateh dimensions. This yields vectors a, b of N elements, where a/n] = Σ; An, The
individual gradients' lg norm is then given by a o b where is element-wise multiplication.
Second moment:
The gradient second moment (or more specifically, the diagonal of the second moment) is the sum of the squared elements of the individual gradients in mini-batch, i.e.
N
V oowe (0)3
j=
(10)
N
can be used to evaluate the variance of individual elements of the gradient (see below). The second moment is of dimension d(0) the same dimension the layer parameter Ø(t), Similarly to the & norm, it can be computed from individual gradients, but is more efficiently computed implicitly.
Revisiting the example of the linear layer from the individual & norm computation, the second moment of the parameters 0(i,ji j] is given by (A(n,, B (n. il)², which can be directly computed by taking the element-wise square of A and B element-wise, A B², and computing A B².
Variance: Gradient variances over mini-batch (or more precisely, the diagonal of the covariance) can be computed using the second moment and the gradient itself,
N
V Von & (0)) Vo [())
j j= ((0)
(11)
N
The element-wise gradient variance of same dimension as the layer parameter (61), i.e. d(0)
A.2 SECOND-ORDER QUANTITIES BASED ON THE GENERALIZED GAUSS-NEWTON
The computation of quantities that originate from the approximations of the Hessian require an additional backward pass (see Dangel et al. (2019)). Most curvature approximations supported by BACKPACK rely on the generalized Gauss-Newtor (GGN) matrix (Schraudolph, 2002)
G(0) = Σ (Jo 0)) V) 0), (J). 0))
(12)
One interpretation the GGN is that corresponds to the empirical risk Hessian when the model is approximated with its first-order Taylor expansion, i.e. by linearizing the network and ignoring
12

---

Published as a conference paper at ICLR 2020
second-order effects. Hence, the effect of module curvature in the recursive scheme of Eq. 4 can be ignored to obtain the simpler expression
N
G(
= Σ
1
JJcc f) V
Vjl(f(an 0), (a
N
==1
(13)
(Joc
"T=
N
for the exact block diagonal of the full GGN. analogy to G(((1)) we have introduced the [dl x dimeenniio quantity
G( (J ()) ( 0), yn yn) _
that needs be backpropagated. The curvature backpropagation also follows from Eq. 4
G( ((=((-1))= = (J (-=
i=
(14a)
and is initialized with the Hessian of the loss function with respect to the network prediction, i.e.
G(2 V)l(f(En, 0), yn)-
(14b)
Although this scheme is exact, it is computationally infeasible as it requires the backpropagation of N [A(*) x A(1 matrices between module + 1 and i. Even for small N, this is not possible for networks containing large convolutions.
As an example, the first layer of the ALL-CNN-C network outputs 29 x 29 images with 96 channels, which already gives (+) = 80,736, which leads to half Gigabyte per sample. Moreover, storing all the [d()) X ddmmmmmniioo blocks ())))) is not possible. BACKPACK implements different approximation strategies, developed by Martens & Grosse (2015) and Botev al. (2017) that address both of these complexity issues from different perspectives.
Symmetric factorization scheme: One way to improve the memory footprint of the backpropa- gated matrices in the case where the model prediction's 's dimension C' (the number of classes in an image classification task) is small compared to all hidden features (((i) is to propagate a symmetric factorization of the GGN instead. It relies on the observation that if the loss function itself is convex,
even though its composition with the network might not be, its Hessian with respect to the network output can be decomposed as
S(z ((2(")))(()) ((2(0))((()))
(15)
with the [C x di-diennsional matrix factorization of the loss Hessian, S(z! 1.), for sample n. Con- sequently, the GGN in Eq. 12 reduces to an outer product,
N
G(0) = N ==1
Σ (Jaf)
S(z(1)) (laf))(((=)1)))
(Jaf) s(=())))
(16)
The analogue for diagonal blocks follows from Eq. 13 and reads
G( Σ
(J (Jaca 2 TT)())))))))
S(==(())
(17)
where we defined the (M(0) -i-dimensional matrix square root S(z(²) (J (L Sz21
Instead of having layer backpropagate N objects of shape (h()) h00 according to Eq. 14, we instead backpropagate the matrix square root via
S(z((-1) (J_
S( ((-
i= i=,,..,, L.
(18)
starting with Eq. 15. This reduces the backpropagated matrix of layer i (h(*) CJ for each sample.
13

---

Published as a conference paper at ICLR 2020
A.2.1 DIAGONAL CURVATURE APPROXIMATIONS
Diagonal of the GGN (DiagGGN): The factorization trick for the loss Hessian reduces the size of the backpropagated quantities, but does not address the intractable size of the GGN diagonal blocks G( J. In BACKPACK, we can extract diag (G(O())) given the backpropagated quantities S(zn"), ); I,...,N, N, without building up the matrix representation of Eq. 17. In particular, we compute
diag (G((())
N Σ diag ([Ubco [((())))))()( [(0))(())))
(19)
Diagonal of the GGN with MC sampled loss Hessian (DiagGGN-MC); We use the same back- propagation strategy of Eq. 18, replacing the symmetric factorization of Eq. 15 with an approxima- tion by a smaller matrix (z)11) of size |C' x C1 and Ö <C
vil(f(zn, 0), yn) so ((x(")) (5(=(()))
(20)
This further reduces the size of backpropagated curvature quantities. Martens & Grosse (2015) in- troduced such sampling scheme with KFAC based on the connection between the GGN and the Fisher. Most loss functions used in machine learning have probabilistic interpretation as negative log- log-likelihood of probabilistic model. The squarred error of regression is equivalent to a Gaussian noise assumption and the cross-entropy is linked to the categorical distribution. In this case, the loss Hessian with respect to the network output is equal, in expectation, to the outer products of gradients if the outpul of the network is sampled according to a particular distribution, P (2x), defined by the network output f (z). Sampling outputs P, have that
E)-ryy
Vol(s(x 0), 9) Vo T)
Vje(s(x 0). y)
(21)
Sampling one such gradient leads to a rank-1 1 MC approximation of the loss Hessian. With the sub- stitution S ++ S, we compute an MC approximation of the GGN diagonal in BACKPACK
diag (G(O(())
diag (Ieenz (00) (((((")))) (Joco 2 [0)
(22)
A.2.2 KRONECKER-FACTORIE CURVATURE APPROXIMATIONS
A different approach to reduce memory complexity of the GGN blocks G( apart from diagonal curvature approximations, is representing them as Kronecker products (KFAC for linear and con- volution layers by Martens & Grosse (2015); Grosse & Martens (2016) KFLR and KFRA for linear layers by Botev et al. (2017)).
G( = A (1) ② B
(23)
For both linear and convolution layers, the first Kronecker factor A is obtained from the inputs 2n (i-1) to layer i. Instead of repeating the technical details of the aforementioned references, we will focus on how they differ in (i) the backpropagated quantities and (ii) the backpropagation strategy.
As a result, we will be able to extend KFLR and KFRA to convolutional neural networks?.
KFAC and KFLR: KFAC uses an MC-sampled estimate of the loss Hessian with square root fac- torization Š (2)) like in Eq. 20. The backpropagation is equivalent to the computation the GGN diagonal. For the GGN of the weights of linear layer i, the second Kronecker term is given by
(i) KFAC
1 1 N
+ 5(x (3(=
We keep the PYTORCH convention that weights and bias are treated as separate parameters. For the bias terms, we can store the full matrix representation of the GGN. This factor reappears in the Kronecker factoriza- tion of the GGN with respect the weights.
14

---

Published as a conference paper at ICLR 2020
which the same time corresponds the GGN of the layer's bias
In contrast to KFAC, the KFLR approximation backpropagates the exact square root factorization ((z)) (L)), i.e. for the weights of a linear layer* (see Botev et al. (2017) for more details)
N
B BBd KILK
Σ n=1
(())))
(S(=!)))*
N
KFRA: The backpropagation strategy for KFRA eliminates the scaling of the backpropagated cur- vature quantities with the batch size N in Eq. 14. Instead of having layer receive the N exact [M () matrices G(zn"), as an approximation. In particular, the recursion changes to
,....,,, NN only a single averaged object, denoted (()) used
N 1 n=1
= N
( (1))
(0)
i= 1=,..... L,
(24a)
and is initialized with the batch- loss Hessian
N
G (1)
Σ
Van (all/an:0)) 0), ym)
(24b)
N
7m
For a linear layer, KFRA uses (see Botev et al. (2017) for more details)
B KFRA
(
A.3 THE EXACT HESSIAN DIAGONAL
For neural networks consisting only of piecewise linear activation functions, computing the diagonal of the Hessian equivalent to computing the GGN diagonal. This is because for these activations the second term in the Hessian backpropagation recursion (Eq. 4) vanishes.
However, for activation functions with non-vanishing second derivative, these residual terms have to be accounted for in the backpropagation. The Hessian backpropagation for module i reads
V2 oco (() = (Jac) =/ (00
(O)) x(00) + R (0()),
(25a)
2
{()) = (J (1-1) (8) (V (()) --- R (2 (-)),,
(25b)
for n = 1,... N. Those (1( x h (1) -dimensional residual terms are defined as
R (0(:)) = Σ 12 [ (())
(=((--) Σ 2 (--
(11
[ V 2 {(0)
For common parameterized layers, such linear and convolution transformations, (0(1) 0.
the activation function is applied element-wise, (((i) (zn (1-1)) are diagonal matrices.
Storing these quantities becomes very memory-intensive for high-dimensional nonlinear activation
layers. BACKPACK, this complexity reduced by application of the aforementioned matrix square root factorization trick. To do so, express the symmetric factorization of R (z! (2()-)))
(2(() -
(2 ( (z((-1)) p (=
((i)
(N (2( (2(1-11))
(26)
where P. (z
( ), N (zf D) represent the matrix square root of R! (z! -)) projected on its positive and negative eigenspace, respectively.
'In the case of convolutions, one has sum over the spatial indices of single channel of (4) as the bias is added entire channel, Grosse & Martens (2016) for details.
15

---

Published as a conference paper at ICLR 2020
the This composition allows for the extension of the GGN backpropagation: In addition S(zn ). decompositions ( (=1 ), " (z)--1)
for the residual parts also have to be backpropagated according Eq. 18. All diagonals are extracted from the backpropagated matrix square roots (see Eq. 19). All diagonals stemming from decompositions in the negative residual eigenspace have to be weighted by factor of before summation.
In terms of complexity, one backpropagation for R. (2 changes the dimensionality as follows
R ))
(h) ! + (f(1-1) 111111 + (1(1-2) x 11-211
With the square root factorization, one instead obtains
(=((-11) (z) ))
(h(0) 10 + (f-1) 1 - [h(--2) h1(1
(1 x | + [1 [1 [/ h0
Roughly speaking, this scheme is more efficient whenever the hidden dimension of nonlinear activation layer deceeds the largest hidden dimension of the network.
Example: Consider one backpropagation step of module Assume R. (0()) 0, i.e. linear, convolution, or non-parameteriize layer. Then the following computations are performed in the protocol for the diagonal Hessian:
Receive the following quantities from the child module + (for n .,...,,))
p((+1) (2
(((1-) (2 (z())),
(J 20 2/ (i+1)) p((+2) (2 (1+1)).
(J zn
(1+1))
N
(1+2) (2 (111))
(J ( 2n
((+1))
(1+2))T
(J [[ (c- (b (b-2)) p(b-11) (1(-1)(2(1-- (2(1--1))
(J m (1+1))
) (J (1_ 2)} (L--11 1= 211-2))
Extract the module parameter Hessian diagonal, diag (V p (())
For each quantity E extract the diagonal from the square root factorization and sum over the samples, i.c. compute
N diag ([]000210) A. (O) )..)))
Multiply the expression by 1 if A stems from backpropagation of residual's nega-
tive eigenspace's factorization.
Sum all expressions to obtain the block Hessian's 's diagonal diag (V2 (. C(0))
Backpropagate the received quantities to the parent module
For each quantity A, E h apply (J [1(1))))))) ) A.,
Append ((+
(zn
( and N
to +
B ADDITIONAL DETAILS ON BENCHMARKS
KFAC YS. KFLR: As the KFLR of Botev et al. (2017) orders of magnitude more expensive to compute than the KFAC of Martens & Grosse (2015) CIFAR-100, it not included in the
16

---

Published as a conference paper at ICLR 2020
main plot. This is not implementation error; follows from the definition of those methods. To approximate the GGN, G( = [ofnl V ln [Jofu] KFAC uses a rank-1 approximation for each of the inner Hessian V in 8,, , and needs propagate vector through the computation graph for each sample. KFLR uses the complete inner Hessian instead. For CIFAR-100, the network has 100 output nodes-one each class and the inner Hessians are [100 x 100] matrices. KFLR needs to propagate matrix through the computation graph for each sample, which is 100x more expensive as shown Fig.
CIFARI AlI-CNN-C, B=16
300
Time [tns)
200
100
2nd Moment
DiagCGN-MC
Indiv. Grad
KFLR
Grad Batch L2
KFAC
Variance
DiagGGN
Figure 8: KFLR and DiagGGN are more expensive rum on large networks. The gradient takes less than 20 ms to compute, but KFLR and DiagGGN are approximately 100x more expensive.
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 piecewise linear. Be- cause of this, the diagonal of the GGN is equivalent to the diagonal of the Hessian (Martens, 2014). However, for networks that use non piecewise linear activation functions like sigmoids or tanh, com- puting the Hessian diagonal can be much more expensive than the GGN diagonal. To illustrate this point, we modify the smaller network used in our benchmarks to include a single sigmoid activation function before the last classification layer. The results in Fig. 9 show that the computation of the diagonal of the Hessian already an order of magnitude more expensive than for the GGN.
CIFARIO on 3C3D with one sigmoid
Figure 9: Diagonal of the Hessian the GGN. If the network contains single sigmoid activation function, the diagonal of the Hessian is an order of magnitude more computationally intensive than the diagonal of the GGN.
Betch
Tima (mos]
10a
Grad (ref.)
Diag GGN
Diag Hessian
C ADDITIONAL DETAILS ON EXPERIMENTS
C.I
PROTOCOL
The optimizer experiments are performed according to the protocol suggested by DEEPOBS:
Train the neural network with the investigated optimizer and vary its hyperparameters on specified grid. This training performed for single random seed only.
DEEPOBS evaluates metrics during the training procedure. 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 highest final accuracy the validation set.
For better understanding of the optimizer performance with respect randomized rou- tines in the training process, DEEPÔBS reruns the best hyperparameter setting for ten dif- ferent 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 provided by DEEPOBS.
17

---

Published as a conference paper at ICLR 2020
The optimizers built upon BACKPACK's curvature estimates were benchmarked on the DEEPOBS image classification problems summarized in Table
Table Test problems considered from the DEEPOBS library (Schneider et al., 2019).
Codename
Deseription
Dataset
# Parameters
LoGREG
Linear model
MNIST
7,850
2c2p
2 convolutional and dense linear layers
FASHION-MNIST
3,274,634
3C3D
3 convolutional and dense linear layers
CIFAR- 10
895,210
ALL-CNN-C
9 convolutional layers (Springenberg et al.. 2015)
CIFAR- 100
1,387,108
C.2 GRID SEARCH AND BEST HYPERPARAMETER SETTING
Both the learning rate and damping tuned over the grid
a {10- 4 10-3 10 2 100000'1 .1}. X (10 10-3, 100 2,0 1,0',,11000 .1,10]
We use the same batch size (N = 128 for all problems, except N 256 for ALL-CNN-C on CIFAR-100) as the base lines and the optimizers run for the identical number of epochs.
The best hyperparameter settings are summarized in Table 3.
C.3 UPDATE RULE
We use a simple update rule with a constant damping parameter X. Consider the parameters 0 of single module in neural network with fgrregularization of strength 7). Let G(O,) denote the curvature matrix and (e((O,) the gradient at step t. One iteration of the optimizer applies
0 + 0, + (G((,) + (A + n)) [Vo[(0,) r0,1
(27)
The inverse cannot be computed exactly (in reasonable time) for the Kronecker-factoree curvatures KFAC, KFLR, and KFRA. We use the scheme of Martens & Grosse (2015) to approximately invert G(O,) & It replaces the expression (A + O) by diagonal terms added to each Kronecker factor. In summary, this replaces
(A0))+(A-m) if G(O,) is Kronecker-factoree
[A(),) @ B + (A + M) by [A(0.) + T. / nI x B + V nI
(28)
A Mo B(,) for an arbitrary matrix norm ||. We follow Martens & Grosse (2015) and choose the trace norm,
A principled choice for the parameter ne is given by ग.
tr,) dim(B) dim A) Ø tr(B(O,))
Wf
(29)
Table 3: Best hyperparameter settings for optimizers and baselines shown in this work. In the Mo- mentum baselines, the momentum was fixed to 0.9. Parameters for computation the running aver- ages in Adam the default values (31, }22) (0.9, 0.999). The symbols / and X denote whether the hyperparameter setting is an interior point of the grid or not, respectively.
mnist Logreg
Emnist_2e2d
cifar10_3034
cifer100_a coas
Curvature
A
int
X
int
int
int
DiagUGK
10
10
10
10
X
10
10"
DiagGON MC
10
10
x
10
/
10
10 3
KFAC
10 2
10-2
3
10
10
10
KFRA
10-2
10-3
Baseline
Momentum
z 2.07 10
R5 2.07 10
Rs 3.79
10
RS 4.83. 10
Adam
2.98
1.27
2.98
6.95
18

---

Published as a conference paper at ICLR 2020
C.4 ADDITIONAL RESULTS
This section presents the results for MNIST using a logistic regression in Fig. 10 and FASHION- MNIST using the 2C2D network, composed of two convolution and two linear layers, Fig. 11.
MNIST: LOGREG
DiggGGN DiggGNNMMC
Test loes
0.35
0.35
Train loss
0.3
Momentur
0.3
0.25
0.25
Test necuracy
0.94
0.94
Train accuracy
0.93
0.93
0.91
0.91
0
10
20
30
40
50
10
20
30
40
50
Epoch
Epoch
Figure 10: Median performance with shaded quartiles of the best hyperparameter settings chosen by DEEPOBS for logistic regression (7,850 parameters) MNIST. Solid lines show well-tuned baselines momentum SGD and Adam that are provided by DEEPOBS.
FASHION-MNIST: 2C2D
0.4
Train loss
DaggoN
DiggGON-MCC
Test
0.2
0.2
KHAC
KFLR
Adam
Momentam
1
Test uccuracy
Train accuracy
0.95
0.95
0.9
0.9
20
60
100
20
80
Epoch
Epoch
Figure 11: Median performance with shaded quartiles of the best hyperparameter settings chosen by DEEPOBS for the 2c2D network (3,274,634 parameters) on FASHION-MNIST. Solid lines show well-tuned baselines of momentum SGD and Adam that provided by DEEPOBS.
19

---

Published as a conference paper at ICLR 2020
D BACKPACK CHEAT SHEET
Assumptions
Feedforward network
T:
(=
7: o(2) (=
Tile) 001) (=
((=)"),) .v)
2
(())
d Dimension of parameter (()
Empirical risk
((0) [ ((((0,)n))
Shorthands
(0) = ((f(0, xn), )n))))n)
n= ,,...,N: N.
(0) f(0, = 2, ())
1 N
Generalized Gauss-Newton matrix
G(0) Σ (Jafx) (
N ==1
Approximative GGN via MC sampling
N
Σ (Jofa) G = N
9.) (nmppscom) (Jofn)
Table 4: Overview of the features supported in the first release of BACKPACK. The quantities are computed separately for all module parameters, i.e. i 1....,L L.
Feature
Details
Individual gradients
N
V
(0),
n
N [ (())}
Batch variance
Σ
N
oww (0))
j= 1, d()
2nd moment
Tom En (0)
j j= 1, 1...., d()
.().
m=
Indiv. gradient & norm
V (n))
n= I....,N N
N
DiagGGN
diag
(G((()))
DiagGGN-MC
diag
((((1))
Hessian diagonal
diag
(2:0 ((0)
(((()) ~ A @ B KEAC
KFAC
KFLR
G(0()) (6)
X B KELK
KFRA
G(0 ~ (4)
B KFRA
20


Reviews:
[{"review_id": "SJxJeLvmcS", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper adds a very interesting and useful feature to existing autodifferentiation for training neural networks. The second-order information can be backprogated just as the first-order ones, which can be used to accelerate training.\n\nThis idea, although according to the paper, is developed upon existing works, still, strongly attracts as the second-order information is crucial for training and perhaps visualizing the landscape of neural networks. I vote for an acceptance as this brings a significantly important feature to PyTorch, and the author's good experiments results and open-sourced code.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "SJeQXw-85r", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper presents a Pytorch framework for experimenting with first and second order extensions to standard gradient updates via backpropagation.  At the time of writing, the implementation supports feed-forward networks where there is a strict module ordering (by which I mean residual connections are not possible).\n\nThe framework enables researchers to easily experiment with techniques which involve modifying the update according to quantities computed over the batch of gradients (i.e. before they are summed to form the standard SGD update)—these are ‘first-order’ extensions—and it also makes use of the block-diagonal factorisation of the Hessian outlined in Mizutani & Dreyfus as well as Dangel & Hennig to enable the computation of second order quantities via ~ ‘hessian prop’.\n\nI think the paper is a strong accept: the framework has some limitations in the current form (mostly in terms of what architectures are supported), however it still provides a very useful and extensible tool for researchers to efficiently experiment with a variety of more complex optimisation architectures.  This is (as the paper states) a large bottleneck for much optimisation research in deep learning.\n\nIn section 2.3 you state that generalised Gauss-Newton (GGN) is guaranteed positive semi-definite.  It would also be nice to add a sentence as to when (even intuitively) the Fisher information coincides with GGN; (in practice, as the GGN uses a (possibly rank-bounded) sample size of ‘N’, while the Fisher is the expectation under the data generating distribution, one could argue that even when they should be ==, it would only be as N->\\infty).\n\n\n\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "S1eDIl2NqH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This is a good paper. The authors present a software implementation which allows one to extend PyTorch to compute quantities that are irritatingly difficult to compute in PyTorch directly, or in other automatic differentiation frameworks, particularly if efficiency is a concern. Issues of this kind have been discussed at length within the community, particularly on GitHub, and related issues with optimization of automatic differentiation code have motivated other software developments, such as Julia's Zygote package. Having wasted a large amount of my time implementing the WGAN gradient penalty from scratch - which, to implement scalably, requires one to use both forward-mode and reverse-mode automatic differentiation simultaneously - I appreciate what the authors are trying to do here to make research that requires more sophisticated automatic differentiation more accessible. \n\n---\n\nDetailed remarks below.\n\n* The paper's title is not very informative about what the paper is about. The authors should choose a different title - perhaps something like \"BackPACK: user-friendly higher-order automatic differentiation in deep networks\" or something similar but less long.\n\n* The authors focus on KFAC-type methods as their key illustrated use case, but actually software frameworks like this are also helpful for certain GAN losses - WGAN-GP in particular. These losses require one to compute a term that involves the norm of a certain gradient of the output. The gradient of this gradient can be handled efficiently with Hessian-Vector products, which in principle are computable efficiently via automatic differentiation, but in practice a huge pain because of the need to use both forward-mode and reverse-mode automatic differentiation and lack of first-class support from automatic differentiation packages. Provided I've not misunderstood BackPACK and that it would help in making such computations less tedious (and I can't verify this myself, as my own implementation of WGAN-GP was not in PyTorch), I would highly recommend the authors to add an extra paragraph discussing this particular use case, because this would increase the paper's impact on the community by connecting it to another literature which is not mentioned in the paper.\n\n* The entire paper could do with talking about backpropagation less, and automatic differentiation more, because it illustrates that the concerns addressed are not solely limited to deep networks, even if the package does primarily target them.\n\n* P2: these issues are not limited to the Python community, and specialized automatic differentiation software has also been developed for Julia. The authors should cite Innes [1,2] and related papers from that literature.\n\n* Figure 1: from the sample code, I worry about how generic BackPACK is. I think the package authors should be careful not to specialize too much to particular kinds of deep networks, particularly since a much wider variety of models and network architectures are starting to be used with automatic differentiation.\n\n* P2: capital \\Omega notation is confusing, please replace with capital \\Theta.\n\n* P2: L_2 regularization should more technically be \\ell^2 instead.\n\n* P4: please cite Baydin [3] who provides a very nice review of automatic differentiation. It may help explanation to introduce dual numbers, which make forward-mode much easier to understand.\n\n* P6: please write out \"with respect to\" for \"w.r.t.\".\n\n* P7: I really liked this section. The simplicity of implementing the example method using the authors' software framework feels compelling to me. However, genericness is still a concern: by analogy, every deep learning framework can do MNIST easily, but some of them make it much harder to do customized or advanced implementation than others. The latter cases are often the ones that matter to practitioners. It's hard to tell how easy it will be to implement something the authors did not foresee or consider - but this will necessarily be the case in any software paper.\n\n* P8: \"and in part driven\" - missing a comma.\n\n* P8: please spell out \"Table\" in \"Tab. 1\".\n\n[1] M. Innes. Don't Unroll Adjoint: Differentiating SSA-Form Programs. NeurIPS, 2018.\n[2] M. Innes. Flux: Elegant Machine Learning with Julia. Journal of Open Source Software, 2018.\n[3] Baydin, Atilim Gunes and Pearlmutter, Barak A and Radul, Alexey Andreyevich and Siskind, Jeffrey Mark. Automatic differentiation in machine learning: a survey. JMLR, 2017.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "N/A"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper efficiently computes quantities, such as variance estimates of the gradient or various Hessian approximations, jointly with the gradient, and the paper also provides a software package for this. All reviewers agree that this is a very good paper and should be accepted.
Decision: Accept (Talk)