Search papers, labs, and topics across Lattice.
The paper introduces LITE, a novel optimization strategy for accelerating LLM pre-training by enhancing training dynamics along flat directions in the loss landscape. LITE builds upon a unified Riemannian ODE framework that reveals how preconditioners and momentum operate synergistically in adaptive optimizers. By applying larger Hessian damping coefficients and learning rates specifically along flat trajectories, LITE improves convergence speed.
By strategically amplifying updates along flat directions in the loss landscape, LITE unlocks faster LLM pre-training with existing matrix-based optimizers like Muon and SOAP.
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.
