Title: BACKPACK: PACKING MORE INTO BACKPROP

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 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 BACKPACK 1 , 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.

1. INTRODUCTION:

The success of deep learning and the applications it fuels can be traced to the popularization of automatic differentiation frameworks. Packages like TENSORFLOW (Abadi et al., 2016) , CHAINER (Tokui et al., 2015) , MXNET (Chen et 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 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 identification of free parameters, and to push the development of more efficient 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 scientific 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 benefits. 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 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 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 firstorder 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 features2 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 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 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.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 backpropagation modules; we provide an implementation on top of PYTORCH, coined BACKPACK, available at https://f-dangel.github.io/backpack/. The initial release supports efficient computation of individual gradients from a mini-batch, their 2 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 flexibility 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 PYTORCH . . . To illustrate the capabilities of BACKPACK, 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 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 a 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 as the variance of the gradients within the mini-batch or the 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 -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 defined, 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 fit 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, BACKPACK 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 LSTMs (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 ŷn using some parameters θ ∈ Θ. The predictions are evaluated with a loss function : Y × Y → R, for example the cross-entropy, which compares them to the ground truth y n . This leads to the objective function L : Θ → R,

{z (i-1) n } N n=1 {z (i) n } N n=1 {∇ z (i-1) n n} N n=1 {∇ z (i) n n} N n=1 ∇ θ (i) L θ (i) T (i) θ (i) z (i-1) n = z (i) n

L(θ) = 1 N N n=1 (f (θ, x n ), y n ) .

(1)

As a shorthand, we will use n (θ) = (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 { n } N n=1 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 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

where L) ]. The loss function can also be seen as another transformation, appended to the network. Let z

EQUATION

(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

z (0) n T (1) θ (1) (z (0) n ) --------→ z (1) n T (2) θ (2) (z (1) n ) --------→ . . . T (L) θ (L) (z (L-1) n ) --------→ z (L) (z (L) n ,yn) --------→ n (θ) .

To compute the gradient of n with respect to the θ (i) , one can repeatedly apply the chain rule,

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 : ∇ z (i-1) n n (θ) = (J z (i-1) n z (i) n ) (∇ z (i) n n (θ)

). 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 n (θ). It then extracts the gradient with respect to its parameters and inputs, ∇ θ (i) n (θ) and ∇ z (i-1) n 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) n (θ), can be obtained by the recursion

EQUATION

and a similar relation holds for the Hessian with respect to each module's output, ∇ 2

z (i) n 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 flexibility 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. The plot shows computation time, comparing to a traditional gradient computation, on the 3C3D network (See §4) for the CIFAR-10 dataset (Schneider et al., 2019) . 

{z (i-1) n } N n=1 {z (i) n } N n=1 {∇ z (i-1) n n} N n=1 {∇ z (i) n n} N n=1 {∇ θ (i) n n} N n=1 ∇ θ (i) L θ (i) T (i) θ (i) z (i-1) n = z (i) n

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 are not directly accessible. The naïve way to compute N 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 approach batches computations to obtain large efficiency gains, as illustrated by Fig. 3 .

As the quantities necessary to 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 to this information, before it is cleared for memory efficiency, BACKPACK 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 PY-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.

This scheme for individual gradient computation is the basis for all first-order extensions. In this direct form, however, it is expensive in memory: if the model is D-dimensional, storing O(N D) elements is prohibitive for large batches. For the variance, 2 nd moment and 2 norm, BACKPACK 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.3. SECOND-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-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 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. (2019) .

For an objective function that can be written as the composition of a loss function and a model f , such as Eq. 1, the GGN of

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 } N n=1 {z (i) n } N n=1 {∇ z (i-1) n n} N n=1 {∇ z (i) n n} N n=1 {[Jz (i) n fn] Sn} N n=1 {[Jz (i-1) n fn] Sn} N n=1 {[Jθ(i) fn] Sn} N n=1 ∇ θ (i) L θ (i) T (i) θ (i) z (i-1) n = z (i) n

G(θ) = n [J θ fn] SnS n [J θ fn]

alongside the gradient at the ith module, for N samples.

for example using a Kronecker factorization. The approach used by BACKPACK for their computation is a refinement of the Hessian Backpropagation equations of Dangel et al. (2019) . 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 n = ∇ 2 f (f (θ, x n ), y n ), it is sufficient to compute [J θ f n ]

S n and square the result. A network 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 [D×C] matrix, which makes it possible to efficiently compute GGN block diagonals. Also, the computation is very similar to that of a gradient, which computes

[J θ f n ] ∇f n 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. 5 . 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 specific information such as the diagonal. If B is the symmetric factorization for a GGN block, the diagonal can be computed as

[BB ] ii = j [B] 2 ij

, where [•] ij denotes the element in the ith row and jth 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 n = ∇ 2 fn n , yields the KFLR approximation. KFAC uses an MC-approximation by sampling a vector s n such that E sn [s n s n ] = ∇ 2 fn n . KFLR is therefore more precise but more expensive than KFAC, especially for networks with high-dimensional outputs, which is reflected 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 BACKPACK extensions are documented in Appendix A.2.

3. EVALUATION AND BENCHMARKS:

We benchmark the overhead of BACKPACK on the CIFAR-10 and CIFAR-100 datasets, using the 3C3D network3 provided by DEEPOBS (Schneider et al., 2019) and the ALL-CNN-C4 network of Springenberg et al. (2015) . The results are 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 2 norm, 2 nd moment and variance can be extracted efficiently. Regarding second-order 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 two backward passes. This last point is encouraging, as our optimization experiment in Section 4 suggest that this approximation is reasonably accurate. 

4. EXPERIMENTS:

To illustrate the utility of BACKPACK, 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-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(θ t

(i)

), which could be a diagonal or Kronecker factorization of the GGN blocks, and a damping parameter λ to precondition the gradient:

EQUATION

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 KFRA5 (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 DEEPOBS (Schneider et al., 2019) . 6 We use the setup (batch size, number of training epochs) of DEEPOBS' 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 final 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 3C3D 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 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 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 DiagGGN-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.

Published as a conference paper at ICLR 2020 

5. 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 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 BACKPACK, an efficient implementation in PYTORCH of recent conceptual advances and extensions to backpropagation (Tab. 1 lists all features). BACKPACK enriches the syntax of automatic differentiation packages to offer additional observables to optimizers beyond the batch-averaged gradient. Our experiments demonstrate that BACKPACK's implementation 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.

Table 1 : Overview of the features supported in the first release of BACKPACK.

Feature Details:

Individual gradients 

1 N ∇ θ (i) n (θ), n = 1, . . . , N Batch variance 1 N N n=1 [∇ θ (i) n (θ)] 2 j -[∇ θ (i) L(θ)] 2 j 2 nd moment 1 N N n=1 [∇ θ (i) n (θ)] 2 j , j = 1, . . . , d (i) . Indiv. gradient 2 norm 1 N ∇ θ (i) n (θ) 2 2 , n = 1, . . . , N DiagGGN diag G(θ (i) ) DiagGGN-MC diag G(θ (i) ) Hessian diagonal diag ∇ 2 θ (i) L(θ) KFAC G(θ (i) ) ≈ A (i) ⊗ B (i) KFAC KFLR G(θ (i) ) ≈ A (i) ⊗ B (i) KFLR KFRA G(θ (i) ) ≈ A (i) ⊗ B (i)

G(θ) = n [J θ f n ] ∇ 2 fn n [J θ f n ]

, KFAC uses a rank-1 approximation for each of the inner Hessian ∇ 2 fn n = s n s n , and needs to propagate a 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 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 . 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. Because 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, computing 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 is already an order of magnitude more expensive than for the GGN. If the network contains a single sigmoid activation function, the diagonal of the Hessian is an order of magnitude more computationally intensive than the diagonal of the GGN.

C ADDITIONAL DETAILS ON EXPERIMENTS C.1 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 a specified grid. This training is performed for a 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 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 DEEPOBS.

The optimizers built upon BACKPACK's curvature estimates were benchmarked on the DEEPOBS image classification problems summarized in Table 2 . Both the learning rate α and damping λ are tuned over the grid α ∈ 10 -4 , 10 -3 , 10 -2 , 10 -1 , 1 , λ ∈ 10 -4 , 10 -3 , 10 -2 , 10 -1 , 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 λ. Consider the parameters θ of a single module in a neural network with 2 -regularization of strength η. Let G(θ t ) denote the curvature matrix and ∇ θ L(θ t ) the gradient at step t. One iteration of the optimizer applies

EQUATION

The inverse cannot be computed exactly (in reasonable time) for the Kronecker-factored curvatures KFAC, KFLR, and KFRA. We use the scheme of Martens & Grosse (2015) .

A principled choice for the parameter π t is given by π t = A(θt)⊗I B I A ⊗B(θt) for an arbitrary matrix norm • . We follow Martens & Grosse (2015) and choose the trace norm,

EQUATION

)

Table 3 : Best hyperparameter settings for optimizers and baselines shown in this work. In the Momentum baselines, the momentum was fixed to 0.9. Parameters for computation of the running averages in Adam use the default values (β 1 , β 2 ) = (0.9, 0.999). The symbols and denote whether the hyperparameter setting is an interior point of the grid or not, respectively. 

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, in Fig. 11 . 

:

See, e.g., the Github issues github.com/pytorch/pytorch/issues/1407, 7786, 8897 and forum discussions discuss.pytorch.org/t/1433,8405, 15270, 17204, 19350, 24955.

3C3D is a sequence of 3 convolutions and 3 dense linear layers with 895,210 parameters.

ALL-CNN-C is a sequence of 9 convolutions with 1,387,108 parameters.

KFRA was not originally designed for convolutions; we extend it using the Kronecker factorization ofGrosse & Martens (2016). While it can be computed for small networks on MNIST, which we report in Appendix C.4, the approximate backward pass of KFRA does not seem to scale to large convolution layers.

6 https://deepobs.github.io/. We cannot run BACKPACK on all test problems in this benchmark due to the limitations outlined in Section 2. Despite this limitation, we still run on models spanning a representative range of image classification problems.

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 factorization of the GGN with respect to the weights.

Figures and Tables:

Figure 2: Schematic representation of the standard backpropagation pass for module i with N samples.

Figure3: Computing individual gradients in a batch using a for-loop (i.e. one individual forward and backward pass per sample) or 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 dataset(Schneider et al., 2019).

Figure 4: Schematic representation of the individual gradients' extraction in addition to the standard backward pass at the ith module for N samples.

Figure 5: Schematic of the additional backward pass to compute a symmetric factorization of the GGN,

Figure 7: Median performance with shaded quartiles of the DEEPOBS benchmark for (a) 3C3D network (895,210 parameters) on CIFAR-10 and (b) ALL-CNN-C network (1,387,108 parameters) on CIFAR-100. Solid lines show baselines of momentum SGD and Adam provided by DEEPOBS.

Figure 8: KFLR and DiagGGN are more expensive to run on large networks. The gradient takes less than 20 ms to compute, but KFLR and DiagGGN are approximately 100× more expensive.

Figure 9: Diagonal of the Hessian vs. the GGN.If the network contains a single sigmoid activation function, the diagonal of the Hessian is an order of magnitude more computationally intensive than the diagonal of the GGN.

Figure 10: Median performance with shaded quartiles of the best hyperparameter settings chosen by DEEPOBS for logistic regression (7,850 parameters) on MNIST. Solid lines show well-tuned baselines of momentum SGD and Adam that are provided by DEEPOBS.

Overhead benchmark for computing the gradient and first-or second-order extensions on real networks, compared to just the gradient. Most quantities add little overhead. KFLR and DiagGGN propagate 100× 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 diagonal, in Appendix B.

Table content: <table><tr><td>50</td><td>Batch = 32</td><td colspan="2">3C3D on CIFAR10 48 64</td><td/><td>100</td><td>All-CNN-C on CIFAR100 Batch = 32 48 64</td></tr><tr><td>Time [ms]</td><td/><td/><td/><td/><td>50</td><td/></tr><tr><td>0</td><td colspan="2">Grad (ref.) 2nd Moment Batch L2 DiagGGN-MC KFAC</td><td>Indiv. Grad KFLR</td><td>DiagGGN</td><td>0</td><td>Grad (ref.) 2nd Moment Batch L2 DiagGGN-MC KFAC</td><td>Indiv. Grad</td></tr><tr><td colspan="2">Figure 6:</td><td/><td/><td/><td/><td/></tr></table>

KFRAPerspectives for Science", EXC 2064/1, project number 390727645; the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039A); and funds from the Ministry of Science, Research and Arts of the State of Baden-Württemberg. F. D. is grateful to the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for support. main plot. This is not an implementation error; it follows from the definition of those methods. To approximate the GGN,

Table content: <table/>

Test problems considered from the DEEPOBS library(Schneider et al., 2019).

Table content: <table><tr><td>Codename</td><td>Description</td><td>Dataset</td><td># Parameters</td></tr><tr><td>LOGREG</td><td>Linear model</td><td>MNIST</td><td>7,850</td></tr><tr><td>2C2D</td><td>2 convolutional and 2 dense linear layers</td><td>FASHION-MNIST</td><td>3,274,634</td></tr><tr><td>3C3D</td><td>3 convolutional and 3 dense linear layers</td><td>CIFAR-10</td><td>895,210</td></tr><tr><td colspan="3">ALL-CNN-C 9 convolutional layers (Springenberg et al., 2015) CIFAR-100</td><td>1,387,108</td></tr><tr><td colspan="2">C.2 GRID SEARCH AND BEST HYPERPARAMETER SETTING</td><td/><td/></tr></table>

Peer review (iclr2020):
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. This 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.