Title: MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates
Venue: NeurIPS 2023
Accepted: yes
Keywords: machine learning

deep learning

optimizers

distributed training;second-order optimization;
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-Factor-Based Optimizer Using Rank-1 Updates

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. 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 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 . 57 × and 1 . 85 × respectively on BERT-Large-Uncased on 64 GPUs.

1 Introduction

Second-order optimization methods have recently gained popularity in the training process of deep neural networks (DNNs) due to their higher convergence rate in comparison to their first-order counterparts. One of the well-known second-order methods, Newton's method, uses the inverse of the Hessian of the objective function as a preconditioner to 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 most common approximations of Newton's method is Natural Gradient Descent (NGD) [1], where the Hessian is substituted with the Fisher Information Matrix (FIM) [4] to reduce the cost

Sikan Li

Maryam Mehri Dehnavi

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 to the design of algorithms based on block-diagonal approximations of FIM; each block corresponds to a layer in the neural network (NN). Inverting the diagonal blocks is also not practical due to the large number of parameters in a 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 [25; 21]. KAISA [21] is a framework with state-of-the-art distributed KFAC implementation. The computational complexity of KFAC-based methods is O ( d 3 ) , where d is the layer dimension. These methods work well for small models with a 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-GoldfarbShannon (BFGS) [13]-based method for computing and updating the Kronecker factor and has O ( bd 2 ) complexity, where b is the batch size. The class of SNGD (Sherman-Morrison-Woodburybased NGD) methods [23; 31; 17] uses the Sherman-Morrison-Woodbury (SMW) identity to calculate the FIM blocks with a complexity of O ( b 3 ) , making the complexity independent of the model size. HyLo [17] 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 [33] attempts to reduce this cost further to O ( d 2 ) 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 M omentum-Enabled K ronecker-Factorizatoin-Based O ptimizer with R ank-1 Updates. (1) MKOR approximates the inverse of covariance matrices using rank-1 updates in the Sherman-Morrison-Based (SM-Based) Matrix Inversion , reducing the inversion computation complexity from O ( d 3 ) to O ( d 2 ) . As a 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 O ( d 2 ) to O ( d ) . Also, MKOR uses half-precision 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, that combines 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, using a loss-reduction-rate-based switching mechanism.

MKOR outperforms state-of-the-art distributed second- and first-order methods by up to 2 . 57 × , 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-art metrics on the GLUE dataset, where other second-order methods such as KFAC fail to converge to the baseline.

Activation

Covariance

Gradient

Covariance

SM-Based Inversion

a. Rank-1 Approximation & Synchronization -

c. SM-Based Inversion -

d. Precondition & Rescale Gradients

2 Background

Training a neural network involves solving an optimization problem to find the optimal values for a set of weights W = { W m } M m =1 , where M is the number of layers in the network and W m is a matrix in R d × d . 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 1, where w m ∈ R d 2 is the vector representation of W m , F m is the block corresponding to that layer and L is the loss function. [15] Shows that the FIM matches the Gauss-Newton matrix under certain conditions.

KFAC-based methods reformulate the FIM block as the Kronecker product of two matrices. Equation 2 shows the update rule in KFAC, where L is the loss function and ( L t m ) -1 and ( R t m ) -1 are the inverse of the left and right factors, respectively.

( L t m ) -1 and ( R t m ) -1 in Equation 2 are computed using equations 3 and 4, respectively, where a m is the activation value of a sample at layer m , and g m = ∇ a m -1 L and γ incorporate the momentum feature to avoid extreme changes in the factors.

3 MKOR: Momentum-Enabled Kronecker-Factorization-Based Optimizer with Rank-1 Updates

In this section we first present the MKOR algorithm, its 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

3.1 The MKOR Algorithm

Algorithm 1 summarizes the MKOR optimizer for a single layer and Figure 1 shows the workflow. For each layer (line 1 in Algorithm 1) 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 14 in Algorithm 1).

Rank-1 Approximation. For the rank-1 approximations of the covariance matrices, we use the average of the values across all the samples, i.e. a m -1 t = E [ a m -1 t ] and g t m = E [ g t m ] (lines 2 and 3 in Algorithm 1 and Figure 1-a). ( A m -1 t ) -1 : ,i and ( G t m ) -1 : ,i show the i th column of ( A m -1 t ) -1 and ( G t m ) -1 respectively, where A m -1 t and G t m are the activations and the gradients of layer m respectively.

Norm-Based 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, it can cause oscillations or even divergence. MKOR uses a normbased stabilizer to detect this and addresses it by modifying the inverse of the factors accordingly (lines 5 and 6 in Algorithm 1 and Figure 1-b). More detail on the norm-based stabilizer are in Section3.3.

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 E [ g m g m T ] is approximated using a rank-1 matrix g m g m T and using the Sherman-Morrison identity, Equation 5 in obtained (line 7 in Algorithm 1 and Figure 1-c).

Furthermore, if Equation 4 is approximated using E [ a m -1 t a m -1 t T ] ≈ a t m a m T t with a similar derivation, Equation 6 is obtained (line 8 in Algorithm 1 and Figure 1-c).

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 1 and Figure 1-d).

Complexity Analysis. MKORreduces the memory, communication, and computation costs for factor inversion. Table 1 compares the overheads of different optimizers. (1) Computation Complexity. MKOR inverts the left and right factors in Equation 2 using equations 5 and 6, both of which can be computed using matrix-vector multiplications, and have O ( d 2 ) computation complexity, in contrast to KFAC and SNGD methods that need O ( d 3 ) and O ( b 3 ) complexity to invert matrices in R d × d and R b × b respectively. (2) Communication Complexity. The only data that is synchronized among different workers in MKOR is the two rank-1 approximations that have 2 d 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 4 d 2 data transfers. In SNGD , the activations and gradients are synchronized, leading to 2 bd data transfers and the inverted kernels are broadcast, resulting in b 2 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 2 d 2 +2 d 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 4 d 2 memory overhead. SNGD stores the activations, the gradients, and the kernels they use as second-order information, leading to 2 bd + b 2 memory complexity.

The low memory overhead of MKOR allows us to use larger batch sizes compared to other secondorder 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 fairness, 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-order information of the loss functions approach identity near convergence points. Thus we designed a hybrid second- and firstorder optimizer with a loss decrease rate-based switching method (MKOR-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.

3.3 MKORConvergence and Stability

Inversion Frequency Due to 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. f is the reciprocal of inversion frequency and varies from a few 100s to a 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 HyLo is 150 × 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-order 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 implementation 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 µI to each of them to improve numerical stability ( µ > 0 is called the damping factor), but MKOR doesn'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 R b × b of rank min ( 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 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/underflow possibility and the need for damping factors (required by other second-order methods for computational stability).

Lemma 3.1. The factors computed using Equation 5 and 6 are all positive-definite.

[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 2 × while using cheaper computation blocks for half-precision operations. Lemma 3.2 shows an upper bound on the quantization error effect in the MKOR updates.

Lemma 3.2. Assuming that the maximum quantization error is ϵ , the maximum number in matrices and vectors is m , and the dimension of the vectors and matrices are d and d × d respectively, the quantization error of formulas 5 and 6 is O (( γ +4 (1 -γ ) γ 2 m 3 d 2 ) ϵ )

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 firstorder methods, by choosing a learning rate that doesn'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 second-order methods and make them almost as 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 8, where ζ is a hyperparameter that controls the amount of information from the original factors that needs to be saved in the new factors.

By expanding Equation 2 with the new factors, we will get Equation 9, which reduces the loss-based on lemma 3.3. The first term in the right-hand-side of 9 is the KFAC term, the second and third terms are the left and right preconditioned versions, and the last term is the SGD term.

Lemma 3.3. Given a differentiable function L ( w ) with first-order Taylor series approximation ˆ L ( w -∆ w ) = L ( w 0 )∆ w T ∇ w L ( w 0 ) around point w 0 , assuming that at point w 0 the second-order derivative of the function L ( w ) is given as ∇ 2 w L ( w 0 ) = H = L ⊗ R , where L and R are positivesemi-definite matrices, for a value of ∆ w = (( ζL -1 +(1 -ζ ) I ) ⊗ ( ζR -1 +(1 -ζ ) I )) ∇L ( w 0 ) , the inequality ˆ L ( w 0 -∆ w ) < L ( w 0 ) holds.

While this modification can avoid exploding gradients, overusing it with small values of ζ will convert MKORto SGD. MKOR uses a factor norm-based metric that observes the infinity norm of the factors, and if they are greater than a 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 questionanswering 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-training, are used; the latter dataset is not thoroughly available which results in a small reduction in the baseline accuracies achieved in our experiments from the original BERT results. Following [21], 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-based 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 A100 GPU. In addition, HyLo doesn'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 2 shows 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 an end-to-end speedup. MKOR-H will converge in 600 steps, reducing the number of steps in LAMB by 2 . 6 × , while achieving the same accuracy. In addition, it achieves 2 . 57 × speedup over the LAMB optimizer and 1 . 75 × speedup over KAISA. As another second-order baseline, we consider Eva, which converges in 1000 iterations, and MKOR achieves 1 . 69 × speedup over it.

For classification tasks, we fine-tune BERT on the GLUE [29] dataset. Table 3 compares the results for different classification tasks in the GLUE dataset. MKOR with 1500 steps achieves a new stateof-the-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 a factor of 2 . 6 × . MKOR and MKOR-H both achieve 2 . 57 × 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 . 89 × . Eva requires 1000 steps to converge to the target average metric, being 1 . 69 × slower than MKOR-H with 600 steps and 1 . 24% less accurate than MKOR with 1500 steps.

BERT-Large-Uncased Training Loss

0

200

400

600

800

1000

1200

1400

1600

Iteration

2

4

6

8

10

Loss

LAMB

MKOR

K-FAC

MKOR-H

Eva

0

1

2

3

4

5

6

7

8

Time (h)

2

4

6

8

10

Loss

LAMB

MKOR

K-FAC

MKOR-H

Eva

Per figure 2, which shows the training error during the training of BERT, MKOR decreases the error in less iterations in comparison to KAISA, Eva, and LAMB, leading to faster convergence. From Tables 2 and 3, MKOR-H converges in only 600 steps.

Inversion Frequency. Due to the low computation complexity of the updates on MKOR, the factor inversion frequency ( f ) in MKOR is in the range of 10. Figure 4-a shows that while the average iteration cost in KAISA is heavily dependent on the inversion frequency, MKOR's cost is almost independent of the inversion frequency. Also 4-b shows that increasing the inversion frequency leads to higher convergence rate. In addition, using stale factors may result in converging to a 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 use a simple autoencoder [24] on CIFAR-100 [11] in this experiment.

Performance Analysis. We compare the performance of different parts of the optimizers to illustrate the bottlenecks and advantages of 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-Uncased, a transformer-based 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 to the batch size, the factor computation and inversion is more expensive for KAISA compared to HyLo. This cost is significantly reduced in MKOR.

For BERT-Large-Uncased, because of the large size of the model, the factor inversion time for KAISA is large. Also, due to the large sequence length value in this model, the kernel inversion time for

Factor

Computation

Precondition

Update

Weights

10

-

1

10

0

Time (s)

MKOR

Eva

LAMB

Factor

Computation

Precondition

Update

Weights

10

-

2

Time (s)

ResNet-50 - ImageNet

MKOR

HyLo

Eva

SGD

ADAM

(a)

(b)

10

0

10

1

10

2

10

3

0

50

100

150

200

250

Time/Iteration (ms)

Autoencoder

MKOR

10

0

10

1

10

2

10

3

0

1000

2000

3000

4000

5000

6000

7000

Time/Iteration (ms)

BERT-Large-Uncased

MKOR

(a) - Average Time per Iteration

0

5

10

15

20

25

Epochs

0

.

4

0

.

5

0

.

6

0

.

7

0

.

8

0

.

9

Loss

Autoencoder

MKOR - Reuse 10

MKOR - Reuse 100

0

25

50

75

100

125

150

175

Iteration

2

3

4

5

6

Loss

MKOR - Reuse 50

MKOR - Reuse 10

(b) - Test Loss

HyLo is comparable to KAISA's inversion time. But as expected, because of its low computational complexity, the aforementioned cost in our method is much smaller than the total training time, leading to speedups. It is important to 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 matrices, MKOR utilizes rank-1 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 the covariance matrices can be approximated with rank-1 matrices with low error and higher rank approximations are unnecessary in practice. Figure 5 shows the error distribution of the optimal rank-1 approximation methods of the covariance matrices in ResNet-50 and BERT-Large-Uncased pre-training. Our extensive tests

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Activation Error

0

.

00

0

.

05

0

.

10

0

.

15

Probability Density

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Input Gradient Error

0

.

000

0

.

025

0

.

050

0

.

075

0

.

100

0

.

125

0

.

150

Probability Density

(a)

(b)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Activation Error

0

.

0

0

.

1

0

.

2

0

.

3

0

.

4

Probability Density

ResNet-50 Activation Error Distribution

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Input Gradient Error

0

.

00

0

.

02

0

.

04

0

.

06

0

.

08

0

.

10

Probability Density

ResNet-50 Input Gradient Error Distribution

(c)

(d)

on well-known benchmarks show this property holds for all models and we have not come across a 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 = XX T , where C ∈ R d × d is the covariance matrix and X ∈ R d × 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 ) . If the per-GPU batch sizes are small, the covariance matrices in each GPU will be low-rank. Rank-1 approximation methods can work well in these scenarios. If 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.

Extending MKOR to Higher Ranks: Furthermore, one can extend MKOR to use higher-rank covariance matrices. Let's assume that C = ∑ r i =1 c i c T i where r is the rank of the covariance matrix C . We can apply SMW identity to compute C new 1 = ( C old + c 1 c T 1 ) -1 with O ( d 2 ) computational complexity. Then we can compute C new 2 = ( C new 1 + c 2 c T 2 ) -1 using SMW identity with O ( d 2 ) computational complexity. We can continue the same pattern by computing C new i = ( C new i -1 + c i c T i ) -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 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.

5 Conclusion

Wepropose MKOR, a Momentum-Enabled Kronecker-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-art first- and second-order optimizers on large language models.

6 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, it should be acknowledged 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 (RGPIN06516, DGECR00303, RGPIN-2023-04897, DGECR-2023-00133), National Science Foundation (NSF CCF-2106621), the Canada Research Chairs program, the Ontario Early Researcher Award, the Digital Research Alliance of Canada ( https://www.alliancecan.ca ) and Texas Advanced Computing Center ( https://www.tacc.utexas.edu ). Work of Zhao Zhang was supported by OAC-2106661.

References

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. arXiv preprint arXiv:2002.09018 , 2020.

Argonne Leadership Computing Facility. Polaris. https://www.alcf.anl.gov/ polaris .

Nihat Ay. On the locality of the natural gradient for learning in deep bayesian networks. Information Geometry , pp. 1-49, 2020.

Jimmy Ba, Roger Grosse, and James Martens. Distributed second-order optimization using kronecker-factored approximations. 2016.

Compute Canada. Compute canada. https://computecanada.ca/ .

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pp. 248-255. Ieee, 2009.

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.

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.

Jeremy Howard and Sebastian Ruder. Universal language model fine-tuning for text classification. arXiv preprint arXiv:1801.06146 , 2018.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Communications of the ACM , 60(6):84-90, 2017.

Dong C Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical programming , 45(1):503-528, 1989.

Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983 , 2016.

James Martens. New insights and perspectives on the natural gradient method. The Journal of Machine Learning Research , 21(1):5776-5851, 2020.

James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning , pp. 2408-2417. PMLR, 2015.

Baorun Mu, Saeed Soori, Bugra Can, Mert Gürbüzbalaban, and Maryam Mehri Dehnavi. Hylo: a 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.

NVIDIA Corporation. Nvidia deep learning examples. https://github.com/NVIDIA/ DeepLearningExamples .

Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and Satoshi Matsuoka. Large-scale distributed second-order optimization using kronecker-factored approximate curvature for deep convolutional neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 12359-12367, 2019.

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.

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.

Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. arXiv preprint arXiv:1606.05250 , 2016.

Yi Ren and Donald Goldfarb. Efficient subsampled gauss-newton and natural gradient methods for training neural networks. arXiv preprint arXiv:1906.02353 , 2019.

Jürgen Schmidhuber. Deep learning in neural networks: An overview. Neural networks , 61: 85-117, 2015.

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.

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:2011.04006 , 2020.

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.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461 , 2018.

Wikipedia. Wikipedia corpus. https://meta.wikimedia.org/wiki/Data_dump_ torrents#English_Wikipedia .

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.

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:1904.00962 , 2019.

Lin Zhang, Shaohuai Shi, and Bo Li. Eva: A general vectorized approximation framework for second-order optimization, 2023.

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 IEEE international conference on computer vision , pp. 19-27, 2015.

8 Supplementary Material

In this section, we start by discussing an experiment on 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 report the GLUE results achieved on each task in 8.2. Next, we discuss the derivation of KFAC and SNGD approximation methods in 8.3. We then discuss some of the features of MKOR an other optimizers and back them up with quantitative data in 8.4 and 8.5. We will provide scalability results of MKOR in 8.6 and provide more data to back up the low-rank features of the covariance matrices in 8.7. In 8.9 and 8.10, and 8.11 we discuss the training setup used in 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.

8.1 ResNet-50

We train ResNet-50, a convolutional neural network with more than 25M parameters, on ImageNet, a text classification task with more than 1.2M samples. The same setup is used in [21], and SGD is used as the first-order baseline.

The target accuracy in this experiment is 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 . 49 × speedup over SGD. KAISA, the second-order baseline converges in 54 epochs, but due to its expensive steps, MKOR still converges 1 . 04 × 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 1 for comparison of complexity between methods.

We could not reproduce the ResNet-50 results of Eva [33] on ImageNet because the hyperparameters are not reported. We tried to tune Eva on multiple settings and none converged to desired accuracy. 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 [16], dated to 2015. A number of followup works, mentioned in 1, have provided faster implementations of KFAC. We use KAISA [21], 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.

ResNet-50 Test Accuracy

0

20

40

60

80

Epochs

0

20

40

60

80

88

57

54

ResNet-50

SGD

MKOR

KFAC

0

5000

10000

15000

20000

25000

30000

35000

Time (s)

0

20

40

60

80

32815

22002

22978

SGD

MKOR

KFAC

(a)

(b)

ζΓ↽ΓDΓFΓWΓRΓUΓHΓGΓΛΓϱΓDΓVΓHΓGΓΛ

Γ5ΓHΓVΓKΓDΓSΓH

ΓDΓηΓΛΓϱΓOΓRΓFΓNΓΛΓφΓLΓDΓJΓRΓQΓDΓOΓΛ

ΓEΓηΓΛΓ▷ΓUΓRΓQΓHΓFΓNΓHΓU

ΓΛΓ↪ΓGΓHΓQΓWΓLΓWΓ♮

ΓϖΓSΓSΓUΓR ♭ΓLΓPΓDΓWΓRΓLΓQ

8.2 GLUE Results

Wediscussed the speedup achieved using MKOR on the GLUE dataset in Section 4. For completeness, Table 4 shows the metrics achieved in each of the different GLUE tasks on BERT-Large-Uncased trained on different optimizers.

8.3 Derivation of NGD Approximations

Natural Gradient Descent (NGD) In NGD, which is a second-order method, we use the inverse of the Fisher Information Matrix (FIM) as a preconditioner to the gradients as shown in Figure 7-a. Equation 1 shows the update rule of NGD , where F m is the FIM block corresponding to block m . Equation 10 shows the definition of FIM for an arbitrary layer in our model, where x i is the i th sample in the batch.

Kronecker Factorization (KFAC). KFAC methods reformulate the FIM block as the Kronecker product of two matrices as shown in Equation 11 where g m = ∇ x m L and a m is the vector form of the activation output of layer m and E is the expectation operator and x m is the input matrix of layer m . Please note that we have used the mixed-product property of Kronecker multiplication for getting the right hand value.

Furthermore, we assume that E [( g m g mT ) ⊗ ( a m -1 a m -1 T )] ≈ E [ g m g mT ] ⊗ E [ a m -1 a m -1 T ] , which is 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.

By using the mixed Kronecker matrix-vector product property, we can get the update value in Equation 2, which is illustrated in Figure 7-b. We refer to L m and R m 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 3 and 4.

Sherman-Morrison-Woodbury-Based Natural Gradient Descent (SNGD). In this method, the SMW identity is used for approximating the inverse of ( F m + µI ) ∈ R d 2 × d 2 , where µ is a damping factor used in the preconditioning. Equation 13 shows the process of computing the inverse of the FIM for a single layer in the network, where A m ∈ R d × b is the batch of activations of layer l and G m ∈ R d × b is the batch of gradients of the loss function with respect to the inputs of that layer and U = [ ∇ W m L ( W , x 1 ) , ..., ∇ W m L ( W , x b )] T ∈ R d 2 × b 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 b × b is inverted, as shown in Figure 7-c.

8.4 Numerical Instability of Second-order Methods

In Section 3.3, we 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 on 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.

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

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 a lot of time tuning this hyperparameter. More specifically, in first-order methods, a large learning rate can easily lead to divergence and numerical instability, and in second-order methods large learning rates can lead to exploding gradient descent as discussed in Section 3.3. Using small learning rates can lead to slow convergence in both first- and second-order 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 is its robustness against a wide range of learning rates. As Table 5 shows, first-order methods are extremely sensitive to learning rate, and the secondorder 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 a high

ResNet-50 - CIFAR-10 - Eigenvalues of Right Factor

0

100

200

300

Iteration (x1000)

10

-

3

10

-

2

10

-

1

10

0

Eigenvalue

Minimum Eigenvalues

Maximum Eigenvalues

(a)

ResNet-50 - CIFAR-10 - Condition Number of Right Factor

0

100

200

300

Iteration (x1000)

10

0

10

1

10

2

10

3

Condition Number

Condition Number

(b)

convergence rate for a wide range of learning rates, and by directly modifying the inverse of factors as discussed in 3.3 can find a 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 9 shows the strong scalability of MKOR on BERT-Large-Uncased on up to 64 GPUs.

0

10

20

30

40

50

60

#GPUs

20

30

40

50

60

Speedup

BERT Large Uncased Scalability

Num GPUs

MKOR

0

200

400

600

800

1000

1200

Iteration

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

Average Covariance Rank-1 Error

ResNet-50 Average Covariance Rank-1 Error

Gradient

Activation

8.7 Decaying Eigenvalues and Rank-1 Approximations

Figure 10 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 work independently in the beginning of the training, but as 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 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 m -1 t -1 -1 and R t m -1 -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.

8.8 Memory Overheads

The memory overheads of MKOR in comparison to other optimizers are reported in Table 6. 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 × .

8.9 Hyperparameters

BERT-Large-Uncased. For the BERT-Large-Uncased pre-training, we use the same hyperparameters used in [18]. The factors in KAISA are updated every 50 iterations, and the factors in MKOR and MKOR-H are updated every 10 iterations.

ResNet-50. For SGD and KAISA, we use 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 in MKOR decays by a factor of 2 at the end of epochs 25, 35, 40, 45, 50, 55, and 56.

8.10 Experiment Settings

For the BERT-Large-Uncased pre-training and fine-tuning experiments, we have used up to 64 A100 GPUs on the Polaris [3] cluster which has 560 nodes, each with 4 NVIDIA A-100 GPUs with NVLink interconnects.

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 our training and fine-tuning experiments. In addition, the results in Figures 3 and 4-a are gathered on a single 4-GPU node of V100s. The results of Figure 4-b are gathered on 64 A100 GPUs.

8.11 Reproduction

Our code base is publicly available on https://github.com/Mohammad-Mozaffari/ mkor , and the instructions for running each experiment are available there.

8.12 Training Accuracy Experiments

To evaluate MKOR as an optimizer that tries to minimize a specific objective function, we 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 a non-zero weight decay adds a 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 as residual networks [9]. In our training accuracy experiment, we use another CNN benchmark, AlexNet [12] with more than 20M parameters on CIFAR-100 [11] consisting of 50K training and 10K validation images of 100 classes. Figures 11-a and 12-a show the convergence properties of different optimizers. MKOR is 1 . 26 × , 1 . 31 × , and 1 . 58 × 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 is a large language model with two variants, BERT-Base with more than 108M parameters and Bert-Large with more than 335M parameters. As shown in Figures 11-c and 12-c, we have fine-tuned BERT-Large on the IMDB dataset which is a text classification task with 25K training and 25K test samples. MKOR outperforms SGD and HyLo-KIS by a speedup factor of 1 . 22 × and 1 . 43 × respectively. We have also fine-tuned BERT-Base on SQuAD dataset, which is a question answering task with 87.6K training and 10.6K test samples. MKOR achieves 1 . 26 × and 1 . 56 × speedup over SGD and HyLo respectively. We couldn't find a learning rate with which KAISA would converge on any of our BERT experiments which are based on the HuggingFace implementation of BERT, and the reason for lack of convergence of KAISA is exploding gradients.

8.13 Knee-Point Learning Rate Scheduler

While using large learning rates is crucial for utilizing the higher convergence rate of second-order methods, it is necessary to reduce the convergence rate after some iterations so that the model

0

10000 20000 30000 40000 50000 60000 70000 80000

Time (s)

0

20

40

60

80

100

78474

51275

63918

BERT Base Cased Question Answering - SQuAD

HyLo

MKOR

SGD

0

10000

20000

30000

40000

50000

Time (s)

0

20

40

60

80

100

15639

22323

19154

BERT Large Cased Text Classification - IMDB

MKOR

HyLo-KIS

SGD

(a)

(b)

0

250

500

750

1000

1250

1500

1750

2000

Time (s)

0

20

40

60

80

100

2119

1400

1509

1562

AlexNet - CIFAR-100

MKOR

HyLo-KIS

SGD

(c)

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 necessarily work ideally in all optimization problems. We used knee-point learning rate scheduler in Section 8.12.

To fully utilize the potentials of the optimizers, we have designed a learning rate scheduler that monitors the rate of improvement in accuracy or decrease in the loss function value, and based on that decides when to decrease the learning rate. The scheduler detects knee-points in the accuracy/loss and decreases the learning rate when a knee-point is observed.

By definition, knee-points are defined as the points where the average accuracy/loss rate is less than β times the increment/decrement in the accuracy/loss since using the current learning-rate. For averaging the accuracy/loss rate, we use an exponential moving average, and β is a hyperparameter that we can choose to show how much the scheduler can tolerate lack of improvement to detect the accuracy/loss.

0

5

10

15

20

25

30

35

40

Epochs

0

20

40

60

80

100

11

20

22

BERT Large Cased Text Classification - IMDB

MKOR

HyLo-KIS

SGD

0

10

20

30

40

50

60

Epochs

0

20

40

60

80

100

60

37

57

BERT Base Cased Question Answering - SQuAD

HyLo

MKOR

SGD

(a)

(b)

0

50

100

150

200

250

300

350

400

Epochs

0

20

40

60

80

100

395

261

280

293

AlexNet - CIFAR-100

MKOR

HyLo-KIS

SGD

(c)

8.14 Proofs

Lemma 3.1 :

Proof. Given a matrix positive definite matrix J t -1 and a vector j and scalar 0 < γ < 1 , we show that Equation 14 results in a positive-definite matrix.

Since γ > 0 and J t -1 is a positive-definite matrix, γJ t -1 is also positive-definite.

̸

Also, since J t -1 is a positive-definite matrix, ∀ x = 0 : x T J t -1 x > 0 .

Now we show that J t -1 -1 jj T J t -1 -1 is also positive-definite.

̸

Since both matrices on the right-hand-side of Equation 14 are positive-definite and the sum of two positive-definite matrices is a positive-definite matrix, the left-hand-side of it will be also positive definite.

Lemma 3.3:

Proof. By defining P = ( ζL -1 +(1 -ζ ) I ) ⊗ ( ζR -1 +(1 -ζ ) I ) , Equation 17 will hold.

Now, we will show that P is a positive-semi-definite matrix. By using the associativity property of Kronecker multiplication, we can represent P as in Equation 18. Please note that different identity matrices in 18 have different shapes.

Since matrices L and R are positive-semi-definite, the Kronecker products L -1 ⊗ R -1 , L -1 ⊗ I , and I ⊗ R -1 are also positive-semi-definite. As a result, for any non-zero vector x , x T Px > 0 . So, based on Equation 17, -∇L ( w 0 ) T P ∇L ( w 0 ) . As a result, Equation 19 holds.

Lemma 3.2

Proof. Considering the quantization of matrix J and vector j in Equation 20 and assuming the maximum quantization error is ϵ and the maximum values in vector j and matrix J is m , we can consider one of the three possible cases:

Vector-Vector Dot Product : The resulting error is O (2 ϵm 2 d ) because the error of each multiplication can at most be 2 mϵ + ϵ 2 and by adding the d multiplications, the maximum error can be (2 m +1) dϵ .

Vector-Matrix Product : The resulting error is O (2 ϵm 2 d ) because a vector-matrix product can be considered as multiple independent vector-vector dot products.

Vector-Vector Outer Product : The resulting error is O (2 ϵm ) because each element in the resulting matrix is computed using a multiplication with the maximum error of 2 ϵm 2 + ϵ 2 .

The error imposed by quantizing γJ is at most γϵ . The quantization error in the denominator of the fraction in 20 is negligible in comparison to 1 and won't change the growth order of the final error term. The quantization error of Jj is O (2 m 2 dϵ ) as discussed earlier in the case of vector-matrix product. Jjj T J T = ( Jj )( Jj ) T is a vector-vector outer product, resulting in O (4 m 3 d 2 ) error.

So the final error quantization error is O (( γ +4 (1 -γ ) γ 2 m 3 d 2 ) ϵ )

Peer reviews:
comments: 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.

STRENGTHS:
1. Efficient second-order optimization algorithm is important for the training of DNNs.

2. Experimental results show great speed-ups compared to the baseline algorithms.

3. 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?

QUESTIONS:
n/a

LIMITATIONS:
n/a
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-06-30

comments: 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.

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.

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.

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.
[1] SeKron: A Decomposition Method Supporting Many Factorization Structures arXiv:2210.06299
 
- Does adding a learning rate scalar in Equation (2) make the method more flexible to converge?  
- 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

comments: 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.

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.

The 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).

While this additional experiment would greatly improve the paper, the experimental methodology in this paper is already relatively robust.

QUESTIONS:
No questions in particular.

LIMITATIONS:
Limitations are not discussed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nSee Questions.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-10"}