[ICLR 2025] LITE: Speeding Up LLM Pre-training via Flat-Direction Dynamics Enhancement

TL;DR

Training Large Language Models (LLMs) is a race against computational constraints. While matrix-based optimizers like Muon and SOAP have pushed the boundaries of efficiency, they often treat all directions in the optimization landscape equally. LITE breaks this isotropy by selectively accelerating the "slow but crucial" flat directions using the $2^{nd}$-order insights of a Riemannian ODE framework. The result? A 2x speedup in token efficiency and consistently lower terminal loss across models up to 1.3B parameters.

The Motivation: Escaping the Isotropic Trap

The loss landscape of a deep neural network is notorious for its ill-conditioning. If you visualize the Hessian (the matrix of second derivatives), you'll see a few "sharp" directions with massive eigenvalues and a vast "flat" sea of near-zero or negative eigenvalues.

Previous research suggests a time-scale separation:

1. Fast Dynamics (Sharp Directions): Rapid oscillations that determine stability but barely reduce loss.
2. Slow Dynamics (Flat Directions): Steady, slow progress that accounts for the bulk of final performance.

Standard optimizers (like AdamW or even matrix-based ones) tend to produce isotropic updates. They are effectively "speed-limited" by the sharpest directions to avoid divergence, meaning they crawl through the flat valleys where they should be sprinting.

Figure: The "River-Valley" landscape (left) and the conceptual mechanism of LITE acceleration (right).

Methodology: The Riemannian ODE Framework

The authors don't just add a heuristic; they build a unified ODE framework that treats the parameter space as a Riemannian manifold.

In this view:

• The Preconditioner induces the geometry (the "shaping" of the hills).
• The Momentum acts as a Riemannian damping term (the "friction" preventing overshooting).

By analyzing the continuous-time limit of optimizers, LITE proposes a decoupled update rule. It keeps the sharp directions stable while injecting "extra fuel"—higher learning rates ($\chi$) and higher Hessian damping (${\beta }_2$)—into the flat directions.

Practical Implementation

Identifying these directions could be a computational nightmare if it required calculating the full Hessian. LITE avoids this by using the preconditioners of Muon/SOAP as a proxy for the Hessian eigenspace. Using an efficient Composite Newton-Schulz iteration, it projects the update into flat/sharp subspaces with almost zero overhead (~1% throughput drop).

Note: The algorithm uses an amplification ratio $\chi \ge 1$ and increased damping ${\beta }_2$ to drive the flat-direction traversal.

Experiments: Breaking the Scaling Laws

The authors tested LITE across LLaMA (Dense) and Qwen (MoE) architectures.

Key Findings:

1. Terminal Loss: In every case—C4, Pile, Cosine, or WSD schedules—LITE-accelerated optimizers reached lower losses than well-tuned baselines.
2. 2x Speedup: In long-horizon training (200B tokens), MUON-LITE achieved the same loss as Muon in roughly half the time.
3. Scaling Performance: As the model size grew from 130M to 1.3B, the performance gap widened, suggesting that larger models (which are often more ill-conditioned) benefit more from this directional acceleration.

Figure: LITE (solid lines) consistently stays below the Muon/SOAP baselines (dotted lines) across different datasets.

Critical Insights & Conclusion

LITE is a sophisticated yet practical upgrade to the modern optimizer stack. Its brilliance lies in selective aggression: it recognizes that "stability" is a local requirement, not a global one. By allowing the optimizer to move faster where the ground is flat, it effectively "whitens" the training dynamics more completely than second-order methods alone.

Limitations: Currently, the sharp subspace dimension $d_s$ is a hyper-parameter (though the paper provides a robust heuristic, $r_s\approx 0.1$). Future work on automatically tuning this ratio per layer could make LITE even more plug-and-play.

For practitioners looking to shave weeks off their LLM pre-training runs, LITE proves that sometimes, the fastest way to the bottom of the valley is to know exactly where you can afford to run.

References

• Jordan et al., 2024. Muon: An optimizer for hidden layers.
• Vyas et al., 2025. SOAP: Improving and stabilizing shampoo.
• Wen et al., 2025. The River Valley View of LLM landscapes.