Source-linked AI summary
How Does Local Landscape Geometry Evolve in Language Model Pre-Training?
Zhanpeng Zhou, Yuhan Sun, Bingrui Li, Jinbo Wang, Huaijin Wu, Lei Wu, Junchi Yan
TL;DR
Efficient language-model pre-training needs principled hyperparameter tuning, but guidance remains unclear. The paper analyzes local landscape geometry and identifies an early sharp-to-flat phase and a later noise-governed phase, connecting them to warmup and batch-size schedules.
Problem
Efficiently tuning language-model pre-training remains difficult because principled guidance is unclear and grid search is costly and unreliable at scale.
Source: §1 · p. 1Method
Results
The study identifies an early sharp-to-flat transition associated with LR warmup and a late noise-governed regime in which preserving η/B yields nearly identical behavior across schedulers.
Source: §1 · pp. 1–2, §6 · p. 9Takeaways & Limitations
The paper suggests warming up longer for larger peak LRs and starting with a small batch size before increasing it late in training.
Source: §7 · p. 9
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper studies how local loss-landscape geometry evolves during language-model pre-training and how that evolution can guide hyperparameter tuning. It identifies an early sharp-to-flat phase linked to LR warmup and a later noise-governed phase motivating batch-size scheduling.
Our contributions
- The paper frames its systematic study against prior work on blockwise sharpness, WSD landscape geometry, and fine-tuned-model landscapes.Those earlier analyses addressed language models or fine-tuning, while this paper focuses on landscape evolution during pre-training.Source: §1 · p. 1
Phase I: From Sharp to Flat Landscapes
- Early pre-training shifts models from sharp toward flatter regions, and this sharp-to-flat transition informs the paper’s hyperparameter-tuning strategies.The introduction contrasts this trajectory with progressive sharpening reported in prior work (Cohen et al., 2021; Song & Yun, 2023; Cohen et al., 2025).Source: §1 · p. 1
- High early sharpness makes large peak LRs without sufficient warmup prone to instability, including loss spikes and plateaus.The paper argues that warmup keeps LR small until sharpness has decayed; within a reasonable range, larger peak LRs call for proportionally longer warmup.Source: §1 · p. 2
Phase II: Basin Selection Governed by Noise Scale
- In Phase II, smaller batches widen the loss basin, while larger batches deepen it through a noise-driven depth–flatness trade-off.The theory attributes the trade-off to gradient noise: increased noise tends to regularize sharpness, while reduced noise tends to minimize loss.Source: §1 · p. 2
- A batch-size scheduler that starts small and ramps up late in training is motivated as a practical strategy for this later regime.The paper reports that the schedule supports steady loss reduction with few token consumption.Source: §1 · p. 2
2 RELATED WORKS
Prior landscape studies span small-network optimization, language-model sharpness, warmup, and batch-size scheduling. This paper positions its contribution as a systematic account of landscape evolution in large-scale language-model pre-training and its connection to tuning.
Local landscape geometry evolution
- Earlier landscape studies reported that GD often moves from flatter to sharper regions, while SGD sharpness can increase or decrease by setting; these results were largely limited to small networks.The later phase can also be governed by optimizer properties, with stochastic noise cited as a bias toward flat minima (Zhu et al., 2019; Xie et al., 2021; Wu et al., 2022).Source: §2 · p. 2
Large-scale pre-training: learning rate warmup
- Warmup is standard in large-scale pre-training, but its mechanism lacks a unified explanation; this paper offers a geometric account linking peak LR to warmup duration.Prior explanations include preventing large early updates, bias reduction, and mitigating momentum bias correction and correlated gradients (Gotmare et al., 2019; Bergsma et al., 2025; Kosson et al., 2024).Source: §2 · pp. 2–3
Large-scale pre-training: batch size schedules
- Prior batch-size work emphasizes critical batch size, while scheduling has received less systematic analysis; this paper links when to increase batch size to evolving landscape geometry.Existing language-model training used stage-wise batch-size schedules without systematic analysis (Brown et al., 2020; Touvron et al., 2023; Liu et al., 2024; Li et al., 2025).Source: §2 · p. 3
3 PRELIMINARIES
The paper sets up a preconditioned-SGD analysis and experiments on small LLaMA-2 models, varying learning rate, warmup, and batch size. A warmup-stable schedule is used to separate batch-size effects from LR decay.
Basic notations
- The notation section defines vector and matrix conventions, three matrix norms, trace, and the Hadamard product symbol.The Hadamard product is denoted by ⊙.Source: §3 · p. 3
Theoretical setup
- The theoretical setup uses preconditioned SGD with a fixed positive-definite preconditioner and models batch-size-dependent gradient noise.The analysis assumes the late-training iterate remains close to a global minimum, where the loss is approximated quadratically.Source: §3 · p. 3
Experimental setup
- Experiments use 93M- and 170M-parameter LLaMA-2 models, and vary LR and batch size under a warmup-stable schedule.The study fixes batch size while examining LR and warmup, then fixes LR while examining batch size; the LR is held constant after warmup to decouple batch-size ramping from decay.Source: §3 · p. 3
4 PHASE I: FROM SHARP TO FLAT LANDSCAPES
In Phase I, local landscapes move from sharp regions toward flatter ones, while sharpness constrains stable learning rates. The resulting analysis connects early loss spikes and plateaus to insufficient warmup and motivates longer warmup for larger peak LRs within a tested range.
Motivating observations: instability and loss plateaus early in training
- With warmup fixed at 16 iterations and peak learning rate varied across models of different sizes, loss plateaus reliably appear near warmup’s end across all settings, while larger learning rates produce higher spikes.Source: §4 · p. 4
Lyapunov stability analysis: sharpness matters
- For preconditioned GD, stability requires the LR to remain below a boundary set by the largest eigenvalue of the preconditioned curvature matrix.When curvature along the sharpest direction is too large, only a sufficiently small LR can prevent divergence.Source: §4 · p. 4
- As LR approaches the stability boundary 2/Λk, the guaranteed per-step loss decrease tends to zero, helping explain spikes and plateaus near that boundary.The stated condition bounds curvature along the update segment by Λk.Source: §4 · pp. 4–5
The early dynamics: from sharp to flat landscapes
- Hessian top eigenvalues decrease sharply and random-direction loss landscapes widen as early training proceeds, indicating a shift from sharp to flat regions.The reported trajectories use the same random direction across iterations for fair comparison.Source: §4 · p. 5
A Tuning Recipe: Larger Peak LR, Longer Warmup
- If LR rises too quickly while the model remains in sharp regions, shortened warmup leads to loss spikes and plateaus.As training proceeds and the landscape flattens, the same LR no longer threatens stability in the same way.Source: §4 · p. 5
- The recommended practice is to keep LR small until sharpness decays, allowing iterates to transition into flatter landscapes.The paper presents this warmup strategy as a way to prevent loss spikes and plateaus.Source: §4 · p. 5
- Within peak LRs 2^-8 to 2^-11, larger peak LRs require proportionally longer warmup to achieve optimal validation loss.This proportionality is not universal: at η = 2^-7, the optimal warmup remains 210 iterations, as for η = 2^-8.Source: §4 · pp. 5–6
5 PHASE II: LOCAL LANDSCAPE GOVERNED BY NOISE SCALE
In Phase II, batch size controls the noise scale and shifts basin selection between wider, shallower minima and deeper minima. The authors use this trade-off to motivate a small-to-large batch-size schedule, ramped late in training to improve data efficiency.
The effect of BS: local landscapes late in training
- Larger batches yield lower terminal loss and deeper basins, while smaller batches produce flatter, wider basins late in training.The comparison includes validation-loss curves and loss landscapes across batch sizes.Source: §5 · p. 6, Figure 5 · p. 7
Widen or deepen: noise scale governs basin selection
- The theory attributes basin selection to free energy: early on loss dominates, while later the flatness penalty becomes comparable and selection grows more sensitive to noise scale τ ∝ η/B.The stationary probability of residing in a basin is governed by the free-energy function.Source: §5 · pp. 6–7
BS scheduler: design principle II
- Runs with different ramp times collapse onto the same loss trajectory, but earlier ramping consumes more data by a given iteration.The experiments compare schedules starting at 0.49M and ramping to 2× or 4× the initial batch size.Source: §5 · p. 7, Figure 6 · p. 8
Scope
- The analysis concerns how reduced noise can move optimization toward deeper minima, a perspective distinct from flat-minima arguments in fully interpolating regimes.The authors motivate scheduling particularly for data-limited pre-training, where excessively large batches increase data consumption.Source: §5 · p. 7
BS scheduler: design principle II
- Ramping batch size late—when loss reduction becomes marginal—is most effective; delaying it preserves progress along the shared loss trajectory under a data-limited budget.Ramping too early consistently harms final performance, while late ramping sharply reduces noise and drives convergence toward deeper minima.Source: §5 · p. 8
Efficient pre-training: a BS scheduler in data-limited regime
- About 1/4 of the tokens and a ~4× speedup suffice for batch-size ramping to 7.8M to achieve nearly the final validation loss of a fixed 7.8M batch.The schedule starts at 0.49M and ramps by 4× when loss minimization slows.Source: §5 · p. 8
Phase II implications
- The authors predict that batch-size ramping and learning-rate decay yield comparable performance because both reduce the noise scale; Figure 8 compares the schedules.The stated proportionality is τ ∝ η/B.Source: Figure 8 · p. 9
6 MORE DISCUSSIONS: LR DECAY AND BS RAMPING
The analysis treats LR decay and BS ramping as ways to reduce the noise scale η/B. Experiments report similar loss behavior when schedules make comparable changes to that scale.
6 MORE DISCUSSIONS: LR DECAY AND BS RAMPING
- BS ramping and LR decay produce remarkably similar validation loss curves when their changes are aligned across positions.The comparison uses a one-time 16× batch-size increase or 16× learning-rate decrease.Source: §6 · p. 9
- Three schedules with the same η/B yield nearly identical loss curves, while the 4× LR-drop schedule follows a noticeably different trajectory.The matching schedules are a 16× LR drop, a 16× BS ramp, and a simultaneous 4× LR drop plus 4× BS ramp.Source: §6 · p. 9
- Late-phase training dynamics are governed by η/B, and schedulers preserving it exhibit nearly identical behavior.The authors report that LR decay and BS ramping reduce the noise scale in the same manner.Source: §6 · p. 9
7 CONCLUSION AND LIMITATIONS
The conclusion presents pre-training as an early sharp-to-flat transition followed by a late noise-governed regime, with different tuning implications for each. It also states theoretical limitations and identifies an unresolved question about batch-size schedules.
7 CONCLUSION AND LIMITATIONS
- The paper’s two-phase account links the early sharp-to-flat transition to LR warmup and the late noise-governed regime to increasing BS late in training.It suggests proportionally longer warmup for larger peak LRs and a schedule that starts with small BS.Source: §7 · p. 9
- The theory relies on strong assumptions, including an infinitesimal-LR SDE, and does not fully explain loss-curve collapse under different BS schedules.The conclusion identifies learning dynamics under different BS schedules as an open question.Source: §7 · p. 9
A TERMINOLOGIES
- Table 1 presents terminology and its usage in the paper.Source: Table 1 · p. 14
B.1 PHASE I: LYAPUNOV STABILITY ANALYSIS
The appendix derives Phase I stability and loss-change conditions, then decomposes the discrete-time dynamics into eigenbasis modes to characterize their stationary variances.
B.1 PHASE I: LYAPUNOV STABILITY ANALYSIS
- Requiring every eigenvalue of the update matrix to have magnitude below 1 gives the linear-system stability condition.The eigenvalues are expressed as 1 − ηλ_j(S_k).Source: §B.1 · p. 15
- Under the segment curvature bound Λ_k, η ≤ 2/Λ_k guarantees non-increasing loss, while the guaranteed decrease per step tends to zero as η approaches 2/Λ_k.Source: §B.1 · p. 16
B.2.1 DISCRETE-TIME SOLUTION
- The eigenbasis transformation decouples the recursion into independent scalar modes and yields a stationary variance for each mode.The variance expression depends on η, G_jj, and λ_j.Source: §B.2.1 · pp. 16–17
B.2.2 CONTINUOUS-TIME LIMIT
The continuous-time analysis specifies a noise-covariance scaling under which the discrete process converges to an Itô SDE. It then gives the SDE’s stationary distribution and a result for basin probabilities across multiple local minima.
B.2.2 CONTINUOUS-TIME LIMIT
- Under the specified noise-covariance scaling, the discrete process converges weakly to an Itô SDE.The scaling is parameterized by temperature τ.Source: §B.2.2 · pp. 17–18
- The limiting SDE has a Gibbs stationary distribution.Source: §B.2.2 · p. 18
- For empirical risk with multiple local minima, the theorem gives the stationary probability that training resides in each basin.The theorem is titled “Free Energy Minimization.”Source: §B.2.2 · p. 18
C.1 EXPERIMENTAL SETUPS
The experiments use GPT-2 and LLaMA models, FineWeb-Edu data, and four optimizers with optimizer-specific settings. Model configurations are detailed in Table 2.
C.1 EXPERIMENTAL SETUPS
- Experiments use GPT-2 (small), implemented via nanoGPT, and LLaMA, a decoder-only Transformer implemented with HuggingFace Transformers; configurations for both models are detailed in Table 2.GPT-2 uses GELU activations and standard LayerNorm; LLaMA incorporates RoPE, SwiGLU, and RMSNorm.Source: §C.1 · p. 20
- Training uses a randomly sampled FineWeb-Edu subset of around 100B GPT-2 tokens.The passage notes that FineWeb-Edu is widely used in LLM pre-training research.Source: §C.1 · p. 20
- The study uses AdamW, Muon, Adam-mini, and Lion with optimizer-specific momentum, weight-decay, and learning-rate settings.Lion’s learning rate is divided by 10× relative to AdamW, while its weight decay is ramped up to 10× to keep effective LR λη = 0.1.Source: §C.1 · p. 20
- All optimizers use gradient clipping at 1.0 for stability.Source: §C.1 · p. 20
C.2 MORE RESULTS UNDER VARIOUS SETUPS.
The experiments extend the sharp-to-flat dynamics and batch-size scheduling studies across architectures, optimizers, and larger models. They also explore the warmup-tuning recipe on additional architectures.
C.2 MORE RESULTS UNDER VARIOUS SETUPS.
- The extensions focus primarily on validating early sharp-to-flat dynamics and the proposed batch-size scheduling principle across settings.The warmup-tuning recipe is also explored on additional architectures.Source: §C.2 · p. 20
- GPT-2 experiments extend the architecture studies, with Figures 9 and 10 presenting the results.Source: §C.2 · p. 20
- The extensions include Adam-mini and Lion optimizers, with results presented in Figures 11 and 12.Source: §C.2 · p. 20
- The larger-model experiments include models with 270M and 530M parameters.Source: §C.2 · p. 20
C.3 ABLATION STUDIES ON EARLY INSTABILITIES.
The ablations link early training instability to sharp landscape regions and inadequate warmup rather than AdamW startup issues. They also examine how peak learning rate and batch size relate to instability and scheduling.
C.3 ABLATION STUDIES ON EARLY INSTABILITIES.
- Figure captions report early sharp-to-flat dynamics and data-efficiency gains from batch-size scheduling across GPT-2, Adam-mini, Lion, and larger-scale experiments.Source: Figure 9 · p. 21, Figure 11 · p. 21, Figure 12 · p. 22
- Figure 10 indicates that larger peak LRs require longer warmup, selecting the optimum by validation loss across GPT-2 runs at 100 TPP.Source: Figure 10 · p. 21
- Loss spikes and plateaus persist with Muon when warmup is shortened or peak LR increased, ruling out AdamW startup issues as their source.Source: §C.3 · p. 21, Figure 14 · p. 22
- Shorter warmup raises the possibility of loss spikes; the authors explain that warmup gradually raises LR until training enters flatter regions that tolerate larger updates.Source: §C.3 · p. 22, Figure 15 · p. 23
- With no warmup, small-batch runs show more significant spikes; small batch size does not replace warmup because sharpness remains high.The passage explains that an immediate jump to a large LR while sharpness is high causes a spike almost at initialization.Source: §C.3 · pp. 22–23, Figure 15 · p. 23