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

cs.LGSep 30, 2026

How Does Local Landscape Geometry Evolve in Language Model Pre-Training?

Authors: Zhanpeng Zhou, Yuhan Sun, Bingrui Li, Jinbo Wang, Huaijin Wu, Lei Wu, Junchi Yan

Organizations: Shanghai Jiao Tong University · Tsinghua University · Peking University

Abstract

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.

Figures & tables

Figure 1: The evolution of local loss landscape throughout pre-training. We train LLaMA-2 models with

170

M parameters using different BSs (

0.49

M and

7.8

M), and visualize the one-dimensional loss landscape at iterate

\bm{\theta}_{t}

along a random direction

\bm{\delta}

, i.e., plot

L(\bm{\theta}_{t}+\alpha\bm{\delta})

vs. the perturbation coefficient

\alpha

. The landscapes are shown across different training iterations

t

. Phase I. The landscapes gradually widen for both training runs. Phase II. Training with smaller BS produces wider landscapes than training with larger BS.

Figure 2: Loss spikes and plateaus in Phase I. We train a series of LLaMA-2 models with $93$ M and $170$ M parameters. We adopt a warmup-stable schedule, where the warmup length is shortened to $16$ iterations and the peak LR is varied, $\eta\in\{2^{-11},2^{-10},2^{-9},2^{-8},2^{-7}\}$ . (Left). LR schedule: $\eta_{t}$ vs. training iteration $t$ . (Middle, right). Training loss curves for different model sizes: $L(\bm{\theta}_{t})$ vs. training iteration $t$ . The vertical dashed line marks the end of the warmup phase.

Figure 3: Early pre-training shifts iterates from sharp to flat regions. We visualize the local landscape geometry evolution of training runs in Figure 2 . For each model size, we select the training run with LR $2^{-10}$ . (Top). Evolution of the top eigenvalues of the Hessian across iterations: $\lambda_{i}(\mathbf{H}(\bm{\theta}_{t}))$ vs. iteration $t$ . (Bottom). One-dimensional loss landscape along a random perturbation direction: the perturbed loss $L(\bm{\theta}_{t}+\alpha\bm{\delta})$ vs. perturbation coefficient $\alpha$ , shown across early training iterations $t$ .

Figure 4: Larger peak LR, longer warmup. We plot the best validation loss $L(\bm{\theta}_{\text{bst}})$ vs. warmup $T_{w}$ for different $\eta$ . Optimal $T_{w}$ is highlighted with a star.

Figure 5: Large BS deepens the basin, small BS widens the basin. We train a series of LLaMA-2 models (170M) for $T=20{,}480$ iterations, using BSs $B\in\{0.49\text{M},0.98\text{M},1.9\text{M},3.9\text{M},7.8\text{M}\}$ . (Top left). Validation loss curves for different BSs: $L(\bm{\theta}_{t})$ vs. training iteration $t$ . (Top right). One-dimensional loss landscapes at the final iterates $\bm{\theta}_{T}$ along a random perturbation direction: perturbed loss $L(\bm{\theta}_{T}+\alpha\bm{\delta})$ vs. perturbation coefficient $\alpha$ , visualized across different BSs. (Bottom). One-dimensional loss landscape: the perturbed loss $L(\bm{\theta}_{t}+\alpha\bm{\delta})$ vs. perturbation coefficient $\alpha$ , shown across late training iterations $t$ for $B=0.49$ M and $B=7.8$ M.

Figure 6: (Left) Collapse of loss curves under different BS schedules. Validation loss curves for training with different BS scheduling. In all runs, BS starts at $0.49$ M. For blue curves, BS is ramped up to $4\times$ its initial value; for red curves, BS is ramped to $2\times$ . The ramping times, $T_{2\times}$ or $T_{4\times}$ , are varied across different positions. (Middle, Right) Ramping BS is more efficient late in training. We evaluate a two-stage BS-ramping schedule with ramp times $T_{1}$ and $T_{2}$ . For the red curves, we fix $T_{2}=10$ B and vary $T_{1}$ ; for the blue curves, we fix $T_{1}=10$ B and vary $T_{2}$ . (Middle). Illustration of BS schedulers. (Right). Final validation loss vs. the relative ramping time, i.e., $(T_{1})/T,(T_{2}-T_{1})/T\in[0,0.5]$ , where $T$ denotes the total training tokens.

Figure 7: BS scheduling improves data efficiency. We train LLaMA-2 models with $93$ M and $170$ M parameters, using a BS schedule that starts at $0.49$ M and increases by $4\times$ at each ramp. Models are trained with $1,2$ , or $3$ ramping steps, while models without ramping serve as the baseline. Vertical gray dashed lines indicate ramping positions. (Left, Middle). The validation curves for each run. (Right). Comparison between training with BS ramping to $7.8$ M and training with a fixed $7.8$ M BS.

Figure 8: (Left) BS ramping performs similarly to LR decay. Validation loss curves for training with either BS ramping or LR decay. For BS ramping, BS increases to $16\times$ its initial value; for LR decay, the LR drops to $1/16$ of its initial value. Each method applies a single step at varying positions. (Right) Interplay between BS ramping and LR decay. We evaluate four different scheduling strategies. At the 10B tokens, (a) LR drops to $1/4$ of its initial value; (b) LR drops to $1/16$ ; (c) BS ramps to $16\times$ . (d) LR drops to $1/4$ and BS ramps to $4\times$ . In all runs (both left and right), BS starts at $0.49$ M and LR begins at $2^{-10}$ (after linear warmup).

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Terminology	General Meaning	Usage in This Paper
Sharpness	A measure of curvature in the loss landscape, often characterized via the Hessian. Different works may define it differently.	We define sharpness as the curvature along the sharpest direction of the loss landscape. Mathematically, it is presented as the largest eigenvalue of the Hessian $\lambda_{\max}(\mathbf{H}(\bm{\theta}_{t}))$ or of the preconditioned curvature matrix $\lambda_{\max}(\mathbf{S}(\bm{\theta}_{t}))$ .
Flat/sharp minimum	A minimum is a point where the gradient vanishes $\nabla L(\bm{\theta})=0$ and the loss does not decrease in a small neighborhood. A sharp minimum has large curvature; a flat minimum has small curvature.	We use these terms sparingly and follow the standard definitions from the sharpness/flat-minima literature.
Wide/deep basin	A loss basin is a region of the landscape surrounding a minimum. A wide basin rises loss slowly in most directions, whereas a deep basin has a significantly lower minimum value compared to its surroundings.	We use these terms to establish the depth–flatness trade-off: large noise scales tend to find wide basins, while small noise scales tend to find deeper regions with lower loss.

Appendix

Table 1: Terminology and usage in this paper.

Acronym	Size	$d_{\mathrm{model}}$	$d_{\mathrm{FF}}$	n $\_$ head	depth
GPT-2 (small)	124M	768	3072	12	12
LLaMA (93M)	93M	512	2048	16	8
LLaMA (170M)	170M	768	3072	12	8
LLaMA (270M)	270M	1024	4096	16	8
LLaMA (530M)	530M	1536	6144	24	8

Appendix

Table 2: Model configurations.

Figure 9: Extensive to GPT-2 architectures. (Left). Early sharp-to-flat dynamics. Evolution of the top eigenvalues of the Hessian across iterations: $\lambda_{i}(\mathbf{H}(\bm{\theta}_{t}))$ vs. iteration $t$ . (Right). BS scheduling improves data efficiency. BS scheduling improves data efficiency. We use a BS schedule that starts at $0.49$ M and increases by $4\times$ at each ramp. Models are trained with $1,2$ , or $3$ ramping steps, while models without ramping serve as the baseline. Vertical gray dashed lines indicate ramping positions.

Figure 10: Extensive to GPT-2 architectures. Larger Peak LR, Longer Warmup. We train a series of GPT-2 models with 100 TPP. We vary the peak LRs $\eta$ and warmup lengths $T_{w}$ . We plot the best validation loss $L(\bm{\theta}_{\text{bst}})$ vs. $T_{w}$ for different $\eta$ . The optimal $T_{w}$ is highlighted with a star.

Figure 11: Extensive to Adam-mini optimizer. (Left). Early sharp-to-flat dynamics. Evolution of the top eigenvalues of the Hessian across iterations: $\lambda_{i}(\mathbf{H}(\bm{\theta}_{t}))$ vs. iteration $t$ . (Right). BS scheduling improves data efficiency. BS scheduling improves data efficiency. We use a BS schedule that starts at $0.49$ M and increases by $4\times$ at each ramp. Models are trained with $1,2$ , or $3$ ramping steps, while models without ramping serve as the baseline. Vertical gray dashed lines indicate ramping positions.

Figure 12: Extensive to Lion optimizer. (Left). Early sharp-to-flat dynamics. Evolution of the top eigenvalues of the Hessian across iterations: $\lambda_{i}(\mathbf{H}(\bm{\theta}_{t}))$ vs. iteration $t$ . (Right). BS scheduling improves data efficiency. BS scheduling improves data efficiency. We use a BS schedule that starts at $0.49$ M and increases by $4\times$ at each ramp. Models are trained with $1,2$ , or $3$ ramping steps, while models without ramping serve as the baseline. Vertical gray dashed lines indicate ramping positions.

Figure 13: Extensive to Larger Scale. (Left). Early sharp-to-flat dynamics. Evolution of the top eigenvalues of the Hessian across iterations: $\lambda_{i}(\mathbf{H}(\bm{\theta}_{t}))$ vs. iteration $t$ . (Right). BS scheduling improves data efficiency. BS scheduling improves data efficiency. We use a BS schedule that starts at $0.49$ M and increases by $4\times$ at each ramp. Models are trained with $1,2$ , or $3$ ramping steps, while models without ramping serve as the baseline. Vertical gray dashed lines indicate ramping positions.

Figure 14: Muon consistently shows loss spikes and plateaus early in training. We train a series of LLaMA-2 models with $170$ M parameters. We adopt a warmup-stable schedule, where the warmup length is shortened to $16$ iterations and the peak LR is varied, $\eta\in\{2^{-11},2^{-10},2^{-9},2^{-8},2^{-7}\}$ . (Left). LR schedule: $\eta_{t}$ vs. training iteration $t$ . (Middle, Right). Training loss curves for different model sizes: $L(\bm{\theta}_{t})$ vs. training iteration $t$ . The vertical dashed line marks the end of the warmup phase.

Figure 15: (Left.) Shorter warmup, more loss spikes. We train a series of LLaMA-2 models with $170$ M parameters. We adopt a warmup-stable schedule, where the warmup length varies from $\{2^{5},2^{6},2^{7},2^{8},2^{9},2^{10}\}$ iterations and the peak LR is fixed $\eta=2^{-7}$ . (Right). Zero warmup and small BS leads to larger loss spikes. We train a series of LLaMA-2 models with $170$ M parameters. We adopt a constant LR schedule, where no warmup and the peak LR is varied, $\eta\in\{2^{-11},2^{-10},2^{-9},2^{-8},2^{-7}\}$ . We also use the $0.49$ M BS.

Explore similar work

Jul 30, 2026cs.LG

Towards joint scaling laws with optimal batch size schedules

Modern deep learning typically keeps the batch size static throughout training, thus overlooking the joint effect of learning rate and batch size on the training dynamics. In this paper, we study the deep learning dynamics through the lens of convex optimization and derive a joint characterization of loss in terms of both schedules, applicable to general optimizers and model architectures. This characterization yields a closed-form optimal batch size schedule for any prescribed learning rate schedule, and further leads to joint scaling laws that consistently outperform static batch size baselines, highlighting the significance of dynamic batch size schedule in large language model training.

Jiaxiang Li, Zhiqi Bu, Shiyun Xu

Meta · Independent researcher

Batch Large Language Model Training Scaling Laws Learning Dynamics Schedule Convex Optimization

Ask PDF

Aug 28, 2026cs.LG

Deriving Scaling Laws for OpenEuroLLM Models: Learning Rate, Batch Size and Loss

We study the scaling behavior of learning rate and batch size in pretraining dense large language models on English-prevalent corpora. Beyond scaling jointly optimal learning rates and batch sizes, we investigate their marginal evolution with model capacity and data scale and develop a model that captures these relationships. As we employ a Warmup-Stable-Decay learning rate schedule, we further investigate the gains from learning rate annealing over a broad range of hyperparameters settings, models and data budgets, and whether the optimal learning rate and batch size transfer between the stable and decay phases. Finally, we characterize the dependence of loss on model capacity and dataset size, evaluating recently proposed scaling forms that explicitly model their interaction. We find these approaches particularly effective at capturing both undertraining and overtraining regimes across our experiments. This study establishes a first baseline and scaling procedure for the development of future OpenEuroLLM models. We open-source the complete collection of pretraining runs used in this study.

Niccolò Ajroldi, Diana Alexandra Onutu, Haider Al-Tahan +4

ELLIS Institute Tübingen · OpenEuroLLM · Eindhoven University of Technology +6

Batch Scaling Laws Capacity

Ask PDF

May 26, 2026cs.LG

The Stability of Singular Distribution: A Spectral Perspective on the Two-Phase Dynamics of Language Model Pre-training

Large language model pre-training typically exhibits a two-phase trajectory: a fast initial loss drop followed by a prolonged slow improvement. We identify an underlying spectral phenomenon, Stability of Singular Distribution (SoSD), where the trace-normalized singular value spectrum stabilizes early, even as parameter matrices continue to evolve. We demonstrate that synchronization between SoSD and the slow-descent regime is widely observed across diverse architectures (GPT-2, LLaMA) and settings, including various schedules (Step-wise, WSD, Cosine Decay), weight decays, and optimizers (AdamW, Muon). By analyzing a simplified Transformer, we prove that growing weight norms inevitably precipitate an early SoSD threshold, after which the rate of loss decrease becomes theoretically bounded by the variation in the singular distribution. We further interpret strategies like WSD and Muon through their ability to modulate the SoSD scale, offering a spectral lens for understanding efficient pre-training dynamics.

Hongtao Zhang, Wenjie Zhou, Chenxi Jia +2

State Key Laboratory of AI Safety, Institute of Computing Technology, Chinese Academy of Sciences, Beijing, China · School of Advanced Interdisciplinary Sciences, University of Chinese Academy of Sciences, Beijing, China · University of Chinese Academy of Sciences, Beijing, China +1

Large Language Model Pretraining Pretraining Singular Value Decomposition Spectral Descent Distributions

Ask PDF