SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 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 a high
ResNet-50 - CIFAR-10 - Eigenvalues of Right Factor
(a)
ResNet-50 - CIFAR-10 - Condition Number of Right Factor
(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 are 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 CIFAR10. MKOR is the least sensitive optimizer to learning rate, converging in almost the same number of iterations for a wide range of learning rate, while other optimizers either diverge (D) or converge to a local-minimum (* superscript).
Learning Rate Optimizer	10	1	0.1	0.01
MKOR	94	79	78	76
KAISA	112	100	90	89*
HyLo	D	123*	98	150*
SGD	D	D	108	145*
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.
Figure 9: Scalability of MKOR.
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 on BERT-Large-Uncased pre-trainingResNet-50 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 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_{t - 1}^{m - 1^{-1}}$  and  $R_{t - 1}^{m}$  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 \times$ .
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 as 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.
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.1M	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-Cased	108.9M	SQuAD v1.1	87.6K	10.6K	V100	4
BERT-Large-Cased	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 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 \times$ ,  $1.31 \times$ , and  $1.58 \times$  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 \times$  and  $1.43 \times$  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 \times$  and  $1.56 \times$  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
(c)
Figure 11: Training time for distributed first- and second-order optimizers SGD, MKOR, KAISA, and HyLo on BERT-Large-Cased on IMDB (a), BERT-Base-Cased on SQuAD (b), and AlexNet 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 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  $\beta$  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  $\beta$  is a hyperparameter that we can choose to show how much the scheduler can tolerate lack of improvement to detect the accuracy/loss.
(c)
Figure 12: Training accuracy vs. the number of epochs for distributed first- and second-order optimizers SGD, MKOR, KAISA, and HyLo on BERT-Large-Cased on IMDB (a), BERT-Base-Cased on SQuAD (b), and AlexNet on CIFAR-100 (c). In all the experiments, MKOR outperforms other optimizers in convergence rate.
8.14 Proofs
Lemma 3.1:
Proof. Given a matrix positive definite matrix  $J_{t-1}$  and a vector  $j$  and scalar  $0 < \gamma < 1$ , we show that Equation 14 results in a positive-definite matrix.
$$
J _ {t} ^ {- 1} = \gamma J _ {t - 1} ^ {- 1} + \frac {(1 - \gamma)}{\gamma^ {2} (1 + \gamma (1 - \gamma) j ^ {T} J _ {t - 1} ^ {- 1} j)} J _ {t - 1} ^ {- 1} j j ^ {T} J _ {t - 1} ^ {- 1} \tag {14}
$$
Since  $\gamma > 0$  and  $J_{t-1}$  is a positive-definite matrix,  $\gamma J_{t-1}$  is also positive-definite.
Also, since  $J_{t-1}$  is a positive-definite matrix,  $\forall x \neq 0: x^T J_{t-1}x > 0$ .
$$
j ^ {T} L _ {t - 1} ^ {- 1} j > 0 \xrightarrow {0 <   \gamma <   1} 1 + \gamma (1 - \gamma) j ^ {T} J _ {t - 1} ^ {- 1} j > 0 \tag {15}
$$
Now we show that  $J_{t - 1}^{-1}jj^{T}J_{t - 1}^{-1}$  is also positive-definite.
$$
\forall x \neq 0: x ^ {T} J _ {t - 1} ^ {- 1} j j ^ {T} J _ {t - 1} ^ {- 1} x = \left(j ^ {T} J _ {t - 1} ^ {- 1} x\right) ^ {T} \left(j ^ {T} J _ {t - 1} ^ {- 1} x\right) = \| \left(j ^ {T} J _ {t - 1} ^ {- 1} x\right) \| ^ {2} > 0 \tag {16}
$$
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 = (\zeta L^{-1} + (1 - \zeta)I) \otimes (\zeta R^{-1} + (1 - \zeta)I)$ , Equation 17 will hold.
$$
\hat {\mathcal {L}} \left(w _ {0} - \Delta w\right) - \mathcal {L} \left(w _ {0}\right) = - \nabla \mathcal {L} \left(w _ {0}\right) ^ {T} P \nabla \mathcal {L} \left(w _ {0}\right) \tag {17}
$$
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.
$$
P = \zeta^ {2} L ^ {- 1} \otimes R ^ {- 1} + \zeta (1 - \zeta) L ^ {- 1} \otimes I + \zeta (1 - \zeta) I \otimes R ^ {- 1} + I \tag {18}
$$
Since matrices  $L$  and  $R$  are positive-semi-definite, the Kronecker products  $L^{-1} \otimes R^{-1}$ ,  $L^{-1} \otimes I$ , and  $I \otimes R^{-1}$  are also positive-semi-definite. As a result, for any non-zero vector  $x$ ,  $x^T P x > 0$ . So, based on Equation 17,  $-\nabla \mathcal{L}(w_0)^T P \nabla \mathcal{L}(w_0)$ . As a result, Equation 19 holds.
$$
\hat {\mathcal {L}} \left(w _ {0} - \Delta w\right) <   \mathcal {L} \left(w _ {0}\right) \tag {19}
$$
Lemma 3.2
Proof. Considering the quantization of matrix  $J$  and vector  $j$  in Equation 20 and assuming the maximum quantization error is  $\epsilon$  and the maximum values in vector  $j$  and matrix  $J$  is  $m$ , we can consider one of the three possible cases:
1. Vector-Vector Dot Product: The resulting error is  $O(2\epsilon m^2 d)$  because the error of each multiplication can at most be  $2m\epsilon + \epsilon^2$  and by adding the  $d$  multiplications, the maximum error can be  $(2m + 1)d\epsilon$ .
2. Vector-Matrix Product: The resulting error is  $O(2\epsilon m^2 d)$  because a vector-matrix product can be considered as multiple independent vector-vector dot products.
3. Vector-Vector Outer Product: The resulting error is  $O(2\epsilon m)$  because each element in the resulting matrix is computed using a multiplication with the maximum error of  $2\epsilon m^2 + \epsilon^2$ .
$$
\gamma J + \frac {(1 - \gamma)}{\gamma^ {2} (1 + \gamma (1 - \gamma) j ^ {T} J j)} J j j ^ {T} J ^ {T} \tag {20}
$$
The error imposed by quantizing  $\gamma J$  is at most  $\gamma \epsilon$ . 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(2m^2 d\epsilon)$  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(4m^3 d^2)$  error.
So the final error quantization error is  $O((\gamma + 4 \frac{(1 - \gamma)}{\gamma^2} m^3 d^2) \epsilon)$