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

Towards joint scaling laws with optimal batch size schedules
Jul 30, 2026Jiaxiang Li, Zhiqi Bu, Shiyun XuBatch Large Language Model Training
Meta · Independent researcher
Deriving Scaling Laws for OpenEuroLLM Models: Learning Rate, Batch Size and Loss
Aug 28, 2026Niccolò Ajroldi, Diana Alexandra Onutu, Haider Al-Tahan +4Batch Scaling Laws
ELLIS Institute Tübingen · OpenEuroLLM · Eindhoven University of Technology +6
The Stability of Singular Distribution: A Spectral Perspective on the Two-Phase Dynamics of Language Model Pre-training
May 26, 2026Hongtao Zhang, Wenjie Zhou, Chenxi Jia +2Large Language Model Pretraining Pretraining
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

Jul 30, 2026, cs.LGJ/K move · Enter open · S save

Towards joint scaling laws with optimal batch size schedules

Jiaxiang Li, Zhiqi Bu, Shiyun Xu

Meta · Independent researcher

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Towards joint scaling laws with optimal batch size schedules

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

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