# Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement

arXiv: https://arxiv.org/abs/2602.22681 · HTML: https://mlsanity.com/paper/2602.22681/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 56 references that are in the map.
- First posted on arXiv 26 Feb 2026.

## Builds on (most cited first, 30 of 41)

- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [The Pile: An 800GB Dataset of Diverse Text for Language Modeling](https://mlsanity.com/paper/2101.00027/) (2020)
- [Qwen2 Technical Report](https://mlsanity.com/paper/2407.10671/) (2024)
- [Optimizing Neural Networks with Kronecker-factored Approximate Curvature](https://mlsanity.com/paper/1503.05671/) (2015)
- [MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies](https://mlsanity.com/paper/2404.06395/) (2024)
- [New insights and perspectives on the natural gradient method](https://mlsanity.com/paper/1412.1193/) (2014)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [An Investigation into Neural Net Optimization via Hessian Eigenvalue Density](https://mlsanity.com/paper/1901.10159/) (2019)
- [PyHessian: Neural Networks Through the Lens of the Hessian](https://mlsanity.com/paper/1912.07145/) (2019)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [A Differential Equation for Modeling Nesterov's Accelerated Gradient Method: Theory and Insights](https://mlsanity.com/paper/1503.01243/) (2015)
- [Training Deep Learning Models with Norm-Constrained LMOs](https://mlsanity.com/paper/2502.07529/) (2025)
- [Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models](https://mlsanity.com/paper/2208.06677/) (2022)
- [Scaling Laws and Compute-Optimal Training Beyond Fixed Training Durations](https://mlsanity.com/paper/2405.18392/) (2024)
- [Towards Theoretically Understanding Why SGD Generalizes Better Than ADAM in Deep Learning](https://mlsanity.com/paper/2010.05627/) (2020)
- [Why Transformers Need Adam: A Hessian Perspective](https://mlsanity.com/paper/2402.16788/) (2024)
- [The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm](https://mlsanity.com/paper/2505.16932/) (2025)
- [A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale](https://mlsanity.com/paper/2309.06497/) (2023)
- [Understanding the Acceleration Phenomenon via High-Resolution Differential Equations](https://mlsanity.com/paper/1810.08907/) (2018)
- [The alignment property of SGD noise and how it helps select flat minima: A stability analysis](https://mlsanity.com/paper/2207.02628/) (2022)
- [Understanding Optimization in Deep Learning with Central Flows](https://mlsanity.com/paper/2410.24206/) (2024)
- [MARS: Unleashing the Power of Variance Reduction for Training Large Models](https://mlsanity.com/paper/2411.10438/) (2024)
- [A New Perspective on Shampoo's Preconditioner](https://mlsanity.com/paper/2406.17748/) (2024)
- [The AdEMAMix Optimizer: Better, Faster, Older](https://mlsanity.com/paper/2409.03137/) (2024)
- [PolarGrad: A Class of Matrix-Gradient Optimizers from a Unifying Preconditioning Perspective](https://mlsanity.com/paper/2505.21799/) (2025)
- [Cautious Optimizers: Improving Training with One Line of Code](https://mlsanity.com/paper/2411.16085/) (2024)
- [Muon Outperforms Adam in Tail-End Associative Memory Learning](https://mlsanity.com/paper/2509.26030/) (2025)

## Built on by (most cited first, 8 of 8)

- [Why Muon Outperforms Adam: A Curvature Perspective](https://mlsanity.com/paper/2606.04662/) (2026)
- [AMUSE: Anytime Muon with Stable Gradient Evaluation](https://mlsanity.com/paper/2605.22432/) (2026)
- [Negligible in Size, Significant in Effect: On Scale Vectors in Large Language Models](https://mlsanity.com/paper/2605.26895/) (2026)
- [Towards Understanding the Power and Limits of the Muon Optimizer: A River-Valley Perspective](https://mlsanity.com/paper/2606.21514/) (2026)
- [Rethinking Bregman Divergences in Kronecker-Factored Optimizers](https://mlsanity.com/paper/2606.00542/) (2026)
- [A Geometric Perspective on Stabilizing Value Conflict Resolution](https://mlsanity.com/paper/2607.17946/) (2026)
- [Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining](https://mlsanity.com/paper/2608.28442/) (2026)
- [SOLAR: A State-Driven Online Learning Rate Scheduler for LLM Pretraining](https://mlsanity.com/paper/2609.34681/) (2026)

## Abstract (arXiv)

Pre-training Large Language Models requires immense computational resources, making optimizer efficiency essential. The optimization landscape is highly anisotropic, with loss reduction driven predominantly by progress along flat directions. While matrix-based optimizers such as Muon and SOAP leverage fine-grained curvature information to outperform AdamW, their updates tend toward isotropy -- relatively conservative along flat directions yet potentially aggressive along sharp ones. To address this limitation, we first establish a unified Riemannian Ordinary Differential Equation (ODE) framework that elucidates how common adaptive algorithms operate synergistically: the preconditioner induces a Riemannian geometry that mitigates ill-conditioning, while momentum serves as a Riemannian damping term that promotes convergence. Guided by these insights, we propose LITE, a generalized acceleration strategy that enhances training dynamics by applying larger Hessian damping coefficients and learning rates along flat trajectories. Extensive experiments demonstrate that LITE significantly accelerates both Muon and SOAP across diverse architectures (Dense, MoE), parameter scales (130M--1.3B), datasets (C4, Pile), and learning-rate schedules (cosine, warmup-stable-decay). Theoretical analysis confirms that LITE facilitates faster convergence along flat directions in anisotropic landscapes, providing a principled approach to efficient LLM pre-training. The code is available at https://github.com/SHUCHENZHU/LITE.