Organizations: Stanford University · University of California, San Diego · University of Chicago · Yale University · Toyota Technological Institute at Chicago
The empirical success of pretraining large language models has inspired a deeper investigation into the underlying loss landscapes and the optimization dynamics. Recent empirical and theoretical study suggest that the training loss landscape often exhibits a "river-valley" structure, which features a low-loss manifold (river) flanked by sharp orthogonal directions with higher loss (mountains). In the long term, the optimization progress is determined primarily by the progress along the river. Within such a landscape, gradient descent with large learning rates can move faster along the river despite high apparent loss due to vertical oscillations, while a subsequent sharp decay in the learning rate suppresses these oscillations, revealing genuine optimization progress. This explains the recent success of warmup-stable-decay (WSD) learning rate scheduler which, unlike cosine scheduling, keeps stable high learning rate and decays before producing intermediate checkpoints. Building on this foundation, in this work we take a step further and study the role of momentum within such a loss landscape. We establish theoretical analysis that characterizes how momentum accelerates optimization by stabilizing large learning rates that can not be tolerated by vanilla GD without deviating significantly from the river. The enabled large learning rate in-turn gives greater speed along the river and makes faster essential progress in the long run. Another intriguing observation from theory is that for a river-valley landscape with very flat and slow-spinning river, the momentum itself does not contribute directly to acceleration in terms of the speed of tracking the river, while the main acceleration comes from the admissible larger learning rate.
Figures & tables
Figure 1: GD-M stabilizes a river-valley trajectory. At η=0.4 , GD-M ( β=0.9 , red) damps transverse oscillations and moves along the black river, whereas GD ( β=0 , blue) keeps oscillating. Opacity increases with iteration.
Figure 2: Numerical verification on L(x1,x2)=2(x2−sin(x1))2−0.1x1 , initialized at w0=(0,1) . (a) During the first 160 iterations at η=0.4 , GD-M with β=0.9 damps transverse oscillations and advances along the river, whereas GD remains in persistent oscillations and makes no sustained progress. (b) After T=1000 , stable runs attain nearly the same loss at fixed η , while larger β permits larger stable learning rates and hence lower loss. Color encodes β ; arrows point to categorical post-threshold markers, not loss values. (c) Gradient-colored circles and labels give empirical thresholds; the solid curve is ηmaxGD(1+β)/(1−β+0.05) .
Figure 3
Figure 5: GPT-Neo training on TinyStories ( Eldan and Li, 2023 ) over the first 15,000 steps; columns correspond to η∈{0.0001,0.0006,0.0008} . Top: per-batch training losses recorded every 100 steps (faint) and their 900 -step moving averages (solid). Bottom: full-validation losses evaluated at steps 5,000 , 10,000 , and 15,000 .
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Proof-dependency graph for Theorem C.1 . Lemma C.7 supplies (I1); Lemmas C.8 and C.12 support Lemma C.9 , after which Lemmas C.10 and C.11 establish (I2) and (I4) in sequence and close the endpoint induction. Lemma C.4 extends these bounds to the interpolated trajectory. Lemma C.13 then proves proximity to the river, and Lemma C.14 converts the flat-direction control into near-linear reference-flow time. Arrows inside the central block show the one-step induction dependencies, while the vertical spine records how the resulting endpoint bounds are converted into the two conclusions of Theorem C.1 .
Figure 7: Training and validation losses for η=0.0001 and β1∈{0,0.85,0.9,0.95} .
Figure 8: Training and validation losses for η=0.0002 and β1∈{0,0.85,0.9,0.95} .
Figure 9: Training and validation losses for η=0.0004 and β1∈{0,0.85,0.9,0.95} .
Figure 10: Training and validation losses for η=0.0006 and β1∈{0,0.85,0.9,0.95} .
Figure 11: Training and validation losses for η=0.0008 and β1∈{0,0.85,0.9,0.95} .
Figure 12: Training and validation losses for β1=0 and learning rates {1,2,4,6,8}×10−4 .
Figure 13: Training and validation losses for β1=0.85 and learning rates {1,2,4,6,8}×10−4 .
Figure 14: Training and validation losses for β1=0.9 and learning rates {1,2,4,6,8}×10−4 .
Figure 15: Training and validation losses for β1=0.95 and learning rates {1,2,4,6,8}×10−4 .
Recently, Muon has gained substantial attention as an appealing alternative to Adam-like optimizers, with many works highlighting its advantages through spectral normalization and improved conditioning. Yet this positive theoretical narrative contrasts with its empirical performance in large language model (LLM) training, where Muon's gains over Adam/AdamW are often mixed, schedule-sensitive, and not uniformly superior. To address this gap, we develop a trajectory-level theory characterizing both the strengths and limitations of Muon. We introduce a mixed-spiked matrix sensing model whose sensing operator decomposes into signal, spike, and bulk components, capturing a mixture of anisotropic structure and long-tail information reminiscent of LLM training. On top of it, we adopted a river-valley perspective in which we view the landscape as composed of a river direction flowing to the desired solution and hill directions encoding nuisance or task-irrelevant information. In the momentum-free setting, we show that Muon moves faster along the information-bearing river direction during early optimization, but can converge much more slowly near the river bottom than gradient descent. We then extend the river-valley perspective to general nonconvex objectives with momentum by studying points on the spectral river. There, while Muon converges faster early on, its orthogonalized update removes residual scale information, making it prone to overshooting and oscillation near the target solution. Together, these results suggest that our characterizations extend beyond spiked matrix sensing and motivate switching to GD-like refinement optimizers in the final phase, rather than relying only on a fixed learning-rate schedule for Muon. We also provide preliminary evidence supporting this two-stage approach in language model training experiments.
Tianqi Shen, Jinji Yang, Runze Shi +3
Department of Computer Science City University of Hong Kong (CityUHK) · Department of Statistics and Data Science University of Pennsylvania · School of Statistics and Data Science Shanghai University of Finance and Economics
Learning-rate warmup is a standard technique in language-model training, yet its duration remains largely heuristic. Common approaches use either a fixed number of updates or a fixed fraction of the training horizon, two choices that imply very different scaling as training gets longer. When should warmup stay fixed, and when should it grow with the horizon? We address this question with a quadratic model whose modes respond differently to the peak learning rate. Warmup slows progress in directions that already contract well at the peak rate, but can remove persistent error in directions near the stability edge, with higher peak rates shifting the balance toward longer warmup durations. This yields a compact horizon scaling law that captures regimes ranging from essentially no warmup, through fixed-duration warmup, to durations that grow with the training horizon, and explains how the preferred regime changes with peak learning rate. Because the law captures the tradeoff between giving up early progress and improving the trajectory that follows, it can be fit using shorter runs and used to predict warmup at substantially longer horizons. Together, our results explain several familiar properties of warmup through a single tradeoff and suggest treating warmup duration as a horizon-dependent hyperparameter rather than a fixed training heuristic.
The scale and expense of pre-training language models make efficient hyperparameter tuning essential, yet a principled guidance is still missing. In this work, we analyze language model pre-training dynamics from a local landscape geometry perspective. Our study reveals two distinct phases. In Phase I, sharpness of the local landscape is initially high, leading to instability and loss plateaus under large learning rates (LRs). The landscape shifts from sharp to flatter regions early in training. This dynamic explains the necessity of LR warmup and further suggests that larger peak LRs require proportionally longer warmup periods. In Phase II, the local landscape is governed by the gradient noise scale. Our theory identifies a depth flatness trade-off: high noise from smaller batches widens the loss basin, whereas reduced noise from larger batches deepens it. This theory motivates a dynamic batch-size (BS) scheduler that begins with a small BS and increases it late in training. Together, we provide a unified view of loss landscape evolution, which translates into actionable tuning strategies for large-scale pre-training.
Zhanpeng Zhou, Yuhan Sun, Bingrui Li +4
Shanghai Jiao Tong University · Tsinghua University · Peking University