Paper deep dive
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
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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.
Tags
Links
- Source: https://arxiv.org/abs/2609.39767v1
- Canonical: https://arxiv.org/abs/2609.39767v1
Trouble viewing inline? Open PDF directly →
Full Text
82,124 characters extracted from source content.
Expand or collapse full text
How Does Local Landscape Geometry Evolve in Language Model Pre-Training? Zhanpeng Zhou †thanks: Equal contribution. Correspondence to Zhanpeng Zhou (zzp1012@sjtu.edu.cn). Yuhan Sun Bingrui Li Affiliation: Tsinghua University Jinbo Wang Affiliation: Peking University Huaijin Wu Affiliation: Shanghai Jiao Tong University Lei Wu Affiliation: Peking University Junchi Yan Affiliation: Shanghai Jiao Tong 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 I, 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 Training language models efficiently requires carefully tuned hyperparameters, yet a principled guidance for tuning remains unclear. While practitioners often rely on grid search, these approaches are costly and unreliable at scale. Recent research (Foret et al., 2021; Cohen et al., 2021; Gilmer et al., 2022) has highlighted that the geometry of the local loss landscape offers fundamental insights into optimization, revealing how factors such as sharpness11 1 To avoid misunderstanding, we clarify the terminology in Table 1. (Keskar et al., 2017; Zhang et al., 2017; Jiang et al., 2020) shape training dynamics. Leveraging insights from the local loss landscape presents a promising path toward principled hyperparameter tuning for language model pre-training. Several pioneering works have already attempted to study language models from the local loss landscape perspective. Zhang et al. (2024a); Wang et al. (2025) identified blockwise sharpness patterns in language models through Hessian-based analyses. Wen et al. (2024) introduced the “river-valley” landscape to explain the effectiveness of Warmup-Stable-Decay (WSD) schedules (Hu et al., 2024). Peng et al. (2024); Chen et al. (2025) further visualized the loss landscapes of finetuned language models, offering geometric insights into the safety alignment. To this end, we pose the central questions of this paper: 1. How does the local landscape geometry evolve in language model pre-training? 2. What implications does this evolution have for principled hyperparameter tuning? Our contributions. In this work, we present the systematic study of the evolution of local landscape geometry during language model pre-training. As illustrated in Figure 1, our analysis reveals two distinct phases, each with implications for hyperparameter tuning. • Phase I: From Sharp to Flat Landscapes. In Phase I, we observe that the model shifts from sharper regions of the loss landscape toward flatter ones, contrary to the progressive sharpening phenomenon in prior works (Cohen et al., 2021; Song and Yun, 2023; Cohen et al., 2025). Lyapunov stability analysis in Section 4 shows that the maximum stable learning rate (LR) is inversely proportional to sharpness. Since sharpness is extremely high early in pre-training, using large peak LRs without sufficient warmup leads to instabilities, such as loss spikes and plateaus (see Figure 2). Implications. The sharp-to-flat transition explains the necessity of LR warmup: LR should remain small until sharpness has sufficiently decayed, preventing training instabilities. This further provides a tuning recipe: within a reasonable range, larger peak LRs require proportionally longer warmup, to safely navigate the sharpest stage of training. • Phase I: Basin Selection Governed by Noise Scale. In Phase I, the local loss landscape is largely governed by the noise scale during training, with batch size (BS) serving as its primary controller. Our analysis shows that smaller BS widens the loss basin, while larger BS deepens it. Theoretically, we analyze a continuous setup of preconditioned SGD, which uncovers a depth flatness trade-off: reduced gradient noise tends to minimize the loss, leading to deeper minima; whereas increased noise tends to regularize the sharpness of landscape, moving toward wider ones. Implications. The trade-off, together with the ramping-time experiment in Figure 6, motivates a BS scheduling strategy: begins with a small BS and ramps it until the late phase of training. Our scheduling ensures steady loss reduction with few token consumption. Moreover, since the noise scale is proportional to η/Bη/B in our theory, we predict that BS ramping and LR decay reduce the noise scale in similar ways and thus yield comparable performance (see Figure 8). In summary, our work provides a two-phase picture of pre-training: an early sharp-to-flat transition that necessitates LR warmup, and a late noise-driven regime that motivates BS scheduling. This unified view advances our understanding of pre-training dynamics. Figure 1: The evolution of local loss landscape throughout pre-training. We train LLaMA-2 models with 170170M parameters using different BSs (0.490.49M and 7.87.8M), and visualize the one-dimensional loss landscape at iterate t θ_t along a random direction δ, i.e., plot L(t+α)L( θ_t+α δ) vs. the perturbation coefficient α. The landscapes are shown across different training iterations t. Phase I. The landscapes gradually widen for both training runs. Phase I. Training with smaller BS produces wider landscapes than training with larger BS. 2 Related Works Local landscape geometry evolution. Understanding how local landscape geometry, particularly sharpness, evolves during training has drawn attention before the success of large language models. Wu et al. (2018); Cohen et al. (2022); Song and Yun (2023); Cohen et al. (2025) showed that initially gradient descent (GD) tends to move from flatter to sharper regions of the landscape. In addition, Jastrzębski et al. (2019); Jastrzebski et al. (2020) argued that in SGD, sharpness also changes monotonically but either increase or decrease depending on the setting. In the later phase, however, sharpness is largely governed by the properties of the optimizer (Zhou et al., 2025). One notable example is that the stochastic noise of SGD and its variants implicitly biases training toward flat minima (Wu et al., 2018; Zhu et al., 2019; Xie et al., 2021; Wu et al., 2022). Yet, these findings are largely restricted to small networks; In comparison, our work presents a systematic study of how local landscape geometry evolves in large-scale language model pre-training. Large-scale pre-training: learning rate warmup. Learning rate warmup, first introduced in large-batch ResNet (He et al., 2016; Goyal et al., 2017) and Transformer training (Vaswani et al., 2017), is now standard in large-scale pre-training (Shoeybi et al., 2019; Zhang et al., 2022; Hu et al., 2024). Its mechanism, however, remains poorly understood. Gotmare et al. (2019) showed that warmup prevents excessively large early parameter updates; Bergsma et al. (2025) attributed the early updates to bias reduction. Kosson et al. (2024) showed in language model pre-training that warmup mitigates momentum bias correction and correlated gradients that drive unstable representation shifts. Yet no unified explanation exists. In comparison, our work views warmup from a geometric perspective, suggesting that larger peak LRs require proportionally longer warmup. Large-scale pre-training: batch size schedules. Batch size is another critical hyperparameter in large-scale pre-training, controlling the trade-off between time efficiency and data efficiency. Prior works (McCandlish et al., 2018; Kaplan et al., 2020; Gray et al., 2023; Gray et al., 2024; Zhang et al., 2025) have focused on the critical batch size (CBS), the point where further increasing BS yields diminishing returns. However, CBS is usually considered as constant throughout training, and less attention has been given to BS scheduling. Early works on adaptive sampling proposed gradually increasing BS to balance efficiency and noise reduction (De et al., 2017; Lau et al., 2024b; Lau et al., 2024a; Lau et al., 2025; Ostroukhov et al., 2024). Advanced language models (Brown et al., 2020; Touvron et al., 2023; Liu et al., 2024; Li et al., 2025) employed stage-wise BS schedules, but without systematic analysis. In contrast, our work connects BS scheduling to the evolving local loss landscape, providing a principled way for when to increase the BS. 3 Preliminaries Basic notations. We use bold lowercase letters (e.g., =(xi) x=(x_i)) to denote vectors and bold uppercase letters (e.g., =(aij)A=(a_ij)) to denote matrices. For a matrix A, let ‖2 _2, ‖F _F, and Tr() (A) denote its spectral norm, Frobenius norm and trace, respectively. The Hadamard product is denoted by ⊙ . Theoretical setup. Our theory focuses on the preconditioned stochastic gradient descent (PSGD). We consider a model with parameters ∈ℝp θ ^p and a training set of n examples. Let Li()L_i( θ) be the fitting error evaluated at the i-th example and L()=1n∑i=1nLi()L( θ)= 1n _i=1^nL_i( θ) be the empirical risk. We analyze the preconditioned SGD with a fixed positive-definite22 2 Most practical preconditioners are positive-definite: =IM=I for SGD, diagonal M for AdaGrad (Duchi et al., 2011), RMSProp (Tieleman and Hinton, 2012), Adam, etc. preconditioner ≻0M 0. At iteration k, the update rule gives: k+1=k−η(∇L(k)+k), θ_k+1= θ_k- (∇ L( θ_k)+ ξ_k), (1) where η>0η>0 is the LR and k\ ξ_k\ are i.i.d. random noise vectors with [k]=,[kk⊤]=(k)/B. [ ξ_k]= 0, [ ξ_k ξ_k ]= ( θ_k)/B. (2) Note that (k)=1n∑i=1n∇Li(k)∇Li(k)⊤−∇L(k)∇L(k)⊤ ( θ_k)= 1n _i=1^n∇ L_i( θ_k)∇ L_i( θ_k) -∇ L( θ_k)∇ L( θ_k) is the gradient covariance at k θ_k, and B denotes the BS. During the late phase of training, the model remains close to some global minimum ⋆ θ and the loss can be approximated quadratically: L() L( θ) =L(⋆)+12(−⋆)⊤(⋆)(−⋆),(⋆):=∇2L(⋆)≻0. =L( θ )+ 12( θ- θ ) H( θ )( θ- θ ), ( θ ):=∇^2L( θ ) 0. (3) Similar formulations have been widely used in dynamical stability analyses (Wu et al., 2018; Cohen et al., 2021; Zhou et al., 2025) and theoretical advances on BS scaling (McCandlish et al., 2018). Experimental setup. Our experiments are mainly conducted on LLaMA-2 architecture (Touvron et al., 2023) models with 9393M and 170170M parameters. Training is performed on the FineWeb-Edu dataset (Penedo et al., 2024), with sufficient training budgets ranging from 5050 to 10001000 tokens-per-parameter (TPP)33 3 At least 10× over Chinchilla-optimal tokens (Hoffmann et al., 2022). and a context length of 10241024. We adopt AdamW (Kingma and Ba, 2014) with hyperparameters β1=0.95 _1=0.95, β2=0.95 _2=0.95, and weight decay 0.10.1, together with gradient clipping at 1.01.0 for stability. The evaluation is conducted on a held-out validation split of approximately 5050M tokens. More experiments on larger scales, other architectures, and optimizers are deferred to Appendix C. Our experiments vary the LRs and BSs. In Section 4, we primarily study the role of LR and warmup length, fixing BS at 7.87.8M. In Section 5, we focus on the effect of BS, with LR fixed at 2−102^-10. To decouple BS ramping from LR decay, we adopt a warmup-stable schedule: after linear warmup to the peak value, the LR remains constant (similar to WSD (Hu et al., 2024), but without decay phase). Figure 2: Loss spikes and plateaus in Phase I. We train a series of LLaMA-2 models with 9393M and 170170M parameters. We adopt a warmup-stable schedule, where the warmup length is shortened to 1616 iterations and the peak LR is varied, η∈2−11,2−10,2−9,2−8,2−7η∈\2^-11,2^-10,2^-9,2^-8,2^-7\. (Left). LR schedule: ηt _t vs. training iteration t. (Middle, right). Training loss curves for different model sizes: L(t)L( θ_t) vs. training iteration t. The vertical dashed line marks the end of the warmup phase. 4 Phase I: From Sharp to Flat Landscapes In this section, we provide evidence that, during Phase I, the local landscape of language models evolves from sharp regions toward flatter ones. We first observe that training with large LRs and insufficient warmup often leads to instability and early loss plateaus. Via Lyapunov stability analysis, we attribute these behaviors to the sharp-to-flat dynamics. This finding explains the necessity of LR warmup and suggests that larger peak LRs require proportionally longer warmup. Motivating observations: instability and loss plateaus early in training. The loss curves are typically smooth initially; the model escapes from random initialization and the loss decreases rapidly. Yet, surprisingly, when the warmup length is extremely shortened, we consistently observe loss spikes and plateaus near the end of the warmup phase. To demonstrate this, we train models of different sizes with a fixed warmup length of 1616 iterations while varying the peak LR. As shown in Figure 2, a loss plateau reliably appears around the end of the warmup phase across all settings. Additionally, larger LRs produce higher spikes, which mark a characteristic feature of early training instability. Given these results, two natural questions arise: Q1. Why does shortened warmup induce training instability? Q2. Why do spikes and plateaus occur only at the very beginning of training? To shed light on these questions, we analyze the dynamics of PSGD via Lyapunov stability analysis. Lyapunov stability analysis: sharpness matters. Let k θ_k, ~k θ_k be two nearby trajectories, and define their difference as k:=~k−k e_k:= θ_k- θ_k. When the noise term ξ is set to zero, the evolution of k e_k satisfies: k+1=k−η(∇L(k+k)−∇L(k))=(Linearization)(−η(k))k, e_k+1= e_k- (∇ L( θ_k+ e_k)-∇ L( θ_k)) (Linearization)=(I- ( θ_k)) e_k, (4) The dynamics in Equation 4 describe the local sensitivity of the iteration: if matrix (−η(k))(I- ( θ_k)) repeatedly expands k e_k, small perturbations grow exponentially and the iterates are linearly unstable. Intuitively, the LR η interacts directly with the curvature of the landscape: if η is too large relative to the sharpest direction, the update rule amplifies perturbations and leads to loss spikes. The following lemma formalizes this stability condition for preconditioned GD. Lemma 4.1 (Stability condition for preconditioned GD). Define the preconditioned curvature matrix (k):=1/2(k)1/2S( θ_k):=M^1/2H( θ_k)M^1/2, and let λi=1p\λ\_i=1^p be the eigenvalues of (k)S( θ_k). The linear system in Equation 4 is asymptotically stable (i.e., limk→∞k= _k→∞ e_k= 0) if η satisfies 0<η<2λmax((k),∀k≥00<η< 2 _ max(S( θ_k),∀ k≥ 0. Lemma 4.1 shows that the Lyapunov stability is governed by the largest eigenvalue of S. If the curvature along the sharpest direction is too large, only a sufficiently small LR can prevent divergence. 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−102^-10. (Top). Evolution of the top eigenvalues of the Hessian across iterations: λi((t)) _i(H( θ_t)) vs. iteration t. (Bottom). One-dimensional loss landscape along a random perturbation direction: the perturbed loss L(t+α)L( θ_t+α δ) vs. perturbation coefficient α, shown across early training iterations t. We next characterize the one-step loss change as η approaches the stability boundary 2/λmax(k)2/ _ (S_k). Lemma 4.2 (One-step loss change). Let k:=k+1−k δ_k:= θ_k+1- θ_k. Suppose that along the segment k+αk:α∈[0,1]\ θ_k+α δ_k:α∈[0,1]\, we have 0≤λmin((k+αk))≤λmax((k+αk))≤Λk0≤ _ (S( θ_k+α δ_k))≤ _ (S( θ_k+α δ_k))≤ _k. Then, L(k+1)−L(k)≤−η(1−12ηΛk)(∇L(k))⊤∇L(k). L( θ_k+1)-L( θ_k)≤-η(1- 12η _k)(∇ L( θ_k)) M∇ L( θ_k). In particular, if η↑2/Λkη 2/ _k, the guaranteed decrease per step (L(k)−L(k+1))/η→0 (L( θ_k)-L( θ_k+1))/η→ 0. Lemma 4.2 states that when η is close to 2/Λk2/ _k, each update yields only a marginal decrease in loss. Together with Lemma 4.1, it is clear that training near the stability boundary naturally leads to characteristic loss spikes and plateaus. Importantly, the stability boundary is determined by the sharpness of the loss landscape. To further address Q1-2, we analyze how sharpness evolves during the early phase. The early dynamics: from sharp to flat landscapes. We study how the local landscape geometry, particularly the sharpness, evolves for training runs in Figure 2. Specifically, we track the evolution of the top eigenvalues of the Hessian44 4 Following Cohen et al. (2021), we use the Lanczos algorithm to calculate top eigenvalues of Hessian. (t)H( θ_t). For some early checkpoints t θ_t, we also visualize the one-dimensional loss landscape along a random direction by plotting the function ℒ(α):=L(t+α)L(α):=L( θ_t+α δ) with ∼(0,) δ (0,I). Li et al. (2018) showed that these random-direction visualizations reliably capture intrinsic properties of the loss landscape properties, such as sharpness. To ensure fair comparison across iterations, we fix the same random vector δ for all t θ_t. In Figure 3 (top), the largest eigenvalues of the Hessian (t)H( θ_t) start at high values55 5 In fact, at initialization, sharpness is extremely low but rises sharply after the first update. The sharpness curves reported in Figure 3 therefore start from the first iteration. and then decrease sharply. In Figure 3 (bottom), the loss landscape along a random direction progressively widens as training proceeds, confirming that the model shifts from sharp to flat regions. A Tuning Recipe: Larger Peak LR, Longer Warmup. We have seen that training stability depends on sharpness: when the landscape is shar, only a sufficiently small LR can keep training stable; and pre-training initially traverses from sharp landscapes to flatter ones. Now let us return to Q1 and Q2: A1. If the warmup phase is shortened, the LR rises too quickly while the model is still in sharp regions, leading to loss spikes and plateaus. A2. As training progresses, the landscape becomes flatter and the same LR no longer threatens stability, which explains why instability is confined to the very beginning. Therefore, in practice, we need a sufficiently long warmup phase to keep the LR small until sharpness has decayed, thereby preventing loss spikes and plateaus. This further suggests a practical tuning recipe: the larger the peak LR, the longer the warmup should be, ensuring iterates safely transition into flatter landscapes. To validate this, we train models with varied peak LRs η and warmup lengths TwT_w (in iterations). In Figure 4, within a LR range of 2−82^-8 to 2−112^-11, larger peak LRs require proportionally longer warmup to achieve the optimal validation loss L(best)L( θ_best). However, this proportionality does not hold universally. When η=2−7η=2^-7, the optimal warmup length remains 2102^10 iterations, the same as for η=2−8η=2^-8. Thus, the proportional relationship applies only within a certain range. 5 Phase I: Local Landscape Governed by Noise Scale Figure 4: Larger peak LR, longer warmup. We plot the best validation loss L(bst)L( θ_bst) vs. warmup TwT_w for different η. Optimal TwT_w is highlighted with a star. In this section, we turn to the local landscape geometry in late phase. We observe that BS plays a central role: training with a large BS tends to find a deeper basin of the landscape, whereas a small BS favors a wider basin. Theoretically, we prove that this trade-off between widen or deepen is governed by the noise scale. Building on this, we propose a BS scheduler for the data-limited regime: use small BS early and ramp the BS late, which consumes fewer tokens to achieve the same loss. The effect of BS: local landscapes late in training. We conduct experiments to systematically investigate the role of BS in shaping the local loss landscape during the late phase of training. Specifically, we train models with different BSs for T=20,480T=20,480 iterations. Figure 5 (top left) shows the validation loss curves for each run. Larger BS consistently leads to lower terminal loss and faster convergence in term of iterations66 6 In terms of processed tokens, small BS training converges faster.. We then visualize the loss landscape around the final iterate T θ_T. In Figure 5 (top right), it is clear that small BS produces flatter basins, whereas large BS yields deeper ones. Furthermore, Figure 5 (bottom) compares the landscape evolution of runs with B=0.49B=0.49M and B=7.8B=7.8M, indicating that in the late phase, larger BS tends to deepen the basin, while smaller BS shifts toward wider basins. Despite these results, two key questions remain: Q3. Why is there a trade-off between widening and deepening the basin? Q4. Which factor underlying the hyperparameter BS governs this trade-off? Figure 5: Large BS deepens the basin, small BS widens the basin. We train a series of LLaMA-2 models (170M) for T=20,480T=20,480 iterations, using BSs B∈0.49M,0.98M,1.9M,3.9M,7.8MB∈\0.49M,0.98M,1.9M,3.9M,7.8M\. (Top left). Validation loss curves for different BSs: L(t)L( θ_t) vs. training iteration t. (Top right). One-dimensional loss landscapes at the final iterates T θ_T along a random perturbation direction: perturbed loss L(T+α)L( θ_T+α δ) vs. perturbation coefficient α, visualized across different BSs. (Bottom). One-dimensional loss landscape: the perturbed loss L(t+α)L( θ_t+α δ) vs. perturbation coefficient α, shown across late training iterations t for B=0.49B=0.49M and B=7.8B=7.8M. To delve into Q3-4, we revisit the stochastic differential equation (SDE) in Jastrzębski et al. (2017). Widen or deepen: noise scale governs basin selection. Following Jastrzębski et al. (2017), we take the continuous-time limit of Equation 1. Suppose that the noise covariance satisfies77 7 The assumption of noise covariance is justified in Section B.2. ηB(⋆)⊤=2τ+(η) ηBM ( θ )M =2 +O(η) for some temperature τ>0τ>0. As η→0η→ 0, the scaled discrete process ⌊t/η⌋ θ_ t/η converges weakly to the Itô SDE: dt d θ_t =−∇L(t)dt+2τ1/2dWt =-M∇ L( θ_t)dt+ 2τM^1/2dW_t (5) where WtW_t is standard Brownian motion and the noise scale τ is proportional to η/Bη/B. Building on Equation 5 and the local quadratic model in Equation 3, we establish the trade-off between deepening and widening the loss basin. Theorem 5.1 (Depth-flatness trade-off). Let the empirical risk L()L( θ) admit multiple local minima i⋆i=1m\ θ_i \_i=1^m with Hessians (i⋆)≻0H( θ_i ) 0. Under the SDE in Equation 5 with temperature τ, the stationary probability that training resides in basin i is given by: Pτ(basin i) P_τ(basin i) =exp(−Fi(τ)/τ)∑jexp(−Fj(τ)/τ),Fi(τ):=L(i⋆)+τ2logdet(i⋆). = (-F_i(τ)/τ) _j (-F_j(τ)/τ), F_i(τ):= cyan!20$L( θ_i )$+ green!8$τ$2 orange!20$ ( θ_i )$. Theorem 5.1 states that the basin selection is controlled by the free energy function F(τ)=L(⋆)+τ2logdet(⋆)F(τ)=L( θ )+ τ2 ( θ ). In early training, the loss term L(⋆)L( θ ) dominates, so the model primarily seeks regions of lower loss. In later training, L(⋆)L( θ ) is comparable to the flatness penalty logdet(⋆) ( θ ), and basin selection becomes increasingly sensitive to the noise scale τ∝η/Bτ η/B. Efficient pre-training: a BS scheduler in data-limited regime. Turning back to Q3 and Q4, the trade-off arises because basin selection balances loss minimization against curvature regularization ( A3), with the governing factor being the noise scale τ ( A4). Since the primary objective of pre-training is to minimize the training loss88 8 Note that our analysis focuses on how reduced noise (e.g., via larger BS) helps the optimizer move into deeper minima, conceptually different from the flat-minima perspective in fully interpolating regimes., this balance naturally favors largest BS available (small τ). In practice, however, data availability is limited, and excessively large BS substantially increase data consumption99 9 For example, in Figure 5 (top right), when B=7.8B=7.8M, the run consumes approximately 160160B tokens.. Thus, scheduling BS in pre-training is crucial, particularly in the data-limited regime. • BS scheduler: design principle I. Inspired by our theory, loss reduction dominates early in training, during which large BS yields limited benefit. This suggests the first design principle. Design principle I. Start the training process with a small BS before increasing it later. Related ideas were noted by Li et al. (2025); Merrill et al. (2025), often referred to as BS warmup. However, a key difference in our design lies in when the BS should be increased. Surprisingly, we find that ramping BS later in training yields consistently greater performance. 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.490.49M. For blue curves, BS is ramped up to 4×4× its initial value; for red curves, BS is ramped to 2×2×. The ramping times, T2×T_2× or T4×T_4×, 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 T1T_1 and T2T_2. For the red curves, we fix T2=10T_2=10B and vary T1T_1; for the blue curves, we fix T1=10T_1=10B and vary T2T_2. (Middle). Illustration of BS schedulers. (Right). Final validation loss vs. the relative ramping time, i.e., (T1)/T,(T2−T1)/T∈[0,0.5](T_1)/T,(T_2-T_1)/T∈[0,0.5], where T denotes the total training tokens. • BS scheduler: design principle I. To study this, we train models with different BS schedulers while keeping total training iterations fixed. In Figure 6 (left), all runs begin with an initial BS of 0.490.49M and ramp up to either 4×4× or 2×2× that value at different training iterations. Remarkably, all loss curves eventually collapse onto the same trajectory, regardless of when the BS ramping occurs. Note that when measured at the same training iteration, ramping the BS earlier results in higher data consumption. This indicates that early BS ramping offers no efficiency advantage, achieving the same loss but consuming more data. We next evaluate BS schedules under a fixed token budget. Specifically, we consider a two-stage BS-ramping scheduler characterized by ramp times T1T_1 and T2T_2. To isolate the effect of each stage, we vary either T1T_1 or T2T_2 while keeping the other fixed. See Figure 6 (middle) for an illustration. In Figure 6 (right), a clear trend emerges: BS ramping is most effective when applied late in training (i.e., with T1T_1 and T2T_2 large), whereas ramping too early consistently harms final performance. Together, since BS ramping ultimately leads all runs onto the same trajectory, delaying it allows maximal progress (lower loss) along that trajectory under a data-limited budget. This behavior also aligns with our theory. In the late phase of training, the flatness penalty becomes comparable to the loss term, and BS ramping sharply reduces the noise scale, driving rapid convergence toward deeper minima. This consistency between theory and practice leads to our second principle. Design principle I. Ramp the batch size late in training—when loss reduction becomes marginal. Figure 7: BS scheduling improves data efficiency. We train LLaMA-2 models with 9393M and 170170M parameters, using a BS schedule that starts at 0.490.49M and increases by 4×4× at each ramp. Models are trained with 1,21,2, or 33 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.87.8M and training with a fixed 7.87.8M BS. To further validate our design principle, we train models using a BS schedule that starts at 0.490.49M and ramps by 4×4× whenever loss minimization slows. In Figure 7 (left, middle), models with 1,21,2 or 33 BS ramping steps achieves significant lower validation loss. While additional ramping steps provide diminishing returns, each step still offers a measurable improvement. Moreover, Figure 7 (right) highlights the data-efficiency of the BS scheduling: ramping the BS up to 7.87.8M achieves nearly the same final validation loss as training with a fixed 7.87.8M BS, but requires only about 14 14 of the tokens (i.e., a ∼4× 4× speedup). These results confirm that our BS scheduling design preserves the benefits of large BS while substantially reducing data consumption. Comparison with McCandlish et al. (2018); Merrill et al. (2025). McCandlish et al. (2018) linked BS scaling to the gradient noise and introduced the notion of CBS. Merrill et al. (2025) explored the BS scheduling (BS warmup), doubling BS once the CBS exceeds the current BS. We differs from CBS-related works: without estimating CBS on the fly, we propose to ramp the BS late in training. 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×16× its initial value; for LR decay, the LR drops to 1/161/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/41/4 of its initial value; (b) LR drops to 1/161/16; (c) BS ramps to 16×16×. (d) LR drops to 1/41/4 and BS ramps to 4×4×. In all runs (both left and right), BS starts at 0.490.49M and LR begins at 2−102^-10 (after linear warmup). 6 More Discussions: LR Decay and BS Ramping We have excluded LR decay in our experiments to isolate the effect of BS ramping. Yet, recall that the noise scale τ is proportional to η/Bη/B. Our theory suggests that decaying the LR and ramping the BS both reduce the noise scale, and thus may have similar effects on basin selection. Comparing BS ramping with LR decay. We first train models using either BS ramping or LR decay. Both methods apply a one-time step change: BS ramping multiplies the BS by 1616 at TrampT_ramp, while LR decay divides the LR by 1616 at TdecayT_decay. We align TrampT_ramp and TdecayT_decay so that the changes occur at the same positions. In Figure 8 (left), BS ramping and LR decay produce remarkably similar validation loss curves across all change positions, consistent with the idea that both reduce the noise scale in comparable ways. Interacting BS ramping with LR decay. Furthermore, we study the combined effect of using both BS ramping and LR decay. Specifically, we decay the LR by 4×4× and simultaneously ramp the BS by 4×4× at 10B tokens. We compare this hybrid schedule with three baselines: at the same point, we (a) drops the LR by 4×4×, (b) drops the LR by 16×16× and (c) ramps the BS by 16×16×. We denote the hybrid schedule by (d). In Figure 8 (right), three of the schedules (b c d) produce nearly identical loss curves. Crucially, these three configurations yield the same noise scale η/Bη/B. In contrast, schedule (a) results in a noticeably different trajectory. In summary, our results confirm our theoretical prediction: training dynamics in the late phase are governed by the noise scale η/Bη/B. LR decay reduce the noise scale in the same manner as BS ramping, and any scheduler that preserves η/Bη/B will exhibit nearly identical behavior. 7 Conclusion and Limitations In conclusion, we present a study of how local landscape geometry evolves during language model pre-training. Our analysis reveals two phases: an early sharp-to-flat transition and a late noise-governed regime. Phase I explains the necessity of LR warmup, suggesting that larger peak LRs require proportionally longer warmup lengths. Phase I motivates a BS scheduling that starts with small BS and increases the BS late in training. Limitations. The current theory primarily relies on strong assumptions, such as infinite-small LR in SDE. A natural future direction is to generalize the theory to more realistic settings. Additionally, the current theory cannot fully explain the collapse of loss curves under different BS schedules. Understanding the learning dynamics under different BS schedules remains an open question. References Bergsma et al. (2025) S. Bergsma, N. S. Dey, G. Gosal, G. Gray, D. Soboleva, and J. Hestness Straight to zero: why linearly decaying the learning rate to zero works best for LLMs. In The Thirteenth International Conference on Learning Representations, Cited by: §2. Brown et al. (2020) T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei Language models are few-shot learners. Advances in neural information processing systems 33, p. 1877–1901. Cited by: §2. Chen et al. (2025) H. Chen, Y. Dong, Z. Wei, Y. Huang, Y. Zhang, H. Su, and J. Zhu Understanding pre-training and fine-tuning from loss landscape perspectives. arXiv preprint arXiv:2505.17646. Cited by: §1. Chen et al. (2024) X. Chen, C. Liang, D. Huang, E. Real, K. Wang, H. Pham, X. Dong, T. Luong, C. Hsieh, Y. Lu, et al. Symbolic discovery of optimization algorithms. Advances in Neural Information Processing Systems 36. Cited by: 4th item. Cohen et al. (2025) J. Cohen, A. Damian, A. Talwalkar, J. Z. Kolter, and J. D. Lee Understanding optimization in deep learning with central flows. In The Thirteenth International Conference on Learning Representations, Cited by: §1, §2. Cohen et al. (2022) J. M. Cohen, B. Ghorbani, S. Krishnan, N. Agarwal, S. Medapati, M. Badura, D. Suo, D. Cardoze, Z. Nado, G. E. Dahl, et al. Adaptive gradient methods at the edge of stability. arXiv preprint arXiv:2207.14484. Cited by: §2. Cohen et al. (2021) J. M. Cohen, S. Kaur, Y. Li, J. Z. Kolter, and A. Talwalkar Gradient descent on neural networks typically occurs at the edge of stability. International Conference on Learning Representations. Cited by: §1, §1, §3, footnote 4. De et al. (2017) S. De, A. Yadav, D. Jacobs, and T. Goldstein Automated inference with adaptive batches. In Artificial Intelligence and Statistics, p. 1504–1513. Cited by: §2. Duchi et al. (2011) J. Duchi, E. Hazan, and Y. Singer Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research 12 (61), p. 2121–2159. Cited by: footnote 2. Foret et al. (2021) P. Foret, A. Kleiner, H. Mobahi, and B. Neyshabur Sharpness-aware minimization for efficiently improving generalization. In International Conference on Learning Representations, Cited by: §1. Gilmer et al. (2022) J. Gilmer, B. Ghorbani, A. Garg, S. Kudugunta, B. Neyshabur, D. Cardoze, G. E. Dahl, Z. Nado, and O. Firat A loss curvature perspective on training instabilities of deep learning models. In International Conference on Learning Representations, Cited by: §1. Gotmare et al. (2019) A. Gotmare, N. S. Keskar, C. Xiong, and R. Socher A closer look at deep learning heuristics: learning rate restarts, warmup and distillation. In International Conference on Learning Representations, Cited by: §2. Goyal et al. (2017) P. Goyal, P. Dollár, R. Girshick, P. Noordhuis, L. Wesolowski, A. Kyrola, A. Tulloch, Y. Jia, and K. He Accurate, large minibatch sgd: training imagenet in 1 hour. arXiv preprint arXiv:1706.02677. Cited by: §2. Gray et al. (2024) G. Gray, S. Bergsma, J. Hestness, et al. Normalization layer per-example gradients are sufficient to predict gradient noise scale in transformers. Advances in Neural Information Processing Systems 37, p. 93510–93539. Cited by: §2. Gray et al. (2023) G. Gray, A. Samar, and J. Hestness Efficient and approximate per-example gradient norms for gradient noise scale. In Workshop on Advancing Neural Network Training: Computational Efficiency, Scalability, and Resource Optimization (WANT@ NeurIPS 2023), Cited by: §2. He et al. (2016) K. He, X. Zhang, S. Ren, and J. Sun Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, p. 770–778. Cited by: §2. Hoffmann et al. (2022) J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556. Cited by: footnote 3. Hu et al. (2024) S. Hu, Y. Tu, X. Han, C. He, G. Cui, X. Long, Z. Zheng, Y. Fang, Y. Huang, W. Zhao, et al. Minicpm: unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.06395. Cited by: §1, §2, §3. Jastrzebski et al. (2020) S. Jastrzebski, M. Szymczak, S. Fort, D. Arpit, J. Tabor, K. Cho*, and K. Geras* The break-even point on optimization trajectories of deep neural networks. In International Conference on Learning Representations, Cited by: §2. Jastrzębski et al. (2017) S. Jastrzębski, Z. Kenton, D. Arpit, N. Ballas, A. Fischer, Y. Bengio, and A. Storkey Three factors influencing minima in sgd. arXiv preprint arXiv:1711.04623. Cited by: §5, §5. Jastrzębski et al. (2019) S. Jastrzębski, Z. Kenton, N. Ballas, A. Fischer, Y. Bengio, and A. Storkey On the relation between the sharpest directions of DNN loss and the SGD step length. In International Conference on Learning Representations, Cited by: §2. Jiang et al. (2020) Y. Jiang, B. Neyshabur, H. Mobahi, D. Krishnan, and S. Bengio Fantastic generalization measures and where to find them. In International Conference on Learning Representations, Cited by: §1. Kaplan et al. (2020) J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei Scaling laws for neural language models. arXiv preprint arXiv:2001.08361. Cited by: §2. Karpathy (2022) A. Karpathy NanoGPT. GitHub. Note: https://github.com/karpathy/nanoGPT Cited by: 1st item. Keller et al. (2024) J. Keller et al. Muon optimizer. Note: https://github.com/KellerJordan/Muon?tab=readme-ov-file Cited by: 2nd item. Keskar et al. (2017) N. S. Keskar, D. Mudigere, J. Nocedal, M. Smelyanskiy, and P. T. P. Tang On large-batch training for deep learning: generalization gap and sharp minima. In International Conference on Learning Representations, Cited by: §1. Kingma and Ba (2014) D. P. Kingma and J. Ba Adam: a method for stochastic optimization. arXiv preprint arXiv:1412.6980. Cited by: 1st item, §3. Kosson et al. (2024) A. Kosson, B. Messmer, and M. Jaggi Analyzing & reducing the need for learning rate warmup in GPT training. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, Cited by: §2. Lau et al. (2024a) T. T. Lau, W. Li, C. Xu, H. Liu, and M. Kolar Communication-efficient adaptive batch size strategies for distributed local gradient methods. arXiv preprint arXiv:2406.13936. Cited by: §2. Lau et al. (2025) T. T. Lau, W. Li, C. Xu, H. Liu, and M. Kolar Adaptive batch size schedules for distributed training of language models with data and model parallelism. In Proceedings of Conference on Parsimony and Learning, Cited by: §2. Lau et al. (2024b) T. T. Lau, H. Liu, and M. Kolar AdAdaGrad: adaptive batch size schemes for adaptive gradient methods. arXiv preprint arXiv:2402.11215. Cited by: §2. Li et al. (2025) A. Li, B. Gong, B. Yang, B. Shan, C. Liu, C. Zhu, C. Zhang, C. Guo, D. Chen, D. Li, et al. Minimax-01: scaling foundation models with lightning attention. arXiv preprint arXiv:2501.08313. Cited by: §2, §5. Li et al. (2018) H. Li, Z. Xu, G. Taylor, C. Studer, and T. Goldstein Visualizing the loss landscape of neural nets. In Advances in Neural Information Processing Systems, S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (Eds.), Vol. 31, p. . Cited by: §4. Liu et al. (2024) A. Liu, B. Feng, B. Xue, B. Wang, B. Wu, C. Lu, C. Zhao, C. Deng, C. Zhang, C. Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437. Cited by: §2. McCandlish et al. (2018) S. McCandlish, J. Kaplan, D. Amodei, and O. D. Team An empirical model of large-batch training. arXiv preprint arXiv:1812.06162. Cited by: §2, §3, §5. Merrill et al. (2025) W. Merrill, S. Arora, D. Groeneveld, and H. Hajishirzi Critical batch size revisited: a simple empirical approach to large-batch language model training. arXiv preprint arXiv:2505.23971. Cited by: §5, §5. Ostroukhov et al. (2024) P. Ostroukhov, A. Zhumabayeva, C. Xiang, A. Gasnikov, M. Takáč, and D. Kamzolov AdaBatchGrad: combining adaptive batch size and adaptive step size. arXiv preprint arXiv:2402.05264. Cited by: §2. Penedo et al. (2024) G. Penedo, H. Kydlíček, L. B. allal, A. Lozhkov, M. Mitchell, C. Raffel, L. V. Werra, and T. Wolf The fineweb datasets: decanting the web for the finest text data at scale. In The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track, Cited by: §C.1, §3. Peng et al. (2024) S. Peng, P. Chen, M. D. Hull, and D. H. Chau Navigating the safety landscape: measuring risks in finetuning large language models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, Cited by: §1. Radford et al. (2019) A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever Language models are unsupervised multitask learners. OpenAI blog 1 (8), p. 9. Cited by: 1st item. Shoeybi et al. (2019) M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro Megatron-lm: training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053. Cited by: §2. Song and Yun (2023) M. Song and C. Yun Trajectory alignment: understanding the edge of stability phenomenon via bifurcation theory. arXiv preprint arXiv:2307.04204. Cited by: §1, §2. Su et al. (2024) J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Y. Liu Roformer: enhanced transformer with rotary position embedding. Neurocomputing 568, p. 127063. Cited by: 2nd item. Tieleman and Hinton (2012) T. Tieleman and G. Hinton Rmsprop: divide the gradient by a running average of its recent magnitude. coursera: neural networks for machine learning. COURSERA Neural Networks Mach. Learn 17, p. 6. Cited by: footnote 2. Touvron et al. (2023) H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. Llama 2: open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288. Cited by: 2nd item, §2, §3. Vaswani et al. (2017) A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin Attention is all you need. In Advances in neural information processing systems, Vol. 30. Cited by: §2. Wang et al. (2025) J. Wang, M. Wang, Z. Zhou, J. Yan, W. E, and L. Wu The sharpness disparity principle in transformers for accelerating language model pre-training. In Forty-second International Conference on Machine Learning, Cited by: §1. Wen et al. (2024) K. Wen, Z. Li, J. Wang, D. Hall, P. Liang, and T. Ma Understanding warmup-stable-decay learning rates: a river valley loss landscape perspective. arXiv preprint arXiv:2410.05192. Cited by: §1. Wolf et al. (2020) T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, J. Davison, S. Shleifer, P. von Platen, C. Ma, Y. Jernite, J. Plu, C. Xu, T. L. Scao, S. Gugger, M. Drame, Q. Lhoest, and A. M. Rush Transformers: state-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, Online, p. 38–45. Cited by: 2nd item. Wu et al. (2018) L. Wu, C. Ma, and W. E How SGD selects the global minima in over-parameterized learning: a dynamical stability perspective. Advances in Neural Information Processing Systems 31, p. 8279–8288. Cited by: §2, §3. Wu et al. (2022) L. Wu, M. Wang, and W. J. Su The alignment property of SGD noise and how it helps select flat minima: a stability analysis. In Advances in Neural Information Processing Systems, A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho (Eds.), Cited by: §2. Xie et al. (2021) Z. Xie, I. Sato, and M. Sugiyama A diffusion theory for deep learning dynamics: stochastic gradient descent exponentially favors flat minima. In International Conference on Learning Representations, Cited by: §2. Zhang et al. (2017) C. Zhang, S. Bengio, M. Hardt, B. Recht, and O. Vinyals Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations, Cited by: §1. Zhang et al. (2025) H. Zhang, D. Morwani, N. Vyas, J. Wu, D. Zou, U. Ghai, D. Foster, and S. Kakade How does critical batch size scale in pre-training?. International Conference on Learning Representations. Cited by: §2. Zhang et al. (2022) S. Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin, et al. Opt: open pre-trained transformer language models. arXiv preprint arXiv:2205.01068. Cited by: §2. Zhang et al. (2024a) Y. Zhang, C. Chen, T. Ding, Z. Li, R. Sun, and Z. Luo Why transformers need adam: a hessian perspective. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, Cited by: §1. Zhang et al. (2024b) Y. Zhang, C. Chen, Z. Li, T. Ding, C. Wu, Y. Ye, Z. Luo, and R. Sun Adam-mini: use fewer learning rates to gain more. arXiv preprint arXiv:2406.16793. Cited by: 3rd item. Zhou et al. (2025) Z. Zhou, M. Wang, Y. Mao, B. Li, and J. Yan Sharpness-aware minimization efficiently selects flatter minima late in training. In The Thirteenth International Conference on Learning Representations, Cited by: §2, §3. Zhu et al. (2019) Z. Zhu, J. Wu, B. Yu, L. Wu, and J. Ma The anisotropic noise in stochastic gradient descent: its behavior of escaping from sharp minima and regularization effects. In Proceedings of the 36th International Conference on Machine Learning, K. Chaudhuri and R. Salakhutdinov (Eds.), Proceedings of Machine Learning Research, Vol. 97, p. 7654–7663. Cited by: §2. Appendix A Terminologies 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 λmax((t)) _ (H( θ_t)) or of the preconditioned curvature matrix λmax((t)) _ (S( θ_t)). Flat/sharp minimum A minimum is a point where the gradient vanishes ∇L()=0∇ L( θ)=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. Table 1: Terminology and usage in this paper. Appendix B Missing Proof B.1 Phase I: Lyapunov Stability Analysis Lemma B.1 (Stability Condition for Preconditioned GD). Define the preconditioned curvature matrix (k):=1/2(k)1/2S( θ_k):=M^1/2H( θ_k)M^1/2, and let λi=1p\λ\_i=1^p be the eigenvalues of (k)S( θ_k). The linear system in Equation 4 is asymptotically stable (i.e., limk→∞k= _k→∞ e_k= 0) if η satisfies 0<η<2λmax((k)CLOSE,∀k≥0.0<η< 2 _ max(S( θ_k),∀ k≥ 0. (6) Proof. Since k+1=(−ηk)k e_k+1=(I- _k) e_k, the linear system is asymptotically stable if all eigenvalues of −ηkI- _k have magnitude less than 1. Note that: −ηk=1/2(−ηk)−1/2,I- _k=M^1/2(I- _k)M^-1/2, (7) so the eigenvalues are 1−ηλj(k)1-η _j(S_k). The stability condition |1−ηλj|<1|1-η _j|<1 for all j is equivalent to: 0<η<2λmax(k).0<η< 2 _ (S_k). (8) ∎ Lemma B.2 (Exact one-step loss change). Define: () ( θ) :=1/2()1/2, :=M^1/2H( θ)M^1/2, k g_k :=1/2∇L(k), :=M^1/2∇ L( θ_k), k δ_k :=k+1−k=−η∇L(k). := θ_k+1- θ_k=- ∇ L( θ_k). Then the true loss change can be written exactly as L(k+1)−L(k)=−η‖∇L(k)‖2+η2∫01(1−t)k⊤(k+tk)kt,L( θ_k+1)-L( θ_k)=-η\|∇ L( θ_k)\|_M^2+η^2 _0^1(1-t)\, g_k S( θ_k+t δ_k) g_k\,dt, (9) where ∥∇L(k)∥2:=(∇L(k))⊤∇L(k)\|∇ L( θ_k)\|_M^2:=(∇ L( θ_k)) M∇ L( θ_k). Proof. Let k:=k+1−k=−η∇L(k) δ_k:= θ_k+1- θ_k=- ∇ L( θ_k) and define the scalar function ϕ(t):=L(k+tk),t∈[0,1].φ(t):=L( θ_k+t δ_k), t∈[0,1]. Then L(k+1)−L(k)=ϕ(1)−ϕ(0).L( θ_k+1)-L( θ_k)=φ(1)-φ(0). Compute the derivatives: ϕ′(t) φ (t) =(∇L(k+tk))⊤k, =(∇ L( θ_k+t δ_k)) δ_k, ϕ′(t) φ (t) =k⊤(k+tk)k. = δ_k H( θ_k+t δ_k) δ_k. By Taylor’s theorem with integral remainder: ϕ(1)−ϕ(0)=ϕ′(0)+∫01(1−t)ϕ′(t)t.φ(1)-φ(0)=φ (0)+ _0^1(1-t)φ (t)\,dt. Now evaluate at t=0t=0: ϕ′(0)=(∇L(k))⊤k=−η(∇L(k))⊤∇L(k)=−η∥∇L(k)∥2.φ (0)=(∇ L( θ_k)) δ_k=-η(∇ L( θ_k)) M∇ L( θ_k)=-η\|∇ L( θ_k)\|_M^2. For the second derivative term: ϕ′(t)=k⊤(k+tk)k=η2k⊤(k+tk)k,φ (t)= δ_k H( θ_k+t δ_k) δ_k=η^2 g_k S( θ_k+t δ_k) g_k, since k=−η∇L(k) δ_k=- ∇ L( θ_k) and k=1/2∇L(k) g_k=M^1/2∇ L( θ_k), and thus k⊤(⋅)k=η2k⊤(⋅)k. δ_k H(·) δ_k=η^2 g_k S(·) g_k. Substituting both terms yields the result. ∎ Lemma B.3 (One-step Loss Change). Let k:=k+1−k δ_k:= θ_k+1- θ_k. Suppose that along the segment k+αk:α∈[0,1]\ θ_k+α δ_k:α∈[0,1]\, we have 0≤λmin((k+αk))≤λmax((k+αk))≤Λk0≤ _ (S( θ_k+α δ_k))≤ _ (S( θ_k+α δ_k))≤ _k. Then, L(k+1)−L(k)≤−η(1−12ηΛk)(∇L(k))⊤∇L(k). L( θ_k+1)-L( θ_k)≤-η(1- 12η _k)(∇ L( θ_k)) M∇ L( θ_k). In particular, if η≤2/Λkη≤ 2/ _k, each update is guaranteed to non-increasing in loss, i.e., L(k+1)≤L(k)L( θ_k+1)≤ L( θ_k). Instead, if η↑2/Λkη 2/ _k, the guaranteed decrease per step (L(k)−L(k+1))/η→0 (L( θ_k)-L( θ_k+1))/η→ 0. Proof. From Lemma B.2, we have k⊤(k+tk)k≤Λk‖k‖2 g_k S( θ_k+t δ_k) g_k≤ _k\| g_k\|^2 for all t, since (⋅)S(·) is symmetric. Therefore, L(k+1)−L(k) L( θ_k+1)-L( θ_k) ≤−η‖∇L(k)‖2+η2Λk‖k‖2∫01(1−t)t ≤-η\|∇ L( θ_k)\|_M^2+η^2 _k\| g_k\|^2 _0^1(1-t)\,dt =−η‖∇L(k)‖2+12η2Λk‖k‖2. =-η\|∇ L( θ_k)\|_M^2+ 12η^2 _k\| g_k\|^2. Note that ‖k‖2=‖∇L(k)‖2\| g_k\|^2=\|∇ L( θ_k)\|_M^2, yielding the result. ∎ B.2 Phase I: SDE Analysis B.2.1 Discrete-Time Solution Lemma B.4 (Eigenbasis Decomposition). Let :=1/2(⋆)1/2S:=M^1/2H( θ )M^1/2 with eigendecomposition =⊤S=Q Q , =diag(λ1,…,λd) =diag( _1,…, _d). Define :=⊤1/2(⋆)1/2/BG:=Q M^1/2 ( θ )M^1/2Q/B. In coordinates k:=⊤−1/2k w_k:=Q M^-1/2 e_k, the recursion gives: k+1 w_k+1 =(−η)k+ηk,[kk⊤]= =(I-η ) w_k+η ζ_k, [ ζ_k ζ_k ]=G The stationary covariance w _w has diagonal elements: (w)jj ( _w)_j =η2jj1−(1−ηλj)2=ηjj2λj−ηλj2 = η^2G_j1-(1-η _j)^2= _j2 _j-η _j^2 (10) Proof. First, we verify that =1/2(⋆)1/2S=M^1/2H( θ )M^1/2 can be eigendecomposed. Since both M and (⋆)H( θ ) are positive definite matrices, S is also positive definite matrix. By the spectral theorem, S admits the eigendecomposition =⊤S=Q Q , where Q is orthogonal and =diag(λ1,…,λd) =diag( _1,…, _d) with λi>0 _i>0. Starting from k+1=k+ηk e_k+1=A e_k+ ξ_k with =−η(⋆)A=I- ( θ ), we change variables to k=⊤−1/2k w_k=Q M^-1/2 e_k. k+1 w_k+1 =⊤−1/2k+1=⊤−1/2(k+ηk) =Q M^-1/2 e_k+1=Q M^-1/2(A e_k+ ξ_k) =⊤−1/2(−η(⋆))k+η⊤1/2k =Q M^-1/2(I- ( θ )) e_k+ M^1/2 ξ_k =⊤−1/2k−η⊤1/2(⋆)k+η⊤1/2k =Q M^-1/2 e_k- M^1/2H( θ ) e_k+ M^1/2 ξ_k =k−η⊤1/2(⋆)1/2k+η⊤1/2k = w_k- M^1/2H( θ )M^1/2Q w_k+ M^1/2 ξ_k =k−η⊤k+η⊤1/2k = w_k- SQ w_k+ M^1/2 ξ_k =(−η)k+η⊤1/2k =(I-η ) w_k+ M^1/2 ξ_k Defining k:=⊤1/2k ζ_k:=Q M^1/2 ξ_k, we get: [kk⊤] [ ζ_k ζ_k ] =⊤1/2[kk⊤]1/2=⊤1/2(⋆)1/2/B=: =Q M^1/2E[ ξ_k ξ_k ]M^1/2Q=Q M^1/2 ( θ )M^1/2Q/B=:G As matrix (−η)(I-η ) is diagonal, the recursion now decouples into independent scalar equations for each component j: (k+1)j ( w_k+1)_j =(1−ηλj)(k)j+η(k)j. =(1-η _j)( w_k)_j+η( ζ_k)_j. For each component j, the stationary variance satisfies: (w)jj ( _w)_j =(1−ηλj)2(w)jj+η2jj =(1-η _j)^2( _w)_j+η^2G_j (11) Solving for (w)jj( _w)_j: (w)jj ( _w)_j =η2jj1−(1−ηλj)2=η2jj1−(1−2ηλj+η2λj2) = η^2G_j1-(1-η _j)^2= η^2G_j1-(1-2η _j+η^2 _j^2) =η2jj2ηλj−η2λj2=ηjj2λj−ηλj2 = η^2G_j2η _j-η^2 _j^2= _j2 _j-η _j^2 ∎ B.2.2 Continuous-Time Limit We now take the continuous-time limit (η→0η→ 0) to derive a simpler universal theory. The exact solution for the variance in the eigenbasis from Lemma B.4, i.e., (w)jj=ηjj/(2λj−ηλj2)( _w)_j= _j/(2 _j-η _j^2), guides the necessary scaling for the continuous-time limit. Because (w)jj( _w)_j converges to a finite non-zero value as η→0η→ 0, the numerator ηjj _j must remain finite. This suggests defining a quantity τ such that for each mode j: ηjj→2τasη→0.η\,G_j→ 2τ η→ 0. We strengthen this to : η→2τasη→0.η\,G→ 2 η→ 0. Recalling that =⊤1/2(⋆)1/2/BG=Q M^1/2 ( θ )M^1/2Q/B, this condition in the original coordinate system translates to the required scaling for the noise covariance: ηB(⋆)⊤→2τ. ηBM ( θ )M → 2 . Proposition B.1 (Convergence to SDE). Consider the scaled discrete process ⌊t/η⌋ θ_ t/η as η→0η→ 0. Suppose the noise covariance satisfies ηB(⋆)⊤ ηBM ( θ )M =2τ+O(η), =2 +O(η), (12) for some temperature τ>0τ>0. Then the process converges weakly to the Itô SDE: dt d θ_t =−∇L(t)dt+2τ1/2dWt =-M∇ L( θ_t)dt+ 2τM^1/2dW_t (13) where WtW_t is standard Brownian motion. Proof. Consider the discrete preconditioned SGD update: k+1=k−η(∇L(k)+k), θ_k+1= θ_k- (∇ L( θ_k)+ ξ_k), Define the scaled process (η)(t)=⌊t/η⌋ θ^(η)(t)= θ_ t/η . The increment Δk=k+1−k θ_k= θ_k+1- θ_k satisfies: [Δk∣k=]=−η∇L(),E[ θ_k θ_k= θ]=- ∇ L( θ), Cov(Δk∣k=)=η2B(⋆)⊤.Cov( θ_k θ_k= θ)= η^2BM ( θ )M . Given the scaling condition Equation 12, the covariance is O(η)O(η). The generator ℒ(η)L^(η) of the discrete process for a smooth function f is: ℒ(η)()=1η[(k+1)−(k)∣k=].L^(η) f( θ)= 1ηE[ f( θ_k+1)- f( θ_k) θ_k= θ]. Using a Taylor expansion and taking conditional expectation: [(k+1)−(k)∣]=−η∇()⊤∇L()+12[(Δ)⊤∇2()Δ]+O(η3/2).E[ f( θ_k+1)- f( θ_k) θ]=-η∇ f( θ) M∇ L( θ)+ 12E[( θ) ∇^2 f( θ) θ]+O(η^3/2). For the second term, with Δ=−ηM(∇L()+k) θ=-η M(∇ L( θ)+ ξ_k): [(Δ)⊤∇2()Δ] [( θ) ∇^2 f( θ) θ] =η2[(∇L()+k)⊤⊤∇2()(∇L()+k)] =η^2E[(∇ L( θ)+ ξ_k) M ∇^2 f( θ)M(∇ L( θ)+ ξ_k)] =η2[k⊤⊤∇2()k]+O(η2) =η^2E[ ξ_k M ∇^2 f( θ)M ξ_k]+O(η^2) =η2Tr(⊤∇2()[kk⊤])+O(η2) =η^2Tr(M ∇^2 f( θ)ME[ ξ_k ξ_k ])+O(η^2) =η2BTr(⊤∇2()(⋆))+O(η2) = η^2BTr(M ∇^2 f( θ)M ( θ ))+O(η^2) =η2BTr((⋆)⊤∇2())+O(η2) = η^2BTr(M ( θ )M ∇^2 f( θ))+O(η^2) where we used [⊤]=Tr(A[⊤])E[ ξ A ξ]=Tr(AE[ ξ ξ ]) and trace cyclicity Tr()=Tr()Tr(ABC)=Tr(CAB). Therefore: 12[(Δ)⊤∇2()Δ]=η22BTr((⋆)⊤∇2())+O(η2). 12E[( θ) ∇^2 f( θ) θ]= η^22BTr(M ( θ )M ∇^2 f( θ))+O(η^2). Using the scaling condition Equation 12, we have : η22BTr((⋆)⊤∇2())=η2Tr(2τ∇2())+O(η2)=ητTr(∇2())+O(η2). η^22BTr(M ( θ )M ∇^2 f( θ))= η2Tr (2 ∇^2 f( θ) )+O(η^2)=η (M∇^2 f( θ))+O(η^2). Thus, ℒ(η)()=−∇()⊤∇L()+τTr(∇2())+O(η).L^(η) f( θ)=-∇ f( θ) M∇ L( θ)+ (M∇^2 f( θ))+O(η). As η→0η→ 0, ℒ(η)()L^(η) f( θ) converges to: ℒ()=−∇()⊤∇L()+τTr(∇2()),L f( θ)=-∇ f( θ) M∇ L( θ)+ (M∇^2 f( θ)), which is the generator of the Itô SDE: dt=−∇L(t)dt+2τ1/2dWt.d θ_t=-M∇ L( θ_t)dt+ 2τM^1/2dW_t. By the weak convergence theory (e.g., via the martingale problem or generator convergence), the process (η)(t) θ^(η)(t) converges weakly to the solution of this SDE. ∎ Proposition B.2 (Gibbs Stationary Distribution). The SDE in Equation 13 has stationary distribution: p∞() p_∞( θ) ∝exp(−L()/τ) (-L( θ)/τ) (14) Proof. The generator of the SDE (5) is ℒf=−∇L⋅∇f+τtr(∇2f)Lf=-M∇ L·∇ f+ (M∇^2f). The Fokker-Planck equation for the probability density p(t,)p(t, θ) is: ∂tp _tp =ℒ∗p=∇⋅(∇Lp)+τ∇⋅(∇p) =L^*p=∇·(M∇ L\,p)+τ∇·(M∇ p) where ℒ∗L^* is the adjoint operator. Setting ∂tp=0 _tp=0 for stationarity: 0 0 =∇⋅(∇Lp∞)+τ∇⋅(∇p∞) =∇·(M∇ L\,p_∞)+τ∇·(M∇ p_∞) =∇⋅(∇Lp∞+τ∇p∞) =∇·(M∇ L\,p_∞+ ∇ p_∞) This implies the current J=∇Lp∞+τ∇p∞J=M∇ L\,p_∞+ ∇ p_∞ has zero divergence. For a potential-driven system, we require J=J=0: ∇Lp∞+τ∇p∞ ∇ L\,p_∞+ ∇ p_∞ = =0 ∇Lp∞+τ∇p∞ ∇ L\,p_∞+τ∇ p_∞ =(since ≻0) =0 (since $ M 0$) ∇p∞p∞ ∇ p_∞p_∞ =−∇Lτ =- ∇ Lτ Integrating: logp∞=−L/τ+const p_∞=-L/τ+const, which gives Equation 14. ∎ Theorem B.1 (Free Energy Minimization). Let the empirical risk L()L( θ) admit multiple local minima i⋆i=1m\ θ_i \_i=1^m with Hessians (i⋆)≻0H( θ_i ) 0. Under the SDE in Equation 13 with temperature τ, the stationary probability that training resides in basin i is given by: Pτ(basin i) P_τ(basin i) =exp(−Fi(τ)/τ)∑jexp(−Fj(τ)/τ),Fi(τ):=L(i⋆)+τ2logdet(i⋆). = (-F_i(τ)/τ) _j (-F_j(τ)/τ), F_i(τ):=L( θ_i )+ τ2 ( θ_i ). (15) Proof. From the Gibbs distribution Equation 14, the probability mass in basin i is: Pτ(basin i) P_τ(basin i) =∫Bie−L()/τd∫ℝde−L()/τd = _B_ie^-L( θ)/τd θ _R^de^-L( θ)/τd θ where BiB_i is the basin of attraction around minimum i⋆ θ_i . For the numerator, using the quadratic approximation L()=L(i⋆)+12(−i⋆)⊤(i⋆)(−i⋆)L( θ)=L( θ_i )+ 12( θ- θ_i ) H( θ_i )( θ- θ_i ) in basin i we get: ∫Bie−L()/τd _B_ie^-L( θ)/τd θ =∫ℝdexp(−L(i⋆)τ−12τ(−i⋆)⊤(i⋆)(−i⋆)) = _R^d (- L( θ_i )τ- 12τ( θ- θ_i ) H( θ_i )( θ- θ_i ) )d θ =e−L(i⋆)/τ∫ℝdexp(−12τ(−i⋆)⊤(i⋆)(−i⋆))d =e^-L( θ_i )/τ _R^d (- 12τ( θ- θ_i ) H( θ_i )( θ- θ_i ) )d θ The integral is a multivariate Gaussian with covariance τ(i⋆)−1 ( θ_i )^-1. Using the standard formula for Gaussian integrals: ∫ℝdexp(−12y⊤−1y)y _R^d (- 12y ^-1y )dy =(2π)d/2(det)1/2 =(2π)^d/2( )^1/2 With =τ(i⋆)−1 = ( θ_i )^-1, we have det=τd(det(i⋆))−1 =τ^d( ( θ_i ))^-1 and −1=τ−1(i⋆) ^-1=τ^-1H( θ_i ): ∫ℝdexp(−12τ(−i⋆)⊤(i⋆)(−i⋆)) _R^d (- 12τ( θ- θ_i ) H( θ_i )( θ- θ_i ) )d θ =(2π)d/2(τd(det(i⋆))−1)1/2 =(2π)^d/2(τ^d( ( θ_i ))^-1)^1/2 =(2πτ)d/2(det(i⋆))−1/2 =(2πτ)^d/2( ( θ_i ))^-1/2 Therefore: ∫Bie−L()/τd _B_ie^-L( θ)/τd θ =e−L(i⋆)/τ(2πτ)d/2(det(i⋆))−1/2 =e^-L( θ_i )/τ(2πτ)^d/2( ( θ_i ))^-1/2 =(2πτ)d/2exp(−L(i⋆)/τ−12logdet(i⋆)) =(2πτ)^d/2 (-L( θ_i )/τ- 12 ( θ_i ) ) =(2πτ)d/2exp(−1τ(L(i⋆)+τ2logdet(i⋆))) =(2πτ)^d/2 (- 1τ (L( θ_i )+ τ2 ( θ_i ) ) ) =(2πτ)d/2exp(−Fi(τ)/τ) =(2πτ)^d/2 (-F_i(τ)/τ) Similarly, the total partition function is: Z(τ) Z(τ) =∫ℝde−L()/τd=d∑j=1m∫Bje−L()/τd = _R^de^-L( θ)/τd θ=d _j=1^m _B_je^-L( θ)/τd θ =(2πτ)d/2∑j=1mexp(−Fj(τ)/τ) =(2πτ)^d/2 _j=1^m (-F_j(τ)/τ) Therefore: Pτ(basin i) P_τ(basin i) =(2πτ)d/2exp(−Fi(τ)/τ)(2πτ)d/2∑jexp(−Fj(τ)/τ)=exp(−Fi(τ)/τ)∑jexp(−Fj(τ)/τ) = (2πτ)^d/2 (-F_i(τ)/τ)(2πτ)^d/2 _j (-F_j(τ)/τ)= (-F_i(τ)/τ) _j (-F_j(τ)/τ) This completes the proof of the free energy formula Equation 15. ∎ Appendix C Experimental Setups and More Results C.1 Experimental Setups Models. We utilize two popular classes of LLM models for our pre-training experiments: • GPT-2. We use GPT-2 (small) model (Radford et al., 2019), implemented via the nanoGPT code base (Karpathy, 2022). Following nanoGPT, the model employs Gaussian Error Linear Unit (GELU) activations and standard Layer Normalization (LayerNorm). Detailed model configurations are provided in Table 2. • LLaMA. LLaMA (Touvron et al., 2023) is another popular decoder-only Transformer architecture, incorporating Rotary Positional Encoding (RoPE) (Su et al., 2024), Swish-Gated Linear Unit (SwiGLU), and Root mean square layer normalization (RMSNorm). For implementation, we utilize the LLaMA code from HuggingFace Transformers Library (Wolf et al., 2020). Additional model configurations are detailed in Table 2. Datasets. Training is performed on the FineWeb-Edu dataset (Penedo et al., 2024). We adopt the a subset randomly sampled from the whole dataset of around 100B GPT-2 tokens. The same dataset has been widely used in literature on LLM pre-training. Optimizers. To generalize our findings across different optimizers, we choose: • AdamW. AdamW (Kingma and Ba, 2014) is adopted with hyperparameters β1=0.95 _1=0.95, β2=0.95 _2=0.95, and weight decay 0.10.1. • Muon. Muon (Keller and others, 2024) is used with momentum of 0.950.95 and weight decay 0.10.1. • Adam-mini. The hyperparameter of Adam-mini (Zhang et al., 2024b) is the same as AdamW. • Lion. Lion (Chen et al., 2024) is used with hyperparameters β1=0.95 _1=0.95, β2=0.98 _2=0.98. The LR of Lion η is divided by 10×10× compared with the LR of AdamW in the same experiments, and the weight decay λ is ramped up to 10×10× to keep the effective LR λη=0.1λη=0.1. All these optimizers are used with gradient clipping at 1.01.0 for stability. Table 2: Model configurations. Acronym Size dmodeld_model dFFd_F 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 C.2 More Results Under Various Setups. In this section, we extend our findings to other architectures, optimization algorithms, and larger training scales. Due to computational constraints, we primarily focus on validating the sharp-to-flat early dynamics and the proposed BS scheduling principle across these settings. We also explore the warmup–tuning recipe on additional architectures. Figure 9: Extensive to GPT-2 architectures. (Left). Early sharp-to-flat dynamics. Evolution of the top eigenvalues of the Hessian across iterations: λi((t)) _i(H( θ_t)) vs. iteration t. (Right). BS scheduling improves data efficiency. BS scheduling improves data efficiency. We use a BS schedule that starts at 0.490.49M and increases by 4×4× at each ramp. Models are trained with 1,21,2, or 33 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 η and warmup lengths TwT_w. We plot the best validation loss L(bst)L( θ_bst) vs. TwT_w for different η. The optimal TwT_w is highlighted with a star. Extension to GPT-2 Architectures. See Figures 9 and 10 for details. Figure 11: Extensive to Adam-mini optimizer. (Left). Early sharp-to-flat dynamics. Evolution of the top eigenvalues of the Hessian across iterations: λi((t)) _i(H( θ_t)) vs. iteration t. (Right). BS scheduling improves data efficiency. BS scheduling improves data efficiency. We use a BS schedule that starts at 0.490.49M and increases by 4×4× at each ramp. Models are trained with 1,21,2, or 33 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: λi((t)) _i(H( θ_t)) vs. iteration t. (Right). BS scheduling improves data efficiency. BS scheduling improves data efficiency. We use a BS schedule that starts at 0.490.49M and increases by 4×4× at each ramp. Models are trained with 1,21,2, or 33 ramping steps, while models without ramping serve as the baseline. Vertical gray dashed lines indicate ramping positions. Extension to Other Optimizers. See Figure 11 for Adam-mini, and see Figure 12 for Lion. Figure 13: Extensive to Larger Scale. (Left). Early sharp-to-flat dynamics. Evolution of the top eigenvalues of the Hessian across iterations: λi((t)) _i(H( θ_t)) vs. iteration t. (Right). BS scheduling improves data efficiency. BS scheduling improves data efficiency. We use a BS schedule that starts at 0.490.49M and increases by 4×4× at each ramp. Models are trained with 1,21,2, or 33 ramping steps, while models without ramping serve as the baseline. Vertical gray dashed lines indicate ramping positions. Extension to Larger Models. See Figure 13 for models with 270M and 530M parameters. C.3 Ablation Studies on Early Instabilities. Figure 14: Muon consistently shows loss spikes and plateaus early in training. We train a series of LLaMA-2 models with 170170M parameters. We adopt a warmup-stable schedule, where the warmup length is shortened to 1616 iterations and the peak LR is varied, η∈2−11,2−10,2−9,2−8,2−7η∈\2^-11,2^-10,2^-9,2^-8,2^-7\. (Left). LR schedule: ηt _t vs. training iteration t. (Middle, Right). Training loss curves for different model sizes: L(t)L( θ_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 170170M parameters. We adopt a warmup-stable schedule, where the warmup length varies from 25,26,27,28,29,210\2^5,2^6,2^7,2^8,2^9,2^10\ iterations and the peak LR is fixed η=2−7η=2^-7. (Right). Zero warmup and small BS leads to larger loss spikes. We train a series of LLaMA-2 models with 170170M parameters. We adopt a constant LR schedule, where no warmup and the peak LR is varied, η∈2−11,2−10,2−9,2−8,2−7η∈\2^-11,2^-10,2^-9,2^-8,2^-7\. We also use the 0.490.49M BS. In this section, we conduct ablation studies on the root cause of instabilities, such as loss spikes and plateaus, observed in early training. Is it the unstable optimizer? To disentangle optimizer-induced instability from landscape-induced instability, we repeated the experiments using Muon, a substantially more stable optimizer than AdamW. In Figure 14, the loss spikes and plateaus consistently occurs under Muon when warmup is shortened or the peak LR is increased. This rules out the possibility that the behavior stems from AdamW’s startup issues. What if we use longer warmup? We vary only the warmup length while fixing the peak LR at 2−72^-7. In Figure 15 (left), shorter warmup lengths lead to higher possibilities of loss spikes This behavior is consistent with the sharp-to-flat dynamics. Early in training, the model resides in sharper regions of the landscape, where only sufficiently small LRs ensure stable updates. Therefore, warmup is needed to gradually increase the LR until the trajectory enters flatter regions that can tolerate larger LRs. What if zero warmup and small BS? We also conduct experiments with no warmup and small batch size. In Figure 15 (right), the loss spikes become even more significant. This aligns with our explanation: with no warmup, the LR jumps immediately to a large value while sharpness is still extremely high, causing a spike almost at initialization. More importantly, small BS does not replace warmup, and the instability still appears because of the sharpness.