SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\hat {L _ {t} ^ {m}} = \zeta L _ {t} ^ {m} + (1 - \zeta) I \tag {7}
$$
$$
\hat {R} _ {t} ^ {m} = \zeta R _ {t} ^ {m} + (1 - \zeta) I \tag {8}
$$
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.
$$
\begin{array}{l} \hat {L} ^ {m - 1} \nabla_ {W ^ {m}} \mathcal {L} \hat {R} ^ {m - 1} = \zeta^ {2} L ^ {m - 1} \nabla_ {W ^ {m}} \mathcal {L} R ^ {m - 1} \\ + \zeta (1 - \zeta) L ^ {m ^ {- 1}} \nabla_ {W ^ {m}} \mathcal {L} + \zeta (1 - \zeta) \nabla_ {W ^ {m}} \mathcal {L} R ^ {m ^ {- 1}} + (1 - \zeta) ^ {2} \nabla_ {W ^ {m}} \mathcal {L} \tag {9} \\ \end{array}
$$
Lemma 3.3. Given a differentiable function  $\mathcal{L}(w)$  with first-order Taylor series approximation  $\hat{\mathcal{L}}(w - \Delta w) = \mathcal{L}(w_0)\Delta w^T\nabla_w\mathcal{L}(w_0)$  around point  $w_0$ , assuming that at point  $w_0$  the second-order derivative of the function  $\mathcal{L}(w)$  is given as  $\nabla_w^2\mathcal{L}(w_0) = H = L \otimes R$ , where  $L$  and  $R$  are positive-semi-definite matrices, for a value of  $\Delta w = ((\zeta L^{-1} + (1 - \zeta)I) \otimes (\zeta R^{-1} + (1 - \zeta)I))\nabla\mathcal{L}(w_0)$ , the inequality  $\hat{\mathcal{L}}(w_0 - \Delta w) < \mathcal{L}(w_0)$  holds.
While this modification can avoid exploding gradients, overusing it with small values of  $\zeta$  will convert MKOR to 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.
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-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 \times$ , while achieving the same accuracy. In addition, it achieves  $2.57 \times$  speedup over the LAMB optimizer and  $1.75 \times$  speedup over KAISA. As another second-order baseline, we consider Eva, which converges in 1000 iterations, and MKOR achieves  $1.69 \times$  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 state-of-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 \times$ . MKOR and MKOR-H both achieve  $2.57 \times$  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 \times$ . Eva requires 1000 steps to converge to the target average metric, being  $1.69 \times$  slower than MKOR-H with 600 steps and  $1.24 \%$  less accurate than MKOR with 1500 steps.
Table 2: 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 (×)	1.00	1.39	1.51	2.57	1.52
BERT-Large-Uncased Training Loss
Figure 2: The training loss of BERT-Large-Uncased using different optimizers.
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
Table 3: BERT-Large Uncased results on the GLUE classification tasks.
Optimizer	Iterations	Time (h)	Speedup (×)	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
Figure 3: Per-step breakdown of different optimizes on BERT-Large-Uncased (a) and ResNet-50 (b)
(a) - Average Time per Iteration
(b) - Test Loss
Figure 4: The sensitivity of MKOR and KAISA for BERT-Large-Uncased and an Autoencoder model (a) and the effect of inversion frequency on the convergence properties of these models (b).
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
BERT Activation Error Distribution
(a)
BERT Input Gradient Error Distribution
(b)
ResNet-50 Activation Error Distribution
(c)
Figure 5: Rank-1 error for activation and input gradient covariance matrices for BERT-Large-Uncased pre-training (a, b) and ResNet-50 on ImageNet (c, d).
ResNet-50 Input Gradient Error Distribution
(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 = XXT$ , 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)$ . 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 = \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 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.
We propose 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.
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.
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.
[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.alcf.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-order optimization using kronecker-factored approximations. 2016.
[6] Compute Canada. Compute canada. https://computecanada.ca/.
[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-255. IEEE, 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 classification. arXiv preprint arXiv:1801.06146, 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.
[30] Wikipedia. Wikipedia corpus. https://meta.wikimedia.org/wiki/Data_dump_torrents#English_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:1904.00962, 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 IEEE international conference on computer vision, pp. 19-27, 2015.
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 \times$  speedup over SGD. KAISA, the second-order baseline converges in 54 epochs, but due to its expensive steps, MKOR still converges  $1.04 \times$  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
(a)
Figure 6: Test accuracy of ResNet-50 on ImageNet for MKOR, KAISA, and SGD on 64 GPUs.
Table 4: BERT-Large-Uncased Results on the GLUE classification tasks.
Optimizer	Iterati- ons	MNLI (acc)	QQP (F1)	QNLI (acc)	SST- 2 (acc)	COLA (mcc)	STS- B (corr)	MRPC (F1)	RTE (acc)	Avera- 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-H	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
Figure 7: Approximations in second-order methods.
8.2 GLUE Results
We discussed 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.
$$
F ^ {m} = \frac {1}{b} \sum_ {i = 1} ^ {b} \nabla_ {x _ {i}} \ell (\mathcal {W}, x _ {i}) \nabla_ {x _ {i}} \ell (\mathcal {W}, x _ {i}) ^ {T} \tag {10}
$$
Kronecker Factorization (KFAC). KFAC methods reformulate the FIM block as the Kronecker product of two matrices as shown in Equation 11 where  $g^{m} = \nabla_{x^{m}}\mathcal{L}$  and  $a^m$  is the vector form of the activation output of layer  $m$  and  $\mathbb{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.
$$
F ^ {m} = \mathbb {E} \left[ \left(g ^ {m} \otimes a ^ {m - 1}\right) \left(g ^ {m} \otimes a ^ {m - 1}\right) ^ {T} \right] = \mathbb {E} \left[ \left(g ^ {m} g ^ {m T}\right) \otimes \left(a ^ {m - 1} a ^ {m - 1} ^ {T}\right) \right] \tag {11}
$$
Furthermore, we assume that  $\mathbb{E}[(g^m g^{mT})\otimes (a^{m - 1}a^{m - 1^T})]\approx \mathbb{E}[g^m g^{mT}]\otimes \mathbb{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.
$$
F ^ {m - 1} w ^ {m} = \mathbb {E} \left[ g ^ {m} g ^ {m T} \right] ^ {- 1} \otimes \mathbb {E} \left[ a ^ {m - 1} a ^ {m - 1 ^ {T}} \right] ^ {- 1} w ^ {m} \tag {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 + \mu I) \in \mathbb{R}^{d^2 \times d^2}$ , where  $\mu$  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 \in \mathbb{R}^{d \times b}$  is the batch of activations of layer  $l$  and  $G^m \in \mathbb{R}^{d \times b}$  is the batch of gradients of the loss function with respect to the inputs of that layer and  $U = [\nabla_{W^m} \mathcal{L}(\mathcal{W}, x_1), \dots, \nabla_{W^m} \mathcal{L}(\mathcal{W}, x_b)]^T \in \mathbb{R}^{d^2 \times b}$  is the concatenation of the gradients of the loss function with respect to the parameters of that layer and  $\odot$  shows the Hadamard element-wise product. In this method, a kernel matrix in  $\mathbb{R}^{b \times b}$  is inverted, as shown in Figure 7-c.
$$
(F ^ {m} + \mu I) ^ {- 1} = \frac {1}{\mu} (I - U ^ {m} \left(A ^ {m - 1 ^ {T}} A ^ {m - 1} \odot G ^ {m T} G ^ {m} + \mu I\right) ^ {- 1} U ^ {m T}) \tag {13}
$$
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