HesScale: Scalable Computation of Hessian Diagonals
Authors: Mohamed Elsayed
A. Rupam Mahmood
Venue: Submitted to ICLR 2023 (2023)

Abstract: Second-order optimization uses curvature information about the objective function, which can help in faster convergence. However, such methods typically require expensive computation of the Hessian matrix, preventing their usage in a scalable way. The absence of efficient ways of computation drove the most widely used methods to focus on first-order approximations that do not capture the curvature information. In this paper, we develop \textit{HesScale}, a scalable approach to approximating the diagonal of the Hessian matrix, to incorporate second-order information in a computationally efficient manner. We show that HesScale has the same computational complexity as backpropagation. Our results on supervised classification show that HesScale achieves high approximation accuracy, allowing for scalable and efficient second-order optimization.

Under review as a conference paper at ICLR 2023
HESSCALE: SCALABLE COMPUTATION OF HESSIAN DIAGONALS
Anonymous authors Paper under double-blind review
ABSTRACT
Second-order optimization curvature information about the objective fune- tion, which can help in faster convergence. However, such methods typically require expensive computation of the Hessian matrix, preventing their usage in a scalable way. The absence of efficient ways of computation drove the most widely used methods to focus on first-order approximations that do not capture the cur- vature information. In this paper, we develop HesScale, a scalable approach to approximating the diagonal the Hessian matrix, to incorporate second-order in- formation in computationally efficient manner. We show that HesScale has the same computational complexity as backpropagation. Our results on supervised classification show that HesScale achieves high approximation accuracy, allowing for scalable and efficient second-order optimization.'
INTRODUCTION
First-order optimization offers cheap and efficient way of performing local progress in optimization problems by using gradient information. However, their performance suffers from instability or slow progress when used in ill-conditioned landscapes. Such problem is present because first- order methods do not capture curvature information which causes two interrelated issues. First, the updates in first-order have incorrect units (Duchi et al. 2011). which creates a scaling issue. Second, first-order methods lack parameterization invariance (Martens 2020) in contrast to second- order methods such as natural gradient (Amari 1998) or Newton-Raphson methods. Therefore, some first-order normalization methods were developed to address the invariance problem (Ba et al. 2016, loffe & Szegedy 2015, Salimans & Kingma 2016). On the other hand, some recent adaptive step- size methods try to alleviate the scaling issue by using gradient information for first-order curvature approximation (Luo et al. 2019, Duchi al. 2011, Zeiler 2012, Reddi et al. 2018, Kingma & Ba 2015, Tran & Phong 2019, Tieleman et al. 2012). Specifically, such methods use the empirical Fisher diagonals heuristic by maintaining moving average of the squared gradients to approximate the diagonal of the Fisher information matrix. Despite the huge adoption of such methods due to their scalability, they use inaccurate approximations. Kunstner al. (2019) showed that the empirical Fisher does not generally capture curvature information and might have undesirable effects. They argued that the empirical Fisher approximates the Fisher or the Hessian matrices only under strong assumptions that are unlikely to be met in practice. Moreover, Wilson et al. (2017) presented counterexample where the adaptive step-size methods are unable to reduce the error compared to non-adaptive counterparts such as stochastic gradient descent.
Although second-order optimization can speed up the training process by using the geometry of the landscape, its adoption is minimal compared to first-order methods. The exact natural gradient or Newton-Raphan Raphson methods require the computation, storage, and inversion of the Fisher information or the Hessian matrices, making them computationally prohibitive in large- tasks. Accord- ingly, many popular second-order methods attempt to approximate less expensively. For example, type of truncated-Newton method called Hessian-free methods (Martens 2010) exploits the fact that the Hessian-vector product is cheap (Bekas et al. 2007) and uses the iterative conjugate gradient method perform an update. However, such methods might require many iterations per update or some tricks to achieve stability, adding computational overhead (Martens & Sutskever 201 1).
"Code will he available.

---

Under review as a conference paper at ICLR 2023
Some variations try approximate only the diagonals of the Hessian matrix using stochastic timation with matrix-free computations (Chapelle & Erhan 2011, Martens et al. 2012, Yao et al. 2021). Other methods impose probabilistic modeling assumptions and estimate block diagonal Fisher information matrix (Martens & Grosse 2015, Botev et al. 2017). Such methods are invariant to reparametrization but are computationally expensive since they need to perform matrix inversion for each block.
Deterministic diagonal approximations the Hessian (LeCun et al. 1990, Becker & Lecun 1989) provide some curvature information and are efficient to compute. Specifically, they can be imple- mented to be as efficient as first-order methods. We view this category of approximation methods as scalable second-orde methods. In neural networks, curvature backpropagation (Becker & Lecun 1989) can be used to backpropagate the curvature vector. We distinguish this efficient method from other expensive methods (e.g.. Mizutani & Dreyfus 2008, Botev et al. 2017) that backpropagate the full Hessian matrices. Although these diagonal methods show a promising direction for scalable second-order optimization, the approximation quality is sometimes poor with objectives such as cross-entropy (Martens et al. 2012). A scalable second-order method with high quality approxima- tion is still needed.
In this paper, we present HesScale, a high-quality approximation method for the Hessian diagonals. Our method is also scalable and has little memory requirement with linear computational complexity while maintaining high approximation accuracy.
BACKGROUND
In this section, we describe the Hessian matrix for neural networks and some existing methods for estimating it. Generally, Hessian matrices can be computed for any scalar-valued function that are twice differentiable. If f R" + R is such a function, then for its argument &b E R", the Hessian matrix H E R=x= of f with respect to v is given by Ha, 8 (()))aae.aw Here, the ith element of a vector v is denoted by n; and the element at the ith row and jth column of a matrix M is denoted by Ma.-- When the need for computing the Hessian matrix arises for optimization in deep learning, the function f is typically the objective function, and the vector &b is commonly the weight vector of a neural network. Computing and storing n n matrix, where is the number of weights in a neural network, is expensive. Therefore, many methods exist for approximating the Hessian matrix or parts of with less memory footprint, computational requirement, or both. A common technique is to utilize the structure of the function to reduce the computations needed. For example, assuming that connections from a certain layer do not affect other layers neural network allows one to approximate a block diagonal Hessian. The computation further simplifies when we have piece-wise linear activation functions (e.g., ReLU), which result in a Generalized Gauss-Newton (GGN) (Schraudolph 2002) approximation that is equivalent to the block diagonal Hessian matrix with linear activation functions. The GGN matrix is more favored in second-order optimization since it is positive semi-definite. However, computing block diagonal matrix is still demanding.
Many approximation methods were developed to reduce the storage and computation requirements of the GGN matrix. For example, under probabilistic modeling assumptions, the Kronecker-factored Approximate Curvature (KFAC) method (Martens & Grosse 2015) writes the GGN matrix G as a Kronecker product of two matrices of smaller sizes as: G A & B, where A E/hh J. B = Elgg J. h is the activation output vector, and g is the gradient of the loss with respect to the activation input vector. The A and B matrices can be estimated by Monte Carlo sampling and an exponential moving average. KFAC is more efficient when used in optimization since it requires inverting only the small matrices using the Kronecker-product property (A @ B)- A B ! However, KFAC is still expensive due to the storage of the block diagonal matrices and computation of Kronecker product, which prevent from being used as a scalable method.
Computing the Hessian diagonals can provide some curvature information with relatively less com- putation. However, it has been shown that the exact computation for diagonals of the Hessian typ- ically has quadratic complexity with the unlikely existence of algorithms that can compute the ex- act diagonals with less than quadratic complexity (Martens et al. 2012). Some stochastic methods provide way to compute unbiased estimates of the exact Hessian diagonals. For example, the AdaHessian (Yao et al. 2021) algorithm uses the Hutchinson's 's estimator diag H) = E[z o (Hz)), where z is multivariate random variable with Rademacher distribution and the expectation can
2

---

Under review as a conference paper at ICLR 2023
be estimated using Monte Carlo sampling with exponential moving average. Similarly, the GGN- MC method (Dangel et al. 2020) uses the relationship between the Fisher information matrix and the Hessian matrix under probabilistic modeling assumptions to have MC approximation of the diagonal of the GGN matrix. Although these stochastic approximation methods are scalable due to linear or (n) computational and memory complexity, they suffer from low approximation quality, improving which requires many sampling and factors of additional computations.
3 THE PROPOSED HESSCALE METHOD
In this section, we present our method for approximating the diagonal of the Hessian at each layer in feed-forward networks, where a backpropagation rule is used utilize the Hessian of previous layers. We present the derivation of the backpropagation rule for fully connected and convolutional neural networks in supervised learning. Similar derivation for fully connected networks with mean squared error is presented before (LeCun et al. 1990, Becker & Lecun 1989). However, we use the exact diagonals of the Hessian matrix the last layer with some non-linear and non-clement-wis output activations such as softmax and show that it can still be computed in linear computational complexity. We show the derivation for Hessian diagonals for fully connected networks in the following and provide the derivation for the convolutional neural networks in Appendix B.
We use the supervised classification setting where there collection of data examples. These data examples are generated from some target function f mapping the input to the output y, where the k-th onput- pair is (at, 3k). In this task, the learner is required to predict the output class y E {1,2, m) given the input vector E Rd by estimating the target function f*. The performance is measured with the cross-entropy loss, [(p. q) = Ei= Ps log (r. where E IRm the vector of the target one-hot encoded class and q IR" is the predicted output. The learner is required to reduce the cross-entropy by matching the target class.
Consider neural network with L layers that outputs the predicted output q. The neural network is parametrized by the set of weights (Wi,.., W..). J. where W, is the weight matrix at the 1-th layer, and its element the ith row and the jth column is denoted by W1,.j. During learning. the parameters of the neural network are changed to reduce the loss. At each layer 1, we get the activation output h; by applying the activation function a to the activation input ay: h; (a:). We simplify notations by defining hg = æ. The activation output h is then multiplied by the weight matrix W1+1 of layer 1 assume here that the activation function is element-wise activation for all layers except for the final layer L, where becomes the softmax function. The backpropagation equations for the described network given follows Rumelhart et al. (1986):
produce the activation input: 01+1,1 Σ)au Wi We
lat+l
Oc
OC 0ai+1 Ohu = (arss) 001+1k
OC
Σ
W4+1,,
(1)
day,
0a1t1,k Ohi, Sar, day OC
k=1
ac
ac
h4-15-
(2)
OWia. Das, DWiss
dan
In the following, we write the equations for the exact Hessian diagonals with respect to weights 8² e/oW?,,,, which requires the calculation of 8 2/a/a, first:
82C
larn
JL
d' d'(a)))
WILLRA
Da?
Dols
Bo141
kml
lassi lartal
82C
1art H.rl..Waai W4ILK + o"(ar,)
0041
OL
=
W411k
J041,,0a411
Dap,
Day+1,k
1=1
la
laryal
JC
layil
aC
= =
Jast Wit1p,:W44... oa((a))))
W41Lk
Daj+1k
k=
p=1
8C
④
0L
Jui, h1-1.
8LC 1 Oa
hu 1n-1
11
(3)
OW?,
aWiis
Jai,
Jai,
0ay,5
OWi,..
3

---

Under review as a conference paper at ICLR 2023
Since, the calculation of 2 depends on the off-diagonal terms, the computation complexity becomes quadratic. Following Becker and Lecun (1989), approximate the Hessian diagonals by ignoring the off-diagonal terms, which leads to a backpropagation rule with linear computational complexity for our estimates
8 and 3cC: twt.
Buf,
lauil
Jalill
SL d' (ar.))
82C
OC
Σ
WH + to taax))
Σ
W411K3
(4)
007,
dai+1,
dai-1,k
UL h h2
(5)
OW? [.,.j
Doj
However, for the last layer, we use the exact Hessian diagonals 8 2²c since can be Duj,
du
computed in ((n) for the softmax activation function and the cross-entropy loss. More precisely, the exact Hessian diagonals for cross-entropy loss with softmax is simply q- q q. where q is the predicted probability vector and о denotes element wise multiplication. We found empirically that this small change makes large difference in the approximation quality, shown in Fig. 1a. Hence, unlike Becker and Lecun (1989) who use a Hessian diagonal approximation the last layer by Eq. we use the exact values directly to achieve more approximation accuracy. We call this method for Hessian diagonal approximation HesScale and provide its pseudocode for supervised classification in Algorithm HesScale not specific to eross-entropy loss as the exact Hessian diagonals can
Algorithm HesScale: Computing Hessian diagonals of neural network layer classification
Require: Neural network / and a layer number! |
Require: First and second order information
0c a0t+1,
and 8 00141a5
unless
Require: Input-output pair (x: y)
Set loss function L to cross-entropy loss
Compute preference vector a, < f()) and target one-hot-encoded vector +- onehot (y)
Compute the predicted probability vector +- (az) using softmax function Compute the error C(p, q) if L then Compute Dag aC using Eq. 2 JWL
Computing Hessian diagonals exactly for the last layer 0C consists of elements dc
Oc
Bag
Compute
8c consists of elements
8C
OW
JWL.5.
32 da7
PC consists of elements
C
9 9
day
Du7.,
Compute 026 aw? using Eq. 5
Ow?
0P
consists of elements
aW2.s.
02
else if * T. then
Compute SC da: and ac 8c/ow, using Eq. and Eq. 2
Compute and dc using Eq. 4 and Eq. 5
πC dag
ow?
end
return
8C C oc
and
0c
FW OW? Baj day
be calculated in O(n) for some other widely used loss functions as well. We show this property for negative log-likelihood function with Gaussian and softmax distributions in Appendix A. The computations can be reduced further using a linear approximation for the activation functions (by dropping the second term in Eq. 4), which corresponds to approximation of the GGN matrix. We call this variation of our method HesScaleGN.
Based on HesScale, we make stable optimizer, which we call AdaleesScale, given in Algorithm We use the same style introduced in Adam (Kingma & Ba 2015), using the squared diagonal ap- proximation instead of the squared gradients update the moving average. Moreover, another optimizer based on HesScaleGN, which we call AdaHesScaleGNN We refer the reader to the convergence proof for methods with Hessian diagonals, which was presented by Yao al. (2021).
introduce
4

---

Under review as a conference paper at ICLR 2023
Algorithm 2 AdaHesScale for optimization
Require: Neural network f with weights { .... W and dataset D
Require: Small number +- 10-8
Require: Exponential decay rates B1, 82 (0,1)
Require: step size a
Require: Initialize {W1 1, 1.... W
Initialize time step & +-+-0 0. for 2 {L, L
1, ..... 1) do
Set exponential moving averages at time step 0 to zero D Same size as Wi
M.- + 0; Vi 7 00
end for
for (æ,, v) in D do
t +-t+1
TL+1 f $L+1 - 0
and st stand for
oc daj
and 0 dag
respectively
for in {L, L 1,,,-,,. 1,..,,,, 1
Fi, St T1, f- HesScale(f.ess 2, 8411).
Check Algorithm
Mj- M, + F
Fi stands for ac
2W.
Vi < +1-111111 +
Si stands for 22c
ow?
M, + M1/(1-B) BB)
> Bias-corrected estimate for F
Vi t V1/(1 - (3)
> Bias-corrected estimate for S
W. + Wi aMM 2 (V +c)
> Z is element wise division
A is clement- square root of A
end for
end for
4 APPROXIMATION QUALITY & SCALABILITY OF HESSCALE
In this section, we evaluate HesScale for its approximation quality and computational cost and com- pare it with other methods. These measures constitute the criteria we look for in scalable and ef- ficient methods. For our experiments, we implemented HesScale using the BackPack framework (Dangel et al. 2020), which allows easy implementation of backpropagation of statisties other than the gradient.
We start by studying the approximation quality of Hessian diagonals compared to the true values. To measure the approximation quality of the Hessian diagonals for different methods, we use the L distance between the exact Hessian diagonals and their approximations. Our task here is supervised classification, and data examples are randomly generated. We used a network of three hidden layers with lanh activations, each containing 16 units. The network weights and biases are initialized randomly. The network has six inputs and ten outputs. For each example pair, compute the exact Hessian diagonals for each layer and their approximations from each method. All layers' errors summed and averaged over 1000 data examples for each method. In this experiment, 40 different initializations for the network weights, shown as colored dots in Fig. 1a. Fach point represents the summed error over network layers, averaged over 1000 examples for each different initialization. In this figure, we show the average error incurred by each method normalized by the average error incurred by HesScale. Any approximation that incurs an averaged error above has a worse approximation than HesScale, and any approximation with an error less than 1 has better approximation than HesScale. Moreover, we show the layer-wise error for each method in Fig. 1b.
used
Different Hessian diagonal approximations are considered for comparison with HesScale. We in- cluded several deterministic and stochastic approximations for the Hessian diagonals. We also in- clude the approximation of the Fisher Information Matrix done by squaring the gradients and de- noted by g. which is highly adopted by many first-order methods (e.g., Kingma and Ba, 2015). We compare HesScale with three stochastic approximation methods: AdaHessian (Yao et al. 2021), Kronecker-Fe fontecee--acctred approximate curvature (KFAC) (Martens & Grosse 2015), and the Monte-Care (MC) estimate of the GGN matrix (GGN-MC) (Dangel et al. 2020). We also compare HesScale with two deterministic approximation methods: the diagonals of the exact GGN matrix (Schrau- dolph 2002) (diag(G)) and the diagonal approximation by Becker and Lecun (1989) (BL89). In
5

---

Under review as a conference paper at ICLR 2023
:000
8:0 35
10a
Normalizen 1 error
AddHessian MC1
25-
E. Error
GEN-NCI APAC-00C3
20-
107
15-
KFRC- Adaleesten44C5
Acoeessam
Acwl Aeelee
WPMC MC1
KTAC MESS
10-
diss
MCSO
(acll))
10
(agie
BLAY
Methods
Layer Number (b) Layer-wise L error
(a) Normalized L error with respect HesScale
Figure/ 1: The averaged error for each method is normalized by the averaged error incurred by HesScale. We show 40 initialization points with the same colors across all methods. The norm of the vector of Hessian diagonals [diag( (H)l is shown as a reference.
KFAC, we extract the diagonals from the block diagonal matrix and show the approximation error averaged over 1 MC sample (KFAC-MCI) and over 50 MC samples (KFAC-MC50), both per each data example. Since AdaHessian and GGN-MC are already diagonal approximations, we use them directly and show the error with 1 MC sample (GGN-MCI & AdaHessian and with 50 MC samples (GGN-MCS) & AdaHessian-MCS0 We refer the reader to Appendix C for an additional experiment with MNIST data points.
HesScale provides a better approximation than the other deterministic and stochastic methods. For stochastic methods,
many MC samples improve their approximation. However, their approximation quality is still poor. Methods approximating the GGN diagonals do not capture the complete Hessian information since the GGN and Hessian matrices are different when the activation functions are not piece-wii wise linear. Although these methods approximate the GGN diagonals, their approximation is significantly better than the AdaHessian approximation. And among the methods for approximating the GGN diagonals, HesScaleGiN performs the best and is close to the exact GGN diagonals. This experiment clearly shows that HesScale achieves the best approximation quality compared to other stochastic and deterministic approximation methods.
Next, we perform another experiment to evaluate the computational cost of our optimizers. Our Hessian approximation methods and corresponding optimizers have linear computational complex- ity, which can be seen from Eq. 4 and Eq. 5. However, computing second-order information in optimizers still incurs extra computations compared first-order optimizers, which may impact how the total computations scale with the number of parameters. Hence, we compare the computa- tional cost of optimizers with others for various numbers of parameters. More specifically, measure the update time of each optimizer, which is the time needed to backpropagate first-order and second-order information and update the parameters.
We designed two experiments to study the computational cost of first-order and second-order opti- mizers. In the first experiment, we used a neural network with single hidden layer. The network has 64 inputs and 512 hidden units with tanh activations. We study the increase in computational time when increasing the number of outputs exponentially, which roughly doubles the number of parameters. The set of values we use for the number of outputs is {2 25 26, 27, 2a, 29). The re- sults of this experiment are shown in Fig. 2a. In the second experiment, we used neural network with multi-layers, each containing 512 hidden units with tanh activations. The network has 64 in- puts and 100 outputs. We study the increase in computational time when increasing the number of layers exponentially, which also roughly doubles the number of parameters. The set of values we use for the number of layers is {1,2, 4, 8, 16, 32, 64, 128). The results are shown in Fig. 2b. The points in Fig. 2a and Fig. 2b are averaged over 30 updates. The standard errors of the means of these points are smaller than the width each line. On average, we notice that the cost of AdaHessian, AdaHesScale, and A are three, two, and 1.25 times the cost of Adam, respectively.

---

Under review as a conference paper at ICLR 2023
It is clear that our methods among the most computationally efficient approximation method for Hessian diagonals.
SGD
SGD
Adaltssia AJathscca ETCC
Time in seconds
50
AZalioa GGVVC ASahiiical
in
KTAC
18-1
Number parameters (2*)
Number parameters (2")
(a) Increasing number of outputs in neural network
(b) Increasing number of layers in neural network
Figure 2: The average computation time for each step of an update is shown for different optimizers. The computed update time is the time needed by each optimizer to backpropagate gradients or second-order information and to update the parameters of the network. GGN overlaps with H in (a).
5 EMPIRICAL PERFORMANCE OF HESSCALE IN OPTIMIZATION
In this section, we compare the performance of our optimizers-AdaHeeseee AdallesScaleGN-wiit three second-order optimizers: BL89, GGNMC, and AdaHessian. We also include comparisons to two first-order methods: Adam and SGD. We exclude KFAC and the exact diagonals of the GGN matrix from our comparisons due to their prohibitive computations.
and
Our optimizers are evaluated in the supervised classification problem with series of experiments using different architectures and three datasets: MNIST, CIFAR-10, and CIFAR-100, Instead of attempting to achieve tfate-of-the-ar performance with specialized techniques and architectures, follow the DeepOBS benchmarking work (Schneider et al. 2019) and compare the optimizers in their generic and pristine form using relatively simple networks. It allows to perform a more fair comparison without extensively utilizing specialized knowledge for a particular task. In the first ex- periment, we use the MNIST-MLP task from DeepOBS. The images are flattened and used as inputs to a network of three fully connected layers (1000, 500, and 100 units) with tanh activations. We train each method for 100 epochs with a batch size of 128. We show the training plots in Fig. with their corresponding sensitivity plots in Appendix D. Fig. 9a. In the second experiment, we use the CIFAR1O-333D task from the DeepOBS benchmarking tasks. The network consists of three convo- lutional layers with tanh activations, each followed by max pooling. After that, two fully connected layers (512 and 256 units) with tanh activations are used. We train each method for 100 epochs with batch size of 128. We show the training plots in Fig. 7b with their corresponding sensitivity plots in Fig. 9b. In the third experiment, we use the CIFAR100-3C-23 task from DeepOBS, The network is the same the method for 200 epochs with a batch size of 128. We show the training plots in Fig. 8b with their corresponding sensitivity plots in Fig. 10b. In the fourth experiment, we use the RIFARI00-ALL CNN task from DeepOBS with the ALL-CNN-C network, which consists of 9 convolutional layers (Springenberg et al. 2015) with ELU activations. We use tanh and ELU instead of ReLU, which is used in DeepOBS, to differentiate between the performance of AdaHesScale and AdaHesScaleGN. We show the training plots in Fig. 8a with their corresponding sensitivity plots in Fig. 10a.
used in the second task except for the activations are ELU. We train each
In the MNIST-MLP and CIFAR-10-3C3D experiments, we performed a hyperparameter search for each method to determine the best set of 31, f2, and a. The range of 83 is (0.99, 0.999, 0.9999) and the range of 31 is (0.0, 0.9}. The range of step size is selected for each method to create a convex curve. Our criterion was to find the best hyperparameter configuration for each method the search space that minimizes the area under the validation loss curve. The performance of each method was averaged over 30 independent runs. Each independent run had the same initial representation for the algorithms used in an experiment. Using each method's best hyperparameter configuration the validation set, we show the performance of each method against the time in seconds needed to complete the required number of epochs, which better depicts the computational efficiency of the
7

---

Under review as a conference paper at ICLR 2023
methods. Fig. 3a and Fig. 3b show these results MNIST-MLP and CIFAR-10 tasks. Moreover, we show the sensitivity each method to the step size in Fig. 5a and Fig. 5b.
Test Accuracy
0,9850
0.985
GONAC
0.900
0.955
0.9775
0.9750
0.000
500
1005
1500
2000
2500
20
40
150
Time seconds
Epochs
(a) MNIST-MLP
Test Accuracy
GONMC
0.75
0.75
1000
2000
3000
4000
SC00
6202
7002
80
100
Time seconds
Epochs
(b) CIFAR-I0 3C3D
Figure 3: MNIST-MLP and CIFAR-I0 3C3D classification tasks. Each method is trained for 100 epochs. We show the time taken by each algorithm in seconds (left) and we show the learning curves the number of epochs (right). The performance of each method is averaged over 30 independent runs. The shaded area represents the standard error.
Test Accuracy
GGNAC
AdallesSoate
5GD
0.40
Adan!
Addtissiaa
0.40
2000
4000
6000
8000
10000
12000
24000
25
10
75
100
150
175
200
Time seconds
Epochs
CIFAR- 100 3C3D
Test Accuracy
0.35
GONMC
0.30
0.50
Adallasstar
Time seconds
Epochs
(b) CIFAR-100 AlI-CNN-C
Figure 4: CIFAR- 100 3C3D and CIFAR- - ALL-CNN classification tasks. Each method from the first task is trained for 200 epochs and each method from the second task is trained for 350 epochs. We show the time taken by each algorithm in seconds (left) and we show the learning curves in the number of epochs (right). The performance of each method is averaged over 30 independent runs. The shaded area represents the standard error.
In the CIFAR- SEFAR-100-ALL-C -0--ALL-CNN and CIFAR- 100-331 experiments, we used the set of 31 and Bg that achieved the best robustness in the previous two tasks, which were 0.9 and 0.999 respectively. We did a hyperparameter search for each method to determine the best step size using the specified 31 and (3.. The rest the experimental details are same as the first two experiments. Using each method's best hyperparameter configuration the validation set, we show the performance of each method against the time seconds needed to complete the required number of epochs. Fig, 4a and Fig. 4b show these results the results Appendix
and CIFAR-100-300 tasks. We summarize
Our results show that all optimizers except for BL89 performed well the MNIST-MLP task. However, in CIFAR- 10, CIFAR-IO1 100 3c3d, and CIFAR- --0AR--00 ALL-CNN, we notice that AdaHessian performed worse than all methods except BL.99 This result is aligned with AdaHessian's inability
8

---

Under review as a conference paper at ICLR 2023
25
GCAMI
Adt iedicae
3,0
Test Loss
Add beStoasiiN
14-
BL,9
330
2.5
10-
66-
Adarfzzian
60
(a) MNIST
(b) CIFAR-10
Figure 5: Sensitivity of the step size for each method on MNIST-MLP and CIFAR-10 3C3D tasks. We select the best values of 31 and B2 for each step size a.
to accurately approximate the Hessian diagonals, shown in Fig. 1. Moreover, AdaHessian quired more computational time compared to all methods, which is also reflected Fig. 2. While being time-efficient, AdaHesScaleGN consistently outperformed all methods in CIFAR-10-3C31 and CIFAR-100-3C3D, and it outperformed all methods except AdaHesScale in CIFAR-100 ALL- CNN. This result is aligned with methods' accurate approximation of Hessian diagonals. Our experiments indicate that incorporating HesScale and HesScaleGN approximations in optimization methods be of significant performance advantage in both computation and accuracy. AdaHesS- cale and AdaHesScaleGN outperformed other optimizers likely due to their accurate approximation of the diagonals of the Hessian and GGN, respectively.
CONCLUSION
HesScale is scalable and efficient second-order method for approximating the diagonals of the Hessian at every network layer. Our work is based on the previous work done by Becker and Lecun (1989). We performed series of experiments to evaluate HesScale against other other scalable algorithms in terms of computational cost and approximation accuracy. Moreover, we demonstrated how Hes- Scale can be used to build efficient second-order optimization methods. Our results showed that our methods provide more accurate approximation and require small additional computations.
7 BROADER IMPACT
Second-order information is used in domains other than optimization. For example, some works al- leviating catastrophic forgetting use a utility measure for the network's connections to protect them. Typically, auxiliary loss is used between such connections, and their old values are weighted by their corresponding importance. Such methods (LeCun et al. 1990, Hassibi Stork 1993, Dong et al. 2017, Kirkpatrick et al. 2017. Schwarz al. 2018, Ritter et al. 2018) use the diagonal of the Fisher information matrix or the Hessian matrix as a utility measure for each weight. The quality of these algorithms depends heavily the approximation quality of the second-order approximation. Second-ordes information can also be used in neural network pruming. Molchanov et al. (2019) showed that second-order approximation with the exact Hessian diagonals could closely represent the true measure of the utility of each weight.
The accurate and efficient approximation for the diagonals of the Hessian each layer enables Hes- Scale to be used in many important lines of research. Using this second-ord order information provides reliable measure connection utility. Therefore, using HesScale in these types of problems can potentially improve the performance of neural network pruning methods and regularization-basee catastrophic forgetting prevention methods.
9

---

Under review as a conference paper at ICLR 2023
REFERENCES
Amari, S. (1998). Natural Gradient Works Efficiently in Learning. Neural Computation, 10(2), 251-276.
Ba, J. L., Kiros, J. R., & Hinton, G. E. (2016). Layer normalization. arXiv:1607.06450.
arXiv preprint
Becker, S. & Lecun, Y. (1989). Improving the convergence of nack-propagation learning with second-order methods. Proceedings of the 1988 Connectionist Models Summer School (pp- 29-37).
Bekas, C., Kokiopoulou, E., & Saad, Y. (2007). An estimator for the diagonal matrix. Applied Numerical Mathematics, 57(11), 1214-1229.
Botev, A., Ritter, H., & Barber, D. (2017). Practical Gauss-Newtor optimisation for deep learning.
Proceedings of the 34th International Conference on Machine Learning, 70, 557-565.
Chan, A., Silva, H., Lim, S., Kozuno, T., Mahmood, A. R., & White, M. (2022). Greedification operators for policy optimization: Investigating forward and reverse KL divergences. Journal of Machine Learning Research, 23(253), 1-79.
Chapelle, O. Erhan, D. (2011). Improved preconditioner for hessian free optimization. NIPS Workshop Deep Learning and Unsupervised Feature Learning.
Dangel, F., Kunstner, F., Hennig, P. (2020). BackPACK: Packing into backprop. Interna- tional Conference on Learning Representations.
Dong, X., Chen, S., & Pan, S. J. (2017). Learning to prune deep neural networks via layer-wise optimal brain surgeon. Proceedings of the 31st International Conference Neural Information Processing Systems (pp. 4860-4874).
Duchi, J., Hazan, E., & Singer, Y. (2011). Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(61), 2121-2159.
Hassibi, B. & Stork, D. (1993). Second order derivatives for network pruning: Optimal brain surgeon. Advances in Neural Information Processing Systems, 5.
loffe, S. and Szegedy, C. (2015). Batch normalization: Accelerating deep network training by reducing internal covariate shift. International conference on machine learning (pp. 448-456).
Kingma, D. P. Ba, J. (2015). Adam: A method for stochastic optimization. International Conference on Learning Representations.
Kirkpatrick, J., Pascanu, R., Rabinowitz, N.. Veness, J.. Desjardins, G., Rusu, A. A.. Milan, K., Quan, J., Ramalho, T., Grabska- Barrikskka A., et al. (2017). Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 1/4(13), 3521-3526.
Kunstner, F., Hennig, P., & Balles, L. (2019). Limitations of the empirical fisher approximation for natural gradient descent. Advances in Neural Information Processing Systems, 32.
10

---

Under review as a conference paper at ICLR 2023
LeCun, Y., Denker, J., & Solla, S. (1990). Optimal brain damage. Advances in Neural Information Processing Systems,
Luo, L., Xiong, Y., Liu, (2019). Adaptive gradient methods with dynamic bound learning rate. International Conference Learning Representations.
Martens, J. (2010). Deep learning via hessian-free optimization. International Conference Machine Learning (pp. 735-742).
Martens, J. (2020). New insights & perspectives on the natural gradient method. Machine Learning Research, 21(146), 1-76.
Journal of
Martens, & Grosse, R. (2015). Optimizing neural networks with kronecker-factored approximate curvature. International Conference on Machine Learning (pp. 2408-22222 (207).
Martens, J. & Sutskever, I. (2011). Learning recurrent neural networks with hessian-free optimiza- tion, International Conference on Machine Learning (pp. 1033-1040).
Martens, J., Sutskever, & Swersky, K. (2012). Estimating the hessian by Back-propagating curvature. International Conference on International Conference Machine Learning(pp- 963-970).
Mizutani, E. & Dreyfus, S. E. (2008). Second-order stagewise backpropagation for hessian-matrix analyses & investigation of negative curvature. Neural Networks, 21(2), 193-203.
Molchanov, P., Mallya, A., Tyree, S., Frosio, I,, & Kautz, J. (2019). Importance estimation for neural network pruning. JEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 11264-11272).
Reddi, S. J., Kale, S., & Kumar, S. (2018). On the convergence of Adam and beyond. International Conference on Learning Representations.
Ritter, H., Botev, A., and Barber, D. (2018). Online structured Laplace approximations for overcoming catastrophic forgetting. International Conference Neural Information Processing Systems (pp. 3742-3752).
Rumelhart, D. E., Hinton, G. E., and Williams, R. J. (1986). Learning representations by back- propagating errors. Nature, 323(6088), 533-536.
Salimans, and Kingma, D. P. (2016). Weight normalization: A simple reparameterization accelerate training of deep neural networks. Advances in neural information processing systems, 29, 901-909999999
Schneider, F., Balles, L., and Hennig. (2019). DeepOBS: A deep learning optimizer benchmark suite. International Conference Learning Representations.
Schraudolph, N. N. (2002). Fast curvature matrix-vector products for second-order gradient descent. Neural Computation, 14(7), 1723-1738.
Schwarz, J., Czameekk, W., Luketina, J., Grabska-Barvinska, A., Teh, Y. Pascanu, R., and Hadsell, R. (2018). Progress & compress: A scalable framework for continual learning.

---

Under review as a conference paper at ICLR 2023
International Conference Machine Learning (pp.
Springenberg, J., Dosovitskiy, A., Brox, T., and Riedmiller, M. (2015). Striving for simplicity: The all convolutional net. International Conference on Learning Representations [Workshop].
Tieleman, T., Hinton, G., et al. (2012). Lecture 6.5-RMSProp: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2), 26-31.
Tran, TT and Phong. L. T. (2019). On the convergence proof of AMSGrad and a new version. IEEE Access, 7, 61706-61716.
Wilson, A. C., Roelofs, R., Stern, M., Srebro, N., and Recht, B. (2017). The marginal value of adaptive gradient methods in machine learning. International Conference on Neural Information Processing Systems (pp. 4151-4161).
Yao, Z., Gholami, A., Shen, S., Keutzer, K., and Mahoney, M. W. (2021). Adahessian: An adaptive second order optimizer for machine learning. AAAI Conference on Artificial Intelligence, 35((2), 10665-10673.
Zeiler, M. (2012). Adadelta: An adaptive learning rate method. arXiv preprint arXiv: 1212.5701.
12

---

Under review as a conference paper at ICLR 2023
A HESSIAN DIAGONALS OF THE LOG-LIKELIHOOD FUNCTION FOR TWO COMMON DISTRIBUTIONS
Here, we provide the diagonals of the Hcssian matrix of functions involving the log-likelihood of two common distributions: normal distribution and categorical distribution with probabilities represented by a softmax function, which we refer to as a softmax distribution. We show that the exact computations of the diagonal can be computed with linear complexity since computing the diagonal elements does not depend on off-diagonals in these cases. In the following, we consider the softmax and normal distributions, and we write the exact Hessian diagonals in both cases.
A.1 SOFTMAX DISTRIBUTION
Consider a cross-entropy function for a discrete probability distribution as f - -i=1 Ps log qi (0), where q a probability vector that depends on a parameter vector 0, and p one-hot vector for the target class. For softmax distributions, q is parametrized by a softmax function q = c²/ E ea. write the gradient of the cross-entropy function with respect to 0 In this case,
Vof(0) q p.
Next, we write the exact diagonal elements of the Hessian matrix as follows:
diag Ho) diag( Vo pp) q a² -
where q² denotes element-wise squaring of q, and V operator applied vector denotes Jacobian. Computing the exact diagonals of the Hessian matrix depends only on vector operations, which means that we can compute it in Onz). The cross-entropy loss used with softmax distribution in many important tasks, such as supervised classification and discrete reinforcement learning control with parameterized policies (Chan al. 2022).
A.2 MULTIVARIATE NORMAL DISTRIBUTION WITH DIAGONAL COVARIANCE
For a multivariate normal distribution with diagonal covariance, the parameter vector 0 is determined by the mean-varian variance vector pair: A = (n, 2.). The ligeeiiii of random vector æ drawn from this distribution can be written as
1 5( (a-w) u) D(o2)-11a -w) --)- log( ((o") +
log ((:; μ1 2²)
1 - n1) D(0)))-2222 11) log( o +
where D( gives a diagonal matrix with の² in its diagonal, M is the determinant of matrix M and is some constant. We can write the gradients of the log-likelihood function with respect to and 2² as follows:
y logg(e;; u. u, 0²) D()²) ( (x-+)= (x 20003 (x
a (og [-1 -
- 112² 2 0² - -
where an all-ones vector, and @ denotes element-wise division. Finally, we write the exact diagonals of the Hessian matrix as
diag( Hμ,) ) diag V Vu(x 11) 0 -
=11003
diag(H,>) = diag Vo 1 11)² 0² 0 =[0.51 (x ((==-)²0 1)) 0² o
Clearly, the gradient and the exact Hessian diagonals can be computed in O(n). logeliiel functions for normal distributions are used in many important problemm,,,oollms,,, such as variational inference and continuous reinforcement learning control.
13

---

Under review as a conference paper at ICLR 2023
B HESSCALE WITH CONVOLUTIONAL NEURAL NETWORKS
Here, we derive the Hessian propagation for convolutional neural networks (CNNs). Consider a CNN with L layers followed by fully connected layer that outputs the predicted output q- The CNN filters are parameterized by (W,..., W where Wi is the filter matrix at the I-th layer with the dimensions 1 X kr,2, and its element at the ith and the jth column is denoted by Wi,a.j- For the simplicity of this proof, we assume that the number of filters at each layer is one; the proof can be extended easily to the general case. The learning algorithm learns the target functioo f by optimizing the loss L. During learning, the parameters of the neural network are changed to reduce the loss. At the layer I, we get the activation output matrix H; by applying the activation function a to the activation input AY: Hi = a (At). We assume here that the activation function element-wise activation for all layers except for the final layer L, where it becomes the softmax function. We simplify notations by defining Ho = X, where X is the input sample. The activation output H, is then convoluted by the weight matrix W1+1 of layer ( to produce the next activation
11,1-11 -4,-1 W4+.,m, H1, ((+m), We denote the size of the activation input: A1+1,á. = == _n=
output at the 1-th layer by h X wy. The backpropagation equations for the described network are
given following Rumelhart et al. (1986):
R1111 11 1111
OC
Oc
L 0A
Σ
dAi.
m=0
14
1ah1-1kke 2
11-111- Σ
-1441221
dc
0H, (aammmmmm))... 4mm)..immmx''
Σ
Σ
Wi
0 (1
0 A
m=0
R=U
n'=0
8411111
2
Oc
Σ
W WL+1,,,n' (Atis)
a A441,66m)).6-
11
11-22-1
0c
d' (AAag)
Σ
a (1 W1+1,m,n
(6)
wi k1.2
hi-ki wi wi-
7C
OC
0A,m.
OL (itm)))
Σ
Σ
OWi.a
A AWL.s.
0A
In the following, we write the equations for the exact Hessian diagonals with respect to weights 8 2//oow.... which requires the calculation of o 22//oA,.. first:
0
111111
11112211
ac
d'(
Σ
W441m
0A
aAiu
M=D
==U
1141,2
2
OC
a 0A (A--1,(t-))00s)) W441,m
dA(Li)))
Σ
(A++1((-m))6
A
m.p=0
n.q==
1
0C
+ +o" (A(,)))))
Σ
W411m.
a 0 (1t)),(imm).0--
m=0
R=0
hu kin
wi wi-ki
02C
ac
Σ
Σ
H 1.(+m))(tt)
OW2.,
OWias
A
m=0
n=0
ht k1,1 wi -1.2
JC
AA H1-1.644m
Σ
A aA A aW.ss
m.
Since the calculation of 82/0AA,,,, and 8 c/ow2,, depend on the off-diagonal terms, the compu- tation complexity becomes quadratic. Following Becker and Lecun (1989), approximate the
14

---

Under review as a conference paper at ICLR 2023
Hessian diagonals by ignoring the off-diagonal terms, which leads backpropagation rule with linear computational complexity for estimates
BPC aw 3A7
and
SLL = (At,))))
1111.2- 1111,211
JL
W 1+1mm. (i)
0AP
0A
k141.2 2- 1
Oc
+ 0" (Aris)
W Wi+1mm
(8)
n=0 m==0
by h1-
Wi k1.2
0LL
PL
H
(9)
OW
DA
L.m.
m=0
C APPROXIMATION QUALITY WITH MNIST DATA
We repeat the experiment shown in Fig. with MNIST data points instead of random data points. The experimental details are the same except for two changes. First, we used a larger network where changed the number of units in each hidden layer to 32 instead of 16. Second, we performed an optimization update with SGD at each data point. The results shown in Fig. 6 are similar to the results shown Fig. I where HesScale gives better approximation quality than other methods. This experiment shows that our results hold for realistic settings where learning is involved.
:000
10'
Normalizeed/ L error
10'
AdbHierstan-MC]
35-
30
Error
AdaHessden-MCS0
KFRCINCI GEN-MIC1
1>
Ada Martastin--C
KTNC-MCI
<F CPAC-NCO
dagg.
Hertrancc
COLACI
13a
HesStale
HiSScrkGN
Methods
Layer Number (b) Layer-wise L' error
(a) Normalized L error with respect to HesScale
Figure 6: The averaged error for each method is normalized by the averaged error incurred by HesScale for data points coming from MNIST. We show 40 initialization points with the same colors across all methods. The norm the vector of Hessian diagonals [diag [H] is shown as reference.
15

---

Under review as a conference paper at ICLR 2023
D OPTIMIZATION PLOTS IN THE NUMBER OF EPOCHS
We give the training loss, training accuracy, validation loss, validation accuracy, test loss, and test accuracy for each the methods we include in our comparison in Fig. and Fig. 8. Moreover, we give the sensitivity plots for 31, $2, and a for each method in Fig. 9 and Fig. 10.
0.5
6.0
GGNMC
GONMC
AdarlesScalle
4.5
AdaHesScale
AdaHeeS
SGD
4.0
SGD
Test Loss
Test Loss
3.5
0.3
AdeHessian
AdaHessian
Aster
1.5
1.0
5.0
4.5
4.0
Train Loss
Train Loss
3.5
3.0
2.5
1.0
Test Test Accuracy
Test Test Accuracy
0.97
0.90
0.3
Train Accuracy
Train Accuracy
0.7
0.90
0.97
0.4
0.94
80
100
0.3
Epochs
Epochs
(a) MNIST
(b) CIFAR-10
Figure 7: Learning curves of each algorithm two tasks, MNIST-MLP and CIFAR-10 3C3D, for 100 epochs. We show the best configuration for each algorithm on the validation set. The best parameter configuration for each algorithm is selected based on the area under the curve for the validation loss.
16

---

Under review as a conference paper at ICLR 2023
GGNMC ScaeHeeS
AdaHosBcale
3.2
SGD
4.5
S6D
Test Loss
3.0
AdaHessian
Adam
Test Loss
4.0
AdaHeessan
Adam
2.0
2.2
2.5
3.5
3.75
3.0
3.50
Train Loss
Train Loss
2.5
3.00
2.0
2.75
1.5
2.25
1.0
2.00
0.660
0.525
Test Accuracy
0.000
Test Accuracy
0.455
0.45
0.425
0.40
0.400
0.35
0.360
0.30
0.75
0.70
Train Accuracy
0.65
Train Accuracy
0.60
0.60
0.55
0.50
0:45
0.40
0.35
60
100
130
200
250
300
300
0.30
75
100
125
150
176
200
Epochs
Epochs
(a) CIFAR- 100 All CNN
(b) CIFAR-100 3C3D
Figure 8: Learning Curves of each algorithm on CIFAR-100 with AII-CNN and 3C3D architectures, for 100 epochs. We show the best configuration for each algorithm the validation set. The best parameter configuration for each algorithm is selected based on the area under the curve for the validation loss.
17

---

Under review as a conference paper at ICLR 2023
B2 = 0.99
=
B2 = 0.9999
CONMC
Adart
1.0
1.0
1.0
0.5
0.5
0.5
0.0
0.0
0.0
10
10
10
10
10
10
10
10
0.9
15
1.0
1.0
1.0
05
0.5
055
0.0
0.0
10
10
10
10"
10
10
10"
1.0
0.5
0.0
(a) MNIST
B2 = 0.99
B2 0.999
82 0.9999
0.0
DUNIIG Acallaticale KASMMObaaO
B1 =
10
10
10
10"
16
10
10
10"
10
10
10
(b) CIFAR-10
Figure 9: Parameter Sensitivity study for each algorithm on two data sets, MNIST and CIFAR- The range of 32 is {0.99, 0.999, 0.9999} and the range of 31 is (0.0,0.9) Each point for each algorithm represents the average test loss given a set of parameters.
18

---

Under review as a conference paper at ICLR 2023
6
5
GGNMC
AdaHesScale
4
GGNMC
AdaHesScaleGGN
AdaHesScale
SGD
AdaHesSoaleGN
AdaHessian
3
SGD
Adam
AdaHessian
Adam
2 unl
10
10
10
10
10
10
10
10
10
10
10
10
(a) CIFAR-100 AlI-CNN
(b) CIFAR-100 3C3D
Figure 10: Parameter Sensitivity study for each algorithm on CIFAR-100 with All-CNN and 3C3D
architectures. The range of step is {10-1 5, 10 10 10 10 10"). We choose 31 to be equal to 0.9 and 52 to be equal to 0.999. Each point for each algorithm represents the average test loss given set of parameters.
E SUMMARY OF OPTIMIZATION RESULTS
We summarize the final performance of AdaHesScale and AdaHesScaleGN against other optimizers on train sets and test sets, in Table and Table respectively.
Table Performance of optimization methods on the train sets of different problems.
MNIST
CIFAR-10
CIFAR-100 3C3D
CIFAR-100 ALL-CNN-C
AdaHesScalc
99.68%
78.19%
57.29%
71.89%
AdaHesScaleGN
99.64%
82.11%
65.39%
58.10%
GGN-MC
99.64%
78.92%
64.62%
59.14%
SGD
99.65%
77.93%
61.37%
63.37%
Adam
99.67%
80.24%
58.36%
60.82%
AdaHessian
99.71%
74.94%
46.92%
62.43%
BI.89
96.99%
33.89%
Table 2: Performance of optimization methods on the test sets of different problems.
MNIST
CIFAR-10
CIFAR-100 3C3D
CIFAR-100 ALL-CNN-C
AdaHesScale
98.17%
77.98%
49.13%
53.59%
AdaHesScaleGN
98.17%
79.79%
49.99%
50.16%
GGN-MC
97.99%
77.93%
49.41%
48.06%
SGD
98.02%
77.28%
47.65%
49.13%
Adam
98.10%
79.07%
47.30%
47.97%
AdaHessian
98.23%
73.91%
45.10%
47.63%
BL89
95.13%
36.74%
19


Reviews:
[{"review_id": "wxbZaq2f4Z", "reviewer": "Reviewer_hQqe", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "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": "This study proposes HesScale, which can quickly compute the diagonal elements of loss Hessian in deep neural networks by back-propagating only the diagonal component of Hessian w.r.t. activation. For an MLP, it is shown that HesScale estimates a more accurate diagonal Hessian at a lower computational cost than existing methods such as AdaHessian. This study also proposes AdaHesScale, which incorporates the approximate diagonal Hessian calculated by HesScale into an Adam-style update rule. The experiment results show that AdaHesScale (and its variant) can achieve a certain test accuracy faster than SGD and Adam in training CNNs and MLPs for image classification tasks.\n", "summary_of_the_review": "While the computation of fast and accurate second-order information is indeed an important topic, questions remain about the validity of the approximations in the proposed HesScale and AdaHesScale, their superiority over existing methods, and the generality of the proposed method’s effectiveness. The quality of this study could be improved by more profound validation of the approximations used in HesScale (ignoring non-diagonal elements of the back-propagated Hessian w.r.t. activation) and more appropriate comparisons with existing second-order methods. \n", "strength_and_weaknesses": "Strengths\n- The computational cost, training convergence, and sensitivity to hyperparameters of several diagonal second-order optimization and adaptive gradient methods are compared. (The validity of the comparison is questionable, though.)\n\nWeaknesses\n- The proposed HesScale appears to be a very crude approximation without justification.\n    - HesScale ignores the off-diagonal elements of Hessian w.r.t. activation, but no theoretical or intuitive justification for this has been provided. Although the experimental results shown (Figure 1) indicate that HesScale has a relatively small error, it is unclear whether this is true in other tasks.\n    - The recursive computation of the Hessian is already discussed by Botev et al. (2017) (https://arxiv.org/pdf/1706.03662.pdf) and Dangel et al. (2020) (https://openreview.net/forum?id=BJlrF24twB), where they “back-propagate” the Hessian w.r.t. activation including the off-diagonal elements.  \n- The validity of the comparison is questionable.\n    - It is unfair to compare only the diagonal elements of K-FAC in Figure 1. K-FAC approximates the layer-wise Fisher matrix, including off-diagonal elements, by sacrificing the accuracy of the diagonal elements. Therefore, the error of the layer-wise Fisher approximation should be compared. If only the diagonal elements matter, then a computational approach that only calculates the diagonal elements of the K-FAC matrix should be used to compare the computing time. \n    - Although diag(H) is used as the grand truth, H (Hessian) is not necessarily a valuable measure of information for second-order optimization in deep learning because the generalized Gauss-Newton and Fisher matrix is preferred in nonconvex optimization (e.g., https://www.cs.toronto.edu/~jmartens/docs/Deep_HessianFree.pdf).  \n    - Similarly, the error with diag(H) is not a good indicator because what second-order optimization uses is the inverse of the (damped) second-order information: H_inv = (H + dI)^{-1}, where d > 0 is the damping value corresponding to \\epsilon in Algorithm 2. Depending on the choice of d (it could be 1e-1 to 1e-4 in practice https://arxiv.org/pdf/1806.03884.pdf), the importance of the accuracy of estimating diag(H) will be limited.\n    - Based on the above, I suggest reviewing the comparison approaches.\n- The advantage of AdaHesScale over Adam is unclear.\n    - AdaHesScale (Algorithm 2) uses the 'square root of' the inverse of the approximate diagonal Hessian as the preconditioner of the gradient, which does not correspond to the Newton-Raphson method. While the motivation for this study comes from improving approximation accuracy (and computational speed) of the Newton-Raphson method in large settings, it is unclear how much of the second-order information is utilized in AdaHesScale since it takes an update rule almost identical to that of Adam. \n    - According to the leaderboard for the CIFAR-100 classification task with All-CNN-C in DeepOBS (https://deepobs.github.io/leaderboardP6.html), Adam achieves a test accuracy of >61% at about 350 epochs, which is much better than Adam (47%) and AdaHesScale (53%) in Figure 4(b). Furthermore, training with momentum SGD achieves >66%. Even if it is difficult to reproduce the leaderboard results perfectly, I do not believe it is reasonable to use a result about 12% lower than the existing results as a baseline for Adam.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "wb5M1YXBqkW", "reviewer": "Reviewer_i83K", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "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": "The authors present a modification of a method of Becker and Lecun (1989) using a diagonal Hessian approximation to improve the convergence of stochastic gradient descent schemes.  The scheme takes about the same amount of time as standard back-propagation, unlike some other approximation schemes.  Based on this modified Hessian approximation, they introduce the AdaHessian method, which compares favorably to other standard optimizers on a set of test problems.", "summary_of_the_review": "With a small-but-critical change to a 1989 scaling technique, together with adopting modern SGD framework, the authors introduce a new SGD-style optimizer for NN training that includes second-order information \"for cheap\" and seems to lead to better training performance (in terms of test accuracy vs time) than natural competitors.  I am a fan of \"not a big change, but the right change\" work, and would like to see this published.", "strength_and_weaknesses": "To the extent that I can read the figures, the approximation scheme does indeed seem to lead to better convergence results than competitors on the test problems.  The modification compared to B&L 89 is not large, but the authors point out that it makes a big difference in the convergence of the final method.\n\nThere is no new theory behind this, or at least none given in the paper.  The verification of the quality is purely empirical.  It would be interesting if there was a theoretical argument for why the modification makes as much difference as it does.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "HhxSsH3Zkj", "reviewer": "Reviewer_HXwV", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "The authors present a method for approximating the diagonal of the Hessian of neural networks in order to use this approximate curvature information during training.\nThe approximation proposed by the authors is based on considering the exact symbolic derivative for the Hessian diagonal entries and then dropping terms as necessary during backprop in order to make the computation fast (constant factor of the function evaluation like standard backprop gradient).\nThe authors analyze the accuracy of their Hessian approximation compared to other baseline methods for approximating curvature information and are able to demonstrate that their approximation has much smaller approximation error in practice for some small toy networks (3 layeers with 16 unites, 1 layer with 512 both with tiny input size).\nThe authors also present some experimental results demonstrating the usefulness of their Hessian diagonal optimization in terms of reaching high test set accuracy faster.", "summary_of_the_review": "I do not recommend to accept the paper as the paper does not have either realistic experiments that are meaningful to the machine learning practitioner nor does it present a theoretical analysis of approximation bounds and resulting possible statements about optimizer convergence for the proposed Hessian approximation", "strength_and_weaknesses": "Strengths:\n- Interesting approach because it utilizes an exact symbolic differentiation approach in a computationally efficient manner\n- The authors present convergence results over wall-clock time demonstrating that their approximation is in principle computationally feasible in practice\n\nWeaknesses:\n- The background section is potentially a bit extensive on backgroud about block structure that is not necessary for the method proposed in the paper\n- It would be interesting to know whether there can be theoretical bounds on the error made from dropping the terms in the symbolic derivative\n- Tiny toy example network used for accuracy analysis (3 layers with 16 units, 1 layer with 512 units), not even something like MNIST but just random data examples so it is unclear how the accuracy statements hold in more realistic settings\n- Increasing the number of outputs is not a practical application for increasing parameter size since the number of classes does not typically change, furthermore this seems biased as the last layer computation of the Hessian is exact and different from the others. So cost scaling should be done by increasing number of layers and layer width (and much larger input space should be used for the examples)\n- The experimental section could be improved by presenting the results in a more structured way also presenting results numerically in tables, the scaling of the plot axes could be misleading\n- Not sure if the hyperparam tuning for the different optimizers was sufficient to be convincing\n- Need to better motivate the method compared to quasi-Newton type methods for estimating a Hessian diagonal (especially in a full-batch setting)\n- Need to motivate the method in the context of nonconvex stochastic estimation or some other application of machine learning. In the end I am not sure sure what the relevance of curvature information is for machine learning. In theory curvature is useful for different contexts, it can be useful to escape saddle point quicker, but that's not how it seems to be used in this paper. The other purpose can be to converge to a an optimum faster when already in a convex region around the optimum. But getting close to an optimum of an empirical loss is not always beneficial for generalization. The authors show test set plots which is unconventional, but maybe they can be viewed like validation set plots. It would have been interesting as well to see how the training set loss behaves to see what the generalization gap is for these applications.\n\n*Questions and remarks*:\n- How does this related to edge pushing algorithm (see https://arxiv.org/pdf/2007.15040.pdf)?\n- What is the quality of the Hessian approximation when computed on a mini-batch?\n- Not sure about beta_1 search space of just two choices, that seems limited?\n- Why is the standard error for the CIFAR-10 3C3D AdaHessian run so large?\n- Color codes for different optimizers should be identical between plots (Figure 3 a) has AdaHessian in pink and b) has Adam in pink)\n- Why does Figure 5 report test loss while Figure 3 and 4 report test accuracy?\n- Typo: \"sensitivty\"\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "0iOm6mQg8o", "reviewer": "Reviewer_37p6", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "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": "The paper proposes an efficient heuristic approximation to compute the diagonal of the Hessian matrix in neural nets, and empirically shows that it is significantly more accurate than other existing deterministic and stochastic approximations.  The paper then goes on to propose a second-order optimization method that uses these estimated hessian values, and shows on experimental results that the method is competitive and sometimes better than other standard optimizers (SGD, ADAM) in computer vision tasks. ", "summary_of_the_review": "The paper proposed an improved approximate estimator of the diagonal of the Hessian for neural networks, and uses it to propose a second-order method. Experimental results are done well,  but could be more comprehensive.  It'd be very helpful if the authors made an attempt at gaining insight into why the method works as well as it does (and whether there are settings where it works less well) . ", "strength_and_weaknesses": "Strength:  A principled and inexpensive approximate second-order optimization method for learning neural nets would be very impactful -- and the paper is taking a shot at this problem.  The method is a small modification of an existing proposal by Becker and Lecun, that observes that the diagonal of the Hessian of the last layer can be computed exactly (rather than approximately) for common loss functions,  but this seems to make a very significant difference in the approximation error (in experiments).  The paper then uses this Hessiagn-diagonal estimator for optimization. Experimental results with image classification show that the method is competitive with other optimizers, and shows the best performance on some if the instances. \n\nWeaknesses: \n1) The paper goes into a lot of depth describing the hessian-diagonal approximation, but the actual optimization method is just a paragraph + an algorithm.  Some additional explanation of steps / parameters would make it more readable. \n\n2) I was struck by the impressive results in for HessScale in Figure 1,  but I have some questions on them.  a) Do you have any intuition, insight or analysis on how can the method do so much better than other alternatives (order of 10-20x reduction in error w.r.t best other methods).   Are the discarded off-diagonal terms somehow conveniently cancelling for random problems, or what's going on?  Is the setting of random neural nets on random inputs somehow fortuitous for this evaluation, or does this improvement in Hessian approximation also hold for realistic neural nets in the middle of the training? \n\n3) What limitations are there on when the method can be used for more general neural nets? E.g. can it be applied with dropout, batch-norm layers, or for other architectures (ResNets, RNNs, LSTMs, transformer-style architectures)?  What common losses are not supported? \n\n4) Evaluation.  It would be great to include some examples from other (non-vision) domains and with different achitectures -- e.g. natural language, tabular data, recommendation, maybe speech.  It seems that there are differences in what's the default optimizer in each community, and it could be motivated by artifacts of the structure of the problem / architectures commonly used. \n\n5) For second order methods ideally it would be better to compute the diagonal of the inverse Hessian -- rather than inverting the diagonal of the Hessian? Is that something that can be computed efficiently?  Having access to Hessian-vector products it is possible to estimate the diagonal of the inverse -- using Hutchinson's like methods,  or alternatives, like (this is in context of Gaussian graphical models, but it's essentially an estimator for the diagonal of the inverse of the matrix): \nWilsky, Malioutov, Johnson, \"GMRF Variance Approximation using Spliced Wavelet Bases.\" 2007 IEEE ICASSP, 2007. \n\n6) For AdaHessScaleGCN -- it seems that well into convergence the optimizer occasionally has some large instabilities with large oscillations (e.g. top-left and in Figure 1, and some runs in Fig 6).  Why does this happen, and can it be improved somehow? \n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject