Theoretical Analysis of Muon on Transformers for Next-Token Prediction
Authors: Shengwei Wu, Yixuan Yang, Gao Huang, Song Li
Organizations: School of Mathematical Sciences, Zhejiang University, Hangzhou, Zhejiang, China · Department of Mathematics, National University of Singapore, Singapore
Muon has recently shown strong empirical performance in training Transformer-based large language models (LLMs), yet its optimization dynamics remain poorly understood from a theoretical perspective. In this paper, we analyze the training dynamics of Muon for next-token prediction (NTP) in Transformers. The training procedure consists of two phases: the preprocessing phase learns the output-value component, while the main phase learns the key-query component. For the preprocessing phase, we characterize the optimization trajectory exactly and show that the normalized spectral margin is optimal after each update. For the main phase, we establish nonasymptotic convergence rates for the training loss and the normalized margin gap under softmax attention. These results suggest that Muon's update rule can help models efficiently learn which tokens matter for prediction, offering a possible explanation for its strong performance in large language model training. The learned weight structure also yields a generalization guarantee for unseen contexts. Furthermore, we validate our theoretical insights with experiments on next-token prediction tasks.
Figures & tables
Figure 1: Normalized margin behavior of Muon and Adam on the synthetic task. Subfigures (a) and (b) show the relative spectral and Frobenius margin gaps G∞(t) and G2(t) , respectively. Subfigures (c) and (d) show the Frobenius cosine similarities to the unique Frobenius hard margin solution V2 and a selected spectral hard margin solution V∞ , respectively. Each curve shows the mean over ten independent runs.
Figure 2: Dependence of Muon’s relative spectral margin gap G∞(t) on the initial step size and decay exponent. Subfigures (a), (b), and (c) use η2,0=0.5 , 1 , and 2 , respectively, and compare α∈{1/4,1/2,3/4,1} under the step size schedule η2,t=η2,0(t+1)−α . The exponent giving the smallest observed gap varies with both η2,0 and the iteration count.
α
D(t)
1/St
Qt/St
Larger term
1/4
0.086223
0.001774
0.078711
Qt/St
1/2
0.075970
0.011272
0.026953
Qt/St
3/4
0.114558
0.058158
0.037658
1/St
1
0.122928
0.209107
0.085985
1/St
Table 1: Normalized spectral margin gap D(t) and the terms 1/St and Qt/St for η2,0=0.5 at iteration t=8,000 . The last column indicates the larger of the two terms for each decay exponent.
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure A.1 : Effect of the ambient embedding dimension. Subfigures (a) and (b) show the relative spectral and Frobenius margin gaps at t=20,000 , respectively. At every dimension, Muon has the smaller spectral margin gap, while Adam has the smaller Frobenius margin gap. Both gaps remain unchanged for Muon and decrease for Adam as d increases.
Figure A.2 : Effect of vocabulary nonorthogonality. Subfigures (a) and (b) show the relative spectral and Frobenius margin gaps at t=20,000 , respectively, for different perturbation strengths σ . At every perturbation strength, Muon has the smaller spectral margin gap, while Adam has the smaller Frobenius margin gap.
Figure A.3 : Effect of Muon’s momentum parameter. Subfigure (a) shows the relative spectral margin gap over the first 20,000 Phase II iterations for different values of β2 , together with the Adam trajectory. Subfigure (b) shows the gap at t=100,000 for β2=0.99 and at t=20,000 for the other momentum values. The horizontal line in (b) shows Adam’s gap at t=20,000 . The spectral margin gap decreases more slowly for β2=0.99 over the tested iteration range.