🤖 machine learning

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

This paper introduces MKOR, a momentum-enabled second-order optimizer that utilizes rank-1 updates and Sherman-Morrison-based inversion to reduce computational and communication complexity from cubic to quadratic and linear respectively, thereby significantly accelerating the training and convergence of large language models compared to state-of-the-art first- and second-order methods.

Original authors: Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

Published 2024-01-30✓ Author reviewed ⓘ

📖 6 min read🧠 Deep dive

Original authors: Mohammad Mozaffari, Sikan Li, Zhao Zhang, Maryam Mehri Dehnavi

✨ This is an AI-generated explanation of the paper below. It is not written by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the vast landscape of modern artificial intelligence, the most powerful systems are built on deep neural networks, intricate mathematical structures that learn to recognize patterns, translate languages, and generate text. To teach these systems, researchers must solve a massive optimization problem: finding the perfect set of internal numbers, or weights, that allow the network to make accurate predictions. For decades, the standard approach has been to nudge these weights in the direction that most quickly reduces error, a method akin to a hiker feeling their way down a foggy mountain by checking the slope under their feet. This first-order method is reliable and widely used, but it can be slow, often taking thousands of steps to reach the bottom of the valley.

A more ambitious strategy, known as second-order optimization, attempts to give the hiker a map of the entire terrain, including the curvature of the ground, allowing for much faster descents. By understanding how the slope changes, these methods can theoretically reach the optimal solution in far fewer steps. However, this approach has historically been too computationally expensive to be practical for the massive models used today. Calculating the necessary curvature information for a large network requires immense processing power and memory, often making the method slower than the simpler approach it aims to replace. The challenge has been to capture the speed of the map-reading hiker without the heavy burden of carrying the map.

A team of researchers from the University of Toronto, Rutgers University, and the Texas Advanced Computing Center has developed a new tool called MKOR that solves this dilemma. They created a momentum-enabled optimizer that uses a specific mathematical shortcut to approximate the terrain's shape without calculating every detail. Instead of trying to compute the full, complex curvature of the loss landscape, MKOR uses a technique called rank-1 updates. This approach simplifies the calculation by focusing on the most significant direction of change, effectively reducing the computational cost from a cubic relationship to a quadratic one. In practical terms, this means the method scales much better as models grow larger, avoiding the bottlenecks that have previously prevented second-order methods from being used in large language models.

The researchers tested this new optimizer on a large language model known as BERT-Large-Uncased, running experiments on 64 high-performance graphics processors. They compared MKOR against the best existing first-order methods and the most advanced second-order implementations available. The results were striking. In a task involving question answering, the standard first-order method required 1,536 steps to reach a specific level of accuracy. The new MKOR method reached the same or better accuracy in just 600 steps. When measuring the total time required to complete the training, MKOR finished the task in just over three hours, whereas the standard method took nearly eight hours. This represents a speedup of 2.57 times, effectively cutting the training time by more than half.

Beyond raw speed, the study addressed a critical stability issue that often plagues second-order methods. Because these methods adjust the learning process based on complex mathematical factors, they are prone to "exploding gradients," a phenomenon where the adjustments become so large that the training process crashes or diverges. Previous solutions often required slowing down the learning rate significantly to maintain safety, which negated the speed benefits. MKOR includes a built-in stabilizer that monitors the size of these adjustments. If the values grow too large, the system automatically rescales them, preventing the training from blowing up while maintaining the aggressive pace of the second-order approach. This allows the optimizer to move quickly without the risk of instability.

The researchers also observed that the advantage of second-order methods is most pronounced in the early stages of training, when the model is learning the broad contours of the data. As the model gets closer to the optimal solution, the benefits of the complex curvature calculations diminish. To address this, they introduced a hybrid version of the optimizer, called MKOR-H. This version starts with the powerful second-order method to accelerate the initial learning phase but automatically switches to a simpler, first-order method once the training loss stops decreasing rapidly. This hybrid approach allowed the system to converge in even fewer steps, achieving the same final accuracy as the standard method in less than half the time.

The study also examined why this specific mathematical shortcut works so well. The researchers analyzed the data structures within the neural network and found that the covariance matrices, which describe how different parts of the network interact, are often low-rank. This means that despite the high dimensionality of the data, the information is concentrated in a few key directions. By approximating these matrices using rank-1 updates, MKOR captures the essential information needed for optimization without the overhead of full matrix calculations. The team verified through extensive testing that this approximation introduces very little error, confirming that the simplification does not come at the cost of accuracy.

In addition to speed and stability, the new method significantly reduces the amount of data that needs to be transferred between different processors during training. In distributed computing, where thousands of processors work together, communication can become a major bottleneck. MKOR reduces this communication complexity from a quadratic relationship to a linear one, meaning that as more processors are added, the system scales efficiently without being bogged down by data transfer delays. This efficiency allows the method to handle the massive batch sizes required by modern large language models, which was a limitation for previous second-order approaches.

The implications of this work extend beyond just training one specific model. By making second-order optimization practical for large-scale systems, MKOR offers a path toward more efficient training of artificial intelligence. The researchers noted that their method could help save significant amounts of energy and time for machine learning practitioners. While the technology itself is a neutral tool, its application in fields ranging from medical diagnostics to climate modeling could accelerate scientific discovery. However, the authors also acknowledged that efficient optimization algorithms can be applied to any model, including those with negative societal impacts, emphasizing that the responsible use of such technology depends on the choices of the people who wield it.

Ultimately, the paper presents a concrete solution to a long-standing problem in deep learning. By combining momentum, rank-1 approximations, and a hybrid switching mechanism, the researchers have created an optimizer that is faster, more stable, and more scalable than its predecessors. The experiments demonstrate that it is possible to harness the power of second-order methods without the prohibitive computational costs that have kept them on the sidelines. As the field of artificial intelligence continues to push toward larger and more complex models, tools like MKOR provide a way to train these systems more efficiently, turning what was once a theoretical advantage into a practical reality.

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

1. Problem Statement

Second-order optimization methods, such as Natural Gradient Descent (NGD) and Kronecker-Factored Approximate Curvature (KFAC), offer superior convergence rates compared to first-order methods (e.g., SGD, Adam, LAMB) by utilizing curvature information (Hessian or Fisher Information Matrix approximations). However, their adoption in training large-scale Deep Neural Networks (DNNs), particularly Transformers and Large Language Models (LLMs), is hindered by three primary bottlenecks:

Computational Complexity: Traditional KFAC methods require inverting covariance matrices of size $d \times d$ (where $d$ is the layer dimension), resulting in $O(d^3)$ complexity. Similarly, Sherman-Morrison-Woodbury (SMW) based methods (SNGD) often scale with batch size $b$ ( $O(b^3)$ ). In Transformers, where batch sizes scale with sequence length, these complexities become prohibitive.
Communication Overhead: Distributed training requires synchronizing second-order information (covariance matrices and their inverses) across workers. KFAC requires synchronizing $O(d^2)$ data, while SNGD requires $O(bd + b^2)$ , creating significant communication bottlenecks that limit scalability.
Update Frequency and Stability: Due to high inversion costs, state-of-the-art implementations (like KAISA) update second-order information infrequently (e.g., every 100–1000 iterations) using "stale" factors. This reduces the effectiveness of the curvature information. Furthermore, second-order methods are prone to numerical instability (exploding gradients, singular matrices) and often require damping factors or high-precision arithmetic, which further increases computational costs.

2. Methodology: The MKOR Algorithm

The authors propose MKOR (Momentum-Enabled Kronecker-Factor-Based Optimizer Using Rank-1 Updates), which addresses the above issues through four core mechanisms:

Rank-1 Approximation & SM-Based Inversion:
Instead of computing and inverting full covariance matrices, MKOR approximates the covariance of activations and gradients using rank-1 matrices derived from the average of samples ( $E[aa^T] \approx \bar{a}\bar{a}^T$ ). By applying the Sherman-Morrison identity to these rank-1 updates, MKOR updates the inverse of the Kronecker factors directly. This reduces the computational complexity of factor inversion from $O(d^3)$ to $O(d^2)$ .
- Update Rule: The inverse factors are updated using momentum ( $\gamma$ ) and the Sherman-Morrison formula, allowing for frequent updates without the cost of full matrix inversion.
Linear Communication Complexity:
In distributed settings, MKOR only synchronizes the two rank-1 approximation vectors (size $2d$ ) between workers, rather than full covariance matrices ( $4d^2$ ). This reduces communication complexity from quadratic $O(d^2)$ to linear $O(d)$ . Additionally, MKOR utilizes half-precision (FP16) for these vectors, further halving communication costs, a feat difficult for other second-order methods due to numerical sensitivity.
Stabilization Mechanisms:
To prevent numerical instability and exploding gradients common in second-order methods:
- Norm-Based Stabilizer: Monitors the infinity norm of the inverse factors. If they exceed a threshold, the factors are blended with the identity matrix ( $\hat{L} = \zeta L + (1-\zeta)I$ ) to prevent divergence.
- Gradient Rescaling: Preconditioned gradients are rescaled to match the norm of the original gradients, ensuring learning rate schedulers remain effective.
- Theoretical Stability: The authors prove (Lemma 3.1) that the factors remain positive-definite without needing artificial damping factors ( $\mu I$ ), and that the single scalar division in the update rule is guaranteed to be non-zero.
Hybrid MKOR (MKOR-H):
Recognizing that second-order advantages often diminish as training converges (where the loss landscape approaches identity), MKOR-H introduces a switching mechanism. It monitors the loss reduction rate and automatically falls back to a first-order optimizer if second-order updates no longer provide acceleration, optimizing the trade-off between cost and convergence speed.

3. Key Contributions

Complexity Reduction: MKOR achieves $O(d^2)$ computational complexity and $O(d)$ communication complexity for second-order updates, alleviating the bottlenecks of $O(d^3)$ and $O(d^2)$ found in KFAC and $O(b^3)$ in SNGD.
High-Frequency Updates: By reducing inversion costs, MKOR enables second-order updates every iteration (or very frequently), eliminating the need for "stale" factors and fully exploiting curvature information.
Robustness: The introduction of norm-based stabilization and gradient rescaling allows MKOR to operate stably without damping factors or double-precision arithmetic, enabling efficient use of half-precision hardware.
Hybrid Strategy: The MKOR-H variant dynamically switches between second-order and first-order methods to maintain efficiency throughout the entire training lifecycle.

4. Experimental Results

The authors evaluated MKOR on BERT-Large-Uncased using 64 A100 GPUs, comparing it against the state-of-the-art first-order optimizer LAMB and second-order implementations KAISA (KFAC) and HyLo (SNGD), as well as Eva.

Speedup: MKOR achieved a 2.57× speedup over LAMB and a 1.85× speedup over KAISA on the pre-training task. Specifically, it reduced the training time of BERT-Large-Uncased from 8 hours to 3 hours.
Convergence: On the SQuAD v1.1 question-answering task, MKOR-H converged in 600 steps (vs. 1,536 for LAMB) while achieving comparable or better F1 scores (90.64 vs. 90.44).
GLUE Benchmark: MKOR achieved new state-of-the-art metrics on the GLUE dataset with 1,500 steps. Notably, while KAISA failed to converge to the baseline average accuracy even after 1,563 steps, MKOR successfully converged.
Scalability: Unlike HyLo, which diverged or failed to fit in GPU memory due to high memory overhead ( $O(bd + b^2)$ ), MKOR successfully trained with large batch sizes (up to 64k effective batch size via accumulation) without memory issues.
Approximation Validity: Experiments on ResNet-50 and BERT showed that rank-1 approximations of covariance matrices incur negligible error, validating the core approximation assumption.

5. Significance and Claims

The paper claims that MKOR successfully bridges the gap between the theoretical benefits of second-order optimization and the practical constraints of training modern, large-scale DNNs. By reducing the complexity of second-order updates to a level comparable to first-order methods (in terms of scaling with model size) while maintaining the convergence advantages, MKOR enables:

Faster Training: Significant reduction in wall-clock training time for LLMs.
Better Convergence: The ability to converge to higher accuracy metrics where other second-order methods (like KAISA) fail or stagnate.
Scalability: The ability to scale to large batch sizes and model dimensions without the communication and memory bottlenecks that previously limited second-order methods.

The authors emphasize that their work makes second-order optimization viable for the "large model" era, specifically addressing the scalability issues introduced by the sequence-length-dependent batch sizes in Transformer architectures. They conclude that MKOR represents a new state-of-the-art for efficient and effective training of deep neural networks.