Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation
Authors: Lin Zhang
Shaohuai Shi
Bo Li
Venue: ICLR 2023 poster (2023)
Keywords: Deep learning
Second-order optimization
Approximation

Abstract: Second-order optimization algorithms exhibit excellent convergence properties for training deep learning models, but often incur significant computation and memory overheads. This can result in lower training efficiency than the first-order counterparts such as stochastic gradient descent (SGD). In this work, we present a memory- and time-efficient second-order algorithm named Eva with two novel techniques: 1) we construct the second-order information with the Kronecker factorization of small stochastic vectors over a mini-batch of training data to reduce memory consumption, and 2) we derive an efficient update formula without explicitly computing the inverse of matrices using the Sherman-Morrison formula. We further provide a theoretical interpretation of Eva from a trust-region optimization point of view to understand how it works. Extensive experimental results on different models and datasets show that Eva reduces the end-to-end training time up to $2.05\times$ and $2.42\times$ compared to first-order SGD and second-order algorithms (K-FAC and Shampoo), respectively. 

Published as a conference paper at ICLR 2023
EVA: PRACTICAL SECOND-ORDER OPTIMIZATION WITH KRONECKER-VECTORZZZ APPROXIMATION
Shaoouaa Shi ' Li Bo' Lin Zhang', Shaohuai
Honn Kong University of Science and Technology, "Harbin Institute of Technology, Shenzhen 1 zzannbvve connect. wwtnnnetttsttt...... hk, shachuais@hin .ehii.edu.cnttt....... bli0cse. ust.hk
ABSTRACT
Second-order optimization algorithms exhibit excellent convergence properties for training deep learning models, but often incur significant computation and memory overheads. This can result in lower training efficiency than the first-order counterparts such as stochastic gradient descent (SGD). In this work, we present memory- and time-efficient recond-order algorithm named Eva with two novel techniques: 1) we construct the second-order information with the Kronecker fac- torization of small stochastic vectors over mini-batch of training data to reduce memory consumption, and 2) we derive an efficient update formula without ex- plicitly computing the inverse of matrices using the Beeeman-Morrrsiee formula. We further provide theoretical interpretation of Eva from trust-region optimiza- tion point of view to understand how it works. Extensive experimental results different models and datasets show that Eva reduces the end-to-end training time up 2.05 and 2.42x compared to first-order SGD and second-order algorithms (K-FAC and Shampoo), respectively.
INTRODUCTION
While first-order optimizers such stochastic gradient descent (SGD) (Bottou al., 1998 and Adam (Kingma & Ba 2015) have been widely used in training deep learning models (Krizhevsky al. 2012 He al. 2016 Devlin et al. 2019), these methods require large number of iterations to converge by exploiting only the first-order gradient to update the model parameter (Bottou et al.] 2018). To overcome such inefficiency, second-order optimizers have been considered with the po- tential to accelerate the training process with a much fewer number of iterations to converge (Osawa al. 2019 2020 Pauloski al. 2020 202T). For example, our experimental results illustrate that second-order optimizers, e.g., K-FAC (Martens & Grosse 2015). require ~00 fewer iterations to reach the target top-1 validation accuracy 93.5% than SGD, training ResNet- 10 (He al.) 2016) model on the Cifar- 10 dataset (Krizhevsky 2009) (more results are shown in Table 2
The fast convergence property of second-order algorithms benefits from preconditioning the gra- dient with the inverse of matrix C' of curvature information. Different second-order optimizers construct C' by approximating different second-order information, e.g., Hessian, Gauss-Newton. and Fisher information (Amari 1998), to help improve the convergence rate (Dennis & Schnabel) 1983). However, classical second-ord -odder optimizers incur significant computation and memory over- heads in training deep neural networks (DNNs), which typically have large number of model parameters, as they require quadratic memory complexity to store C, and a cubic time complexity to invert C, w.r.t. the number of model parameters. For example, ResNet-50 (He al.) (016) model with 25.6M parameters has to store more than 650T elements in C using full Hessian, which affordable on current devices, e.g., an Nvidia A100 GPU has S0GB memory.
To make second-order optimizers practical in deep learning, approximation techniques have been proposed to estimate C with smaller matrices. For example, the K-FAC algorithm (Martens & Grosse 2015) uses the Kronecker factorization of two smaller matrices to approximate the Fisher information matrix (FIM) each DNN layer, thus, K-FAC only needs store and invert these small matrices, namely Kronecker factors (KFs), to reduce the computing and memory overheads.
"Corresponding author.

---

Published as a conference paper at ICLR 2023
Table 1: Time and memory complexity comparison of different second-orde algorithms. d is the dimension of a hidden layer, L is the number of layers, and m is the number of gradient copies.
Complexity Time
Newton
K-FAC
Shampoo
M-FAC
Hva
Lid
()
L)
L)
L)
Memory ( L)
O(2
O L)
0
Second-order Info.
Hessian
KFs
Gradient Statistics
Gradient Copies
KV's
However, even by doing so, the additional costs of each accond-order update are still significant, which makes it slower than first-order SGD. In our experiment, the iteration time of K- is 2.5x than that of SGD training ResNet-50 (see Table 4 and the memory consumption of storing KFss and their inverse results is 12x larger than that of storing the gradient. Despite the reduced number of iterations, existing second-order algorithms, including K-FAC (Martens & Grosse| 2015), Sham- poo (Gupta et al. 2018), and M-FAC (Frantar et al. 2021), are not time-and-memon
efficient, as shown in Table One limitation in K-FAC and Shampoo is that they typically require dedicated system optimizations and second-order update interval tuning to outperform the first-order counter-
part (Osawa et al. 2019 Pauloski et al. 2020 Anil et al. 2021 Shi et al., 2021
To address the above limitations, we propose novel second-order training algorithm, called Eva, which introduces a matrix-free approximation to the second-order matrix to precondition the gra- dient. Eva not only requires much less memory to estimate the second-order information, but it also does not need to explicitly compute the inverse of the second-order matrix, thus eliminating the intensive computations required in existing methods. Specifically, we propose two novel tech- niques in Eva. First, for each DNN layer, we exploit the Kronecker factorization of two small stochastic vectors, called Kronecker vectors (KVs), over a mini-batch of training data to construct rank-one matrix to be the second-order matrix C' for preconditioning. Note that our constructed second-order matrix is different from the average outer-product of gradient (i.e., Fisher informa- tion (Amari 1998)) that has been used in existing K-FAC related algorithms (Martens & Grosse, 2015 George et al. 2018) ($3.I). KVs require only a sublinear memory complexity w.r.t. the model size. which is much smaller than the linear memory complexity in existing second-order algorithms like storing KFs in K-FAC (Martens & Grosse 2015), gradient statistics in Shampoo (Gupta et al. 2018), or gradient copies in M-FAC (Frantar et al.] (021)) (see Table [].. Second, we derive new update formula to precondition the gradient by implicitly computing the inverse of the constructed Kronecker factorization using the Sherman-Morrison formula (Sherman & Morrison 1950). The new update formula takes only linear time complexity; it means that Eva is much more time- efficient than existing second-orde in inverting matrices (see Table IT. Finally, we provide theoretical interpretation to Eva from trust-region optimization point of view to understand how it preserves the fast convergence property of second-order optimization (Ai & Duchi 2019 Bae et al. 2022).
optimizers which normally take a superlinear time complexity
We conduct extensive experiments to illustrate the effectiveness and efficiency of Eva compared to widely used first-order (SGD, Adagrad, and Adam) and second-order (K-FAC, Shampoo, and M- FAC) optimizers on multiple deep models and datasets. The experimental results show that 1) Eva outperforms first-order optimizers- achieving higher accuracy under the same number of iterations or reaching the same accuracy with fewer number of iterations, and 2) Eva generalizes very closely to other second-order algorithms such as K-FAC while having much less iteration time and memory footprint. Specifically, in terms of per-iteration time performance, Eva only requires an average of 1.14x wall-clock time over first-order SGD, while K-FAC requires 3.47x in each second-order update. In term of memory consumption, Eva requires almost the same memory as first-order SGD, which is up to 31 % and 45% smaller than second-order Shampoo and K-FAC respectively. In term of the end-to-one aod-to-ent end training performance, Eva reduces the training time on different training benchmarks up to 2. 2.05x, 1.58x, and 2.42x compared to SGD, K-FAC, and Shampoo respectively.
In summary. our contributions are as follows: (1) We propose a novel efficient second-order op- timizer Eva via Kronecker-vectorizee approximation, which uses the Kronecker factorization of two small vectors to be second-order information so that Eva has sublinear memory complexity and requires almost the same memory footprint as first-order algorithms like SGD. (2) We derive a new update formula with an implicit inverse computation in preconditioning by exploiting the Sher- man-Morrison formula to eliminate the expensive explicit inverse computation. Thus, Eva reduces (3) We conduct extensive experiments to
each second-order update cost to linear time complexity. validate that Eva can converge faster than SGD, and it is more system efficient than K-FAC and Shampoo. Therefore, Eva is capable of improving end-to-end training performance.
2

---

Published as a conference paper at ICLR 2023
2 BACKGROUND
In this section, we introduce the background of first-order and second-order optimization algo- rithms. We start from the supervised learning, in which DNN model is trained by randomly going through the dataset D many times (i.e., epochs) to minimize loss function E(w; D). The loss function measures the average dis- tance between model predictions and ground- truth labels, and w is the model parameter should be trained. Given a DNN with L learnable layers, the model parameter typically consists of set of parameter matrices (Wi)=
Loss
Loss
Loss
FF2
BP
A1 FF2
7b
BP
G2
31
FF2
BP
92
FF1
21
BP
80
PF
(4)
(b)
(c)
Figure 1: Examples of different optimization al-
gorithms in two-layer DNN model: (a) SGD, (b) K-FAC, (c) Eva.
SGD. The stochastic gradient descent (SGD) algorithm and its variants (e.g., Adam) with first-order gradient information the main optimizers for training DNN models. In each iteration with mini-hatch of data, SGD updates the model parameter layer-wisely as follows:
w(+++) w/(h) a Vl (B)))
(1)
where W (w, E x are model parameter and first-order gradient matrices layer 1, respectively. (1 > is the learning rate at iteration t. The mini-batch of data B(a) is sampled randomly from the training dataset. SGD typically takes a large number of iterations to converge as it only utilizes the first-order gradient to optimize the model (Bottou et al. 2018).
Compared to SGD, second-order algorithms precondition the gradient by the inverse of the curva- ture matrix, i.e., C-- vee(( (V)), for model update. The straight-forward way is to use the Hessian to precondition the gradient using the Newton method, but it takes extremely high time and memory complexity as shown Table | Due to the large model size in DNNs, storing and inverting G expensive, which motivates approximation approaches to alleviate the memory and computa- tion complexity. One of representative methods is K-FAC (Martens & Grosse 2015) that has been applied in large-scale training to achieve comparable performance over SGD Osawa et al.) 2019
2020 Pauloski et al. 2020 2021). Our proposed method in this work is also close to K-FAC 33.1
K-FAC. The K-FAC algorithm (Martens & Crosse 2015) adopts the Fisher information matrix (FIM) as C for layer in a DNN and approximates Ci with the Kronecker product of two smaller For example, for linear layer a = o(Wai-=)), in which @ is an element-wise non-linear activation function, the FIM can be
matrices (Martens Grosse 2015 Grosse & Martens 2016) []]
approximated via Fi A/-1 1 @ G₁, where
A = E(ar-1a711 and G: = =gggg]]])
A1-1 E R4-1xd-1 and G; E di are symmetric matrices, called Kronecker factors (KFs). a E IRd is the input vector of layer I (i.e., output of layer 1 1-1) and gi the pre-activation gradient vector of layer I. & denotes the Kronecker product. KFs are used precondition the gradient with a damping parameter 7 > O via (Ar-1 @
G. + - veevvvee (Vt). Formally, we have 71)
w((+1) = w/ - -a (G V7 1- (l((")) (B)) +
(3)
where the scalar T. = (((ii1)/ /TAAII//VTGG and T(A) is the trace of A divided by its dimension. The detailed derivation of K-FAC given in Appendix A.2 An example of 2-layer DNN is shown in Fig. a and Fig. b)) to demonstrate the difference between SGD and K-FAC. Compared to the training process of SGD, K-FAC extra constructs two KFs, which have memory complexity of O(2d)), to approximate FIM and computes its inverse, which has time complexity of O(2d)), to precondition the gradient in each layer.
Additional approaches. There exist some other second-order optimizers for DNN training. For example, Shampoo (Gupta et al. 2018) is a full-matrix adaptive algorithm that constructs gradient statistics matrices (similar to KFs) by building smaller matrices layer-wisely for preconditioning,
In this paper, we focus on the aft-used K-FAC with the empirical FIM, discussed in Appendix A.2
3

---

Published as a conference paper at ICLR 2023
which introduces a time complexity of L) to invert these matrices, and a memory complexity of ((2 L) to store them. d is the dimension matrix-free algorithm that utilizes matrix-vector products with many gradient copies. M- M-FAC does not invert any matrix but requires m copies of gradient to estimate FIM, which brings time and memory costs of ((md² T.). m is typically suggested to be high 1024 (Frantar et al. 2021
hidden layer. M-FAC (Frantar et al. 2021 is a
In summary, existing second-order methods have high time and/or memory complexity in each second-order update as shown in Table In this work, we present a novel second-order optimizer, Eva, with much lower time and memory complexity.
EVA: KRONECKER-FACTOREN APPROXIMATION WITH SMALL VECTORS
Following standard feed-forward (FF) and pack-propagation (BP) processes, shown in Fig. [c). we propose to construct the second-order curvature matrix, denoted as Rz for the sake of clarity, for layer / using the Kronecker product of two small vectors a and gi, i.e.,
R (a1-17711) & (gi) (ā:-1 1 @ (a) 2 gg)
(4)
where
(1)
a
E
and g E
g1
(5)
B
B
iEB
EEB
al- and gi are activation and pre-activation gradient vectors at layer I, respectively, sampled from mini-batch of data B. We call av-1 and g: as Kronecker vectors (KVs), whose dimensions are the
(0)
(6)
same with a and gi respectively and they are irrelevant with the batch size. Thus, we just need to store inexpensive KVs to construct Rú which is much more memory-efficient than storing the whole matrix like K-FAC or Shampoo. The preconditioner of Eva becomes (Ri --) However, it still needs to compute the inverse of the damped Rt, which is computation-eepeensi in training DNNs. From Eq. 4 RS is a rank-one matrix. Using this good property, we can implicitly compute its inverse efficiently using the Sherman-Morrisonn formula (Sherman & Morrison
1950). Specifically,
(a 0 (8(8
(Ri + = ((ar-1 gi (a & g:) + 7 a 1 g:) + +7
(6)
Based on the above preconditioner, we can derive the update formula of Eva to be
w((++) = w/th) (t)) (VE)(
(( B(C) a (a a a) g: + graf
(7)
The detailed derivation of Eva is deferred to Appendix B.2 and the study of its training dynamies is provided in Fig. i Appendix C Compared to SGD, the precondition process of Eva changes the gradient direction with compensating in the direction of g:a"", and scales the step size by 1/7.
In Eva, KVs are calculated over mini-batci of data for constructing the second-order precon- ditioner by their Kronecker product. For a large-scale training dataset, the approximation can be stabilized by long-term run of estimation using the whole data. Thus, we use the running average strategy which is also commonly used in Adam Kingma & Ba (2015) with 1st and 2rd moments, or K-FAC (Martens & Grosse 2015) with KFs, that is
(t) a {a + (1 {)a"">" and
(t)
(1) +- {g) + (1- () (t-1)
(t-1)
(8)
where & E (0. 1] is the running average parameter in iteration t, a and (2) are new KVs caleu- lated during each FF and BP computation and they are used to update the state of KVs.
Complexity analysis. The extra time costs of Eva come from constructing KVs and precondition- ing the gradients with KVs, apart from sharing the same FF, BP, and update computations as SGD. As the overhead of estimating KVs can be ignored, the main time cost of Eva is multiple vector multiplications for preconditioning, which can be denoted as (d. It means Eva has a time com- plexity that is linear to the total number of parameters, which is much smaller than the superlinear complexity in K-FAC and Shampoo as shown in Table Besides, the memory complexity Eva is O(2dL) for storing KVs, which is sublinear to the total number of parameters. In summary, Eva has very little extra time and memory costs in each recond-order update compared to first-order SGD, but it enjoys the fast convergence property of second-order K-FAC.
4

---

Published as a conference paper at ICLR 2023
3.1 THEORETICAL UNDERSTANDING
Trust-region optimization. optimization algorithm (uaan 2015) can be for- mulated as
The trust-region
Wie1
WL+1
WL+1
WI+
w (++1) min ((w),
(9)
(b)
(d)
s.t. p(w, w()) - X,
(10)
Trust regions of different proximal con- straints: (a) ball region, (b) strip region, (c) trust region of Eva, (d) trust region of K-FAC.
Figure 2:
where w()) and ((+1) are current and next pa- rameter points, respectively. p(w, w is the
proximal function to measure the distance between w and w(t) (the smaller the closer). Trust- region optimization treats each update step as finding the next parameter point that minimizes the loss function and close to the current parameter point (i.e., located a trust region of {w p(w, w(?) S )J. where À 0 is the threshold). Different proximal functions correspond to different trust regions. For example, the euclidean distance, p(w, (w()()))= = ||w - w0) gives a ball region shown in Fig. a).
Adopting the same assumption of K-FAC (Martens Grosse 2015 activation gradient vectors are fairly independent, I.e., Ela & ~
where the input and pre-
X E[g]. we have
(11) the current mini-batch of data. The above equation indicates that the second-order matrix in Eva is in fact the approximation
R = [E(a] @ (E(a) @ Eig))* = E(a @ g E/a & g)" T E ((ve][[cvvc
where E. c V() is the estimated gradient vector
the outer Here we ignore the subscrip- tion of the layer index for simplicity. By applying OPAG into the proximal function to measure the parameter closeness, we obtain
product of average gradient (OPAG) (i.e., (Vee
p(w, w(())= (w- (O) w()) ((AA"((vvcCV)))) < X,
(12)
((t) is the parameter update. Therefore, the corresponding trust region is as shown in Fig. 2bb), which limits the update size along the gradient direction. According to
where Aw
Eq. 7 Eva use the preconditioned gradient using the damped OPAG as the parameter update, i.e.,
Aw (R + ? (Ve((V)) It indicates that the trust region of Eva is an intersection between the ball region and the strip region, as shown in Fig. c). More details are given in Appendix [B]
Relation to K-FAC. From the trust-region optimization perspective. the K-FAC algorithm with the
true FIM is known choose p as KI.-divergence between model distributions (Bae et al. 2022
which however is not the case for the K-FAC with empirical FIM that has been widely applied in
large-scale training (Osawa et al.) 2019 2020 [eno al.) 20200] Pauloski et al.]
2020
2021 We
focus on the analysis on oft-used K-FAC with empirical FIM as it is much more practical than true
FIM. We build the connection between K-FAC with empirical FIM (Osawa al.
2019
2020 and
our proposed Eva. Formally, the relationship between the empirical FIM and OPAG
B(vec(V()yec(Ve))
(erecc(V + Cov(ve:(V!))
(13)
empirical FIM
OPAG: outer-product of average gradient gradient covariance matrix
which indicates the empirical FIM is the sum of OPAG and the gradient covariance matrix (Mc- Candlish al. 2018). K-FAC and Eva utilize the approximate empirical FIM and approxi-
mate OPAG as second-order information for calculating preconditioners, respectively. That is F (ereecvVVveeeeVVVV
As Cov(vec(V1)) in Eq. 13 is positive semi-definite, we have v² Fv 2 v Rv for any vector Thus, we have
and It [[vee((()][vvecce))
{w: Aw" + AT) S A} C (w Aw" R R +/) S A},
(14)
for a given threshold A, which indicates the trust region of Eva is larger than K-FAC shown in Fig. 2 d). Therefore, the update step of K-FAC to minimize the loss is more conservative than Eva.
EVALUATION
4.1 CONVERGENCE PERFORMANCE
We evaluate the generalization performance of Eva with three representative models, VGG-19 (Si- monyan & Zisserman 2015). ResNet-110 (He et al. 2016), and WideResNet-28-10 (WRN-28-
5

---

Published as a conference paper at ICLR 2023
Table 2: Validation accuracy (%) comparison between Eva and SGD/K-FAC algorithms for training from scratch with different epoch buckets. indicates training with extra tricks.
Cifar-10)
Cifar- 100
Model
Epoch
SGD
K-FAC
Eva
SGD
K-FAC
Eva
50
90.951..2
92.57.0.3
92.63 =.2
61.69
70.23 70.14...11
VGG-19
100
92.27
93.37
93.20
67.97
72.25 10.1
71.7940.5
200
93.02+0.1
93.464402
93.59
70.98+0.0
72.90+0.3
72.72++0.4
50
90.98 90.9840.6
93.03+0.1
93.02
67.93 67.93+0.9
71.0340.5
71.13+0.4
ResNet- 10
100
92.49
93.764.0.1
93.76
70.74.0.7
72.31
72.38 10.2
200
93.8010.2
94.21
93.99
72.43
72.9610.3
73.29 10.3
50
95. 15.28
96.12.0.2
96.19 0.11
79.03
80.9810.2
81.15 10.2
WRN-28-10
100
96.88
97.21.1.00
97.05
83.0310.3
83.54
83.52.1.2
200
97 97.33+0.1
97.44
97.
84.5440.0
84.5640.1
84.52.0.1
Table 3: Validation accuracy (%) comparison between Eva and 4 more algorithms for training from scratch on Cifar- 10 with 100 epochs. indicates training with extra tricks (except M-FAC).
Model
Adagrad
AdamW
Shampoo
M-FAC
Eva
VGG-I9
92.42+0.0
92.97+0.1
93.38+0.1
92.50+0.1
93. 93.20+0.1
ResNet-110
90.3410.2
92.61
92.4710.2
93.45...1
93.76.10.0
WRN--2-10f
93.72+0.2
96.91
96.9940.1
94.5440.2
97.05
10) Zagoruyko & Komodakis 2016) on Cifar-10 (Krizhevsky) 2009) and Cifar-100 (Krizhevsky 2009) datasets. We compare Eva to the finst-order baseline SGD (with a momentum of 0.9), and the second order baseline K-FAC (Martens& & Grosse 2015). Following the configurations of (Pauloski eet al. 2020, 2021), we set the same hyper-parameterr for all algorithms for a fair comparison, and the details are given in Appendix(C.. Since training with more epochs can generalize better (Hoffer al. 2017), compressed to sufficient training budgets. The validation accuracy comparison is shown in Table 2 and Table 3 where we report the mean and std over three independent runs. We summarize the results in the following three aspects.
each algorithm with 50, 100, and 200 epochs for better comparison from
First, in the models VGG-19 and ResNet- with commonly used settings from their original papers (details in Appendix C.I)) it is seen that Eva performs closely to the second-order K-FAC, and they both consistently outperform SGD under different training budgets. To be specific, under the 50-epoch setting, both Eva and K-FAC outperform SGD by large margin, e.g., +8.5% for training VGG-19 on Cifar- under the setting of sufficient 200 epochs, Eva and K-FAC still achieve slightly better generalization performance than SGD. Due to the page limit, we put the convergence curves under the 50-epoch setting in Fig. 7 in Appendix C which shows that Eva converges similarly K-FAC and learns much faster than SGD. Eva and K-FAC with 50 (and 100) epochs achieve the loss or validation accuracy that SGD needs to take 100 (and 200) epochs. In summary, our proposed second-orde Eva has similar good convergence performance with K- K-FAC, and both of them train the models around 2x faster than SGD in terms of iterations. It validates that second-order algorithms have better convergence performance than SGD and they can reach the same target accuracy in fewer number of training iterations.
Second, considering that existing training paradigms typically use extra tricks like CutMix (Yun) ett al. 2019) and AutoAAgment (Cubuk et al. 2011) to achieve better validation accuracy, we train relatively new model WRN-28-TU (Zaooruyko & Komodakis 2016) on Cifar- 10 and Cifar-100 to compare the performance of different optimizers. Note CutMix (Yun et al.) 2019) and AutoAug- ment (Cubuk et al.] 2019) have been particularly developed and heavily tuned for the first-order SGD optimizer, which could be detrimental to second-order algorithms. However, as shown in Table our Eva and K-FAC still learn faster and generalize better than SGD under the same num- ber of training epochs. Additional results to verify the generalization performance for fine-tuning pretrained models are provided in Table 6 in Appendix [C] In summary, Eva achieves the same gen- eralization performance with less number of iterations than well-tuned SGD it achieves higher generalization performance with the same training number of iterations as SGD.
Third, compare the convergence performance of Eva with 4 more popular optimizers: Ada- grad (Duchi et al.] (2100), AdamW (Loshchilov & Hutter, (2019), Shampoo (Gupta et al.) 2018), and M-FAC (Frantar al. (022). Adagrad and AdamW are common adaptive gradient methods, while

---

Published as a conference paper at ICLR 2023
Shampoo and M-FAC recently proposed second-order algorithms. We tune the learning rate for each algorithm to choose best one for particular algorithms (see Appendix (C.I) to train three models Cifar-10 for 100 epochs. The results are given in Table []] showing that Eva achieves comparable performance to other second-order algorithms Shampoo and M-FAC, and outperforms first-order adaptive methods Adagrad and AdamW on different models. M-FAC, however, results in 2.5% accuracy loss, compared to Eva and Shampoo in training WRN-28-10. Note that we exclude CutMix and AutoAugment in M-FAC as they cause M-FAC divergences training WRN-28-10.
Table 4:
Relative iteration time and memory over SGD. Values in parentheses represent the results with increased second-order update intervals (10 Cifar- and 50 ImageNet).
Shampoo
K-FAC
Eva
Dataset
Model
Time
Mem
Time
Mem
Time
Mem
VGG-19
19.6x (2.89x)
1.01x
5.57x (1.68x)
1.01x
1.13x
1.00x
Cifar-10
ResNet- 110
6.79x (1.90x)
1.00x
1.64x
(1.16x)
1.03x
1.16x
1.00
WRN-28-10
6.69x (1.60x)
1.05x
2.68x (1.25x)
1.38x
1.03x
1.00x
ResNet-50
30.7x (1.71x)
I.07x
2.52x
(1.14x)
1.06x
1.09
1.00x
ImageNet
Inception-v4
75,3x (2.70x)
1.11x
3,95x
(1.28x)
1,42x
1.28x
1.00
ViT-B/16
199.x (6,20x)
1.31x
4.47x (1.43x)
1,45x
1.18x
1.00x
4.2 TIME AND MEMORY EFFICIENCY
As we have shown the good convergence and generalization performance of Eva, we demonstrate its per-iteration training time and memory consumption compared with SGD, Shampoo, and K-FAC. We do not report M-FAC here as needs store extra m gradients for FIM estimation (m 1024
by default), which is very memory inefficient (normally causes out-of -fmooory))) To cover relatively large models, we select three extra popular deep models trained on ImageNet (Deng et al.) 2009): ResNet-50, Inception-v4 (Szegedy et al. 2017), and ViT-B/16 (Dosovitskiy et al., 2021). We run each algorithm on an Nvidia RTX2080TI GPU. Following the training recipes in (Pauloski et al.) 2020). we set the second-order update interval as 10 Cifar-10, and 50 on ImageNet for both Shampoo and K-FAC to reduce their average iteration time. More detailed configurations can be found in Appendix C.1 We report relative time and memory costs over SGD in Table 4
Time efficiency. We can see that Eva has the shortest iteration time among all evaluated second- order algorithms. Eva has only 1.14x longer iteration time average than SGD, but it achieves an average of 3.04x and 49.3x faster than K-FAC and Shampoo, respectively. This is because our Eva does not need to explicitly compute the inverse of second-order matrix while K-FAC requires expen- sive inverse matrix computations on KFss and Shampoo performs even more computations of inverse p-th roots (p 2 2). Therefore, Eva is able to update second-order information iteratively achieve faster convergence, but K-FAC and Shampoo have to update their preconditioners infrequently. For example, the average time can be reduced to 1.32x and 2.83x for K-FAC and Shampoo, respec- tively, when increasing the second-order update interval (10 on Cifar- 10 and 50 on ImageNet). We will show the end-to-end training performance in $4.
Memory efficiency. In terms of memory consumption, Eva takes almost the same memory as SGD, and spends much less memory than Shampoo (with 1.09x) and K-FAC (with 1.22x), since Eva only needs to store small vectors, but Shampoo and K-FAC have to store second-order matrices and their inverse results. We notice that K-FAC consumes more GPU memory than Shampoo as K-FAC requires extra memory on intermediate states to estimate KFs, such as unfolding images.
In summary, Eva is a much more efficient second-order optimizer than K-FAC and Shampoo in terms of the average iteration time and memory consumption.
4.3 END-TO-END TRAINING PERFORMANCE
We further compare the end-to-end wall-clock time reach the target accuracy training with differ- ent optimizers. We train VGG-19, 19, ResNet-110, and WRN-28- 10 on Cifar- 10 with one RTX2080Ti GPU (IIGB memory), and ResNet-50 on ImageNet with 32 RTX2080Ti GPUs. We set the second- order information update interval to 10 on Cifar-10 and 50 on ImageNet for K-FAC and Shampoo. More configurations can be found in Appendix C.T
7

---

Published as a conference paper at ICLR 2023
Vel socuccu
Val Alcuraa 8
1 1
Accuracy
18
2x
8
5
50
K-AA3300
GFACB33
20
150
200
1:0
200
300
Wall clock Time (minutes)
Wall clock Time (minutes)
Wall dock Time (minutes)
Wall clock Time (minntes)
(a) VGG-19, Cifar-10 10
(b) ResNet-- 10, Cifar-10 (c) WRN-28-10. 10, Cifar- 10 (d) ResNet-50, ImageNet
Figure 3: Wall-clock time comparison between Eva and SGD/K FAC/SSapoo algorithms for train- ing multiple DNNs. The inset plot reports relative time-to-solution over Eva.
The results are given in Fig. 3 showing that Eva optimizes faster than all the other evaluated al- gorithms SGD, K-FAC, and Shampoo. On the Cifar-10 dataset, Eva is 1.88x, 1.26x, and 2. x faster on average than SGD, K-FAC, and Shampoo, respectively. This is because Eva requires less training epochs than SGD for convergence, and the iteration time of Eva smaller than K-FAC and Shampoo. We notice that K-FAC is also possible to optimize faster than SGD, but it needs to increase the second-order update interval. Otherwise, the training time of K-FAC would be much more expensive as studied in Fig. 8 in Appendix
Though increasing the interval makes K-FAC and Shampoo computationally efficient, they require more GPU memory than Eva. For fair comparison, on the large-scale ImageNet dataset, we set per-GPU batch size to 96 for SGD and Eva, and 64 for K-FAC and Shampoo, to maximize the GPU utilization. In this setting, the results shown in Fig. d) indicates that Eva achieves 1.74x, 1.16x, 1.86x speedups over SGD, K-FAC, and Shampoo, respectively, to achieve the target accuracy of 75.9% on the validation set according to MLPerf.. The median test accuracy of the final 5 epochs on ImageNet is 76.02%, 76.25% 76.06%, and 75.96% for SGD, Shampoo. K-FAC, and Eva, with 100, 60, 55, and 55 epochs, respectively in training ResNet-50. The throughput improvement of Eva by using larger per-GPU batch size is studied in Table in Appendix C
Vel. Aroncyy
80
SGD
KFF
GEAC
K-FAC
KFAC
86
Eva
Evo
Era
0.004
0.04
0.4
256
512
1024
2046
0.303
0.88
0.3
0.05
0:25
0.85
635
(a) Learning rate
(b) Batch size
(c) Damping
(d) Running average
Figure 4:
study of Eva by training ResNet- 110 on C0 with 100 epochs.
4.4 HYPER-PARAMETER AND ABLATION STUDY
We study the convergence performance of Eva with different hyper-parameters, including learning rate, batch size, damping, and running average. The results are shown in Fig. Din training ResNet- 110 on Cifar-10, and the similar results of training VGG-19 on Cifar- 100 are reported in Fig. P in Appendix C First, the learning rate and batch size are two key Eva, as shown in Fig. aa)) and (b), but Eva and K-FAC can consistently outperform SGD under different settings, and largely outperform SGD in the large batch size regime. However, a large learning rate would cause performance degradation in SGD and Eva, but Eva is more robust the learning rate and batch size than SGD. Second, we tune the damping and running average hyper- parameters introduced in K-FAC and Eva, as shown in Fig. ( and (d). K-FAC and Eva perform very closely and they are both robust to damping and running average values, which implies that one can use default hyper-parameters to achieve good performance. Therefore, the used in SGD and some dedicated typer-parameten used in K-FAC can be directly applied in Eva to achieve good convergence performance instead of tuning them. Finally, we conduct the ablation study on Eva to validate the necessity of using momentum, KL clip (details in Appendix (B22). and KVs. We compare the performance of Eva to three variants without using momentum, KL clip (see Eq. 45), and KVs (see Eq. 38), respectively. The results are demonstrated in Table [5] which show that discarding any part of them would cause performance degradation. Specifically. momentum can improve the generalization of Eva (similar to SGD with momentum). KI. clip is of importance
to both SGD and
8

---

Published as a conference paper at ICLR 2023
Table 5: Ablation study Eva without using momentum, KI. clip, and KVs, respectively.
Dataset
Model
Eva
w/o momentum
w/o KI. clip
w/o KVs
Cifar-10
ResNet-I10
93.86
89.3
90.9540.5
92.62+0.2
Cifar-10
WRN-28-10'
97.03
40.5
67. 67.27
96.67
Cifar-100
VGG-19
72.04
66.00+0.3
60.93+2.3
66.6440.8
to prevent exploding the preconditioned gradients (similar to K-FAC), otherwise, Eva w/ KL clip could cause divergence as shown in training WRN-28-10 on Cifar-10. addition, KVs are required, rather than gradient norm, to construct useful curvature information that helps optimization.
4.5 LIMITATION AND FUTURE WORK
Though we have demonstrated the good performance of our proposed Eva, we would like to dis- cuss several limitations and possible future work: (1) there lacks a solid theoretical analysis on the convergence rate for both Eva and K-FAC with empirical FIM; (2) since most training tricks were initially proposed for first-order algorithms, is of interest to design novel second-order friendly strategies for achieving possibly better performance; (3) we will conduct more experiments to show the effectiveness of Eva on other applications like NLP; (4) our prototype implementation of Eva currently only supports data parallclism for distributed training, which can be further integrated with model parallelism (Huang et al. 2019 Shoeybi et al. 2019) for training very large models.
5 RELATED WORK
Matrix-free methods do not explicitly construct second-order matrix, but they rely on the matrix- vector products to calculate the preconditioned gradients. The very initial work in this line is the Hessian-free method (Martens et al. 2010).. which requires only Hessian-vector products with an iterative conjugate gradient (CG) approach. To reduce the cost of each CG iteration, subsampled mini-batch can be used for Hessian-vector products (Erdogdu & Montanari 2015). Recently, M- FAC (Frantar et al.) (021) is proposed to estimate invsrae vector products with a recursive Woodbury-Shermant
-or-Sheeean--oriioon Morrison formula (Amari) 1998). However, matrix-free free methods forgo second- order matrix at a cost of either performing expensive CG iterations or storing extra sliding gradients.
Approximation methods, on the other hand, construct smaller second-orde matrix with different approximation techniques such as quasi-Newton (Goldfarb et al. 2020),, quantization (Alimisis et al.,
2021 Hessian diagonal (Yao et al.) 2021). and the most relevant one K-FAC (Martens & Grosse, Grosse & Martens 2016), among which K-FAC is a relatively practical one for deep learn- ing. With K-FAC, one only needs to construct and invert KFs for preconditioning, which is much more efficient than inverting large FIM directly. However, inverting and storing KFs are not cheap enough compared with SGD, thus, many recent works attempt accelerate K-FAC with distributed training (Osawa et al.) 2019: 2020: Ueno et al., 22000] Pauloski et al.) 2020; 2021 Shi et al.) 2021 Zhang et al. |022 2023). Besides, full-matrix adaptive methods such Shampoo (Gupta et al., 2018 Anil et al.) 2021) are very similar to K-FAC, which construct second-order preconditioners for gradient tensors on each dimension (coneretely, a variant of full-matrix Adagrad). However, the required inverse p-th root computations are more expensive than inverting KFs. These expensive memory and compute costs make them even slower than first-order SGD.
2015
CONCLUSION
We proposed an efficient second-order algorithm called Eva to accelerate DNN training. We first proposed to use the Kronecker factorization of two small vectors to construct second-order informa- tion, which significantly reduces memory consumption. Then we derived a computational-friet frienfly update formula without explicitly calculating the inverse of the second-order matrix using the Sher- man-Morrisoo formula, which reduces the per-iteration computing time. We also provided a the- oretical interpretation to Eva with trust-region optimization perspective. Extensive experiments were conducted to validate its effectiveness and efficiency. and the results show that Eva outper- forms existing popular first-order and second-order algorithms multiple models and datasets.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
The research was supported in part by a RGC RIF grant under the contract R6021-20, and RGC GRF grants under the contracts 16209120 and 16200221.
REPRODUCIBILITY
We give precise statements for our algorithm in Section and provide all the necessary implemen- tation details in Appendix B.2 We have specified the settings and haper-parameters required to algorithm atop PyTorch framework and provide easy-to-use APIs so that users can adopt it by adding several lines of code in their training scripts. The code is available htttps:: /tip:////gthubbb. com/ zhangbv/eva
reproduce experimental results in Appendix C.1 We implement
REFERENCES
Foivos Alimisis, Peter Davies, and Dan Alistarh. Communication-effiienn distributed optimization with quantized preconditioners. In Proceedings of the 38th International Conference Ma- chine Learning, ICMI,, volume 139 of Proceedings of Machine Learning Research, pp. 196-206. PMLR, 2021.
Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):251- 276, 1998.
Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning, 2021.
Hilal Asi and John C. Duchi. Stochastic (approximate) proximal point methods: Convergence, optimality, and adaptivity. SIAM Optim., 29:2257-2290, 2019.
Jimmy Ba, Roger Grosse, and James Martens. Distributed second-order optimization using krontcken-factored approximations. International Conference on Learning Representations, 2017.
Juhan Bac, Paul Vicol, Jeff Z. HaoChen, and Roger B. Grosse. Amortized proximal optimization. ArXiv, foss/203.000899 2022.
Frederik Benzing. Gradient descent on neurons and its link to approximate second-order optimiza- tion. In ICML, 2022.
Léon Bottou, Frank Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. Siam Review, 60(2):223-31 2213,1 2018.
Léon Bottou et al. Online learning and stochastic approximations. On-line learning neural networks, 17(9):142, 1998.
Kumar Chellapilla, Sidd Puri, and Patrice Simard. High performance convolutional neural networks for document processing. Tenth International Workshop on Frontiers Handwriting Recogni- tion. Suvisoft, 2006.
Mengyun Chen, Kaixin Gao, Xiaolei Liu, Zidong Wang, Ningxi Ni, Qian Zhang, Lei Chen, Chao Ding, Zhenghai Huang, Min Wang, et al. THOR, trace-based hardware drieeedriennnn layer-oriented nat- ural gradient descent computation. Proceedings of the AAAI Conference on Artificial Intelligence, 35(8):7046-7054 2021.
Ekin Dogus Cubuk, Barret Zoph, Dandelion Mané, Vijay Vasudevan, and Quoc V. Le. Autoaugment: Learning augmentation strategies from data. 2019 IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 113-123, 2019.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scall scale hi- erarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-255 leee, 2009.
10

---

Published as a conference paper at ICLR 2023
John E. Dennis and Bobby Schnabel. Numerical methods for unconstrained optimization and non- linear equations. Prentice Hall series in computational mathematics, 1983.
Jacob Devlin, Ming-War Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre- Pre-training training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chupter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Pupers), pp. 4171-4188 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR. OpenReview.net, 2021.
John C. Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. In J. Mach. Learn, Res., 2010.
Murat A. Erdogdu and Andrea Montanari. Convergence rates of sub-sampled newton methods. In NIPS, 2015.
Elias Frantar, Eldar Kurtic, and Dan Alistarh. M-fac: Efficient matrix-free approximations of second-order information. In Advances in Neural Information Processing Systems, volume 34, pp. 14873-14886, 2021.
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approx- imate natural gradient descent in kronecker factored eigenbasis. Advances Neural Information Processing Systems, 31:9550-95566 2018.
Donald Goldfarb, Yi Ren, and Achraf Bahamou. Practical quasi-newton methods for training deep neural networks. In Advances in Neural Information Processing Systems, 2020.
Roger Grosse and James Martens. A kronecker-facttonn approximate fisher matrix for convolution layers. In International Conference on Machine Learning, pp. 573-582, 2016.
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor opti- mization. In ICML, 2018
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 770-778, 2016.
Elad Hoffer, Itay Hubara, and Daniel Soudry. Train longer, generalize better: closing the generaliza- tion gap in large batch training of neural networks. In Advances in Neural Information Processing Systems, pp. 1731-1741, 2017.
Yanping Huang, Youlong Cheng. Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoe V Le, Yonghui Wu, et al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32: 103-112, 2019.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd Inter- national Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Abdoulaye Koroko, Ani Anciaux- Sedrakkin, Ibtihel Ben Gharbia, Valérie Garès, Mounir Haddou, and Quang-Huy Tran. Efficient approximations of the fisher matrix in neural networks using kronecker product singular value decomposition. ArXiv, abs/2201.10088. 2022.
Aravind Krishnamoorthy and Deepak Menon. Matrix inversion using cholesky decomposition. 2013 Signal Processing: Algorithms, Architectures, Arrangements, and Applications (SPA), pp. 70-72,
2013.
Alex Krizhevsky, Learning multiple layers of features from tiny images. Citeseer, 2009.

---

Published as a conference paper at ICLR 2023
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep convo- lutional neural networks. Communications of the ACM, 60:84 90, 2012.
Guan-Horn Liu, Tianrong Chen, and Evangelos Theodorou. Second-order neural ODE optimizer. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
Ilya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with warm restarts. In 5th International Conference Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR, 2019.
Linjian Ma, Gabe Montague, Jiayu Ye, Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael Mahoney. Inefficiency of k-fac for large batch size training. Proceedings the AAAI Conference on Artificial Intelligence, 34(04))5053-5000 5060, 2020.
James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factoree approximate curvature. In International conference machine learning, pp. 2408-2417, 2015.
James Martens, Jimmy Ba, and Matt Johnson. Kronecker-factoree curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018.
James Martens et al. Deep learning via hessian-free optimization. ICML, volume 27, pp. 735-742, 2010.
Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training. ArXiv, abs//881.06612, 2018.
MLPerf. Miperf training v1.1 results.
https://m www.ono.oog//eeo/ooooooooooooe// Ellerf-trraiinggvi11 2021. Accessed: 2022-09-15.
Jorge J Moré. The levenberg-marquardd algorithm: implementation and theory. In Numerical anal- ysis, pp. 105-116. Springer, 1978.
K. Osawa, Y. Tsuji. Y. Ueno, A. Naruse, Foo, and R. Yokota. Scalable and practical natural gradi- ent for large-scale deep learning, IEEE Transactions on Pattern Analysis & Machine Intelligence, 2020.
Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. Large-scale distributed second-order optimization using kronecker-factee approximate curva- ture for deep convolutional neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 13359-12367,2 2019.
J Gregory Pauloski, Zhao Zhang, Lei Huang, Weijia Xu, and Ian T Foster. Convolutional neural network training with distributed K-FAC. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-14, 2020.
J. Gregory Pauloski, Qi Huang, Lei Huang, Shivaram Venkataraman, Kyle Chard, Ian Foster, and Zhao Zhang. Kaisa: An adaptive second-order optimizer framework for deep neural networks. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, 2021.
Yi Ren and Donald Goldfarb. Tensor normal training for deep learning models. In Advances in Neural Information Processing Systems, pp. 26040-26052, 2021.
Alexander Sergeev and Mike Del Balso. Horovod: fast and easy distributed deep learning in tensor- flow. arXiv preprint arXiv: MrXiv:/8000 /00::05799099 2018.
Jack Sherman and Winifred Morrison. Adjustment of an inverse matrix corresponding to a change in one element of given matrix. The Annals of Mathematical Statistics, (1(1):124-12) 1950.
Shaohuai Shi, Lin Zhang, and Bo Li. Accelerating distributed k-fac with smart parallelism of com- puting and communication tasks. In 2021 IEEE 41st International Conference on Distributed Computing Systems (ICDCS). IEEE, 2021.
12

---

Published as a conference paper at ICLR 2023
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model par- allelism. arXiv preprint arXiv: Stt.:::99900008000 2019.
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale im- age recognition. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Christian Szegedy, Sergey loffe, Vincent Vanhoucke, and Alexander A Alemi. Inception-v4, inception-resnet and the impact of residual connections on learning. In Proc. of The 3 AAAI, 2017.
Mingxing Tan and Quoe V. Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In Proceedings of the 36th International Conference Machine Learning, ICML, volume 97 of Proceedings of Machine Learning Research, pp. 6105-6114. PMLR, 2019.
Zedong Tang, Fenlong Jiang, Maoguo Gong, Hao Li, Yue Wu, Fan Yu, Zidong Wang, and Min Wang. Skfac: Training neural networks with faster fronecker-factored approximate curvature. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp- 13479-13487, 2021.
Yuichiro Ueno, Kazuki Osawa, Yohei Tsuji, Akira Naruse, and Rio Yokota. Rich information is affordable: A systematic performance analysis of second-order optimization using K-FAC. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 2145-2153, 2020.
Minghan Yang, Dong Xu, Zaiwen Wen, Mengyun Chen, and Pengxiang Xu. Sketch-based empirical natural gradient methods for deep learning. J. Sci. Comput., 92:94, 2022.
Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney. Adahessian: An adaptive second order optimizer for machine learning. Proceedings of the AAAI Conference Artificial Intelligence, 35(12):10655-006 2021.
Yang You, Zhao Zhang, Cho-Jui Hsieh, James Demmel, and Kurt Keutzer. ImageNet training in minutes. Proceedings of the 47th International Conference on Parallel Processing, pp. 1-10, 2018.
Ya-Xiang Yuan. Recent advances in trust region algorithms. Mathematical Programming, 151: 249-281, 2015.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Young Joon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. 2019 IEEE/CVF International Conference Computer Vision (ICCV), pp.
2019.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Proceedings of the British Machine Vision Conference, BMVA Press, 2016.
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George E. Dahl, Christopher J. Shallue, and Roger Baker Grosse. Which algorithmic choices matter at which batch sizes? insights from noisy quadratic model. In NeurIPS, 2019.
Lin Zhang, Shaohuai Shi, Wei Wang, and Bo Li. Scalable k-fac training for deep neural networks with distributed preconditioning. IEEE Transactions on Cloud Computing. 2022.
Lin Zhang, Shaohuai Shi, and Bo Li. Accelerating distributed k-fac with efficient collective commu- nication and scheduling. IEEE International Conference on Computer Communications, 2023.
13

---

Published as a conference paper at ICLR 2023
A APPENDIX: BACKGROUND OF SECOND-ORDER OPTIMIZATION
A.1 SECOND-ORDER OPTIMIZATION
Second-order optimization algorithms typically utilize the second-order information matrix C to precondition the first-order gradient (i.e., multiplying the gradient by the inverse of C), and then use the preconditioned gradient to update the model parameter as follows:
(f+1)
w(t) C vec(((( B())))
(15)
where w('), vec (el(w() B(O)) E ER are the model parameter vector, and the first-order gradient vector of loss with respect to model parameter, respectively. C' E Raxx is the second-order curva- ture information matrix, and (*) 0 is the learning rate at iteration t. The mini-batch of data B(() is sampled randomly from the training dataset. The matrix-vector multiplication of C-I vec (V)) is called preconditioning, and the second-order information matrix includes Hessian, Newus- Fisher information, their damped versions and so on (Bottou et al. 2018). When C' = I, the update formula becomes the first-order SGD algorithm.
A.2 DETAILED ANALYSIS OF K-FAC
K-FAC has been proven to be one of the most powerful second-order optimizers for training deep models (Osawa et al. 2019 2020, Pauloski et al. 2020 202T), and the idea of Kronecker factorized approximation of curvature information is similar to our algorithm. Here we provide detailed analysis of K-FAC (Martens Grosse 2015), which could help us understand Eva provided in the next section (& B.2). We start from introducing FIM, which is the curvature matrix used in K- FAC, and then establish the Kronecker factorized approximation for FIM, and give the final update formula with the damping technique.
FIM. In the natural gradient descent (NGD) algorithm (Amari 1998; Martens & Grosse 2015), the Fisher information matrix (FIM) is used and inverted to precondition the gradient. The update formula of NGD is:
( w(() Fve B())),
(16)
where F E Rxx" is the FIM, and velvec(VE)) is the natural gradient, which is known to be the steepest direction in the space of model distributions (Amari 1998). The true FIM is derived from
Fime Eg-ppate (/-r(- [vec( (rrecc((/i) (V)vec(Ve)""
(17)
where VL= V log (() is the pseudo- grandent w.r.t. model parameter the expectation E is taken from sampling the input x from data distribution Pduta. and its pseudo-label ğ from the model distribution r(-(z)). As the model distribution is different from the data distribution, the pseudo- gradient is different from the gradient used in SGD. In practice, the one-sample Monte-Carlo (1mc) approach (Martens & Grosse 20155) is used generate a pseudo-label ý for each input from the model output distribution r(-(æ), and then the seeuo-graddeent is calculated via back back-propagation (BP) w.r.t. the loss between model outputs and pseudo-labels, To avoid any confusion, the true FIM is usually denoted as Fine-
However, in many existing K-FAC implementations, is common practice to construct the empiri- cal FIM (Osawa et al. 2019 2020 Ueno et al. 2020 Pauloski et al. 2020 (021) to avoid additional pack-propagation computations with pseudo labels for efficient training. The empirical FIM can be represented as
(18) where V V log r(y(z) is the the log- likeliiid gradient of loss between the input x and its label (sampled from the data distribution). As this gradient is exactly the same one in SGD, F Femp can be efficiently constructed during the same feed-forward and back-propagation process of SGD. In this work, focus the analysis of the empirical FIM as it is more efficient than the true FIM while having the similar convergence performance with the true FIM in practise (Osawa et al. 2019 2020 Ueno et al.] 2020 Pauloski et al. 2020 2021). Thus, throughout the following context, F refers to Femp if not particularly specified.
F Femp = ((z.g))=paan (Ve
K-FAC. Due to the large number of parameters in DNNs, it is impractical to store and invert the whole F. which takes O(n²) memory. Many approximation methods have been proposed to re- duce the quadratic memory cost, such as Matrix-free methods (Martens et al. 2010: Frantar al.
14

---

Published as a conference paper at ICLR 2023
2021
and Kronecker factorization (K-FAC) methods (Martens & Grosse 2015 Grosse & Martens
2016 George et al. 2018). In particular, K-FAC methods provide efficient and effective approxi- mations of F, and have been successfully applied in large-scale distributed DNNs training for faster convergence (Osawa et al. 2019 Ueno et al. 2020 Pauloski et al. 2020 2021).
Specifically, given an L-layer DNN, K-FAC approximates the full F with block-wise diagonal matrix as follows:
F R diag(Fi, Fr), where Pi = (
(19)
where Fi is the FIM block constructed from the gradient information of layer l. For each Fi, K-FAC approximates it as the Kronecker product of two smaller matrices. Take a linear layer as example, where a = (W(a)- 1), Fl can be approximated by
Fi E|vec (Vitveec(e
(20)
= [ ② (ar 2 g:)² = (((ar-1a @ (g/g/))
(21)
2 1aa-aa_s
[ggg A Al-1 & G2.
(22)
The second-line derivation is based on the chain rule of B[V&] = E|gia[11] and the properties of
Kronecker product 8, where V4 is the gradient matrix w.r.t. the parameter matrix W; E Rdixd 1,
a,-1 E 12 is the input vector of layer (i.e., activation of layer 1 - 1), g IR: the gradi- ent vector w.r.t. the pre-activation output of layer 1, and ø is an element-wise non-linear activation function. The third-line approximation is based on the assumption that a and gi are fairly inde- pendent (Martens & Grosse 2015), which gives [B(
E E(gig/") The A E| 1 and Gi E(g/g/ are called Kronecker factors (KFs). We notice that recent work points out K K-FAC may not approximate FIM, but should be linked to gradient descent on neurons (Benzing 2022). We are open to study this new perspective as well in the future work.
(gigi) ]] Efay-1af11
In summary, K-FAC approximates the full FIM diagonally via F' xe diag(I', -... FL), and then use the Kronecker product of two KFss to approximate each FIM block via Fi = A &G, Based on the property of (A @ G))ec(VV = vec A), the preconditioned gradient is computed layer-wisely:
Fe-vec V4) = (47h @ GT vvec (V = vec(G)
(23)
Thus, preconditioning the gradient in K-FAC is simply to left- left-multiply the gradient with G and right-multiply it with A where both KFs are symmetric matrices.
The K-FAC approximation reduces the memory complexity from a quadratic cost storing FIM to a linear cost of storing KFs. Specifically, the total model parameter number is n = (d xd),, while the KF parameter number is (/ + df O(kn), in which & 2 2. In addition, the K- FAC approximation is by no means limited to linear layers, and it has been successfully applied into other types of DNNs, such CNNs, RNNs, and neural ODEs (Grosse & Martens 2016 Martens
let al. 2018 Liu et al. 2021).
Damping technique. To stabilize the second-order optimization, the damping technique (Moré 1978) is commonly applied into the FIM (i.e., Fi + z/) with a damping value > 0, before pre- conditioning. As inverting the large damped FIM directly is very expensive, K-FAC exploits an alternative approximate approach to invert two smaller damped KFs, i.e.,
(Fi + 7T) (Ai1 Gr + (Ai11 + II - s (Gi
(24)
where mi VTTAAIII//VCC T(.). T(-) is the trace of square matrix divided by its dimension, which is derived to minimize the approximation error of inverting damped KFs (Martens & Grosse 2015). To avoid that becomes infinitely large when Gj is nearly zero, restrict n, no larger than 10 in practise. With the damping technique, the K-FAC update formula becomes:
wit++) w/(' (G. + v7 (w (BB()))) MIV
(K-FAC)
where I is an identity matrix. Therefore, K-FAC only needs construct and invert damped KFs for all layers (I = ................... L). The damped KFs are positive definite and they are typically inverted using Cholesky decomposition (Krishnamoorthy & Menon 2013).
15

---

Published as a conference paper at ICLR 2023
A.3 K-FAC VARIANTS
The K-FAC algorithm can potentially improve the training performance by using a much less num- ber of iterations than the SGD counterpart to reach the target accuracy (Osawa et al.) 2019 Ueno et al. 2020). However, the iterative training time of K-FAC is much longer than that of SGD, as needs to construct and invert KFs. which limits the practical usability of K-FAC. Distributed train- ing techniques can also help alleviate the heavy inverse computations with multiple workers (Osawa
al. 2019 2020; Pauloski et al. 2020 2021). As distributed K-FAC introduces new bottleneck of communicating KFs, recent works (Shi et al. [021]] ZZhang et al.] 2022 2023) were proposed to optimize its communication cost to reduce training time. Besides, the eigen-decomposition meth- ods (George et al., 2018 Pauloski et al. 2020 2021) target at providing better FIM approximations, for more stable optimization performance, by eigen-deconposing the KFs. For example, the exact damped preconditioner is given by cigen-decomposing KFFs:
(Fi+t))-1 (A-1)) ( (QA @ QD) @ DE D6+71)-' & Q&), (25)
G1
where A|-1 QADAQ' and G: = QcDaQG are the orthogonal eigen- vigen-decompositions of two KFs. The eigen-decome deeeppeccooooitionn provides the exact preconditioner rather than approximate precondi- tioner with matrix inversion (see Eq. (44). However, eigen- deeendecmoooitionn operations introduce much larger per-iteration computation overheads than matrix inversion operations (Zhang et al. 2022).
In addition, to alleviate the computational bottleneck, SKFAC (Tang et al. 2021 and SENG (Yang) al. 2022) applied the Woodbury formula to invert the Kronecker factors in a smaller dimension of mini-batch size, and KPSVD (Koroko et al., 2022) considered low-rank approximation of the Kronecker factors via expensive singular value decomposition. Unlike Eva, these methods attempt to approximate the low-rank FIM, and they are still compute-inefficiet as they rely either matrix inversion or decomposition operations.
Stale FIM. Existing K-FAC algorithms suffer from the system inefficiency due to the expensive overheads of computing and/or communicating KFs of all layers iteratively. As a compromise, they usually utilize the stale FIM information increase the update interval of KFss to alleviate its inefficiency. However, skipping K-FAC approximations by using the stale statistics could bring potential negative effects on the convergence performance (Ba et al. [017;] Pauloski et 2020 Chen et al. 2021), and increase the complexity of tuning hyper-parameters (Ma et al. 2020))
B APPENDIX: DETAILED DERIVATION OF EVA
We provide the derivations of Eva from the trust-region optimization perspective, and supplement its algorithm and implementation details.
B. TRUST-REGION OPTIMIZATION
The trust-region optimization provides a useful perspective toward understanding many second- order methods. Formally, we have
((+1) = min [(w). s.t., p(w, w (1)) S X,
(26)
(27)
where w!t) is the current model parameter, E the loss function, and is the proximal function to measure the distance between two parameters. Therefore, the objective of trust-region optimization is to find the next model parameter ((+1]) that minimizes the loss function while being close to the current parameter (i.e., in the trust region). It can be well solved by converting into a proximal optimization problem (Bac et al. 2022) as follows:
w (f+1) min (ww + 7 p(w. (())
(28)
Rs mmin((w()) (t) + vec (1)) (w (w--w) (()) + (w - (w0' - w!))
(29)
21
First, the proximal constraint of p(w, w()) < is added into the objective function as a penalty (smaller À means larger penalty). Second, we approximate the left-side loss term with first-order
16

---

Published as a conference paper at ICLR 2023
expansion: f(w) = ((w!") + vec(Ve(((()) ( w(f). and approximate the right-side penalty term with second-order expansion: p(w, w(t)) xe w P(w-w(+) A. where the metric P depends on the choice p.
As the approximated objective function is quadratic to the optimal solution can be obtained as follows:
((+1) = (1) AP veee'veee
(30)
This gives generalized update formula of many second-order methods. For example, NGD uses the true FIM as the metric matrix, i.e., P = Fime, and the proximal constraint with the true FIM is in fact the second-order approximation of the KL-divergence between model output distributions, i.., (ppp))/pww))))))) = (w (v(())TFimell Fim((w-w(())
B.2 EVA
Algorithm. Following the same idea, Eva uses the damped outer-produet of average gradient (OPAG) as the metric, that is
P R+ = Elveec(V + +,I,
(31)
where E[lvec (V))) [ is the OPAG. Put it into the proximal function, we have
p(w, w()) 5 (w- (Ww-wo)TP(ww-w) w() 1 ( [( + -
(32)
where Aw w- (()) is the parameter update. Therefore, the proximal function consists of two terms: the first one constrains the parameter update along the gradient direction, and the second one limits the parameter update size. In other words, is equivalent to combining the strip trust region and ball trust region, as shown in Cc).
Like K-FAC, reduce memory and computation costs, Eva approximates the OPAG with two smaller vectors at each DNN layer. Specifically, given a linear layer = ((Wai-1), the OPAG is approximated as follows:
R₂ EEvvec ([)eec(V = E(ay-1 8 gi
g E(a 8 gi IT.
(33)
(E|a--1) E Elg:)) E
E
(ā:-1 a Oz)
(34)
Similarly, the derivations are based on the chain rule of 44 E alg
1. and the assumption that a and g: are fairly independent (Martens & Grosse 2015 a E[a--]] and g: EEg// are
called Kronecker vectors (KV)), Since R3 viVí where vi = â 2 gr. is a rank-one matrix, the inverse of the damped R2 can be computed efficiently via the Morrion formula (Sherman & Morrison 1950). The Sherman Moorriso formula is
A ' 2 A U
(A + u A-1
(35)
u '
where is an invertible square matrix, u1 and u2 are column vectors. Taking 7II and u1 u2 vi into the formula, we can invert the damped Ri efficiently as follows:
(Ri + - (7) v.v/) - viv?) v/
(a)- a
(Big)
(36)
(af a
Bi) +
Therefore, the preconditioned gradient with damped OPAG in the vector form given by
RRi + -1) - (vec)))) (af ( gi) +
g g'V vec(giaf
(37)
Thus, we can derive the update formula Eva as:
w(((+1) =w())
(1) [14
g V V
8(a
(Eva)
7 (a 1 g g +
17

---

Published as a conference paper at ICLR 2023
Relation to SGD. Assume that and g are fully independent, we have E[vec(VE)) = E
E/a] & Elg). This means one can replace the Kronecker product of KVs by the gradient in Eva's update formula. That is, we can derive simple update formula without using KVs as follows:
w
w/O -
(t)
|vell (IV6 V() w/ (0)
V4
(VA
(38)
This shows that Eva (w/o KVs) uses only gradient norm precondition the gradient in each DNN layer, behaving like a layer-wise adaptive SGD algorithm (You et al. 2018). However, the assump- tion that and g are completely independent is unrealistic, because oi a is rank-one matrix while the expectation of gradient generally not (i.e., E[V&] A g a Therefore, performing Eq. [8 directly will lose second-order information and affect the convergence performance com- pared the Eva algorithm (see Table Nevertheless, the Kronecker factorization of K are able to capture the structure information that helps optimization in deep learning training.
Vss estimation for different types of layers. The key component of Eva is to construct KVs used to precondition the gradient. Due to different structure of different types of layers, how to construct KVs is different. First of all, we consider how to estimate KVss for linear and convolutional layers.
Linear layer without bias: it is straightforward to estimate the KVs during feed-forware (FF) and pockppropagttion (BP) processes with a mini-bateh batch of intermediate states by
1 B
a = F2 al
a
and gi EEgi
8
(39)
B
EEB
BEB
where (1) a and gi input and pre-activation gradient vectors at layer I, sampled from the mini- batch of data B. In practice, the mini-batch size used to estimate KVs can be subset of the current
batch. We find that mini-batch size of 16 works well to estimate KVs.
Linear layer with bias: we can formulate it as matrix-vector multiplication:
S; Waa b, (Wi bi) a/-1
(40)
where Wi and by are weight and bias parameters, respectively. Therefore, to precondition the con- catenated gradient of weight and bias, should estimate â over the input activation associated with one.
Convolutional layer: need convert the convolution operation into linear transformation to timate KVs. Assume that the input size is (Ciny Hin, Win), and the output size is (Cout; Hout, Wout) in conv layer, where C' denotes a number of channels, H is height in pixels, W is width in pixels, the convolution operation is defined as follows:
out Cour]
weight[Cout][es [cin] input(csn). Cout = 1,... . Cout:
(41)
where is the 2D cross-correlation operator. There totally Cin Cout cross-correlation opera- tions, and each * operation involves parameters with kernel size K x Kw. By unfolding the input into a matrix X, one can perform convolution linear transformation (Chellapilla et al.) (0066):
out = reshape(XX
(42)
where X is the unfolding input matrix of shape Hoye W out Kt,KH Kw (each row contains all necessary values within the receptive fields), and W is the weight matrix of shape Cin KH Kw x Cout (each column contains all in-channel kernel weights). To construct KVs used to precondition the gradient of weight matrix, we are interested in the output pixel-wise linear transformation:
y: Wx, i 1, Hout Wouts out;
(43)
where x; and y, are the i-th row of the input matrix and output matrix XWT. In other words, we can treat the pixel-wise out-channel computations as linear transformations on Hout W out vectors. Therefore, the KVs for convolutional layer can be efficiently constructed as follows:
aconv
1
mean (i)) (8) and 8comv
mean (1) out: [1, 21,
(44)
B
B
iBB
18

---

Published as a conference paper at ICLR 2023
where ( is the unfolding input sampled from B, and V is the corresponding pre-activation output gradient tensor. The aconv E IR K and Beor E RRCuus take the average estimations over all rows of all inputs, and all pixels of all outputs, respectively. Like linear layers, the calculated KVs are used to precondition the gradient of the weight matrix. If bias is involved, one can simply add one element into acony as discussed before.
Gradient clipping. After the gradient is preconditioned with KVs, the size of preconditioned gra- dient is typically an order of magnitude larger than the size of the original gradients (in terms of L2-norm), which could cause divergence. To ensure the model update is inside the trust region, different gradient clipping strategies can be applied to prevent exploding the preconditioned gradi- ent. As suggested in (Pauloski et al., [2000, one can clip the preconditioned gradient when the KI. size is higher than a threshold, which is measured by the damped OPAG (Martens & Grosse (0155) vec(Gr)" (R + (f)
(Gi) vec where Gi is the preconditioned gradient at layer Thus, one can scale the preconditioned gradients by a factor of
VKL min
(45)
- vec(())
where > is the threshold for KL clipping. As KL clip is coupled with learning rate schedule algorithms, we find it works well for multi-step learning schedule following the practise of (Pauloski et al. 2020
In addition, to avoid any
for clipping (i.c., k), we propose gradient rescaling to keep the preconditioned gradient in the same size of the original gradient, that is, multiplying pre- conditioned gradient by a factor of
VGR
(46)
EL ||
With gradient rescaling, preconditioning will not change the size of the gradient. We apply it in our finetuning experiments (see Table 6 and the results show that Eva can generalize on par with SGD under the same cosine learning rate schedule. As gradient clipping is of importance for precondi- tioned gradient descent algorithms such as K-FAC and Eva, we leave designing effective gradient clipping or scaling strategies our future work.
Implementation. We implement Eva atop PyTorch. Following (Pauloski et al. 2020 2021), we build preconditioner estimate KVs and precondition the gradients before performing the stan- dard SGD optimizer. Our preconditioner supports Linear and Conv2D layers, and parameters at other layers (e.g., BatchNorm22) will be updated by SGD without preconditioning. To construct KVs at supported layers, we register forward pre-hooks and backward-hooks to capture the acti- vations and pre-activation gradients during the feed-forward and pack-propagation computations, respectively. These intermediate values are used to estimate new KVs and update the running aver- age states. Then we use stored KVs to precondition the gradients layer-wisely, and perform gradient clipping on preconditioned gradients, before they are used to update the model parameters in SGD.
To support data parallelism, we implement our Eva preconditioner to communicate KVs at each worker via all-reduce primitives. The communication of KVs is very efficient as the data vol- ume of KVs sublinear the number of gradients, and small KVs can be merged be commu- nicated together via the tensor fusion technique supported by the distributed training framework Horovod (Sergeev & Del Balso 2018). The aggregated KVs are then used to precondition the ag- gregated gradients distributed Eva does not need to assign
all workers. Unlike distributed K-FAC (Osawa et al. 2019 Shi et al. 2021 tasks at different layers into different work- ers, and it also does need to stale FIM to skip the precondition of many iterations (c.g., update KFs every 50 iterations). Instead, memory- and time-efficient to construct KVs and precondition the gradients all workers during the whole training process.
19

---

Published as a conference paper at ICLR 2023
C APPENDIX: SUPPLEMENT OF EXPERIMENTS
C.I EXPERIMENTAL SETTINGS
Testbed. We conduct our experiment on a 32-GPU cluster. consists of nodes connected 10Gb/s Ethernet, and each node has 4 Nvidia RTX2080Ti GPUs connected by two Intel(R) Xeon(R) Gold 6230 CPUs, 512GB memory, and PC1e3.0x 16. We use some common software including PyTorch- 1.10.0, Horovod-0.22 CUDA-10.2, cuDNN-7.6, and NCCL-2.6.4.
Datasets and Models. We conduct our experiments on three commonly used datasets: Cifar- 10 (Krizhevsky 2009), Cifar-100 (Krizhevsky (2009), and ImageNet (Deng et al. 2009). The Cifar- 10/100 has 50,000 training images and 10,000 validation images. The ImageNet has 1133 train- ing images and 50,000 validation images. On Cifar- 10 and C0 datasets, we choose three representative models: VGG-19 (Simonyan & Zisserman, 2015), ResNet-110 11 (He et al. 2016), and WRN-28-10 10 Zagoruyko & Komodakis 2016). On the ImageNet dataset, we select ResNet-S1 (He)
et al.
2016 Inception-v4 SSzegeyy et al. 2017), and ViT-B/16 (Dosovitskiy et al. 2021 models.
We compare our Eva to the first-order baseline SGD, and second-order baseline K- FAC (Martens& & Grosse, 2015). Besides SGD and K-FAC, we select two first-order adaptive gra- dient algorithms Adagrad (Duchi et al.) (2000) and AdamW (Loshchilov & Hutter (209)), and two other second-order algorithms M-FAC (Frantar et al.) 2021 and Shampoo (Anil et al.) (021). We run each algorithm for 3 runs to compute the average metric (e.g., top-1 validation accuracy).
Baselines.
We provide hyper-parameter parameter configurations for reproducibility as below.
In training VGG-19, ResNet-110, 10, WRN-28-10 on Cifar-10 and Cifar-100 with SGD, K- FAC, and Eva, following (Pauloski al. 2020). we the mini-batch size 512, learning rate to 0.4, and weight decay to 5c-4. We apply the multi-step learning rate schedule (a linear warmup at the first 5 epochs and learning rate decays by a factor of 10 at 35%, 75%, and 90% epochs). For K-FAC and Eva, we set damping to 0.03, running average to 0.95, and KL-cllip to 0.001. The second-order update interval of K-FAC is 10.
In training VGG-19, ResNet- 10, WRN-28-10 on Cifar- 10 with 4 more algorithms (Ada- grad, AdamW, Shampoo, and M-FAC), we set the mini-batch size to 512, and tune the learning rate (from 10e-4 to 4) and choose the best one for each algorithm. Specifically, we set the learning rate as 0.4, 0.06, 0.04, and 0.004 for Shampoo. M-FAC, Adagrad, and AdamW, respectively. For AdamW, we set weight decay to 0.05 (5e-4 for others). For M-FAC, we set the number of gradient copies to 32. For Shampoo, we set the second-order update interval to 10, and apply the same multi-step learning rate schedule. We use co- sine learning rate schedule (Loshchilov & Hutter 2017) in other 3 algorithms (Adagrad, AdamW, and M-FAC) for better results.
In time and memory efficiency comparison, we run Shampoo. K-FAC, and Eva on an Nvidia RTX2080TI GPU. The second-order update interval of Shampoo and K-FAC is 10 on Cifar- (and 50 on ImageNet). On Cifar- we set the batch size as 512 for training VGG-19 and ResNet-110, 10, and 256 for WRN-28-10. 10. On ImageNet, we set the batch size as 64, 32, 16 for training ResNet-50, Inception-ve v4, and ViT-B/16, respectively. We measure the average time over 250 iterations, and obtain the GPU memory consumption by calling the "nvidia-smi" command.
In the end-to-end performance comparison, we train VGG-19, ResNet-110, 110, WRN-28-10 on Cifar-10 10 with one GPU. We set learning rate to 0.4, batch size to 512 for VGG-19, 110 and 256 for WRN-28-100 10. We run SGD for 200 epochs, and run Shampoo, K- FAC, and Eva for 100 epochs. We use the same multi-step learning rate schedule as before. For Shampoo and K-FAC, the update interval is 10. For K-FAC and Eva, the damping is 0.03 and the running average is 0.95. The target validation accuracy is 93% for training VGG-19 and R1, and 96.5% for training WRN-28-10.
In end-to-end performance comparison, we train ResNet-50 on ImageNet-Ik Ik with 32 GPUs. To maximize the GPU memory, we set the per-GPU batch size to 64 for Sham- poo and K-FAC, and to 96 for SGD and Eva. We set the learning rate to 0.05 x 32 = 1.6. Following (Pauloski et al., 2020). we train K-FAC and Eva for 55 epochs with learning rate decays at 25, 35, 40, 45, 50 epochs, and we set the damping to 0.001 and the running aver-
20

---

Published as a conference paper at ICLR 2023
age to 0.95. We train Shampoo for 60 epochs with tuned decays at 30, 45, 55 epochs, and SGD for 100 epochs. with tuned decays at 30, 60, 90 epochs. For Shampoo and K-FAC, the update interval is 50. The target validation accuracy is 75.9% (MLPerf).
C.2 ADDITIONAL EXPERIMENTAL RESULTS
Autoencoder:
Following Martens Grosse, 2013 Goldfarb et al., 2020; Ren & Goldfarb 2021 I)
compare the optimization performance of each algorithm by training 8-layer autoencoder (with hidden dimensions of [1000, 500, 250, 30, 250, 500, 1000)) on three datasets: MNIST, FACES, and CURVES. We set batch size to 1000, and each algorithm for 100 epochs with linear decay learning rate schedule. We tune the learning rate in the range of [0.001, 0.5] for different cases. The results are given in Fig. 5 The Fig. 5 shows that second-order methods K-FAC and Eva optimize much faster than SGD in this task, and our Eva can optimize the autoencoder model at the same con- vergence speed as K-FAC. In addition, one can see that Eva performs closely or better than another second-orde method Shampoo, and Shampoo is faster than Adagrad except on CURVES. Note that Shampoo a full-matrix version of Adagrad. We also compare the generalization performance of each algorithm in this task, and the results are very close to those of optimization comparison, that is, Eva performs similarly to K-FAC and outperforms other first-order counterparts on three datasets.
1.0
SGD
SGD
SGID
Train Less
0.8
K-TAC
Adaggad
Train Loss
0.8
KFFCC
Adagrad
Adagrad
Shampoo
Shampon
Shampoo
0.4
Eva
1.8
0.2
0.2
1.4
20
60
100
40
80
100
20
60
100
Epochs
Epochs
Epochs
(a) Train loss on MNIST
(b) Train loss on FACES
(c) Train loss on CURVES
1.0
Test Loss
0.8
0.6
Shampon
SGD
Adagrad
Test Log
0.8
0.6
FAC
SGD
Shampoo
Adagrad
2.6
2.2
Adagrad
Shampoo
K-FAC
0.4
0.4
Test
Eva
0.2
1.4
20
60
80
100
40
80
100
20
40
60
100
Epochs
Epochs
Epochs
(d) Test loss on MNIST
(c) Test loss on FACES
(f) Test loss on CURVES
Figure 5: Optimizing an autoencoder on MNIST/FACCES/CURVSS with different algorithms.
Training dynamics study. To understand the training dynamics of Eva, we provide case study by setting Eva's 's learning rate as & = 7 X @SGDD and then update formula Eva is given by
w(++1) w:(O) aSGD
(1)
V4 B.. ggaf
(47)
VVaa 8
where B1 ()_, g:) +7
(48)
In this case, 3r has played an important role Eva, making its training dynamics different from SGD (which can be viewed as 31 0). Here we optimize the autoencoder following the original setting, except that we set asSD 0.6, O'Fea 0.018, 7 = 0.03, and disable the KL.-clip for Eva.
The optimization performance comparison is given in Fig. O(a), showing that Eva can still outper- form SGD under the condition that &'Ev = 7 X asGD. This implies that Bi gif can help precondition the gradient information. Therefore, we also plot the values of 3r in each layer during the training process in Fig. bb), showing that 3 are very small (< 10-3) 9) but not zeros, and 3 are gradually decreasing. This implies that Eva preconditions more in the early stage, but acts more like SGD in the late stage. Besides, 3 values are adaptive to different layers, e.g., layer 7 seems the most ill-conditioned it requires the largest values for preconditioning.
21

---

Published as a conference paper at ICLR 2023
SGD
8
layer 1
layer
Train Loss
0.6
Eva
Beta Value
6
layer 2
layer
layer 3
layer
0.4
4
layer 4
layer
2
0.2
20
40
100
20
40
60
80
100
Epochs
Epochs
(a) SGD vs. Eva
(b) 9c dynamics
Figure 6: Training dynamics study of Eva optimizing an autoencoder on MNIST.
Convergence results with 50 epochs. To verify that second-order algorithms K-FAC and Eva can learn faster than first-order algorithms such as SGD, especially in the early epochs (Frantar) et al. 2021), we report the training loss and validation error curves (log-scale), in Fig. 7 for training ResNet- 110 on Cifar-10 and VGG-19 on Cifar-100 with epochs. The results show that K-FAC and Eva converge very closely, and they indeed outperform SGD both optimization and general- ization abilities. As studied in Table 2 K-FAC and Eva can outperform SGD under different epoch budgets.
0.8
2.0
SGD
SGD
Train Loss
K-FAC
Val. Error
0.4
K-FAC
0.5
Eva
Eva
0.2
0.1
0.1
10
20
30
40
50
10
20
30
40
50
Epochs
Epochs
(a) ResNet-1 10 on Cifar-10
4.0
0.9
Train Loss
2.0
Val. Error
0.7
K-FAC
Eva
1.0
SGD
0.5
0.4
K-FAC
0.4
Eva
0.2
0.3
10
20
30
40
50
10
20
30
40
50
Epochs
Epochs
(b) VGG-19 on Cifar-100
Convergence performance (log-scale) comparison between Eva and SGD/K-FAC algo- rithms for training ResNet-110 110 on Cifar-10, 10, and VGG-19 on Cifar-100 in compressed 50 epochs.
Figure/ 7:
Pairwise t-test comparison. In addition to Table we perform pairwise t-test with Bonfereonni correction for the accuracy metric among SGD, K-FAC, and Eva algorithms. The results show that Eva performs closely to K-FAC in all cases (i.e., they have no significant difference with p 0.05). K-FAC and Eva significantly outperform SGD with p < 0.05 in many cases (12 out of 18 cases), including all models with 50 epochs, VGG-19 and ResNet- 10 with 100 epochs, and VGG-19 with 200 epochs two datasets. K-FAC and Eva slightly better than SGD in other cases while SGD trained with sufficient epochs and/or extra tricks.
Finetuning pretrained models. Pretraining model on a large dataset and then finetuning the downstream task (i.c., transfer learning) common practice to produce better model accuracy small datasets. To further verify the effectiveness of Eva in this task, choose two repre-
22

---

Published as a conference paper at ICLR 2023
sentative pretrained models, EfficientNet-b0 (Tan & Le 2019) (pretrained ImageNet-1K3 and ViT-B/16 (Dosovitskiy et al. 202I) (pretrained on ImageNet-21k), We load their weights from pub- licly available checkpoints. except the last classification layer, which is randomly initialized. As suggested in (Tan & Le 2019 Dosovitskiy et al. 2021). we scale down initial learning rate to 0.04 and 0.004 for EffiientNet-Hi and ViT-B/16, set the batch size 96, and use the cosine learning rate schedule (Loshchiloy & Hutter 2017). The weight decay is set to 5e-5 for EfficientNet-bO, and O for ViT-B/16, and the input images are resized 224-pixel. For K-FAC and Eva, we damping value to 0.03 and running average to 0.95. Gradient rescaling is also applied. We finetuning the models with 20 epochs. The results are given in Table [ showing that second-order optimization methods such as K-FAC and Eva can generalize as well as SGD finetuning pretrained models, even though they were pretrained with first-order algorithms (SGD for EfficientNet--b and Adam W for ViT-B/16). Specifically, Eva can achieve very competitive performance, i.e., 98.88 % and 92.88% % top-1 accuracy on Cifar- 10 and Cifar- 100, respectively, when finetuning ViT-B/16.
Table 6: Validation accuracy (%) comparison between Eva and SGD/K-FAC algorithms for finetun- ing pretrained models with 20 epochs. EffNet stands for EfficientNet.
Cifar-10
Cifar- 100 K-FAC
Model
SGD
K-FAC
Eva
SGD
Eva
Bo
97.39+0.1
97.37
97.43+0.1 98.88 10.0
85.41
85.32=0.2
85.38+0.0 92.88 10.0
ViT-B/16
98.87
98.87 10.0
92.79
92.68
The effects of update interval. As second-order algorithms such K-FAC have expensive second- order update costs, it is necessary to increase the second-order update interval to achieve faster convergence (Martens Grosse 2015). Here we compare the end-to-end performance of K-FAC with different update intervals.
Val. Accuracy
MMM
Val. Accuracy
80
80
1x
1.5x
2x
1x
3x
5X
60
K-FAC@1
K-FAC@1
K-FAC@S
K-FACB5
K-FAC@10
K-FAC@10
40
K-FAC015
40
K-FAC@15
10
20
30
40
50
20
40
60
Wall-clock Time (minutes)
Wall-clock Time (minutes)
(a) ResNet-110, 100 K-FAC
(b) VGG-19, K-FAC
Val. Accuracy
80
Val. Accuracy
80
MMW
1x
3x 5x
1x 5x 10x
15x
60
Shampoo@1 Shampoog5
Shampoogs Shampoo@ p
40
Shampoo@10
sham poog 10
20
Shampoo@15
20
Shampoo@ 15
50
100
150
200
200 Wall-clock Time (minutes)
50
100
150
Wall-clock Time (minutes)
(c) ResNet-110, Shampoo
(d) VGG-19, Shampoo
Figure 8: Wall-clock time comparison of K-FAC and Shampoo with different update intervals on 10. The inset plot reports relative time-to-solution over Eva.
As shown in Fig. 8 K-FAC@ I takes much longer training time than Eva, i.e., 1.45x and 5.27x than Eva for ResNet- and VGG-19, respectively. Thus, it is often required to increase the second-order update interval to make K-FAC more affordable. For example. the training speed of K-FAC@ 10 (update KFs and their inverses every 10 iterations) is comparable to Eva for training ResNet- 100, and it is 1.58x slower than Eva for training VGG-19 to achieve the target accuracy.
23

---

Published as a conference paper at ICLR 2023
Unlike K-FAC, our Eva converge quickly using an update interval of second-order informa- tion of 1, i.., updating second-order information iteratively, which avoids the efforts of tuning the interval and the danger of performance degradation with stale information.
We also include the effects of update intervals on Shampoo for training ResNet- 10 and VGG-19 on Cifar-10, in Fig. ( and (d), respectively. The results show that increasing the update intervals of Shampoo can accelerate the training process with little performance degradation, but keep increasing the update interval only brings marginal performance improvement.
The effects of batch size on throughput. As Eva is more memory-efficiet than other second- order algorithms, it is of interest to investigate the effects of batch size on system throughput. For that purpose, we train the ResNet-50 model with SGD, K-FAC@ 50, Shampoo 50, and Eva on ImageNet with 32 GPUs. We set the Per-GPU batch size to maximally utilize GPU memory, i.c., 64 for K-FAC 50 and Shampoo 50, and 96 for SGD and Eva (64 for SGD and Eva is also conducted for comparison). The system throughput is given in Table 7 It shows that Eva, K-FAC@ 50, and Shampoo @ 50 can achieve 90%, 86%, and 68% throughput over SGD, when they adopt the same batch size of 64. However, as we scale the batch size of SGD and Eva to 96 (as they are more memory efficient than K-FAC and Shampoo), Eva achieves a much closer throughput (92%) with SGD, while the performance gap between other second- -oder algorithms (K-FAC@ 50 and Shampoo@ 50) and SGD becomes larger.
Table 7: Throughput (samples per second) comparison between Eva and other algorithms
Algorithm
SGD
Shampoo@50
K-FAC@50
Eva
SGD
Eva
Batch Size
64
64
64
64
96
96
Throughput 6420.1
4366.7
5520.2
5801.7
7420.3
6857.1
In addition to Fig. Bjd, we have added Table S to indicate the required number of epochs to achieve the target accuracy of ResNet-50 on ImageNet as follows.
Table 8: Number of epochs achieve the target accuracy of ResNet-50 on ImageNet.
Algorithm
SGD
Shampoo
K-FAC
Eva
Number of Epochs
91
55
46
46
The effects hyper-parameters on Cifar-100. We conduct hyper-parameter study of Eva, includ- ing learning rate, batch size, damping, and running average, by training VGG-19 on Cifar- The results are given in Fig. 9 showing that Eva can consistently outperform SGD under different learn- ing rate and batch size, and Eva performs closely to K-FAC under different learning rate, batch size, damping and rumning average settings. Besides, both SGD and Eva performs poorly with a large learning rate, but Eva performs much better than SGD with large batch size. We leave the study on the effects of large-batch training with second-order optimization methods (Zhang et al., (019) as our future work. As for damping and running average syper-parameters, we found that damping is more sensitive than running average, so that one can tune the damping value to further boost the performance of Eva.
Vel. Aecurnry
70
SGD
K-FAC Ero
KFAC
K-FAC
KFFAC
Eva
-1- Exo
Bro
0.004
2.64
84
512
1624
204h
0.303
0.00
0.3
0.05
0.35
0.36
0.5
(a) Learning Rate
(b) Batch Size
(c) Damping
(d) Running Average
Figure 9: Hyper-parameter study of Eva by training VGG-19 on 100 with 100 epochs.
24


Reviews:
[{"review_id": "wkGO9MeRXHQ", "reviewer": "Reviewer_1mYD", "summary": "", "questions": "", "limitations": "", "rating": 5, "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 presents an optimizer in the K-FAC family. K-FAC approximates the Fisher information matrix as the Kronecker product of two smaller matrices, by assuming independence --- they approximate E((a_{l−1} ⊗ g_l)(a_{l-1} ⊗ g_l)^T] as E[a_{l-1}a_{l-1}^T] ⊗ E[g_l g_l^T ]. The proposed algorithm, Eva, goes one step further, by approximating E[a_{l-1}a_{l-1}^T] as E[a_{l-1}] ⊗ E[a_{l-1}]^T and E[g_l g_l^T] as\nE[g_l]  ⊗ E[g_l]^T. The result is a rank one approximation to the Kronecker factors, and the authors use damping and the Sherman-Morrison-Woodbury formula to compute the preconditioner. Note that Eva approximates the empirical Fisher information matrix, as opposed to the true Fisher recommended by the K-FAC authors.\n\nThe algorithm uses significantly less memory than K-FAC, in fact uses less memory than Adam or Adagrad --- its memory consumption is similar to Adafactor or SM3. The paper has several comparisons on image datasets, showing that Eva is competitive with K-FAC and Shampoo in terms of accuracy, while achieving performance similar to SGD.", "summary_of_the_review": "The paper presents an interesting and fast approximation to K-FAC. I would like to see some theoretical justification for the approximation, and also some experimental results as detailed above for a higher score.", "strength_and_weaknesses": "Strengths: The paper is clearly written, and the proposed algorithm is very simple and easy to implement. The update rule is very fast to compute, and the algorithm uses very little memory, so is very fast. The experiments for CIFAR-10 are quite comprehensive. \n\nWeaknesses: The approximations made in the algorithm are not motivated in the paper. Why is it acceptable to approximate E[a_{l-1}a_{l-1}^T] as E[a_{l-1}] ⊗ E[a_{l-1}]^T? K-FAC's approximations are justified by independence assumptions, but Eva's further approximations are not justified anywhere. There is no convergence analysis to show that this update is mathematically well grounded.\n\nFor Imagenet, please include a table showing the number of epochs needed to achieve target accuracy. Also for K-FAC and Shampoo, please include results of increasing the update interval further, since their authors say that larger intervals do not degrade accuracy very much. Also for the autoencoders, please include train error and test accuracy for the 3 common datasets --- MNIST, Faces and Curves.\n\nIt is strange that you claim that \"K-FAC and Shampoo are not time-and-memory efficient,  which severely hinders their feasibility in practise (sic)\" --- these optimizers are used to train some of the largest networks today, with billions of parameters.\n\nK-FAC and Eva are both presented for fully connected and convolutional networks --- does Eva extend easily to other networks such as transformers, and to networks with features like batch norm or layer norm?", "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": "Not applicable"}, "extra_text": {}}, {"review_id": "_p5ImLJQSU", "reviewer": "Reviewer_p2Vi", "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 describes what I would call a (simple but clever) evolution of K-FAC style fisher information matrix approximation (FIM) by approximating the FIM with  kronecker products of activations and gradients (i.e. an outer product over vectors) where K-FAC uses a kronecker product of *matrices* derived from the activation and gradient. While this is an additional step of approximation, it appears to work well in practice and the vector parametrization of the FIM allows efficient inversion via the Sherman-Morrison formula, leading to sizeable speedups, evaluated on real world architectures.", "summary_of_the_review": "I think this is a strong paper that has the potential to help bring 2nd order optimisation to the masses, especially for industrial use cases where it is not about individual percentages on benchmarks but about tradeoffs, and the 50 epoch accuracy of Eva seems comparable to the 200 epoch accuracy of SGD, meaing it actually achieves real world speedups (especially for WRN).", "strength_and_weaknesses": "Strengths:\n\n- the explanation of related work, background etc. is well done, I think this paper can serve as a good entrypoint for readers\n- the method is conceptually simple, but clever. I think it is non-obvious that the OPAG approximation would work as well as the K-FAC approximation a priori, and the exploitation of Sherman-Morrison and the rescaling to gradient norm in order to make the method drop-in are nice touches\n- the experimental evaluation is top notch, using the related work and real world datasets and architectures\n- the limitations section is an actual, honest self assessment\n\nWeaknesses\n\n- the authors state most weaknessses in the related work\n- truly a nitpick, but statistical significance tests between results using bonferonni correction would have been a cherry on top", "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": "K1Zbavv66Bl", "reviewer": "Reviewer_2rcV", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "In this paper, a novel second-order technique is proposed for deep neural networks training. ", "summary_of_the_review": "Overall, it is a good try for a cheap second-order method for deep neural network training. I overall give a positive score and if the author could give convincing discussion on the currently weak points, I would like to rise my score.  ", "strength_and_weaknesses": "strength\n+ it is a good try to further reduce the computational and storage burden for second-order method for training DNNs.\n\n+ the method is not complicated and has shown great advantage in numerical experiments\n\nweaknesses\n-  there are two modifications. One is to use a mini-batch and the other is to use Sherman-Morrison formula. It is not clear how they two can coordinate each other and to achieve very good performance, while only one of them does not work.\n\n- the authors have conducted experiments on ImageNet, however, for the most important comparison on test accuracy, the results on ImageNet are missing.\n\n- one question is about the learning rate. I saw still Eva needs a delicate rate schedule but in theory , second-order methods are less sensitive to learning rate. Did Eva demonstrate advantages on this aspect? ", "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": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "-RvgKGH8L-3", "reviewer": "Reviewer_j8t4", "summary": "", "questions": "", "limitations": "", "rating": 6, "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 Eva, a second-order optimization method for training deep neural networks that approximates the Kronecker factor of K-FAC using empirical FIM with a rank one matrix. Averaging batch activations or errors construct the rank one matrix. Experimental results on various computer vision tasks show that Eva achieves similar convergence (number of steps vs. accuracy) as K-FAC, Shampoo, and M-FAC. Since Eva's computational cost per step is considerably smaller than other second-order optimization methods and is close to SGD’s cost, Eva achieves relatively fast training times in wall-clock time.\n", "summary_of_the_review": "Although there are no solid theoretical guarantees, a simple, low-cost method like Eva can achieve the same level of convergence as a high-cost method like K-FAC is a critical data point for future research on efficient training methods in deep learning. However, due to its simplicity, it is unpredictable that Eva is effective in a wide range of tasks. I believe that comparisons and discussions focusing on predictable aspects such as trade-offs between the amount of information (e.g., number of matrix elements, matrix rank to be used for gradient preconditioning) and computational/memory cost would enhance the value of this study.", "strength_and_weaknesses": "Strengthes\n- The paper clearly describes the literature, the differences between the proposed and existing methods, and the experimental details, which are easy to follow.\n- The proposed Eva method is straightforward (in a good sense). It boosts the practicality of second-order optimization in deep learning regarding computational and memory cost and ease of implementation and integration (as long as the fast convergence holds for other tasks). Thanks to the low computational cost, there is no need to use stale second-order information or to tune the second-order information update interval.\n- Experiments are well designed to validate the effectiveness of second-order optimization methods.\n    - (There are still missing points as I will describe in “Weaknesses”, though.) \n    - Several numbers of epochs are considered. A longer training (200 epochs) gives the baseline first-order optimization methods a chance to achieve the same accuracy as the second-order methods. Yet, second-order optimization methods (K-FAC, Eva) show a faster convergence (It would be better if the same comparison (with 50, 100, and 200 epochs) is made for Shampoo and M-FAC, though.)\n    - Hyperparameters specific to second-order optimization are explored. Since damping, second-order information update interval, and running average coefficient largely affect the convergence of mini-batch second-order optimization, a comparison with various values of them is valuable in verifying Eva’s robustness.\n\nWeaknesses\n- As the authors mention, there is no theoretical analysis in this study, and the effectiveness of Eva is supported only by the experimental results. \n    - Eva has shown good experimental performance (as convergent as K-FAC). I appreciate the detailed experiments and analysis in this work. However, there is still no guarantee that good performance will hold for other tasks due to the algorithm's simplicity.\n    - (Here, I am not requesting additional experiments but pointing out that the “no guarantee” is almost inevitable in deep learning, which involves many ‘unpredictable factors.’)\n    - Therefore, I do not think comparing only the computation and memory costs of existing methods and Eva (as in Table1) is fair. The amount of information (e.g., number of matrix elements and matrix rank) that each method uses for preconditioning gradient should be included as a comparison item. Eva’s computational and memory cost is lower than K-FAC simply because it calculates/stores much less amount of information —- it approximates the Kronecker factor matrices with rank-1 matrices (outer products of the Kronecker vectors). I think readers would appreciate it if this work also shows the trade-off between the amount of information and the computational/memory cost (a ‘predictable factor’) to decide which algorithm to use.\n- Comparison between K-FAC and Eva with different mini-batch sizes is missing. \n    - (Here, I do request for additional experiments)\n    - The sum of the outer product of average gradient (OPAG), aka mini-batch empirical FIM, is known to have less second-order information when the mini-batch size is larger (https://www.cs.toronto.edu/~rgrosse/courses/csc2541_2022/readings/L05_normalization.pdf, equation 4) unlike the (per-example) empirical FIM.  \n    - And again, since Eva approximates the Kronecker factor matrices with rank one, it should miss lots of information in large-batch settings (whether and how much the information is beneficial for training convergence and generalization is unpredictable, though).\n    - Therefore, Eva, which uses OPAG or mini-batch empirical FIM, might underperforms K-FAC w/ empirical FIM in large-batch settings so I would appreciate more empirical comparison with different mini-batch sizes.\n    - Regardless of the empirical results, it would be valuable to indicate the amount of information Eva captures/misses clearly.\n- Comparison with several existing studies in the literature are missing.\n    - SKFAC (https://openaccess.thecvf.com/content/CVPR2021/papers/Tang_SKFAC_Training_Neural_Networks_With_Faster_Kronecker-Factored_Approximate_Curvature_CVPR_2021_paper.pdf) and SENG (https://link.springer.com/article/10.1007/s10915-022-01911-x) also apply Sherman-Morrison-Woodbury (SMW) formula to invert the Kronecker factors. Both approaches approximate layer-wise Fisher matrix with low-rank matrix. So one can think of them as higher-rank counterpart of Eva.\n    - KPSVD (https://hal.archives-ouvertes.fr/hal-03541459/document) also consider rank-one approximation of the Kronecker factor. Unlike Eva, it tries to find the optimal rank-one matrices that minimizes the Frobenius norm of the difference between the layer-wise Fisher matrix and the rank-one approximation. Compared to Eva, it might have more computational costs, but it sounds more theoretically-grounded. Again, whether this makes a change is unpredictable in practice, but showing the information-cost trade-off is useful.", "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": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster