Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement

TL;DR

A unified Riemannian Ordinary Differential Equation framework is established that elucidates how common adaptive algorithms operate synergistically and proposes LITE, a generalized acceleration strategy that enhances training dynamics by applying larger Hessian damping coefficients and learning rates along flat trajectories.

Abstract

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.

Figures (7)

• Figure 2: (left) Hessian eigenvalue distribution of an up_proj FFN block in a toy LLaMA model; (right) Schematic illustration of the time-scale separation property and the acceleration mechanism of the LITE approach.
• Figure 3: Coverage of the top eigenspaces of row (column) Hessians by those of $G^\top G$ ($G G^\top$) for the up_proj block of an FFN layer. A higher score indicates a greater degree of containment.
• Figure 4: Performance comparison of Muon and Muon-LITE on LLM pre-training tasks. We evaluate LLaMA2 models across various sizes and datasets using a cos learning rate schedule. The suffixes L and H denote two ablation variants of LITE: in flat directions, L increases only the learning rate ratio $\chi\ge1$, while H increases only the Hessian damping coefficient $\beta_2$. The token batch size is approximately 2M for the 1.3B model, and 1M for others.
• Figure 5: Performance comparison of SOAP and SOAP-LITE on LLM pre-training tasks. Same experimental setup and notations are employed as in \ref{['fig:cosine-llama-muon']}.
• Figure 6: Muon-LITE outperforms Muon in QwenMoE pre-training tasks.

Theorems & Definitions (39)

• Proposition 4.1: Adapted from Proposition \ref{['transform-2nd-to-1st-system']}
• Theorem 7.3
• Theorem 7.4
• Lemma 7.1
• proof
• Proposition 7.2
• Remark 7.3
• proof
• Remark 7.4
• Remark 9.1