Does LLM Pre-Training Typically Occur at the Edge of Stability?

Summary (Overview)

• Two Distinct Failure Modes of Quadratic Approximations: The paper identifies and systematically studies two ways quadratic approximations fail to predict neural network training dynamics: (1) Edge of Convexity (EoC) — an LR-independent failure caused by persistent negative Hessian eigenvalues that makes quadratic dynamics diverge regardless of learning rate, and (2) Edge of Stability (EoS) — an LR-dependent failure where quadratic dynamics diverge only when the LR exceeds a critical threshold.

• Generalized EoS Definition: The authors propose a generalized definition of EoS for stochastic training with adaptive optimizers (including AdamW), based on whether increasing the LR induces significant divergence in quadratic dynamics while leaving true dynamics largely unchanged. For full-batch GD, this definition recovers the original EoS definition of Cohen et al. (2021).

• Scalable Detection Methods: Efficient algorithms are designed to detect both EoC and EoS at scale with theoretical guarantees, using Hessian-vector products (approximately 9 GEMMs per linear layer) rather than forming the full Hessian.

• Empirical Findings on LLM Pretraining: Across GPT-style models from 43M to 787M parameters trained on multiple datasets (FineWeb, C4, DCLM, OpenWebText, OpenWebMath):

  • EoC is always observed during LLM pretraining
  • EoS is prevalent but not universal — it disappears when LR becomes sufficiently small (e.g., after decay) or when batch size falls below a critical threshold (EBZ)

• Key Correlation: The EoS batch size threshold (EBZ) shows a strong positive correlation (Pearson correlation 0.95) with the critical batch size (CBZ) across five datasets.

Introduction and Theoretical Foundation

Background

The paper builds on the classical optimization theory framework where gradient-based methods are analyzed through a local quadratic approximation of the loss function via second-order Taylor expansion:

$$L(w_{t+1}) \approx L(w_t) + (w_{t+1} - w_t)^\top \nabla L(w_t) + \frac{1}{2}(w_{t+1} - w_t)^\top \nabla^2 L(w_t)(w_{t+1} - w_t)$$

This yields the descent lemma: loss decreases when the learning rate satisfies $\eta < 2/\lambda_{max}$, where $\lambda_{max}$ is the maximum Hessian eigenvalue.

The Edge of Stability (EoS) Paradox

Cohen et al. (2021) observed that full-batch gradient descent operates at the Edge of Stability: the largest Hessian eigenvalue $\lambda_{max}$ concentrates near $2/\eta$, placing dynamics at the stability boundary predicted by quadratic models. Yet the loss continues to decrease over long time scales, contradicting the descent lemma.

Gaps in Prior Work

The paper identifies three key limitations of existing research:

1. Full-batch focus: Most EoS definitions assume full-batch training, while stochastic training dominates practice
2. Positive curvature only: Studies focus on failures from large positive Hessian eigenvalues, ignoring negative eigenvalues
3. Small-scale measurements: Direct Hessian analysis is limited to CIFAR-10 and ImageNet scale

Theoretical Foundation: Stochastic Quadratic Stability

The paper analyzes SGD on stochastic quadratics:

$$x_{t+1} = (I - \eta A_t)x_t - \eta b_t$$

with error dynamics:

$$\epsilon_{t+1} = M_t \epsilon_t - \eta r_t, \quad M_t := I - \eta A_t$$

Definition 4.1 (Mean-Square Stability): The system is mean-square stable if $\sup_{t \geq 0} \mathbb{E}[\|\epsilon_t\|^2] < \infty$ for every initialization.

Proposition 4.2: For each fixed $\eta$, there exist matrices $S$ and vector $b_{quad}$ such that:

$$\mathbb{E}[\epsilon_{t+1}] = (I - \eta A)\mathbb{E}[\epsilon_t]$$ $$\text{vec}(\mathbb{E}[\epsilon_{t+1}\epsilon_{t+1}^\top]) = \mathbb{E}[M_t \otimes M_t] \text{vec}(\mathbb{E}[\epsilon_t\epsilon_t^\top]) + S\mathbb{E}[\epsilon_t] + b_{quad}$$

The system has bounded second moments if $\rho(\mathbb{E}[M_t \otimes M_t]) < 1$.

Two Instability Mechanisms

1. Negative Curvature Divergence: If $\lambda_{min}(A) < 0$, then along the corresponding eigenvector:

$$a_{t+1} = (1 - \eta\lambda_{min}(A))a_t = (1 + \eta|\lambda_{min}(A)|)a_t$$

This grows unboundedly for every $\eta > 0$ — the instability is LR-independent.

2. Positive Curvature Divergence: On the positive curvature subspace (projected via $\Pi_+$), the critical LR is:

$$\eta^\star := \sup\left\{\eta > 0 : \sup_{t \geq 0} \mathbb{E}\|\epsilon_t^+\|^2 < \infty \text{ for every } \epsilon_0 \in \mathbb{R}^d\right\}$$

For full-batch GD with $\lambda_{max}(A) > 0$: $\eta^\star = 2/\lambda_{max}(A)$.

For stochastic training, the second-moment transition decomposes as:

$$\mathbb{E}[M_t^+ \otimes M_t^+] = \mathbb{E}[M_t^+] \otimes \mathbb{E}[M_t^+] + \text{Cov}(M_t^+)$$

The covariance term can shift the stability boundary away from $2/\lambda_{max}(A)$.

Methodology

Key Definitions

Definition 5.1 (Saddleness):

$$\nu(w_{t_0}) := \frac{\max\{0, -\lambda_{min}(H_{t_0})\}}{\lambda_{max}(H_{t_0})}$$

This measures negative curvature relative to positive curvature. For LLM pretraining, $\nu \leq 0.1$ after warm-up.

Edge of Convexity (EoC): At checkpoint $w_{t_0}$, training is at EoC when $\lambda_{min}(H_{t_0}) < 0$ but saddleness $\nu(w_{t_0})$ is small.

Stochastic EoS: Training is at EoS when the current LR is close to the stochastic critical LR on the positive-curvature subspace: $\eta \approx \eta^\star(w_{t_0})$.

EoC Detection: Stochastic Power Iteration with Bisection

Algorithm 1 runs the shifted iteration:

$$x_{t+1} = (I - \eta_p(A_{B_t} + \lambda I))x_t, \quad t = 0, \ldots, T_p - 1$$

with a fresh mini-batch at each step. Since $\mathbb{E}[x_{t+1}] = (I - \eta_p(A + \lambda I))\mathbb{E}[x_t]$, if $\lambda_{min}(A) + \lambda < 0$, the mean iterate grows geometrically along the bottom eigenvector direction.

Theorem A.4 (Informal): The stochastic power iteration outputs $\hat{\lambda}_{min}$ satisfying:

$$\hat{\lambda}_{min} - \lambda_{min}(A) \leq O\left(\eta_p \log(2N)\right) + O\left(\frac{1}{\eta_p T_p}\right) + \Lambda 2^{-N}$$

EoS Detection: Compensated Local Simulation

Algorithm 3 compensates for negative curvature by adding penalty $\frac{\lambda_P}{2}\|\epsilon\|_P^2$ with $\lambda_P = \max\{0, -\hat{\lambda}_{min}\}$, then simulates local AGM dynamics:

$$m_{t+1} = \beta_1 m_t + (1-\beta_1)(g_{B_t} + (H_{B_t} + \lambda_P P)\epsilon_t)$$ $$\epsilon_{t+1} = \epsilon_t - \alpha\eta P^{-1}m_{t+1} - \alpha\eta\lambda_{WD}(\epsilon_t + w_{t_0})$$

The estimate $\hat{\eta}^\star$ is the smallest LR in a grid for which displacement reaches a cutoff $M = 10^5$. EoS is declared when $q = \hat{\eta}^\star/\eta \in [1 - \tau_{eos}, 1 + \tau_{eos}]$ with $\tau_{eos} = 0.1$.

Adaptive Gradient Methods (AGMs)

For AGMs with momentum and preconditioner, the effective curvature is:

$$A := P^{-1/2}H_{t_0}P^{-1/2} + \lambda_{WD}I$$

The critical LR is characterized by the spectral radius of the second-moment operator:

$$K^+(\eta) := \mathbb{E}[M_t^+(\eta) \otimes M_t^+(\eta)]$$

Proposition A.3: The critical LR satisfies $0 < \eta^\star = \sup\{\eta > 0 : \rho(K^+(\eta)) < 1\} < \infty$.

Computational Cost

Hessian-vector products cost approximately 9 GEMMs per linear layer (versus 3 for one gradient evaluation), computed via:

$$\text{HVP}(v) = \nabla_w(g(w)^\top v) = \nabla_w^2 L(w)v$$

Empirical Validation / Results

Experimental Setup

• Models: GPT-extra-small (43M), GPT-small (104M), GPT-medium (344M), GPT-large (787M)
• Datasets: FineWeb-10B, OpenWebMath, OpenWebText, C4, DCLM
• Optimizer: AdamW (WD=0.1, $\beta_1=0.9$, $\beta_2=0.95$)
• LR Schedule: Warm-Stable-Decay (WSD) with 4% warmup and 20% decay steps
• Evaluation: Over 3,000 checkpoints from more than 100 pretraining runs

Key Results

1. EoS and EoC Are Widely Observed

Across model sizes from 43M to 787M, the bottom heatmaps in Figure 2 show a sharp transition from bounded displacement to explosion near LR factor 1.0, indicating the training LR sits at the critical threshold. Saddleness remains below 0.1 reference line, consistent with EoC.

2. Batch Size Dependence

Figure 4 shows:

• EoS emerges only above the EoS Batch Size (EBZ) threshold
• The critical LR factor $q = \hat{\eta}^\star/\eta$ decreases as batch size increases
• EBZ and CBZ show positive correlation (Pearson 0.95) across datasets

3. LR Dependence

• Training can enter stable regime after LR decay: After abrupt LR drops, the transition moves to larger LR factors; with sufficiently large reductions, EoS does not return
• Training can enter extremely unstable regime with large LR: At peak LR 0.04, the estimated critical LR is only ~1/10 of the training LR on average

4. Model Scale Results

For GPT-small/medium/large:

• At optimal peak LRs: EoS observed at 51/60, 43/50, 29/30 checkpoints (respectively)
• EoC criterion satisfied at 44/60, 38/50, 26/30 checkpoints
• At extreme peak LRs (0.04): no LR factor in grid stays bounded at most checkpoints

Theoretical Guarantees

Theorem C.16 (Accuracy of expected EoS detector):

$$\min\{\Phi(\mu_\star), 0\} - H_T \leq \frac{1}{\bar{\eta}_\star(M,T)} - \frac{1}{\eta_\star} \leq \max\{\Phi(L) + \delta_{cov}, 0\}$$

where $H_T = O((1 + \log M)/T)$ is the finite-horizon error, $\Phi$ controls the compensation error, and $\delta_{cov}$ bounds the covariance mismatch.

Theoretical and Practical Implications

Theoretical Implications

1. Quadratic approximations have two distinct failure modes: The paper formalizes that negative curvature (EoC) and positive curvature (EoS) cause different types of instability, requiring separate treatment.

2. Stochastic EoS differs from full-batch EoS: The critical LR is not determined by $\lambda_{max}$ alone but by the second-moment operator $\mathbb{E}[M_t^+ \otimes M_t^+]$, which includes the covariance of mini-batch Hessians.

3. Weight decay and preconditioning matter: The effective curvature $A = P^{-1/2}H_{t_0}P^{-1/2} + \lambda_{WD}I$ incorporates both the preconditioner and weight decay, generalizing prior definitions.

Practical Implications

1. LR schedule design: LR decay can move training out of the EoS regime into the stable regime, suggesting that aggressive decay schedules may reduce instability risks.

2. Batch size selection: The positive correlation between EBZ and CBZ suggests that batch sizes near the critical batch size for optimal loss also tend to place training at EoS.

3. Detection tools: The scalable detection algorithms (stochastic power iteration + compensated local simulation) provide practical tools for monitoring training stability in large-scale LLM pretraining.

4. Extremely unstable regime: Training with excessively large peak LRs enters a regime beyond EoS where the quadratic approximation predicts critical LR far below the training LR, correlating with worse final evaluation loss.

Conclusion

This paper systematically characterizes when and how quadratic approximations fail in LLM pretraining, identifying two distinct failure modes:

1. Edge of Convexity (EoC): An LR-independent failure caused by persistent negative curvature, always observed in LLM pretraining
2. Edge of Stability (EoS): An LR-dependent failure, prevalent but not universal — disappearing with small LRs (after decay) or small batch sizes (below EBZ)

The generalized definitions extend to stochastic training with adaptive optimizers, and the scalable detection methods enable monitoring at LLM scale. The positive correlation between EBZ and CBZ provides a practical link between stability thresholds and optimal training configurations.

Future Directions

1. Other optimizers: Extending EoS analysis to optimizers like Muon
2. Theoretical analysis: Developing solid theoretical explanations for realistic neural network training dynamics at EoS and EoC
3. Understanding the mechanisms: Why does the loss continue to decrease despite operating at the stability boundary?