MKOR: Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates

1 Introduction

• Second-order methods offer higher convergence rates, but Hessian and inverse computation and storage become bottlenecks as models grow.Natural Gradient Descent substitutes the Fisher Information Matrix for the Hessian, while layer-wise block approximations still require further approximation of their inverses.Source: §1 · pp. 1–2
• KFAC-based methods have O(d3) complexity and scale poorly for large layers, while KBFGS and SNGD address model-size costs but remain limited by large transformer batch sizes.Transformer batch size scales with attention sequence length, which can reach a few thousand; Eva targets O(d2) complexity but uses damping and cannot leverage momentum.Source: §1 · p. 2
• MKOR uses rank-1 updates to reduce inverse-covariance computation from O(d3) to O(d2), synchronizes rank-1 vectors to reduce communication from O(d2) to O(d), and adds stabilization and gradient rescaling.MKOR-H switches to a first-order optimizer when second-order updates no longer accelerate convergence.Source: §1 · p. 2
• Up to 2.57× faster training: MKOR reduces BERT-Large-Uncased training time from 8 hours to 3 hours on 64 A100 GPUs and reports new state-of-the-art GLUE metrics.The paper also reports that other second-order methods such as KFAC fail to converge to the GLUE baseline.Source: §1 · p. 2

2 Background

• Second-order methods precondition network weights with the inverse Hessian, while block-diagonal NGD approximates it with layer-specific Fisher Information Matrix blocks.The FIM matches the Gauss-Newton matrix under certain conditions.Source: §2 · p. 3
• KFAC represents each FIM block as the Kronecker product of two factors and uses their inverses in its update rule.The factors are computed from activation and gradient information; a momentum parameter moderates changes to them.Source: §2 · p. 3

3.1 The MKOR Algorithm

• For each layer, MKOR updates second-order information, preconditions gradients, and passes them to a backend optimizer; rank-1 covariance approximations enable Sherman-Morrison inverse updates with momentum.The rank-1 approximation uses averages across samples.Source: §3.1 · p. 4
• A norm-based stabilizer modifies inverse factors when extreme activations or gradients threaten numerical instability, and gradient rescaling restores the original gradient norm after preconditioning.Rescaling is intended to keep preconditioning from interfering with the learning-rate effect.Source: §3.1 · pp. 4–5
• MKOR reduces factor-inversion computation to O(d2) from O(d3) for KFAC and O(b3) for SNGD, and synchronizes 2d rank-1 values instead of quadratic or batch-dependent data.It requires 2d2 + 2d memory; half-precision can halve communication size and further reduce memory overhead.Source: §3.1 · p. 5, Table 1 · p. 5
• MKOR's lower memory overhead permits larger batch sizes, but the experiments use equal local batch sizes across optimizers and do not exploit that advantage.Gradient accumulation to increase effective batch size reduces training speed, an issue reported to worsen with KAISA and HyLo.Source: §3.1 · p. 5

3.2 Hybrid MKOR

• MKOR-H monitors loss reduction and switches to a first-order optimizer when second-order updates' convergence advantage becomes less noticeable as the loss flattens.The hybrid is designed to trade costly second-order updates against their convergence benefit.Source: §3.2 · p. 5

3.3 MKOR Convergence and Stability

• Stale-factor methods reuse inverses for hundreds to thousands of iterations; this saves cost but assumes little landscape change, and stale factors can harm convergence and generalization.An inversion iteration for KAISA and HyLo costs 150× an SGD iteration, with more than 98% of that iteration's cost spent on inversion.Source: §3.3 · p. 6
• MKOR's factor updates are positive-definite and use a denominator guaranteed to be nonzero, avoiding the cited need for damping factors and instability from large condition numbers.The paper states this avoids numerical overflows and underflows associated with those inversion issues.Source: §3.3 · p. 6
• When exploding gradients are detected, MKOR modifies factors to lean toward SGD; its local analysis states that a specified Kronecker-factorized curvature approximation yields a loss decrease.Using small ζ too often can convert MKOR to SGD.Source: §3.3 · pp. 6–7

Question answering

• On SQuAD v1.1, MKOR-H reaches the target accuracy in 600 steps, achieving 2.57× speedup over LAMB and 1.75× over KAISA.Vanilla MKOR and KAISA converge after 1,000 iterations; LAMB requires 1,536 steps, and Eva converges in 1,000 iterations with MKOR 1.69× faster.Source: §4 · p. 7, Table 2 · p. 8

Inversion frequency

• More frequent factor inversion increases convergence rate, while MKOR's average iteration cost is almost independent of inversion frequency and KAISA's is heavily dependent on it.The experiment uses a simple autoencoder on CIFAR-100; the authors note stale factors may lead to a local minimum.Source: §4 · p. 8, Figure 4 · p. 9

GLUE classification

• MKOR reaches state-of-the-art GLUE accuracy in 1,500 steps, while MKOR and MKOR-H match the baseline average metric in 600 steps and achieve 2.57× end-to-end speedup.KAISA does not reach baseline average accuracy after 1,563 steps; Eva is 1.69× slower than MKOR-H and 1.24% less accurate than MKOR at 1,500 steps.Source: §4 · p. 7, Table 3 · p. 8

Performance analysis

• MKOR reduces factor-computation and inversion cost on ResNet-50 relative to KAISA; first-order optimizers spend optimization time updating weights, without factorization or preconditioning.Figure 3 breaks down per-step time for BERT-Large-Uncased and ResNet-50.Source: §4 · p. 8, Figure 3 · p. 9

Approximation error

• Experiments report low error for rank-1 approximations of activation and input-gradient covariance matrices across BERT-Large pre-training and ResNet-50 on ImageNet.The authors state that their tests found this property across benchmarks and that they had not encountered a benchmark without low-rank covariance matrices.Source: §4 · pp. 9–10, Figure 5 · p. 10

Higher-rank extension

• Higher-rank MKOR updates have O(rd2) computation, but computing the low-rank approximation requires SVD, which the authors say removes the computational advantage.They report finding no cheaper way to compute covariance-matrix low-rank approximations than the paper's rank-1 method.Source: §4 · p. 10

• MKOR is a Momentum-Enabled Kronecker-Factor-Based Optimizer that uses rank-1 updates.Source: §5 · p. 10
• MKOR reduces computation, communication, and memory usage complexity to improve second-order methods’ end-to-end training time and convergence rate.Source: §5 · p. 10
• Experiments show MKOR outperforms state-of-the-art first- and second-order optimizers on large language models.Source: §5 · p. 10

• Faster convergence in LLMs may help practitioners save time and energy.Source: §6 · p. 11
• The authors state that the computations are fully transparent and note DNN applications including medical diagnostics, voice recognition, and climate change.Source: §6 · p. 11
• DNN optimization may also be applied to models with negative implications, such as privacy invasion; responsible use is beyond this research’s scope.Source: §6 · p. 11

8.1 ResNet-50

• The ResNet-50 ImageNet experiment uses a model with more than 25M parameters and more than 1.2M samples, with SGD as the first-order baseline.Source: §8.1 · p. 14
• At 75.9% target accuracy, MKOR converges in 57 epochs, achieving 1.49× speedup over SGD’s 88 epochs and 1.04× speedup over KAISA.Source: §8.1 · p. 14
• The authors could not reproduce Eva’s ResNet-50 ImageNet results because its hyperparameters are unreported; Eva also uses a 2015 KFAC implementation rather than KAISA.Source: §8.1 · p. 14

8.2 GLUE Results

• Table 4 presents BERT-Large-Uncased optimizer metrics for each GLUE classification task.Source: §8.2 · p. 15, Table 4 · p. 15

8.3 Derivation of NGD Approximations

• NGD uses the inverse Fisher Information Matrix as a gradient preconditioner, and Equation 10 defines the FIM for an arbitrary layer using batch samples.Source: §8.3 · p. 15
• KFAC expresses an FIM block as a Kronecker product of gradient and activation factors, approximating the expectation of their product by the product of expectations.The text identifies this factorization assumption as strong but computationally simplifying.Source: §8.3 · p. 16
• Adding momentum to KFAC’s left and right factors yields the factor equations used in the derivation.Source: §8.3 · p. 16
• SNGD applies the Sherman-Morrison-Woodbury identity to approximate the inverse FIM, requiring inversion of a batch-sized kernel matrix.Source: §8.3 · p. 16

8.5 Sensitivity to Learning Rate

• Large learning rates can destabilize first- and second-order methods, while small rates can slow convergence and erode second-order methods’ convergence advantage.Source: §8.5 · p. 16

8.4 Numerical Instability of Second-order Methods

• KFAC right factors’ minimum eigenvalues approach zero, producing large condition numbers and unstable inversion; MKOR avoids such inversion instability, with complexity independent of condition number.Source: §8.4 · p. 16, Figure 8 · p. 17

8.5 Sensitivity to Learning Rate

• MKOR is least sensitive to learning rate, converging in nearly the same number of iterations across a wide range; other optimizers may diverge or converge to a local minimum.Source: §8.5 · pp. 16–17, Table 5 · p. 17

8.6 Scalability

• MKOR shows strong scalability on BERT-Large-Uncased through 64 GPUs.Figure 9 presents the scalability results.Source: §8.6 · p. 17, Figure 9 · p. 17

8.7 Decaying Eigenvalues and Rank-1 Approximations

• As factor eigenvalues decay during convergence, rank-1 approximations become more effective, while identity initialization makes MKOR more robust to early approximation noise.As training samples increasingly shape the factors, the approximations become more reliable.Source: §8.7 · p. 18, Figure 10 · p. 18

8.8 Memory Overheads

• MKOR reduces KFAC/KAISA memory overhead by up to 1.5×, although all second-order methods have significant overhead over first-order methods.Table 6 reports per-GPU memory usage.Source: §8.8 · p. 18, Table 6 · p. 18

8.9 Hyperparameters

• For BERT-Large-Uncased, MKOR and MKOR-H update factors every 10 iterations, versus every 50 iterations for KAISA; ResNet-50 MKOR updates every 10 iterations.For ResNet-50, MKOR uses KAISA’s learning rate, with its learning rate halved at the ends of epochs 25, 35, 40, 45, 50, 55, and 56.Source: §8.9 · p. 18

8.10 Experiment Settings

• The BERT-Large experiments used up to 64 A100 GPUs on Polaris, while other experiments used the Mist cluster’s V100 GPUs.Training experiments ran five times, with median timing and accuracy reported; scalability runs averaged two jobs per setting.Source: §8.10 · pp. 18–19

8.12 Training Accuracy Experiments

• The experiments set optimizer parameters such as weight decay to zero so that optimizers minimize the same loss objective.The authors note that nonzero weight decay adds a quadratic loss term, and differing weight decay would mean optimizing different objectives.Source: §8.12 · p. 19
• On AlexNet/CIFAR-100, MKOR is 1.26× faster than HyLo-KIS, 1.31× faster than SGD, and 1.58× faster than KAISA.KAISA required small learning rates to avoid exploding gradients, which damaged its convergence rate; Figures 11-a and 12-a show the convergence comparisons.Source: §8.12 · p. 19, Figure 11 · p. 20, Figure 12 · p. 21
• On BERT-Large/IMDB, MKOR is 1.22× faster than SGD and 1.43× faster than HyLo-KIS; on BERT-Base/SQuAD, it is 1.26× faster than SGD and 1.56× faster than HyLo.KAISA did not converge on the reported BERT experiments because of exploding gradients; Figures 11 and 12 compare convergence across the tasks.Source: §8.12 · p. 19, Figure 11 · p. 20, Figure 12 · p. 21

8.13 Knee-Point Learning Rate Scheduler

• The knee-point scheduler monitors accuracy improvement or loss decrease and lowers the learning rate when the average rate of progress falls below a β-scaled threshold.It uses an exponential moving average to estimate the rate; the authors provide their code and experiment instructions publicly.Source: §8.11 · p. 19, §8.13 · p. 20

8.14 Proofs

• Equation 14’s right-hand side is positive definite because it sums positive-definite matrices, including γJ_t−1 for γ > 0.The proof assumes J_t−1 is positive definite.Source: §8.14 · p. 21
• The proof expands P from Equation 17 into four Kronecker-product terms in Equation 18, including the identity term.The expansion uses the associativity of Kronecker multiplication.Source: §8.14 · p. 22
• Because the component Kronecker products are positive semidefinite, the proof concludes that P is positive semidefinite and uses it to establish Equation 19.The argument applies to the factors L and R described in the proof.Source: §8.14 · p. 22

Lemma 3.2

• With maximum value m, quantization error ϵ, and vector length d, the lemma gives O(2ϵm^2d) error for vector-vector dot products and vector-matrix products.The vector-matrix case is treated as multiple independent dot products.Source: Lemma 3.2 · p. 22
• For a vector-vector outer product, the lemma reports quantization error O(2ϵm).The bound is based on the maximum error in each element’s multiplication.Source: Lemma 3.2 · p. 22
• For the update terms, quantizing γJ contributes at most γϵ, denominator error is negligible relative to 1, and the passage reports O(4m^3d^2) outer-product error.The stated error for quantizing Jj is O(2m^2dϵ).Source: Lemma 3.2 · p. 22