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 & tables
Figure 1: Muon-LITE exhibits superior scaling behavior across varying token budgets (left) and model sizes (right) compared to Muon.
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.
Algorithms
Matrix Form
Vector Form
F (up to a constant factor)
N-AdamW
VM+βg
diagv−21(m+βg)
(diagE[gg⊤])21
Muon
Mβ(Mβ⊤Mβ)−21
((Mβ⊤Mβ)⊗I)−21(m+βg)
(E[G⊤G])21⊗I
SOAP
QlVrotQl⊤MQrQr⊤
(Qr⊗Ql)(diagvrot)−21(Qr⊤⊗Ql⊤)m
(Qr⊗Ql)(diagE[grotgrot⊤])21(Qr⊤⊗Ql⊤)
Table 1: Update forms (matrix and vector) and their corresponding approximations of the metric F (up to a constant scaling factor) for common adaptive optimizers. m denotes the EMA-type momentum, g represents the stochastic gradient, and Mβ=M+βG . Capital letters denote the matrix forms corresponding to their vector counterparts (lowercase). For AdamW, we present its Nesterov-accelerated variant (N-AdamW). See Appendix I for further discussions and details.
Figure 3: Coverage of the top eigenspaces of row (column) Hessians by those of G⊤G ( GG⊤ ) 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 χ≥1 , while H increases only the Hessian damping coefficient β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 Figure 4 .
Figure 6: Muon-LITE outperforms Muon in QwenMoE pre-training tasks.
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Acronym
Size
dmodel
dFF
n _ head
depth
lr_max of Muon
lr_max of SOAP
LLaMA (0.13B)
134M
768
2048
12
12
5e-3 (on C4)
-
LLaMA (0.25B)
247M
768
2560
16
24
3e-3 (on Pile)
3e-3 (on Pile)
LLaMA (0.5B)
518M
1280
3456
20
22
2e-3 (on Pile)
2e-3 (on Pile)
LLaMA (1.3B)
1339M
2048
5461
32
24
1e-3 (on Pile)
-
Appendix
Table 2: Dense model configurations and optimally-tuned peak learning rates for Muon and SOAP.
Acronym
Size
Activated Size
dmodel
dFF
n _ head
depth
n _ experts
lr_max
QwenMoE (1B)
1040M
297M
768
3072
12
15
32
2e-3
Appendix
Table 3: MoE model configurations and optimally-tuned peak learning rates for Muon on Pile.
Figure 7: Hessian eigenvalue distribution of different blocks in the toy LLaMA2 model. The results of the up_proj are in Figure 2 .
Figure 8: Hessian eigenvalue distribution of different blocks in the toy LLaMA2 model. The results of the up_proj are in Figure 2 .
Method
ARC_E
ARC_C
PIQA
HellaSwag
BOOLQ
WinoGrande
MMLU
T_mc2
AVG
Muon
52.99
22.10
67.41
35.30
52.45
53.35
23.76
40.14
43.44
Muon-LITE
54.38
23.89
67.63
36.72
60.52
54.78
25.15
40.83
45.49
Appendix
Table 4: Evaluation results on downstream tasks (0-shot with lm-evaluation-harness) of LLaMA models (1.3B) pre-trained on Pile using Muon-LITE and Muon. The best scores in each column are bolded. Abbreviations: T_mc2 = TruthfulQA_mc2 , AVG =Average score.
Pretraining accounts for a large fraction of the total computational cost in LLM training. However, noise-dominant gradients and the highly ill-conditioned loss landscape bring severe challenges. Although modern adaptive optimizers such as AdamW and Muon have achieved great success in large-scale pretraining, their reliance on gradient normalization offers limited mitigation of the ill-conditioned curvature. The progress along flat directions (eigen-directions of small eigenvalues), which dominates the final loss reduction, remains relatively slow. To enhance training dynamics along flat directions, we propose a curvature-conditioned multiscale momentum method with sphere constraints, delivering steady acceleration in LLM pretraining. This multiscale momentum, applied only along flat directions, pairs a slow-decay component for noise reduction with a fast-decay component for rapid curvature adaptation, harnessing their complementary strengths. Crucially, we employ a sphere constraint technique to prevent parameter inflation and excessively rapid effective learning rate decay that would otherwise arise from a naive combination. Extensive experiments show that the proposed method significantly accelerates Muon across diverse architectures (dense, MoE) and model sizes (0.12B--2.3B parameters). Theoretically, we verify the acceleration effect and provide insight into the design principles underlying the flat-direction multiscale momentum.
Modern optimizers like Adam and Muon are central to training large language models, but their reliance on first- and second-order momenta introduces significant memory overhead, which constrains scalability and computational efficiency. In this work, we reframe the exponential moving average (EMA) used in these momenta as the training of a linear regressor via online gradient flow. Building on this equivalence, we introduce LoRA-Pre, a novel low-rank optimizer designed for efficient pre-training. Specifically, LoRA-Pre reduces the optimizer's memory footprint by decomposing the full momentum matrix into a compact low-rank subspace within the online linear learner, thereby maintaining optimization performance while improving memory efficiency. We empirically validate LoRA-Pre's efficacy by pre-training models from the Llama architecture family, scaling from 60M to 1B parameters. LoRA-Pre achieves the highest performance across all model sizes. Notably, LoRA-Pre demonstrates remarkable rank efficiency, achieving comparable or superior results using only 1/8 the rank of baseline methods. Beyond pre-training, we evaluate LoRA-Pre's effectiveness in fine-tuning scenarios. With the same rank, LoRA-Pre consistently outperforms all efficient fine-tuning baselines. Specifically, compared to standard LoRA, LoRA-Pre achieves substantial improvements of 3.14 points on Llama-3.1-8B and 6.17 points on Llama-2-7B, validating our approach's effectiveness across both pre-training and fine-tuning paradigms. Our code is publicly available at https://github.com/mrflogs/LoRA-Pre.
Zhengbo Wang, Jian Liang, Ran He +2
University of Science and Technology of China · NLPR & MAIS, Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences +1
The Muon optimizer has recently offered a promising alternative to AdamW for large language model training, leveraging matrix orthogonalization to produce geometry-aware updates. However, like all first-order methods, Muon can become trapped in sharp local minima. In this work, we present MONA, an optimizer that bridges Muon's orthogonalization framework with curvature-aware acceleration. MONA adds an acceleration term directly into Muon's gradient processing pipeline. This term is calculated from the exponential moving average of gradient differences. We provide a detailed convergence analysis for MONA, showing that the acceleration term introduces curvature-sensitive corrections while preserving Muon's spectral-norm regularization. Empirically, MONA achieves better convergence and downstream task performance compared to both Muon and AdamW across three scales of Mixture-of-Experts pretraining, spanning from 1B to 68B parameters, with the largest model trained on 1 trillion tokens. Furthermore, we conduct supervised fine-tuning on the MOE-68B-A3B model and evaluate it on general capability, mathematical reasoning, and code generation benchmarks, where MONA achieves SOTA performance.