SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mohammad Mozaffari
mmozaffari@cs.toronto.edu
Sikan Li
Texas Advanced Computing Server
sli@tacc.utexas.edu
Zhao Zhang
zhao.zhang@rutgers.edu
Maryam Mehri Dehnavi
mmehride@cs.toronto.edu
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 back 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 \times$  and  $1.85 \times$  respectively on BERT-Large-Uncased on 64 GPUs.
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
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 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  $\mathcal{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-Goldfarb-Shannon (BFGS) [13]-based method for computing and updating the Kronecker factor and has  $\mathcal{O}(bd^2)$  complexity, where  $b$  is the batch size. The class of SNGD (Sherman-Morrison-Woodbury-based NGD) methods [23; 31; 17] uses the Sherman-Morrison-Woodbury (SMW) identity to calculate the FIM blocks with a complexity of  $\mathcal{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  $\mathcal{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 Momentum-Enabled Kronecker-Factorizatoin-Based Optimizer with Rank-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  $\mathcal{O}(d^3)$  to  $\mathcal{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  $\mathcal{O}(d^2)$  to  $\mathcal{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\times$  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.
Figure 1: MKOR for layer  $m$  on a single worker. The inputs of MKOR are the activations  $A_{t}^{m}$ , the gradients of the loss function with respect to the inputs  $G_{t}^{m}$ , and the gradients of the loss function with respect to the weights  $\nabla_{W^m}\mathcal{L}$ . The output is the update values  $\Delta W^{m}$ .
Training a neural network involves solving an optimization problem to find the optimal values for a set of weights  $\mathcal{W} = \{W^m\}_{m=1}^M$ , where  $M$  is the number of layers in the network and  $W^m$  is a matrix in  $\mathbb{R}^{d \times 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 \in \mathbb{R}^{d^2}$  is the vector representation of  $W^m$ ,  $F^m$  is the block corresponding to that layer and  $\mathcal{L}$  is the loss function. [15] Shows that the FIM matches the Gauss-Newton matrix under certain conditions.
$$
w ^ {m} := w ^ {m} - \alpha \left(F ^ {m}\right) ^ {- 1} \nabla_ {w ^ {m}} \mathcal {L} \tag {1}
$$
KFAC-based methods reformulate the FIM block as the Kronecker product of two matrices. Equation 2 shows the update rule in KFAC, where  $\mathcal{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.
$$
W ^ {m} := W ^ {m} - \alpha \left(L _ {t} ^ {m}\right) ^ {- 1} \nabla_ {W ^ {m}} \mathcal {L} \left(R _ {t} ^ {m}\right) ^ {- 1} \tag {2}
$$
$(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} = \nabla_{a^{m - 1}}\mathcal{L}$  and  $\gamma$  incorporate the momentum feature to avoid extreme changes in the factors.
$$
L _ {t} ^ {m} = \gamma L _ {t - 1} ^ {m} + (1 - \gamma) \mathbb {E} \left[ g _ {t} ^ {m} g _ {t} ^ {m T} \right] \tag {3}
$$
$$
R _ {t} ^ {m} = \gamma R _ {t - 1} ^ {m} + (1 - \gamma) \mathbb {E} [ a _ {t} ^ {m - 1} a _ {t} ^ {m - 1 ^ {T}} ] \tag {4}
$$
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$
Data:  $A_{t}^{m - 1},G_{t}^{m},W_{t - 1}^{m}$  Result:  $W_{t}^{m}$    
1 if  $m\in$  Second Order Layers then   
2  $\mathbf{a}_{\mathbf{t}}^{\mathbf{m} - 1}\gets \frac{1}{b}\sum_{i = 1}^{b}(A_{t}^{m - 1})_{:,i};$  // Approx:  $A_{t}^{m - 1}A_{t}^{m - 1T}\approx \mathbf{a}_{t}^{m - 1}\mathbf{a}_{t}^{m - 1T}$    
3  $\mathbf{g}_{\mathbf{t}}^{\mathbf{m}}\leftarrow \frac{1}{b}\sum_{i = 1}^{b}(G_{t}^{m})_{:,i};$  // Approx:  $G_{t}^{m}G_{t}^{mT}\approx \mathbf{g}_{\mathbf{t}}^{\mathbf{m}}\mathbf{g}_{\mathbf{t}}^{\mathbf{m}T}$    
4  $\mathbf{a}_{\mathbf{t}}^{\mathbf{m} - 1},\mathbf{g}_{\mathbf{t}}^{\mathbf{m}}\leftarrow$  AllReduce  $(\mathbf{a}_{\mathbf{t}}^{\mathbf{m} - 1},\mathbf{g}_{\mathbf{t}}^{\mathbf{m}})$  // Synchronize Approximations   
// Norm-Based Stabilization   
5  $L_{t - 1}^{\hat{m}}^{-1}\gets$  if  $\| L_{t - 1}^{m}^{-1}\| >\epsilon$  then  $\zeta L_{t - 1}^{m}^{-1} + (1 - \zeta)I$  else  $L_{t - 1}^{m}^{-1}$    
6  $R_{t - 1}^{\hat{m}}^{-1}\gets$  if  $\| R_{t - 1}^{m}^{-1}\| >\epsilon$  then  $\zeta R_{t - 1}^{m}^{-1} + (1 - \zeta)I$  else  $R_{t - 1}^{m}^{-1}$    
// SM-Based Factor Inversion   
7  $L_{t}^{m - 1}\gets \gamma L_{t - 1}^{\hat{m}}^{-1} + \frac{(1 - \gamma)}{\gamma^{2}(1 + \gamma(1 - \gamma)\mathbf{g}_{\mathbf{t}}^{\mathbf{m}T}L_{t - 1}^{\hat{m}}^{-1}\mathbf{g}_{\mathbf{t}}^{\mathbf{m}})}L_{t - 1}^{\hat{m}}^{-1}\mathbf{g}_{\mathbf{t}}^{\mathbf{m}}\mathbf{g}_{\mathbf{t}}^{\mathbf{m}T}L_{t - 1}^{\hat{m}}^{-1}$    
8  $R_{t}^{m - 1} = \gamma R_{t - 1}^{\hat{m}}^{-1} + \frac{(1 - \gamma)}{\gamma^{2}(1 + \gamma(1 - \gamma)\mathbf{a}_{\mathbf{t}}^{\mathbf{m}T}R_{t - 1}^{\hat{m}}^{-1}\mathbf{a}_{\mathbf{t}}^{\mathbf{m}})}R_{t - 1}^{\hat{m}}^{-1}\mathbf{a}_{\mathbf{t}}^{\mathbf{m}}\mathbf{a}_{\mathbf{t}}^{\mathbf{m}T}R_{t - 1}^{\hat{m}}^{-1}$    
9  $\Delta W_t^\dagger m\gets L_t^{m - 1}\nabla_W^m\mathcal{L}R_t^{m - 1};$  // Precondition Gradients   
0  $\Delta W_t^m\gets \frac{\|\nabla_{W^m}\mathcal{L}\|}{\Delta W_t^m}\Delta W_t^m;$  // Rescale Gradients   
else   
|  $\Delta W_t^m\gets \nabla_{W^m}\mathcal{L};$    
end   
 $W_{t}^{m}\gets$  Optimizer step  $(\Delta W_t^m,W_{t - 1}^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.  $\mathbf{a}_{\mathbf{t}}^{\mathbf{m} - 1} = \mathbb{E}[a_{t}^{m - 1}]$  and  $\mathbf{g}_{\mathbf{t}}^{\mathbf{m}} = \mathbb{E}[g_t^m ]$  (lines 2 and 3 in Algorithm 1 and Figure 1-a).  $(A_{t}^{m - 1})_{;i}^{-1}$  and  $(G_{t}^{m})_{;i}^{-1}$  show the  $i^{th}$  column of  $(A_{t}^{m - 1})^{-1}$  and  $(G_{t}^{m})^{-1}$  respectively, where  $A_{t}^{m - 1}$  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 norm-based 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  $\mathbb{E}[g^m g^{m^T}]$  is approximated using a rank-1 matrix  $\mathbf{g}^{\mathrm{m}}\mathbf{g}^{\mathrm{m}}^{\mathrm{T}}$  and using the Sherman-Morrison identity, Equation 5 in obtained (line 7 in Algorithm 1 and Figure 1-c).
$$
L _ {t} ^ {m - 1} = \gamma L _ {t - 1} ^ {m} ^ {- 1} + \frac {(1 - \gamma)}{\gamma^ {2} (1 + \gamma (1 - \gamma) \mathbf {g} _ {\mathbf {t}} ^ {\mathbf {m} T} L _ {t - 1} ^ {m} ^ {- 1} \mathbf {g} _ {\mathbf {t}} ^ {\mathbf {m}})} L _ {t - 1} ^ {m} ^ {- 1} \mathbf {g} _ {\mathbf {t}} ^ {\mathbf {m}} \mathbf {g} _ {\mathbf {t}} ^ {\mathbf {m} T} L _ {t - 1} ^ {m} ^ {- 1} \tag {5}
$$
Furthermore, if Equation 4 is approximated using  $\mathbb{E}[a_t^{m - 1}a_t^{m - 1^T}] \approx \mathbf{a}_t^{\mathbf{m}}\mathbf{a}_t^{\mathbf{m}^\mathrm{T}}$  with a similar derivation, Equation 6 is obtained (line 8 in Algorithm 1 and Figure 1-c).
Table 1: The computation and communication complexity and memory overhead of the state-of-the-art implementations of the first- and second-order optimizers. The division by 2 in MKOR is because MKOR uses half-precision computations. The complexity of KFAC-based methods depends on layer dimensions while SNGD methods mostly depend on the batch size. In transformers, due to the scaling of the batch size by the sequence length, batch sizes and layer dimensions are comparable, making both KFAC- and SNGD-based methods more expensive than SGD.
Optimizer	Computational Complexity	Memory Overhead	Communication Complexity
MKOR	O(d2+ bd)	O(2d2/2)	O(2d/2)
SNGD (HyLo)	O(b3)	O(2bd + b2)	O(2bd + b2)
KFAC (KAISA)	O(d3)	O(4d2)	O(4d2)
Eva	O(d2+ bd)	O(2d)	O(2d)
SGD (Momentum)	-	O(d2)	-
ADAM / LAMB	-	O(d2)	-
$$
R _ {t} ^ {m - 1} = \gamma R _ {t - 1} ^ {m} ^ {- 1} + \frac {(1 - \gamma)}{\gamma^ {2} (1 + \gamma (1 - \gamma) \mathbf {a} _ {\mathbf {t}} ^ {\mathbf {m} T} R _ {t - 1} ^ {m} ^ {- 1} \mathbf {a} _ {\mathbf {t}} ^ {\mathbf {m}})} R _ {t - 1} ^ {m} ^ {- 1} \mathbf {a} _ {\mathbf {t}} ^ {\mathbf {m}} \mathbf {a} _ {\mathbf {t}} ^ {\mathbf {m} T} R _ {t - 1} ^ {m} ^ {- 1} \tag {6}
$$
Rescaling Gradients. Preconditioning the gradients using the computed factors can change gradient norms. Sometimes, these changes interfere with the effect of the learning rate on the training process. To alleviate this and to make learning rate schedulers more effective, the preconditioned gradients are scaled so that their norm matches the original norms (line 10 in Algorithm 1 and Figure 1-d).
Complexity Analysis. MKOR reduces 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  $\mathcal{O}(d^2)$  computation complexity, in contrast to KFAC and SNGD methods that need  $\mathcal{O}(d^3)$  and  $\mathcal{O}(b^3)$  complexity to invert matrices in  $\mathbb{R}^{d\times d}$  and  $\mathbb{R}^{b\times b}$  respectively. (2) Communication Complexity. The only data that is synchronized among different workers in MKOR is the two rank-1 approximations that have  $2d$  elements. With quantization, this size can be halved. In KFAC, the activation and gradient covariance matrices and the inversion of left and right factors need to be synchronized between all the workers, leading to  $4d^2$  data transfers. In SNGD, the activations and gradients are synchronized, leading to  $2bd$  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  $2d^2 + 2d$  memory overhead, and using half-precision computations further reduces this. KFAC stores the activation and gradient covariance matrices and the left and right factors, leading to  $4d^2$  memory overhead. SNGD stores the activations, the gradients, and the kernels they use as second-order information, leading to  $2bd + b^2$  memory complexity.
The low memory overhead of MKOR allows us to use larger batch sizes compared to other second-order methods. In practice, gradient accumulation methods are used to increase the effective batch size in training, which reduces the training speed significantly. This issue worsens with KAISA and HyLo, but MKOR alleviates this. For 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 first-order 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 MKOR Convergence 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 \times$  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  $\mu I$  to each of them to improve numerical stability ( $\mu > 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  $\mathbb{R}^{b \times 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 \times$  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  $\epsilon$ , the maximum number in matrices and vectors is  $m$ , and the dimension of the vectors and matrices are  $d$  and  $d \times d$  respectively, the quantization error of formulas 5 and 6 is  $O((\gamma + 4\frac{(1 - \gamma)}{\gamma^2} m^3 d^2)\epsilon)$
Exploding Gradients Problem In second-order methods, where the gradients are preconditioned by various factors, the exploding gradient problem is worsened. Our experiments show that in first-order methods, by choosing a learning rate that doesn't 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  $\zeta$  is a hyperparameter that controls the amount of information from the original factors that needs to be saved in the new factors.