MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates
Authors: Mohammad Mozaffari
Sikan Li
Zhao Zhang
Maryam Mehri Dehnavi
Venue: NeurIPS 2023 poster (2023)
Keywords: machine learning
deep learning
optimizers
distributed training;second-order optimization;
TL;DR: MKOR is a distributed second-order optimizer with cheap computatoins and communications and a high convergence rate.

Abstract: This work proposes a Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 updates, called MKOR, that improves the training time and convergence properties of deep neural networks (DNNs). Second-order techniques, while enjoying higher convergence rates vs first-order counterparts, have cubic complexity with respect to either the model size and/or the training batch size. Hence they exhibit poor scalability and performance in transformer models, e.g. large language models (LLMs), because the batch sizes in these models  scale by the attention mechanism sequence length, leading to large model size and batch sizes. MKOR's complexity is quadratic with respect to the model size, alleviating the computation bottlenecks in  second-order methods. Because of their high computation complexity, state-of-the-art implementations of second-order methods can only afford to update the second order information infrequently, and thus do not fully exploit the promise of better convergence from these updates. By reducing the communication complexity of the second-order updates as well as achieving a linear communication complexity, MKOR increases the frequency of second order updates. We also propose a hybrid version of MKOR (called MKOR-H) that mid-training falls backs to a first order optimizer if the second order updates  no longer accelerate convergence.  Our experiments show that MKOR outperforms state -of-the-art first order methods, e.g. the LAMB optimizer, and best implementations of second-order methods, i.e. KAISA/KFAC, up to 2.57x and 1.85x respectively on BERT-Large-Uncased on 64 GPUs.

MKOR: Momentum-Enabled Kronecker-Pactor-Basell Optimizer Using Rank-1 Updates
Mohammad Mozaffari Department of Computer Science University of Toronto
Sikan Li
Texas Advanced Computing Server slistace. utexas.
mmozaffart @cs wwwontoo..dooo.oooo edu
Zhao Zhang
Maryam Mehri Dehnavi
Department Electrical and Computer Engineering Rutgers University
Department of Computer Science
University of Toronto mmehride@cas toronto. edu
zhao rutnggu
Abstract
This work proposes a Momentum -nnable Kronecker-Factor Feea-aaatttrraaaee Optimizer Using Rank-1 Updates, called MKOR, that improves the training time and con- vergence properties of deep neural networks (DNNs). Second-order techniques, while enjoying higher convergence rates vs first-order counterparts, have cubic complexity with respect to either the model size and/or the training batch size. Therefore, they exhibit poor scalability and performance in transformer models, e.g. large language models (LLMs), because the batch sizes in these models scale by the attention mechanism sequence length, leading to large model size and batch sizes. MKOR's 's complexity is quadratic with respect to the model size, alleviating the computation bottlenecks in second-order methods. Because of their high computation complexity, state-of-the-an implementations methods can only afford update the second order information infrequently, and thus do not fully exploit the promise of better convergence from these updates. By reducing the communication complexity of the second-order updates, well as achieving linear communication complexity, MKOR increases the frequency of second-order updates. We also propose a hybrid version of MKOR (called MKOR-H) that mid-training falls backs to a first order optimizer the second order updates longer accelerate convergence. Our experiments show that MKOR outperforms ttate-of-the-us first-order methods, e.g. the LAMB optimizer, and best implementations of second-order methods. i.c. KAISA/KFAC, up 2.57 and 1.85 x respectively on Lart-LLarr--ncssee on 64 GPUs.
second-order
Introduction
Second-order optimization methods have recently gained popularity in the training process of deep neural networks (DNNss due to their higher convergence rate comparison their first-order counterparts. One of the well-known second-order methods, Newton's 's method, uses the inverse of the Hessian of the objective function as preconditioner the gradients, capturing more information on the curvature of the loss function. However, since the size of the Hessian scales with the model size, computing and storing the exact Hessian and its inverse is the main bottleneck in these methods, giving rise to different approximation methods.
One of the common approximations of Newton's 's method is Natural Gradient Descent (NGD) (I] where the Hessian substituted with the Fisher Information Matrix (FIM) to reduce the cost
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

of computing the second-order derivatives. However, as the models become larger, it becomes impractical to store and compute the exact inverse of the FIM, leading the design of algorithms based on block-diagonal approximations of FIM; each block corresponds to layer in the neural network (NN). Inverting the diagonal blocks also not practical due to the large number of parameters in NN layer, thus their inverse is also approximated.
The class of Kronecker-Factored Approximate Curvature (KFAC) [16] methods attempt to reduce the computation costs of the the block inversion. They approximate an FIM block for a batch of samples using the Kronecker multiplication of the covariance of the output activations and input gradients. KFAC is implemented on distributed platforms for both linear and convolutional layers [5)27] 20; 19]. and different computation and communication optimization techniques have been applied to its implementations [15 KAISA [2I is a framework with sante-of-the-asse art distributed KFAC implementation. The computational complexity of KFAC-based methods is O(d)) where d is the layer dimension. These methods work well for small models with small layer dimension d, however, in large models, they don't scale well, resulting in poor performance.
To reduce the effect of model size on the computation complexity of second-order updates, KBFGS [8] which does not have an efficient distributed implementation, uses a Broyden-Fletcher--e Goldfrbr [3]- -based method for computing and updating the Kronecker factor and has Obbd²) complexity, where b is the batch size. The class of SNGD (Sherman Morrison- (Sherman-Morrison- -ooddurr- based NGD) methods (23) [1 uses the the FIM blocks with complexity of Ob²), making the complexity independent of the model size. Hyl.o 117 is the state-of-the-art SNGD implementation. It reduces communication for better scalability. KBFGS and SNGD solve the scalability issue of KFAC-based methods for small batches. However, in transformer models (28), the batch size scales with the sequence length of the attention mechanism (which can be as high as a few thousands (26)) thus limiting the scalability of SNGD and KBFGS methods. Recent work Eva
Shannon (BFGS) [3]
-Morrasn--- Woodbury (SMW) identity to calculate
[3 attempts to reduce this cost further to (dd) by storing vectors instead of Kronecker factors in the KFAC formulation. However, similar to KFAC, Eva uses a damping factor that can lead to additional error in the FIM approximation. Also, because Eva stores the Kronecker vectors instead of factors, it can not leverage the benefits of momentum.
This work presents MKOR, a Momentum- Enabled Rank-1 Updates. (/) MKOR approximates the inverse of covariance matrices using rank-1 updates in the Shermann Based (SM-Based) Matrix Inversion, reducing the inversion computation complexity from (d²) to (dd). As result, the second-order information updates can be applied up to 100 times more frequently compared to KAISA/KFAC and HyLo/SNGD. KFAC computes and inverts the covariance matrices precisely; KFAC updates the second-order information infrequently, e.g. every 100-1000 iterations, due to the high cost of inversion and for time efficiency, which damages its convergence rate and generalization capability. SNGD-based methods suffer from similar overhead due to inversion of their kernel matrices. (2) Second-order methods suffer from high communication costs for synchronizing the inverse of factors between workers. MKOR alleviates this by only synchronizing rank-1 approximation vectors among the workers, reducing the communication costs from (d2) to (d). Also, MKOR uses half- preffireciiion to further reduce communication; other second-order methods cannot use low-precision computations due to their complex matrix inversion algorithms. (3) Second-order methods are more prone to the exploding gradients due to the effects of preconditioning on the norm of the gradients and lack of numerical bounds on the inverse of the factors. MKOR uses a Norm-Based Stabilizer and a Gradient Rescaling Mechanism to detect and prevent exploding gradients. (4) The high convergence rate of second-order methods, including MKOR, stagnates after the first few iterations or epochs of training. We propose a hybrid version of MKOR, namely MKOR-H H. that combines the high convergence rate of second-ordes methods in the initial phase of training with the low overhead of first-order methods in the late stages of the training, using a loss-reduction-ratee ratt--bseedeeed switching mechanism.
Optimizer with
MKOR outperforms state-of- -fate-of-the-art distributed second- and first-order methods by up to 2.57 x, reducing the training time of BERT-Large- Uncased from 8 hours to 3 hours on 64 A100 GPUs. MKOR also achieves new state-of-the-arr metrics on the GLUE dataset, where other second-order methods such as KFAC fail to converge the baseline.
2

---

Activation Covariance
47
A7
1m-1
2 2
t
10
m-
SM-Based Inversion
Gradient Covariance
G G Gmm
ge 9t m
Rank-1 Approximation Synchromization (d)
SM-Based Inversion O
| -11 3
L -1
L -1
AW m
w/
(L,²) -1
V
(R -1
? 11
RTT
[
RSmh
Norn-Based Stabilization
d. Precondition & Rescale Gradients
Figure 1: MKOR for layer m on a single worker. The inputs of MKOR are the activations A the gradients of the loss function with respect the inputs G and the gradients of the loss function with respect to the weights Vwm C. The output is the update values AWW.
Background
Training a neural network involves solving an optimization problem to find the optimal values for of weights W (W"m) m==1- where M is the number of layers the network and W is matrix in Rdd Second-order methods precondition the weights of the network with the inverse of the Hessian for better convergence rates. Block-diagonal approximations of NGD methods replace
the Hessian with the block-diagonal FIM as shown in Equation | where w E & is the vector representation of W Fm block corresponding to that layer and L the loss function. Shows that the FIM matches the Gauss-Newton matrix under certain conditions.
w
2
a: F
V w C
KFAC-based methods reformulate the FIM block as the Kronecker product of two matrices. Equation
shows the update rule in KFAC, where is the loss function and (LTT)-² and (Rm)) are the inverse of the left and right factors, respectively.
mh = Wm - a(L VV C(RT)
(2)
(LI"))' and (RTT) in Equation are computed using equations and respectively, where a" the activation value of sample at layer and Jm V L and y incorporate the momentum feature to avoid extreme changes in the factors.
L L + (1>))Eig"" gt m
(3)
a m IT R RR 1 + - 7Baaa"
(4)
MKOR: Nomentum--nabled Kronecker-Factoriiaaiinnn- Optimizer with Rank-1 Updates
In this section we first present the MKOR algorithm, computation and communication complexity, then present hybrid MKOR (MKOR-H), and finally discuss MKOR's convergence and stability.

---

Algorithm 1: MKOR Algorithm for a Single Layer m
Data: AT-1 G
,
Result: W
m E Second Order Layers then
2
at
f- + - (A,
11 Approx:
A Am-1Am-1" &a
a
m m- a g m
f- 8 Ei= (Gt)) :
// Approx: GT
St
at m f AllReduce a m-1 ma St
Synchronize Aporoximations
// Norm-Based Stabilization
i -1 +- if |
then (4
-1 + (1 - 1 clse L -1
R if R > then (R
-1
-1 + (1 - {)I else R
// SM-Based Factor Inversion (1-7)
I.m-1 <- -
ml T L
t-1
Bt
7- 788 1-1 mm
Rm--1 = R -1
(1-7)
4-1 at
a
m Ri-11
211 1aa R
AWT L - Tnn-'wwe
Wm CR.
11 Precondition Gradients // Rescale Gradients
AW & [/r AW!:
10
AW?"
11 else
AW WwmL;
13
end
W + Opwimizer.steppeeeeooeto AW w
3.1 The MKOR Algorithm
Algorithm D summarizes the MKOR optimizer for single layer and Figure I shows the workflow. For each layer (line in Algorithm IT MKOR updates the second-order information and preconditions the gradients, and at the end the backend optimizer updates the weight using the preconditioned gradients (line in Algorithm [
Rank-I Approximation. For the rank-1 approximations of the covariance matrices, we use the average of the values across all the samples, i.e. at am E [[a">+1] and gt Elg,") (lines and in Algo- rithm and Figure a). (A," (Ar=-))_; and (G!" show the 5th column of (Am--I)- and (GTT)-) respectively, where Am- and G are the activations and the gradients of layer mn respectively.
Norm-Basee Stabilizer: The values in the factor inverses in second-order methods can become large or vanish due to extremely large or small values in activations and gradients, leading to numerical instabilities and over/underflows. Since the inverse of the factors are directly multiplied by the gradients to find the update values, can cause oscillations or even divergence. MKOR uses norm- based stabilizer to detect this and addresses it by modifying the inverse of the factors accordingly (lines and 6 in Algorithm and Figure bb). More detail on the norm- corm-based stabilizer are in Section3.33
SM-Based Inverter. MKOR directly modifies the inverse of the left and right factors using rank-1 updates, while using the momentum for better convergence. If "" is approximated using a rank-1 matrix gm and using the Sherman-Morrison identity, Equation 5 in obtained (line 7 in Algorithm and Figure
(1-2) m 72 +1 7(1 7g
-m-1- -- -1
m -1
gt
mmmm
m. -1
L
(5)
L
gt
Furthermore, if Equation is approximated using E a a tion, Equation 6 is obtained (line 8 in Algorithm and Figure
2 a
am
m² with similar deriva-
c).

---

Table 1: The computation and communication complexity and memory overhead of the state-of- tre-art inplementations of the first- and second-order optimizers. The division by 2 in MKOR is because MKOR uses half- half-precision computations. The complexity of KFAC-based methods depends on layer dimensions while SNGD methods mostly depend on the batch size. In transformers, due to the scaling of the batch size by the sequence length, batch sizes and layer dimensions are comparable, making both KFAC- and SNGD-basee methods more expensive than SGD. Optimizer
Computational Complexity
Memory Overhead
Communication Complexity
MKOR
O(d + bd)
O (2)
(2d/22)
SNGD (HyLo)
O(b")
O + bz)
O + b²)
KFAC (KAISA)
O(d))
(4dd))
(4d)))
Eva
O(d' + bdd)
((2d)
((d))
SGD (Momentum)
((d=)
(dd))
(1-7) m
R Rm-1== R 72 + - 7)am" 1 -1 RT-1
-1 T -1 aa ac R
amma
(6)
Rescaling Gradients, Preconditioning the gradients using the computed factors can change gradient norms. Sometimes, these changes interfere with the effect of the learning rate on the training process. To alleviate this and to make learning rate schedulers more effective, the preconditioned gradients are scaled so that their norm matches the original norms (line 10 in Algorithm I and Figure (dd).
Complexity Analysis. MKOR reduces the memory, communication, and computation costs for factor inversion. Table I compares the overheads of different optimizers. (1) Computation Complexity. MKOR inverts the left and right factors in Equation 2 using equations [5) and 66 both of which can be computed using matrix-vector multiplications, and have Odd) computation complexity, in contrast to KFAC and SNGD methods that need (dd) and O(²)) complexity to invert matrices in Raxd and bb respectively. (2) Communication Complexity. The only data that is synchronized among different workers in MKOR is the two rank-1 1 approximations that have 2d elements. With quantization, this size can be halved. In KFAC, the activation and gradient covariance matrices and the inversion of left and right factors need to be synchronized between all the workers, leading to 4d2 data transfers. In SNGD, the activations and gradients are synchronized, leading to 2bd data transfers and the inverted kernels are broadcast, resulting in b2 data transfers. Reducing the communication complexity of MKOR from quadratic to linear results in better performance on large number of workers. (3) Memory Overhead. MKOR needs to store the inverse of the left and right factors and two rank-1 approximation vectors, leading to 2d + 2d memory overhead, and using half-precision computations further reduces this. KFAC stores the activation and gradient covariance matrices and the left and right factors, leading to 4d memory overhead. SNGD stores the activations, the gradients, and the kernels they use as second-order information, leading 26d + 6² memory complexity.
The low memory overhead of MKOR allows us to use larger batch sizes compared to other second- order methods. In practice, gradient accumulation methods are used to increase the effective batch size in training, which reduces the training speed significantly. This issue worsens with KAISA and HyLo, but MKOR alleviates this. For faimess, in our experiments, we set the local batch sizes to the same value in all optimizers and do not leverage this feature of MKOR.
3.2 Hybrid MKOR
We observed that second-order methods, including MKOR, usually accelerate training more during the first iterations of the training time, and as the loss flattens, their advantage over their first-order counterparts becomes less noticeable. This is because the second-ords -oderr information of the loss functions approach identity near convergence points. Thus we designed hybrid second- and first- order optimizer with loss decrease rate-based switching method (MKOR-H) -H). MKOR-H evaluates the changes in the loss function in different iterations and switches back to first-order methods if needed for an efficient trade-off between the costly second-order updates and their benefits for convergence.
5

---

3.3 MKOR Convergence and Stability
Inversion Frequency Due the high factor inversion costs in KFAC- and SNGD-based methods. researchers use the stale factor approach, which updates the inverted factors every f iterations and reuses the results in the other iterations in their preconditioning to reduce the computation and communication costs. I' is the reciprocal of inversion frequency and varies from few 100s to few 1000s. Our experiments show that in average-sized models such as ResNet-50 9 in an iteration that the inversion of factors is executed, the cost of KAISA and HyL.o is 150x more than an SGD iteration. Furthermore, more than 98% of the total cost in those iterations are spent on matrix inversion.
The stale factors approach can lead to good preconditioners if the loss function landscape does not vary significantly in each iteration. However, this is a strong assumption and doesn't necessarily hold in practice. Also, increasing the inversion frequency can benefit the convergence rate of the second-orde methods, In addition, our experiments show that using stale factors. The stale factors can be good preconditioners lead to converging to local minima in the loss function and damage the generalization of the model.
Numerical Stability In second-order techniques, we need to invert or find the roots of matrices of different sizes, which are usually not full-rank, resulting in numerical issues. The KFAC implementa- tion uses singular value decomposition (SVD) of the factors and masks the eigenvalues that are close to zero to deal with singular matrix inversion issues. In practice, the eigenvalues of the left and right factors in KFAC-based methods computed from equations 3 and 4 are increased manually by adding uI to each of them to improve numerical stability (p > is called the damping factor), but MKOR doesn't 't need such numerical fixes. Furthermore. HyLo uses two decomposition methods to sample the batch of inputs, namely KID and KIS. KID requires inverting matrices in 6xb of rank mén(b, d). thus for batch sizes larger than d in a specific layer, the method fails.
Unlike SVD or other iterative methods used for factor inversion, MKOR doesn't 't suffer from numerical instabilities that rise from large condition numbers. MKOR has a single scalar division, in which the denominator is guaranteed to be non-zero based on lemma 3.1 eliminating the numerical over/under- flow possibility and the need for damping factors (required by other second-order methods for computational stability).
Lemma 3.1. The factors computed using Equation Jand are all positive-definte.
2 suggests using double precision representation of numbers to avoid numerical over/under- flow in inverting or computing the roots of matrices. This approach adds more costs to the matrix inversion and increases the time complexity of the main bottleneck in second-order methods.
MKOR does not need higher precision computations, and can use half-precision floating point operations to reduce costs significantly. This will improve the memory utilization and reduce the communication costs in GPUs by 2x while using cheaper computation blocks for half-precision operations. Lemma 3.2 shows an upper bound on the quantization error effect the MKOR updates. Lemma 3.2. Assuming that the maximum quantization error is e, the maximum number in matrices and vectors is m, and the dimension of the vectors and matrices are d and d x d respectively, the quantization error of formulas and O((y + (1-7), m² )c)))c
72
Exploding Gradients Problem In second-order methods, where the gradients are preconditioned by various factors, the exploding gradient problem is worsened. Our experiments show that in first- order methods, by choosing a learning rate that doesn't 't lead to divergence in the first few iterations, explosion in gradients almost never occurs. On the other hand, in second-order methods, we observe that the explosion can occur at any iteration, and both KFAC and SNGD implementations are prone to this problem. This can lead to ripples in accuracy and divergence. One of the main approaches for solving the exploding gradient problem is choosing small values for the learning rate, limiting the convergence rate significantly. In particular, small learning rates damage the accond-order methods and make them almost performant as their first-order counterparts. Considering that SGD is more robust against the exploding gradients and taking advantage of the direct control of MKOR on the inverse of the factors, the factors in MKOR are modified to lean toward SGD once the possibility of exploding gradients is detected using equations 7 and & where Ç is a hyperparameter that controls the amount of information from the original factors that needs be saved in the new factors.
L CL + (1 6)1
6

---

R = C + (1 Ç)1
(8)
By expanding Equation 2 with the new factors, we will get Equation 9 which reduces the loss-based on lemma The first term in the right-hand-side of[9 the KFAC term, the second and third terms are the left and right preconditioned versions, and the last term is the SGD term.
im Ww CRm =(2L" Vwm CRT
+ {(1- Vwm + C(1 - (1-
Lemma 3.3. Given a differentiable function L(w) with first-order Taylor series approximation [(w - Aw)= = ((wo)Aw"' A v around point wo. assuming that at point we the second-order derivative of the function L(w) is given as C(wo) semi- defindefe matrices. for value of Aw = ((SL + (1 () @ (cR +
H @ R. where L Lad R are positive-
the inequality L(wu- -ouu--wi Aw) < ((wu) holds.
While this modification can avoid exploding gradients, overusing it with small values of Ç will convert MKOR to SGD. MKOR uses factor norm-based metric that observes the infinity norm of the factors, and if they are greater than specific threshold, the process of factor modification will be triggered.
4 Experimental Results
In this section, we demonstrate the performance of MKOR on a large language model using different benchmarks, and analyze the timing of different components in different first- and second-order algorithms. For results on more models and training sets, please refer to the supplementary material.
Large Language Models. We pre-train BERT-Large Uncased and fine-tune it for different question- answering and text classification tasks. A setup similar to (21) for pre-training and fine-tuning is used. The first-order baseline used is Fused LAMB (32). Similar to (21) for the pre-training process, the English Wikipedia [30] and the Toronto BookCorpus [34] dataset, which was used in the original BERT pre- pre-training,
used; the latter dataset is not thoroughly available which results in small reduction in the baseline accuracics achieved in our experiments from the original BERT results. Following 211, due to the time-intensive process of hyperparameter tuning for the first phase of pre-training, we report the effectiveness of MKOR in the second phase of pre-training only while using the checkpoints of the first phase generated using LAMB optimizer.
As expected, the computation, communication, and memory complexity of HyLo is high, and the Khatri-Rao-basse Interpolative Decomposition (KID) approximation method, the main idea of HyLo, cannot be executed because a single sample cannot fit into the 40GB memory of an A A100 GPU. In addition, HyLo doesn't 't support gradient accumulation due to its memory complexity, depending on the batch size; in LLMs such as BERT, the batch sizes are as large as 64k.
For the question answering task, we fine-tune the pre-trained BERT checkpoints on the SQuAD v1.1 22 dataset. Table Saaee@sswwwws the F1 Score achieved using different optimizers and compares their convergence rate and speedups. The vanilla MKOR and KAISA both converge after 1000 iterations, while the LAMB optimizer requires 1, 536 steps. Considering that each step in MKOR is faster than KAISA, MKOR achieves end-to-end speedup. MKOR-H will converge in 600 steps, reducing the number of steps in LAMB by 2.6x, while achieving the same accuracy. In addition, achieves 2.57x speedup over the LAMB optimizer and 1.75x speedup over KAISA. As another second-order baseline, we consider Eva, which converges in 1000 iterations, and MKOR achieves 1.69x speedup over it.
For classification tasks, we fine-tune BERT on the GLUE [9] dataset. Table compares the results for different classification tasks in the GLUE dataset. MKOR with 1 1500 steps achieves a new state- of-the-art art accuracy in GLUE dataset on BERT-Large Uncased, and MKOR and MKOR-H with 600 steps achieve the same average metric as the baseline, while reducing the number of steps by factor of 2.6x. MKOR and MKOR-H both achieve 2.57x x end-to-end speedup. After training KAISA for 1,563 steps, the model does not converge to the baseline average accuracy, while slowing down the convergence by 0.89x. Eva requires 1000 steps to converge to the target average metric, being 1.69x slower than MKOR-H with 600 steps and 1.24% less accurate than MKOR with 1500 steps.
7

---

Table BERT- Large Uncased results on SQuAD v1.1 question answering task
Metric
LAMB
KAISA
MKOR
MKOR-H
Eva
F1
90.44
90.44
90.50
90.64
90.55
Iterations
1,536
1,000
1,000
600
1,000
Time (h)
7.97
5.71
5.25
3.10
5.24
Speedup (x)
1.00
1.39
1.51
2.57
1.52
LeeTT-Larg-Unnasee Training Loss
10
Loss
Loss
LAMB K-FAC
MKOR
K-PAC
MKOR H
MKOR-H
Eva
200
1200
1400
1600
1
2
3
5
Iteration
Time (h)
Figure 2: The training loss of BERT-Large- Uncased using different optimizers.
Per figure which shows the training error during the training of BERT, MKOR decreases the error in less iterations in comparison KAISA, Eva, and LAMB, leading to faster convergence. From Tables and 3 MKOR-H converges in only 600 steps.
Inversion Frequency. Due to the low computation complexity of the updates MKOR, the factor inversion frequency (f) MKOR is in the range of 10. Figure a shows that while average iteration cost in KAISA is heavily dependent the inversion frequency, MKOR's cost is almost independent of the inversion frequency. Also b shows that increasing the inversion frequency leads to higher convergence rate. In addition, using stale factors may result in converging to local minima. Hence, in MKOR we increase the convergence rate by updating the factors more frequently, without affecting the per-iteration cost, leading to end-to-end speedups in training. We simple autoencoder 24 on CIFAR-100 [1] this experiment.
Performance Analysis. We compare the performance of different parts of the optimizers to illustrate bottlenecks and advantages different methods. The training process for an optimizer has three steps: factor computation, precondition, and update weights. Figure 3 shows the time spent on each task in different optimizers on two models; ResNet-50, a CNN and BERT-Large-Uneen LLM with large sequence length. Since first-order optimizers such as SGD, ADAM, and LAMB don't require factorization and preconditioning, their optimization time is only spent in updating the weights. In ResNet-50, since the model size is larger compared the batch size, the factor computation and inversion is more expensive for KAISA compared to HyLo. This cost significantly reduced in MKOR.
transformer-brat
For BERT-Large-Uncased because of the large size the model, the factor inversion time KAISA is large. Also, due to the large sequence length value in this model, the kernel inversion time for
Table 3: BERT-Large Uncased results the GLUE elassification tasks.
Optimizer
Iterations
Time (h)
Speedup (x)
Average Metric
LAMB
1,563
7.97
1.00
0.8023
KAISA
1,563
8.93
0.89
0.796
MKOR
1,500
7.88
1.01
0.8214
MKOR
600
3.10
2.57
0.8078
MKOR-H
600
3.10
2.57
0.811
Eva
1000
5.24
1,52
0,809

---

BERT Large Uncased Wikipedia
ResNet-50
ImageNet
10"
KAESA
KAISA
Time (6)
Time (8)
10-22
10
Factor Computation
Precondition
Update Weights
Factor Computation
Precondition
Update Weights
a
(b)
Figure 3: Per-step breakdown of different optimizes on BERT-Large- Uncgsed (a) and ResNet-50 (b)
Autoencoder
BERT-Large- Uncasd
250
Time /Lterattio (ms)
Time IIteration (ms)
200
4000
100
3000
2000
MKOR KAISA
1000
MKOR KAISA
0
10a
10a
10a
102
10a
Factor Rouse Time
Factor Rouse Time
(a) Average Time per Iteration
Autoencoder
BERT Large Uncased
0.0
0.6
0.7
Loss
Loss
0.6
MKOR Reuse
0.5
KAISA Reuse 100
KAISA Rouse 1000
MKOR Reuse
0.4
MKOR Reuse 100
2
MKOR Reuse
5
10
15
20
25
0
25
50
75
100
125
150
175
Epochs
Iteration
(b) Test Loss
Figure 4: The sensitivity of MKOR and KAISA for BERT-Large- BERT-Large-Uncased and an Autoencoder model (a) and the effect of inversion frequency the convergence properties of these models (b).
HyLo is comparable to KAISA's inversion time. But expected, because of its low computational complexity, the aforementioned cost in our method is much smaller than the total training time, leading speedups. It is important note that HyLo diverges in this training process, hence convergence time is not reported for HyLo.
Approximation Error Experimental Results. Due to the low-rank properties of the covariance matri- ces, MKOR utilizes rank- approximations of the covariance matrices to accelerate the computations and communication in KFAC-based optimizers. Here, we aim to theoretically and experimentally support this choice. As shown in Figure 5 our experiments show that covariance matrices can be approximated with rank-1 matrices with low error and higher rank approximations are unnecessary practice. Figure 5 shows the error distribution of the optimal rank- approximation methods of the covariance matrices in ResNet-50 and BERT-Larger Uncased pre-training. Our extensive tests

---

BERT Activation Error Distribution
BERT Input Gradient Error Distribution
Probability Density
Probability Density
0.15
0.10
0.075
0.05
0.050
0.025
0.00
0.000
Activation Error
Input Gradient Error
(a)
(b)
ResNet-50 Activation Error Distribution
ResNet-50 Input Gradient Error Distribution
Probability Density
Probability Density
0.10
0.4
0.08
0.06
0.04
0.1
0.02
0.0
0.00
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Activation Error
Input Gradient Frror
(d)
Figure 5: Rank-1 I eтror for activation and input gradient covariance matrices for LERT-Large-Uncasee pre-training (a, b) and ResNet-50 on ImageNet (c, d).
well-known benchmarks show this property holds for all models and we have not benchmark that does not have low-rank covariance matrices.
Approximation Error Analysis and Extension to Higher Ranks. Small batch sizes and over parameterization of networks will lead to low-rank covariance matrices in DNNs. Let's consider the covariance matrix C = XXT, where C' E R is the covariance matrix and X E Rdx is matrix in which each column corresponds to single sample and d and b are the sample dimension and the per-GPU batch size respectively. Rank of the covariance matrix is min(b, d). If the Per-GPU batch sizes are small, the covariance matrices each GPU will be low-rank. Rank- 1 approximation methods work well in these scenarios. If the batch sizes in each GPU large, observe that the covariance matrices will stay low-rank. The underlying reason for this observation is that current neural networks over- over-paranne parameterized, and result, different features in the covariance matrices of the activations and the output gradients won't be linearly independent, resulting in low-rank covariance matrices.
Extending MKOR to Higher Ranks: Furthermore, one can extend MKOR to use higher-rank covari- ance matrices. Let's assume that C == - c,c where is the rank of the covariance matrix C.. We apply SMW identity to compute Cinew (Cold +c ) with (dd) computational complexity. Then we can compute Cz = (CTee + C2o3 )-1 using SMW identity with (dd) computational complexity. We can continue the same pattern by computing Crew (C + c.c?") The total computation complexity of this process will be O(rd?). We should add this cost to the cost of computing the low-rank approximation of which requires an SVD. Using SVD kills the main advantage of using low-rank computations, since the computational complexity of applying SVD is the same as inverting the factors directly. We could not find any cheaper way to compute low-rank approximations of the covariance matrices, except for the rank-1 approximation used in this paper.
Conclusion
We propose MKOR, Momentum- Eaneend Sannecker-FFactor Factor-Based Optimizer Using Rank- 1 updates that improves the end-to-end training time and convergence rate of second-order methods by reducing their computation, communication, and memory usage complexity. Our experiments illustrate that MKOR outperforms state-of-the-arr first- and second-order optimizers on large language models.
10

---

Broader Impact
The research described in this paper introduces a novel method for training DNNs that achieves faster convergence in LLMs. Using our work can help save a lot of energy and time for machine learning practitioners. The computations in this work are fully transparent. DNNs can be applied to different problems, including medical diagnostics, voice recognition, and addressing climate changes. However, should be acknowledget that optimization algorithms for DNNs can also be applied to models with potentially negative implications such as privacy invasion. It is important to note that the responsible and ethical utilization of any efficient optimization algorithm, including those we propose, extends beyond the scope of this research.
7 Acknowledgments
We thank the reviewers for their constructive feedback. We would like to express our deepest appreciation to Dr. Behrooz Zarebavani for all the time and energy he put into helping with the writing and formatting of this paper. This work is supported by NSERC Discovery Grants (RGPIN065 6,, DGECR00303, RGPIN-2023-04897
National Science Foundation (NSF CCF-2106621), the Canada Research Chairs program, the Ontario Early Researcher Award, the Digital Research Alliance of Canada (https: Computing Center (https: /www//wwwko.ooo/ www.aat..........o.o........... edu). Work of Zhao Zhang was supported by OAC-2106661.
www.anicccann.... ca) and Texas Advanced
References
[1] Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2): 251-276, 1998.
[2] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018, 2020.
[3] Argonne Leadership Computing Facility. Polaris.
https://www.akee www.walllf.an..o..... anl gov/
polaris
[4] Nihat Ay. On the locality of the natural gradient for learning in deep bayesian networks. Information Geometry, pp. 1-49, 2020.
[5] Jimmy Ba, Roger Grosse, and James Martens. Distributed second-ordes optimization using kronecker-factored approximations. 2016.
[6] Compute Canada. Compute canada. htttp::/ccomputeean.dd. ea/
[7] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 248-244822 2255. Icee, 2009.
[8] Donald Goldfarb, Yi Ren, and Achraf Bahamou. Practical quasi-newton methods for training deep neural networks. Advances in Neural Information Processing Systems, 33:2386-2396 2020.
[9] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
[10] Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classifica- tion. arXiv preprint arXiv:1801.00 146, 2018.
[11] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[12] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Communications of the ACM, 60(6):84-90, 2017.
11

---

[13] Dong C Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical programming. 45(1):503-522 1989.
[14] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv: (60iv.6660.0000 (608803983, 2016.
[15] James Martens. New insights and perspectives on the natural gradient method. The Journal of Machine Learning Research, 21(1):5776-5811, 2020.
[16] James Martens and Roger Grosse. Optimizing neural networks with coonecker-factoree ap- proximate curvature. In International conference machine learning, pp. 2408-2417, PMLR, 2015.
[17] Baorun Mu, Saeed Soori, Bugra Can, Mert Gürbüzbalaban, and Maryam Mehri Dehnavi, Hylo: hybrid low-rank natural gradient descent method. In Proceedings of the International Conference on High Performance Computing, Networking, Storage and Analysis, pp. 1-16, 2022.
[18] NVIDIA Corporation. Nvidia deep learning examples. bitps://qithur beepLearningExamp
NV IDIA/
[19] Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. Large-scale distributed second-order optimization using tronecker-factonee approximate curva- ture for deep convolutional neural networks. In Proceedings of the IEEE/CVF Conference Computer Vision and Puttern Recognition, pp. 12359-12367, 2019.
[20] J Gregory Pauloski, Zhao Zhang. Lei Huang, Weijia Xu, and Ian T Foster. Convolutional neural network training with distributed k-fac. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-12. IEEE, 2020.
[21] 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, pp. 1-14, 2021.
[22] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv: 05250, 2016.
[23] Yi Ren and Donald Goldfarb. Efficient subsampled gauss-newton and natural gradient methods for training neural networks. arXiv preprint arXiv: (ttiv:/996600200 1906.02353, 2019.
[24] Jürgen Schmidhuber. Deep learning in neural networks: An overview. Neural networks, 61: 85-117,2015.
[25] Shaohuai Shi, Lin Zhang, and Bo Li. Accelerating distributed k-fac with smart parallelism of computing and communication tasks. In 2021 IEEE 41st International Conference on Distributed Computing Systems (ICDCS), pp. 550 560. IEEE, 2021.
[26] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. arXiv preprint arXiv:201 1.04006, 2020.
[27] 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.
[28] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances neural information processing systems, 30, 2017.
[29] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: multi-task benchmark and analysis platform for natural language understanding. arXiv preprint 1rXivv1804.07441, 2018.
12

---

[30] Wikipedia. Wikipedia corpus.
http:///wwwtaaaaaaaaaaaaoaa/o/// __dmpppppppppp_
Eorrents VEnglishh Wikipedia
[31] Minghan Yang, Dong Xu, Zaiwen Wen, Mengyun Chen, and Pengxiang Xu. Sketchy empirical natural gradient methods for deep learning. arXiv preprint arXiv:2006.05924 2020.
[32] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv: 4,X99:/994.00062,, 2019.
[33] Lin Zhang, Shaohuai Shi, and Bo Li. Eva: A general vectorized approximation framework for second-order optimization, 2023.
[34] Yukun Zhu, Ryan Kiros, Rich Zemel, Ruslan Salakhutdinov, Raquel Urtasun, Antonio Torralba, and Sanja Fidler. Aligning books and movies: Towards story-like visual explanations by watching movies and reading books. In Proceedings of the IEFE international conference on computer vision, pp. 19-27, 2015.
13

---

8 Supplementary Material
In this section, we start by discussing an experiment training the ResNet-50 [9 model on ImageNet [7] dataset and compare MKOR against state-of-the-art second-order and first-order baselines. Then, for completeness, we discuss the derivation of KFAC and SNGD approximation methods in (33) We then discuss some of the features of MKOR an other optimizers and back them up with quantitative data in 8.4 and 8.3 We will provide scalability results of MKOR in 8.6 and provide more data to back up the low-rank discuss the training setup used our experiments and a link to the MKOR repository. In 8. 12 we analyze MKOR and other optimizers on the training tasks as non-convex optimizers, only concerning their performance on training tasks, We conclude this section by proofs of the lemmas used in the paper.
report the GLUE results achieved on each task in 8.2 Next,
features of the covariance matrices in 8.7 In 8.9 and 8. 10 and 8. 11
8.1 ResNet-50
We train ResNet-50, convolutional neural network with more than 25M parameters, ImageNet, text classification task with more than 1.2M samples. The same setup is used in (21), and SGD is used the first-order baseline.
The target accuracy in this experiment 75.9%. MKOR converges to this target accuracy in 57 epochs, while SGD, the first-order baseline, achieves this accuracy in 88 epochs. MKOR achieves 1.49x speedup over SGD. KAISA, the second-order baseline converges in 54 epochs, but due to its expensive steps, MKOR still converges 1.04x faster than KAISA. We do not compare to HyLo because HyLo is not able to achieve the target accuracy for ResNet (it reaches 75.6% with tuning as reported in [17 and our experiments confirm it). As shown, the effect of complexity reduction and improvement in performance in MKOR is less obvious in ResNet because the model dimension (d) is smaller compared to LLMs such as BERT. Please see Table T for comparison of complexity between
methods.
We could not reproduce the ResNet-50 results of Eva 33 ImageNet because the hyperparameters are not reported. We tried to tune Eva on multiple settings and none converged to desired accuracy. is important to note Eva is not comparing results with the most efficient implementation of KFAC. The KFAC version used in Eva is from [16.. dated to 2015. A number of followup works, mentioned in I have provided faster implementations of KFAC. We use KAISA [21]. the state-of-the- implementation of KFAC. Also from discussions with KAISA authors and our own experiments the
optimal inversion frequency for KFAC is 200. Eva uses an inversion frequency of 50 for KFAC, which makes KFAC slower.
ResNet-50 Test Accuracy
ResNet-50
80
80
Accurayy
60
Accuracy
20
SGD
20
MKOR
MKOR KFAC
KFAC
20
40
60
80
5000
10000
15000
20000
25000
30000
35000
Epochs
Time (s)
(a)
(b)
Figure 6: Test accuracy of ResNet-50 on ImageNet for MKOR, KAISA, and SGD on 64 GPUs.
14

---

Table BERT-Large- Uncased Results on the GLUE classification tasks.
SST-
STS-
Optimizer
Iterati- ons
MNLI (acc)
QQP
QNLI (acc)
COLA
MRPC
RTE (acc)
Avera-
(F1)
(acc)
(mcc)
(corr)
(F1)
ge
LAMB
1,563
0,841
0.878
0.913
0,919
0,516
0,875
0,812
0.664
0.8023
KAISA
1,563
0.821
0.854
0.900
0.921
0.489
0.878
0.888
0.617
0.796
MKOR
1,500
0.844
0.879
0.916
0.923
0.523
0.892
0.905
0.690
0.8214
MKOR
600
0.833
0.878
0.904
0.921
0.494
0.886
0.893
0.653
0.8078
MKOR-F
600
0.838
0.877
0.911
0.921
0.502
0.886
0.898
0.657
0.811
Eva
1000
0.839
0.877
0.907
0.914
0,499
0.890
0.904
0.650
0.809
d
da
F
Reshope
t
L
W
R
F2
b. Kronecker-Factorel Based Approximation
33
H-S
a. Block Diagonal Approximation
SMW Idenity
d²
K
SNGD Based Approximatoin
Figure 7: Approximations second-order methods.
8.2 GLUE Results
We discussed the speedup achieved using MKOR on the GLUE dataset Section For completeness. Table 4 shows 4shoes
the metrics achieved in each of the different GLUE tasks on BERT-LLrrge Uncased trained on different optimizers.
8.3 Derivation of NGD Approximations
Natural Gradient Descent (NGD) In NGD, which is second-order method, we use the inverse of the Fisher Information Matrix (FIM) as preconditioner the gradients as shown in Figure a. Equation T shows the update rule of NGD where Fm is the FIM block corresponding to block m. Equation 10 shows the definition of FIM for an arbitrary layer in our model, where 2; the 5th sample in the batch.
(10)
15

---

Kronecker Factorization (KFAC). KFAC methods reformulate the FIM block as the Kronecker product of two matrices shown in Equation Henaatioonnnhheeeeeeee where g" = ,..' .m C and a" is the vector form of the activation output of layer m and E is the expectation operator and xth is the input matrix of layer m. Please note that we have used the mixed-produs property of Kronecker multiplication for getting the right hand value.
F gl & (0" @ a am=111²1-
"0 = Eg( gm²")) @ (a---1"" (am-1am-1'')
(11)
Furthermore, we assume that [gg gm²)) @ aaamaaam-a strong assumption, but helps us simplify the computation further. Using the inversion property of Kronecker multiplication, we can compute the inverse of FIM using Equation 12
R Eig" gm 3 which
Fm pm-1u" w = E @ Ela" a '-'""" ".
(12)
By using the mixed Kronecker matrix-vector product property, we can get the update value in Equation 2 which is illustrated in Figure We refer to L and Rm as the left and right factors respectively. Adding momentum to the left and right factors and denoting the iteration number with a subscript to the factors, we will get equations Band and
Cherman-Morrrise Setrrison-Woodbur--oro Shhodbury-Based Natural Gradient Descent (SNGD). In this method, the SMW identity is used for approximating the inverse of (FTh + μI) Rd xd2 where je is a damping factor used in the preconditioning. Equation [3 the process of computing the inverse of the FIM for a single layer in the network, where A Rd the batch of activations of layer and GTD E bd is the batch of gradients of the loss function with respect to the inputs of that layer and
Rdd xb is the concatenation of the gradients of the loss function with respect to the parameters of that layer and @ shows the Hadamard element-wise product. In this method, a kernel matrix in R inverted, as shown in Figure 7
- V w L(W, .r),,.., V w L(W, LxW,
1 (I- (ATH 1T A O Gm T GTh ul) U
( + wl)-1
(13)
p.
8.4 Numerical Instability of Second-order Methods
discussed that in second-order methods, multiple matrix inversion or root-finding algorithms need to be executed, which make the second-order methods prone to numerical instabilities. Furthermore, we discussed that left and right factors in second-order methods have large condition numbers, resulting in further issues in inversion. Figure 8 shows the eigenvalues of the right factor and its condition number for ResNet-50 model on CIFAR-10 dataset KFAC algorithm. Even when using damping factors and filtering out extremely small eigenvalues, the condition number of these matrices is large, motivating for using double precision computations for avoiding numerical instabilities.
In Section 3.3
MKOR, on the other hand, doesn't suffer numerical instabilities when inverting such matrices, and its computational complexity isn't dependent on the condition number either.
8.5 Sensitivity to Learning Rate
I Learning rate is one of the main hyperparameters in machine learning (ML) that can directly affect the convergence time of optimization, and ML practitioners have to spend lot of time tuning this perparaameeer. More specifically, in first-order methods, a large learning rate easily lead to divergence and numerical instability, and in second-order methods large learning rates can lead to exploding gradient descent discussed in Section 3.3 Using small learning rates can lead to slow convergence in both first- and second-ordes methods, and can even jeopardize the main advantage of second-order methods, which is their faster convergence rate.
One of the main advantages of our method its robustness against a wide range of learning rates. As Table 5 shows, first-order methods are extremely sensitive to learning rate, and the second- order methods are prone to ripples and divergent for a larger range of learning rates, and lose their performance for small learning rates. Our method, on the other hand, will converge with high
16

---

ResNct-50 CIFAR-10 Eigenvalues of Right Factor
ResNet-50 CIFAR-10 Condition Number of Right Factor
Condition Number
100
10a
Eigenvaiue
10
10a
2
104
Minimum Figenvaltaes Eigenvalues
10a
100
200
300
100
200
300
Iteration (x10000
Iteration (x1000)
(a)
(b)
Figure 8: Maximum and minimum eigenvalues (a) and the condition number (b) of the right factors in KFAC when training ResNet-50 on CIFAR-10. As illustrated, the minimum eigenvalues of the factors in KFAC approach zero, meaning that the factors singular, and hence have large condition numbers, making numerical inversion of them complex and numerically unstable.
Table 5: Number of epochs necessary for convergence in different optimizers for ResNet-50 on CIFAR1O. MKOR is the least sensitive optimizer learning rate, converging in almost the same number of iterations for wide range of learning rate, while other optimizers either diverge (D) or converge to a local-minimum (* superscript).
Learning Rate
10
0.1
0.01
Optimizer
MKOR
94
79
78
76
KAISA
112
100
90
89
HyLo
D
123"
98
150*
SGD
D)
D
108
145*
convergence rate for wide range of learning rates, and by directly modifying the inverse of factors as discussed in [.3 can find proper equilibrium between first- and second-order methods. This table shows that our method is the least sensitive to the learning rate values and can make the job of ML practitioners for tuning this hyperparameter extremely easy.
8.6 Scalability
Figure shows the strong scalability of MKOR on BERT-Large-nnaasel on up to 64 GPUs.
BERT Large Uncased Scalability
Speodup
30
Num GPUs
MKOR
#PUss
Figure 9: Scalability of MKOR.
17

---

Average Covariance Rank-1 Error
ResNet-50 Average Covariance Rank-1 Error
Gradient
Activation
0.8
0.6
0.2
200
400
600
800
1000
1200
Iteration
Figure 10: Average covariance rank-1 approximation error for ResNet-50 in different iterations
Table 6: Per-GPU memory usage (in GB) for MKOR, KFAC/KAISA, LAMB, and SGD BERT- Large-Uncased pre-trainingReeNet-550 training on ImageNet.
Model
MKOR
KFAC/KAISA
LAMB
SGD
ResNet-50 3.88
5,83
3.01
BERT
23.34
29.97
12.80
8.7 Decaying Figenvalues and Rank-1 Approximations
Figure [0 shows that the eigenvalues of the factors will decay as the model converges, making rank-1 I approximations more effective. The reason behind the decay in the eigenvalues is that the weights are initialized randomly and the neurons work independently in the beginning of the training, but the model converges, the neurons become more dependent on each other and thus the activation and input gradients will become linearly dependent. This is also reflected by some large error values in the distributions in Figure 5 a, b, c. d]. The factors in MKOR are initialized with identity, starting MKOR from first-order method. As a result, MKOR is more robust against noise in the approximations in the first iterations (the approximation error does not noticeably affect the factors when replacing I 1 and R in equations and 6 with identity). But as the model converges, the factors in MKOR will be mostly shaped by the training samples, making MKOR more reliable on less erroneous approximations, and the decaying eigenvalues of the factors help MKOR with that.
8.8 Memory Overheads
The memory overheads of MKOR in comparison to other optimizers reported in Table be observed that all the second-order methods have significant memory overheads to the first-order methods, but MKOR reduces the overhead of KFAC/KAISA up 1.5x.
6 It can
8.9 Hyperparameters
BERT-Large- Uncased. For the BERT-Large- Een--L-arg-Unesee pre-training, we use the same hyperparameters used [8] The factors in KAISA are updated every iterations, and the factors in MKOR and MKOR-H are updated every 10 iterations.
the same hyperparameters used in 21 The factors in MKOR are updated every 10 iterations, and the learning rate used there is the same is KAISA. The learning rate MKOR decays by factor of the end of epochs 25, 35, 40, 45, 50, 55, and 56.
ResNet-50. For SGD and KAISA,
8.10 Experiment Settings
For the BERT-Large Large- pre-training and fine-tuning experiments, we have used up to 64 A100 GPUs on the Polaris cluster which has 560 nodes, each with NVIDIA A-100 GPUs with NVL.nn interconnects.
18

---

Table 7: List properties of the models, datasets, and settings used in our experiments.
Model
Dataset
GPU
Name
#Parameters
Name
Train
Test
Arch
#
BERT-Large- Uncased
335.IM
Wikipedia BookCorpus
A100
64
ResNet-50
25,5M
ImageNet
1.2M
50k
V100
64
AlexNet
20.3M
CIFAR-100
50K
10K
V100
4
BERT-Base
108.9M
SQuAD vI.1
87.6K
10.6K
V100
4
BE:RT-Large-Case
335.1M
IMDB
25K
25K
V100
4
The rest of the experiments are conducted on the Mist cluster [6]. which has 54 nodes, each with 32 IBM power9 cores with 256GB RAM and 4 NVIDIA V100 GPUs with 32GB memory and NVLink inter-node connections and InfiBand EDF intra-node connections.
Each of the training experiments are conducted 5 times, and their median timing is reported. The reported accuracies are the median of the results we have obtained from multiple runs. For the scalability experiments, we have trained BERT-Large- Uncased for 100 iterations on two separate jobs for each setting and averaged them. For the time breakdown experiment, we have reported the median of 5 experiments.
Table 7 summarizes the models, datasets, and GPU architectures used in training and fine-tuning experiments. In addition, the results in Figures and 4 are gathered on a single 4-GPU node of V100s. The results of Figure b are gathered 64 A A0O GPUs.
8.11 Reproduction
Our code base is publicly available on https: com/ http///ohaamaaaaaaaaaaaaaa mkor and the instructions for running each experiment are available there.
8.12 Training Accuracy Experiments
To evaluate MKOR optimizer that tries minimize a specific objective function, have considered the case of only minimizing the loss function of models on different tasks and set all the other optimizer parameters such as weight decay to zero. Since using non-zero weight decay adds quadratic term to the loss function and using different weight decays for different optimizers leads to optimizing different objective functions using different optimizers which might be considered unfair.
Recent work has shown the advantage of second-order methods over their first-order counterparts on multiple CNN tasks (17] 21), such residual networks [9. In training accuracy experiment, we use another CNN benchmark, AlexNer [12] with more than 20M parameters on CIFAR- 100 [IT) consisting of 50K training and 10K validation images of 100 classes. Figures [1 and 12 a show the convergence properties of different optimizers. MKOR is 1.26x, 1.31 x, and 1.58 x faster than HyLo-KIS, SGD, and KAISA respectively. The reason for the low convergence speed of KAISA is that we needed to use small learning rates for avoiding exploding gradients in it, which has damaged its convergence rate.
BERT a large language model with two variants, BERT-Bas Base with more than 108M parameters and Bert-Large with more than 335M parameters. As shown in Figures I c and 1 we have fine-tuned BERT-Large on the IMDB dataset which a text classification task with 25K training and 25K test samples. MKOR outperforms SGD and HyLo- Lo- by speedup factor of 1.22x and 1.43x respectively. We have also fine-tuned BERT-Base on SQuAD dataset, which is a question answering task with 87. 37.6K training and 10.6K test samples. MKOR achieves 1.26x and 1.56x speedup over SGD and HyLo respectively. We couldn't i find a learning rate with which KAISA would converge on any of our BERT experiments which are based the HuggingFace implementation of BERT, and the reason for lack convergence of KAISA is exploding gradients.
8.13 Knee-Point Learning Rate Scheduler
While using large learning rates crucial for utilizing the higher convergence rate of second-order methods, it is necessary to reduce the convergence rate after some iterations that the model
19

---

BERT Basc Cased Question Answering SQuAD
BERT Large Cased Text Classification IMDB
51275
62915
75474
10630
10164
22023
101
100
100
Accuracy
HyLo MKOR
MKKOR
HyLo KIS
10000 20000 30000 40000 50000 60000 70000
80000
0
10000
20000
30000
40000
50000
Time (8)
Time (s)
(a)
(b)
AlexNet CIFAR-100
14000
12008
2110
80
Accuracy
60
40
KAISA
MKOR
20
HyLo-KIS
250
500
730
1000
1250
1500
1750
2000
Time (8)
(c)
Figure 11: Training time for distributed first- and second-ordes optimizers SGD, MKOR, KAISA. and HyLo on BERT-Large- Cased on IMDB (a), BERT-Bat Base--Casd on SQuAD (b), and AlexNer on CIFAR-100 (c). In all the experiments, MKOR outperforms other optimizers in convergence speed.
converges, and the number of iterations for changing the learning rate is not known in advance. In practice, machine learning practitioners will manually find the number of iterations by trial and error or use predefined functions [10] 14] that don't It necessarily work ideally in all optimization problems. We used knee-point learning rate scheduler in Section 12
To fully utilize the potentials of the optimizers, we have designed a learning rate scheduler that monitors the rate of improvement in accuracy decrease in the loss function value, and based that decides when to decrease the learning rate. The scheduler detects knee-points in the accuracy/loss and decreases the learning rate when knee-point observed.
By definition, knee-points are defined as the points where the average accuracy/loss rate less than 3 times the incremen//decrementt in the accuracy/loss since using the current learning-rate For averaging the accuracy/loss rate, we use an exponential moving average, and 3 is hyperparameter that we can choose to show how much the scheduler can tolerate lack of improvement to detect the accuracy/loss.
20

---

BERT Large Cased Text Classification IMDB
BERT Basc Cased Question Answering SQuAD
58
100
100
Accuracy
MKOR
Lo
HyLo KIS
MkOo
10
15
20
225
340
35
40
10
20
30
40
50
60
Epochs
Epochs
(b)
AlexNet CIFAR-100
2000
2033
100
80
Accuracy
60
40
KAISA
MKOR
20
KIS
SGD
30
100
150
200
250
300
350
400
Epochs
(c)
Figure 12: Training accuracy vs. the number of epochs for distributed first- and second-order optimizers SGD, MKOR, KAISA, and HyLo WERRT-arge-Cased on IMDB (a), BERT-Base-Case on SQuAD (b), and AlexNet on CIFAR-JO0 (c). In all the experiments, MKOR outperforms other optimizers in convergence rate.
8.14 Proofs
Lemma 3.1
Proof.
Given matrix positive definite matrix Jr-1 and a vector j and scalar < 1, we show that Equation 14 results in desitive-defintt matrix.
J. - = 72 + 7(1 1011-100 (11+)11-))" Jn-1 (11 J.
-
(14)
Since 7 0 and J,1 is nositive-definit matrix, 7J1-1 is also positive-defintt. Also, since J a positive-definite matrix, Vr # :0:2 -T J 0.
1 أااا 0< + > ") J
(15)
Now show that J, 1グز J, -1 is also positive-definite
Vz # 0: 0: J0_1 j - (
(16)
Since both matrices on the right-hand-side of Equation 14 are positive-definit and the sum of two positive-definite matrices is positive-definit matrix, the left-hand-side of it will be also positive definite.
Lemma
21

---

Proof. By defining P == P=(5L-1 + (1- (1-5))) @ (CR + +(1 ()1), Equation 17 will hold.
[(wo - Aw) - [(wn) - - PV ((o))
(17)
Now, we will show that P is positive-semi-deiiie matrix. By using the associativity property of Kronecker multiplication, we can represent P as in Equation 18 Please note that different identity matrices have different shapes.
P (
(18)
Since matrices L and It are positive-semi-defniicc the Kronecker products L & 1-11, L- @ I, and I R-1 are also positive-semi-deffiiit. As result, for any non-zero vector x, x² Pr > 0. So, based on Equation 17 -V PV PVL(wo) As a result, Equation 19 holds.
L(wo Aw) < ((wo)
(19)
Lemma 3.2
Proof. Considering the quantization of matrix J and vector j in Equation 20 and assuming the maximum quantization error is e and the maximum values vector and matrix m, we can consider of the three possible cases:
Vector- Vector Dot Product: The resulting error is O(2em because the error of each multiplication can at most be 2me + e² and by adding the d multiplications, the maximum error can be (2m + 1)dc.
2. Vector-Matrix Product:
The resulting error is O22em2d) because vector-matric product can be considered as multiple independent rector-vector dot products.
3. Vector- Vector Outer Product: The resulting error is ((2mm) because each element in the resulting matrix computed using multiplication with the maximum error of 2em2 e?.
(1-?) 7J 72 + +11++(11-0 JJ) Jr
(20)
The error imposed by quantizing 7 is at most ye. The quantization error in the denominator of the fraction in 20 is negligible in comparison to and won't 't change the growth order the final error term. The quantization eπrог of Jj is O(2m2d de) discussed earlier in the case of vector-matris product. Jjj J" (Jj)(Jj) is a vector-vector outer product, resulting in (44 dd) error.
So the final error quantization error is O(() + 4"اد m² dd)
22


Reviews:
[{"review_id": "mbCJNqTx1F", "reviewer": "Reviewer_KDBs", "summary": "The paper proposes a second-order optimizer named MKOR based on K-FAC. Compared with K-FAC, MKOR uses a rank-1 update to construct Kronecker factors enabling the inverse of Kronecker factors to be efficiently computed via the Sherman-Morrison-Woodbury (SMW) identity. In addition, the rank-1 update of factors in MKOR also helps reduce the communication traffic in distributed training with data parallel for aggregating factors. Experiments on the BERT-Large-Uncased model with the GLUE dataset show that MKOR outperforms first-order LAMB and second-order KAISA.", "questions": "1. The key idea of using rank-1 updates to replace the original Kronecker factors in K-FAC has been proposed in Eva [A]. In [A], the authors provide a better derivation of the preconditioner without storing the factor matrices so that they use Kronecker vectors (KVs) to calculate the inverse of KVs using the SMW formula, which is much better than MKOR as MKOR needs to store the matrix form of factors (i.e., $L_t^m$ and $R_t^m$). The similarity between MKOR and Eva and the advantage of Eva over MKOR make the originality of this work weak. It is suggested to clarify the differences between MKOR and Eva, and highlight some advantages of MKOR over Eva.\n2. The main novel idea in MKOR is the balance between second-order and first-order information in updating $L_t^m$ and $R_t^m$ using Eq. (7) and (8) by controlling $\\zeta$. However, how to choose $\\zeta$ (or dynamically changing $\\zeta$ to switch MKOR to SGD just like MKOR-H) is very important. The paper doesn’t provide such an analysis of when and how to use $\\zeta$ though some results show that it can help improve convergence.\n3. Similarly to the previous problem, MKOR-H is a combination between first-order optimizer and MKOR according to the loss during training. However, how to measure whether the loss is changed is unclear (e.g., results in Figure 2 haven’t such information either) and when to switch to the first-order optimizer is unclear.\n4. What are the update frequencies of the second-order information for MKOR and K-FAC in Figure 2 and Table 3? And how do you choose such frequencies for comparison?\n5. In the original paper of KAISA, the update frequency can be set to be very high (e.g., 100-500) for achieving good accuracy on ResNet models with CIFAR10 and ImageNet datasets, is it possible that MKOR could be better than KAISA in these kinds of scenarios?\n\n[A] Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation, ICLR 2023.", "limitations": "See Questions.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. An implicit inverse computation method with the SMW identity in computing the inverse preconditioner factors and a formula to control the balance of exploitation between the first-order and second-order information.\n2. A hybrid method of MKOR with the first-order method, MKOR-H, which allows switchable from MKOR to the first-order optimizer in the later of training.", "weaknesses": "1. An important baseline, Eva [A], which is extremely similar to MKOR, is missing. \n2. Though MKOR-H seems to achieve some improvement, when switching MKOR to the first-order optimizer is unclear.\n3. The comparison between MKOR and K-FAC seems to be not that fair (e.g., Figure 4(b)).\n\n[A] Eva: Practical Second-order Optimization with Kronecker-vectorized Approximation, ICLR 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Comparison to Eva:**\n\nWhat the reviewer has pointed out as the strengths of Eva (using KVs to not store the factor matrices), is at the same time a weakness for Eva, leading Eva to being less accurate and less efficient vs. MKOR. Using KVs instead of factor matrices disables the proper use of momentums and imposes the use of damping factors in Eva, both of which damage the accuracy of FIM approximations. \n\nSimilar to KFAC, EVA’s formulation involves a damping factor (which helps with numerical stability similar to the left-hand-side of equation 13 in the appendix) that introduces additional error to the FIM approximation, however in MKOR we introduce a new formulation that removes the damping factor by storing the inverse of the previous factors and utilizing the SMW identity (equations 5 and 6 in our paper).\n\nIn addition, Eva doesn’t support proper momentum-based training because it doesn’t store the left and right factors of the FIM (but instead averages the KVs in time which does not have theoretical backup). MKOR stores the left and right factors to enable momentum in training, which has been shown to significantly boost the training process (equations 3 and 4 in our paper).\n\nMKOR has two additional contributions that lead to a more stable and efficient implementation of rank-1 updates. (1) It uses a *Norm-Based Stabilizer* and *Gradient Rescaling Mechanism* to detect and avoid exploding gradients. (2) with MKOR-H we combine the high convergence rate of second-order methods in the initial phase of training with the low overhead of first-order methods in the late stages of the training.\n \nEva comparison Results (please also see the attached PDF figure-file): \n\n* BERT-Large-Uncased: Compared to Eva, MKOR-H is faster in pretraining BERT-Large-Uncased by a factor of 1.69. Its accuracy is also better (an average accuracy of 81.1\\% on GLUE and 90.64\\% on SQuAD v.1 for MKOR with only 600 iterations vs 80.9\\% on GLUE and 90.55\\% on SQuAD v.1 for EVA with 1000 iterations), please see Figure 1 and Table 1 and Table 2 in the attached PDF figure-file. Figure 1 in the PDF file shows the convergence results of all the optimizers used in the paper in addition to Eva.\n\n* ResNet-50: We could not reproduce the ResNet-50 results of Eva on ImageNet because the hyperparameters are not reported. We tried to tune Eva on multiple settings and none converged to desired accuracy. \n\n* It is important to note EVA is not comparing results with the most efficient implementation of KFAC. The KFAC version used in EVA is from [1], from 2015. A number of followup work (please see line 41-42 in the paper) have provided faster implementations of KFAC. We use KAISA, the state-of-the-art implementation of KFAC. Also from discussions with KAISA authors and our own experiments the optimal inversion frequency for KFAC is 200, Eva uses an inversion frequency of 50 for KFAC, which makes KFAC slower.\n\n**Dynamic Controlling of $\\zeta$:**\n\nChanging the stabilization frequency with a constant $\\zeta$ (as done in MKOR) is equivalent to dynamically controlling the value of $\\zeta$. In MKOR, instead of changing $\\zeta$ dynamically for each iteration, we keep it constant but we change the frequency that we stabilize the factors (equations 7 and 8 in the paper). The frequency is dynamically controlled using a norm-based criteria (lines 5 and 6 in Algorithm 1). Also, if the inversion frequency is higher than a threshold, which is similar to having a small small $\\zeta$, we switch to SGD to save on computations. \n\n**Difference of MKOR and MKOR-H in Loss and Switching Criteria:**\n\nMKOR combines the first- and second-order optimization methods using the norm-based stabilization. Per our response to the previous question, the frequency of the stabilization dictates how close MKOR is to first-order methods, i.e. the more frequent the stabilization takes place, the closer MKOR gets to first-order methods. Our switching method looks at the average stabilization frequency in the last few iterations using an exponential moving average, and if the frequency is higher than a threshold and the loss change rate is less than a small ratio of the overall loss reduction, the user will be notified to switch to LAMB (the first-order optimizer) if needed. In the BERT-Large-Uncased pretraining, we switch to LAMB on iteration 300.\n\n**Inversion Frequency**\n\nAs reported in section 7.6 in the appendix of the paper, in all our experiments, the frequency (factor reuse time) of inverting the factors in MKOR and MKOR-H is set to every 10 iterations unless stated otherwise (i.e. in Figure 4 in the paper). For KFAC we use the default settings they provide in their papers.\n\n**Infrequent Factor Inversion in KAISA:** \n\nFor our experiments, including the ResNet-50 experiment, we used the default settings in the KAISA paper, which has a factor reuse time of 200. This number can be large in KAISA, since the factor inversions in KAISA are computed from scratch every time. In MKOR, which is an approximation-based method, we cannot use very stale factors, since each of our updates only modifies the factors slightly and stale factors won’t be useful anymore.\n\n**Fairness of Figure 4.b Comparisons:**\n\nFigure 4.b motivates the use of more frequent updates in second-order information, and it does not intend to compare different optimizers. There might be some confusion when interpreting Figure 4, particularly due to the absence of MKOR 1000 and KAISA 10. Kaisa 10 is not in the figure because it was very slow in comparison to others (we will add a sentence in the figure caption to note this), MKOR 1000 updates the factors so infrequently with very insignificant rank-1 updates that it performs similar to SGD and won’t have any benefit compared to other second-order methods.\n\n[1] Optimizing neural networks with kronecker-factored approximate curvature, Martens et. al., 2015\n\n[2] KAISA: An Adaptive Second-Order Optimizer Framework for Deep Neural Networks, Paulostki et. al., 2021"}, {"review_id": "SC0Glx2Eev", "reviewer": "Reviewer_Ar2i", "summary": "The work introduces a second-order optimizer that uses rank 1 covariance activation and gradient statistics, and efficient inversion algorithms to accelerate an approximated estimate of 2nd order information. Another addition is the rescaling of gradients and norm-based stabilization, which help to stabilize the entire optimization procedure. This procedure yields a highly stable 2nd order optimizer that is both memory (communication) and computationally efficient and leads to improved convergence times compared to baselines.\n", "questions": "No questions in particular.", "limitations": "Limitations are not discussed.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "From the work its clear that the method depends on many independent factors. While some reviewers might see this as a weakness, I see this is a strength. Developing efficient 2nd order optimizer is a highly technical and difficult undertaking and this work provides multiple components that make a stable optimizer possible that is both memory and computationally efficient.\n\nThe evaluation could have been extended to regular language modeling since BERT pretraining has become less academically relevant. Nevertheless, a BERT large pretraining is a rather convincing experimental setup. As such, the experimental robustness is also a strength of the paper.", "weaknesses": "I think additional experiments for causal language modeling would have been beneficial since this would be the main advantage for more efficient training (BERT pretraining is fast enough for most researchers, but most researchers are unable to pretrain LLMs due to their computational cost).\n\nWhile this additional experiment would greatly improve the paper, the experimental methodology in this paper is already relatively robust.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "MKOR is currently the most efficient second-order optimizer that can be used on LLMs and CNNs. As the reviewer has pointed out, similar to most state-of-the-art optimizers, MKOR has hyperparameters that users have to tune to use it properly. We have tried to automate as many of these tunings as possible, e.g. the use of norm-based stabilizer for automating stabilization or the stabilization frequency for switching between first- and second-order methods in MKOR-H or knee-point learning rate scheduler. We are optimistic on MKOR’s performance on other LLMs however were not able to try other models due to their computational and resource costs. "}, {"review_id": "Vbb6jzt5R7", "reviewer": "Reviewer_zw2W", "summary": "MKOR introduces Kronecker rank-1 representation of covariance for second-order optimizers. Authors try to solve the complexity of large language models in second-order optimizers while the Fisher information approximation is computed using rank-1 Kronecker matrix factorization. They also propose a hybrid method to cover the first-order method. Their evaluation of large language models is fair and sound.", "questions": "- Rank1 approximation might be restrictive to approximate the Fisher information properly. I wonder in some cases an extension to higher ranks is necessary. I wonder if an extension towards SeKron [1] is natural.\n[1] SeKron: A Decomposition Method Supporting Many Factorization Structures arXiv:2210.06299\n \n- Does adding a learning rate scalar in Equation (2) make the method more flexible to converge?  \n- How well does the second-order MKOR method compare with the first-order with the same training budget in terms of time and memory?", "limitations": "The numerical stability requires SVD decomposition and matrix inversion which are expensive to compute.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The main idea behind this research is simple but solves the important problem of the applicability of second-order optimizers for large language models. The paper is well-motivated and experiments are sound and relevant. \n\n \n", "weaknesses": "- Authors only argue why rank-1 is proposed from a practical standpoint but there is no theoretical justification to support why rank-1 is enough for Fisher information approximation.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "**Higher Rank Approximations:**\n\nAs shown in figure 2 in the attached PDF figure-file, our experiments show that the covariance matrices can be approximated with rank-1 matrices with low error and higher rank approximations are unnecessary in practice. Figure 2 (in the attached PDF figure-file) shows the error distribution of the optimal rank-1 approximation methods of the covariance matrices in ResNet-50 and BERT-Large-Uncased pretraining. Our extensive tests on well-known benchmarks (shown in the paper) show this property holds for all the models and we have not come across a benchmark that does not have low-rank covariance matrices.\n\n*Sekron and other approximation methods:* MKOR  reduces the factor inversion costs in second-order methods by combining SMW identity with low-weight rank-1 approximation methods. If we use any other method except for rank-1 approximations, e.g. Kronecker approximation discussed in SeKron, we need to find proper formulations to incorporate them in the matrix inversion. To our knowledge, no such cheap formulations exist that can be used for inverting matrices.\n\nIt might be possible to extend MKOR to use higher-rank covariance matrices however it will introduce significant computational overheads to MKOR since it requires computing the SVD of the factors. Let’s assume that $C = \\sum_{i=1}^{r}{c_i c_i^T}$ where $r$ is the rank of the covariance matrix $C$. We can apply SMW identity to compute $C_1^{new} = (C^{old} + c_1 c_1^T)^{-1}$ with $O(d^2)$ computational complexity. Then we can compute $C_2^{new} = (C_1^{new} + c_2 c_2^T)^{-1}$ using SMW identity with $O(d^2)$ computational complexity. We can continue the same pattern by computing $C_i^{new} = (C_{i-1}^{new} + c_i c_i^T)^{-1}$. The total computation complexity of this process will be $O(rd^2)$. We should add this cost to the cost of computing the low-rank approximation of $C$ which requires an SVD. Using SVD kills the main advantage of using low-rank computations, since the computational complexity of applying SVD is the same as inverting the factors directly. We couldn’t find a cheaper way to compute low-rank approximations of the covariance matrices, except for the rank-1 approximation used in this paper.\n\n**Learning Rate in MKOR:** \n\nThanks to the reviewer’s note, we realized a typo in equation 2. The learning rate is indeed accounted for in our implementation of equation 2. Equation 2 should be $W^m := W^m - \\alpha {(L_t^m)}^{-1} \\nabla_{W^m} \\mathcal{L} {(R_t^m)}^{-1}$ where $\\alpha$ is the learning rate. We will fix this typo in the final version of the paper.\n\n**Comparison to First-Order Methods:**\n\nAs shown in figure 2 and 6 in the paper (our appendix), MKOR-H converges 2.57$\\times$ faster than LAMB (the state-of-the-art first-order optimizer for pretraining BERT-Large-Uncased) and 1.49$\\times$ faster than SGD (the state-of-the-art first-order optimizer for training ResNet-50 on ImageNet). So given the same time budget, MKOR converges faster than other first-order methods. The memory overheads of MKOR in comparison to other optimizers are also reported in table 3 in the attached PDF figure-file. It can be observed that all the second-order methods have significant memory overheads to the first-order methods, but MKOR reduces the overhead of KFAC/KAISA up to 1.5$\\times$. We will magnify these in the final version. \n\n**SVD in Low-Rank Approximations:** \n\nInstead of using SVD methods that suffer from numerical instability and high computational overheads, MKOR uses a low-weight averaging method to compute rank-1 approximations (lines 2 and 3 in Algorithm 1 and lines 106 to 109 in the paper text). Hence, MKOR is SVD-free and its computational and memory overheads are negligible."}, {"review_id": "zVXnPreSo1", "reviewer": "Reviewer_dQpV", "summary": "This work introduces a new optimization algorithm for deep neural networks. Building upon the baseline Kronecker-factored curvature (K-FAC) algorithm, this new approach approximates each activation and gradient covariance matrix within K-FAC as a rank-one matrix. This approximation enables the efficient calculation of the inverse Hessian. By leveraging this enhanced efficiency, the algorithm allows for more frequent updates of the Hessian. Experimental results demonstrate great speed improvements compared to the state-of-the-art first and second order optimization algorithms.", "questions": "n/a", "limitations": "n/a", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Efficient second-order optimization algorithm is important for the training of DNNs.\n\n2. Experimental results show great speed-ups compared to the baseline algorithms.\n\n3. The paper is in general well-written, and the main contribution is clear.", "weaknesses": "One notable weakness of this research lies in the insufficient justification of the proposed technique, which involves applying rank-1 approximation to the covariance matrices used in K-FAC. While the work addresses the efficiency benefits of this method, it fails to thoroughly discuss the aspect of why employing rank-1 approximation does not impact the final accuracy or convergence of the algorithm. The following questions should be addressed to strengthen the research: What specific scenarios allow for accurate rank-1 approximation, and when might it be ineffective? Does the effectiveness of the approximation depend on the singular value decays of the covariance matrices? Are there any illustrative toy examples that can shed light on these considerations?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "In the below and the attached PDF figure-file we show that the covariance matrices in training neural networks can be approximated by rank-1 matrices both empirically and theoretically. We will include these discussions in our paper to address the reviewers’ concerns. \n\n**Experimental Results:** As shown in figure 2 in the attached PDF figure-file, our experiments show that the covariance matrices can be approximated with rank-1 matrices with low error and higher rank approximations are unnecessary in practice. Figure 2 (in the attached PDF figure-file) shows the error distribution of the optimal rank-1 approximation methods of the covariance matrices in ResNet-50 and BERT-Large-Uncased pretraining. Our extensive tests on well-known benchmarks (shown in the paper) show this property holds for all the models and we have not come across a benchmark that does not have low-rank covariance matrices.\n\n**Decaying Eigenvalues:** Furthermore, figure 2.e in the attached PDF figure-file shows that the eigenvalues of the factors will decay as the model converges, making rank-1 approximations more effective. The reason behind the decay in the eigenvalues is that the weights are initialized randomly and the neurons are working independently in the beginning of the training, but as the model converges, the neurons will be more dependent on each other and hence the activation and input gradients will become linearly dependent. This also is reflected by some large error values in the distributions in figure 3[a, b, c, d] in the attached PDF figure-file. The factors in MKOR are initialized with identity, starting MKOR from a first-order method. As a result, MKOR is more robust against noise in the approximations in the first iterations (the approximation error does  not noticeably affect the factors when replacing ${L_{t-1}^{m-1}}^{-1}$ and ${R_{t-1}^{m}}^{-1}$ in equations 5 and 6 with identity). But as the model converges, the factors in MKOR will be mostly shaped by the training samples, making MKOR more reliable on less erroneous approximations, and the decaying eigenvalues of the factors help MKOR with that.\n\n**Analysis:** Small batch sizes and overparameterization of the networks will lead to low-rank covariance matrices in deep neural networks. Let’s consider the covariance matrix $C=XX^T$, where $C \\in R^{d \\times d}$ is the covariance matrix and $X \\in R^{d \\times b}$ is a matrix in which each column corresponds to a single sample and $d$ and $b$ are the sample dimension and the per-GPU batch size respectively. Rank of the covariance matrix  is $min(b, d)$. In case the per-GPU batch sizes are small, the covariance matrices in each GPU will be significantly low-rank, and rank-1 approximation methods can work well in those scenarios. In case the batch sizes in each GPU are large, we observe that the covariance matrices will stay low-rank. The underlying reason for this observation is that current neural networks are over-parameterized, and as a result, different features in the covariance matrices of the activations and the output gradients won’t be linearly independent, resulting in low-rank covariance matrices. \n\n**Extending MKOR to Higher Ranks:** Furthermore, one can extend MKOR to use higher-rank covariance matrices. Let’s assume that $C = \\sum_{i=1}^{r}{c_i c_i^T}$ where $r$ is the rank of the covariance matrix $C$. We can apply SMW identity to compute $C_1^{new} = (C^{old} + c_1 c_1^T)^{-1}$ with $O(d^2)$ computational complexity. Then we can compute $C_2^{new} = (C_1^{new} + c_2 c_2^T)^{-1}$ using SMW identity with $O(d^2)$ computational complexity. We can continue the same pattern by computing $C_i^{new} = (C_{i-1}^{new} + c_i c_i^T)^{-1}$. The total computation complexity of this process will be $O(rd^2)$. We should add this cost to the cost of computing the low-rank approximation of $C$ which requires an SVD. Using SVD kills the main advantage of using low-rank computations, since the computational complexity of applying SVD is the same as inverting the factors directly. We couldn’t find any cheaper way to compute low-rank approximations of the covariance matrices, except for the rank-1 approximation used in this paper. \n"}]

Author rebuttal:
We thank all reviewers for their very informative feedback. We have provided a separate answer to each reviewer and have also attached a PDF file with figures and tables that add to our rebuttal (which we refer to as the attached PDF figure-file in our per-reviewer responses). 

Meta-review: The submission presents a family of heuristics for K-FAC which improve its memory and computational complexities and allow adaptively falling to first-order optimization. Experiments on a medium-scale setup show promising accuracy and overall performance. 

The reviewers found the author discussion convicing enough to converge to an acceptance decision. The authors are encouraged to include the discussion points into their next revision.
Decision: Accept (poster)