Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement

Shuchen Zhu; Rizhen Hu; Mingze Wang; Mou Sun; Xue Wang; Kun Yuan; Zaiwen Wen

cs.LGFeb 26, 2026

Accelerating LLM Pre-Training through Flat-Direction Dynamics Enhancement

Authors: Shuchen Zhu, Rizhen Hu, Mingze Wang, Mou Sun, Xue Wang, Kun Yuan, Zaiwen Wen

Organizations: Peking University · Zhejiang Lab, Zhejiang · Meituan, Beijing

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 & tables

Figure 1: Muon-LITE exhibits superior scaling behavior across varying token budgets (left) and model sizes (right) compared to Muon.

Panel of Figure 2 — 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	$\frac{M+\beta g}{\sqrt{V}}$	$\operatorname{diag}v^{-\frac{1}{2}}(m+\beta g)$	$(\operatorname{diag}\mathbb{E}[gg^{\top}])^{\frac{1}{2}}$
Muon	$M_{\beta}(M_{\beta}^{\top}M_{\beta})^{-\frac{1}{2}}$	$((M_{\beta}^{\top}M_{\beta})\otimes I)^{-\frac{1}{2}}(m+\beta g)$	$(\mathbb{E}[G^{\top}G])^{\frac{1}{2}}\otimes I$
SOAP	$Q_{l}\frac{Q_{l}^{\top}MQ_{r}}{\sqrt{V_{\text{rot}}}}Q_{r}^{\top}$	$(Q_{r}\otimes Q_{l})(\operatorname{diag}v_{\text{rot}})^{-\frac{1}{2}}(Q_{r}^{\top}\otimes Q_{l}^{\top})m$	$(Q_{r}\otimes Q_{l})(\operatorname{diag}\mathbb{E}[g_{\text{rot}}g_{\text{rot}}^{\top}])^{\frac{1}{2}}(Q_{r}^{\top}\otimes Q_{l}^{\top})$

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_{\beta}=M+\beta 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.

	$\displaystyle u_{k}=$	$\displaystyle\color[rgb]{0.5,0,0.5}{\displaystyle(I-Q_{k})}F(w_{k})^{-1}(m_{k}+{\color[rgb]{0.5,0,0.5}\beta_{1}}\nabla f(w_{k}))$
		$\displaystyle+{\color[rgb]{0.5,0,0.5}\chi Q_{k}}F(w_{k})^{-1}(m_{k}+{\color[rgb]{0.5,0,0.5}\beta_{2}}\nabla f(w_{k})).$

Algorithm 1 LITE Strategy

Figure 3: Coverage of the top eigenspaces of row (column) Hessians by those of

G^{\top}G

(

GG^{\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\geq 1

, 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 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	$d_{\mathrm{model}}$	$d_{\mathrm{FF}}$	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	$d_{\mathrm{model}}$	$d_{\mathrm{FF}}$	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.

Panel of Figure 7 — Figure 7: Hessian eigenvalue distribution of different blocks in the toy LLaMA2 model. The results of the up_proj are in Figure 2 .

Panel of Figure 8 — 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.

Accelerated Gradient Methods

Accelerated gradient methods use momentum, curvature, or adaptive updates to reduce the iterations needed for optimization.

See all 41 papers

Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining
Aug 28, 2026Shuchen Zhu, Yuxin Fang, Mingze Wang +1Curvature-Aware Optimization Momentum Methods
ByteDance Seed · Peking University
Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation
Feb 27, 2026Zhengbo Wang, Jian Liang, Ran He +2Memory-Efficient Training Low-Rank Compression
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
MONA: Muon Optimizer with Nesterov Acceleration for Scalable Language Model Training
May 26, 2026Jiacheng Li, Jianchao Tan, Hongtao Xu +5Language Model Pretraining Gradient Descent
Meituan, Beijing, China

Aug 28, 2026J/K move · Enter open · S save

Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining

Shuchen Zhu, Yuxin Fang, Mingze Wang, Kun Yuan

ByteDance Seed · Peking University

Curvature-Aware Optimization Momentum Methods Accelerated Gradient Methods Language Model Pretraining

Open paper AskPDF arXiv

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.

Abstract

Accelerated Gradient Methods

Explore similar work

Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining

Taming Momentum: Rethinking Optimizer States Through Low-Rank Approximation

MONA: Muon Optimizer with Nesterov Acceleration for Scalable Language Model Training