PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yuhang Cai, Haodong Wen, Haofeng Huang, Deyi Liu, Yiyuan Ma, Kaifeng Lyu

arXiv:2610.36778v1cs.LG

TL;DR

Quadratic approximations can misrepresent neural-network training, while existing EoS evidence is limited in its focus on full-batch training, positive curvature, and small-scale settings. The paper studies these failure modes in stochastic, adaptive, large-scale training and finds EoC and EoS widely observed, with EoS depending on batch size and learning-rate schedule.

  • Problem

    Existing EoS definitions and measurements focus largely on full-batch training, positive-curvature failures, and small-scale settings, although stochastic training is more common in practice.

    Source: §1 · p. 2

  • Method

    The paper studies quadratic-approximation failures in stochastic, adaptive, large-scale training, distinguishing LR-independent EoC from LR-dependent EoS after accounting for EoC.

    Source: §1 · p. 2

  • Results

    EoC and EoS are widely observed across more than 3,000 checkpoints from over 100 pretraining runs, while EoS is less evident below a batch-size threshold and can recede after LR decay.

    Source: §6.2 · p. 10, §6.3 · p. 10, §6.4 · p. 11

  • Takeaways & Limitations

    The study extends EoS to LLM training and clarifies EoC as a second failure mode of the quadratic approximation.

    Source: §7 · p. 11

    The stochastic EoS threshold depends on training noise as well as λmax, and separating the EoC and EoS curvature subspaces is infeasible for large neural networks.

    Source: §5.2 · p. 8

Abstract

from arXiv · show

Quadratic approximations provide a common lens for studying neural network optimization, but recent evidence challenges their predictive validity. In full-batch gradient descent with learning rate (LR) $\eta$, Cohen et al. (2021) observed the Edge of Stability (EoS), where the largest Hessian eigenvalue concentrates near $2/\eta$, seemingly at odds with classical stability conditions predicted by quadratic approximations. In this work, we revisit the fidelity of quadratic approximations as a model of neural network training dynamics, with a particular focus on their failure modes in LLM training. We first highlight a distinct failure mode of the quadratic approximation, where persistent negative curvature causes the quadratic dynamics to diverge regardless of the LR, even though the real dynamics remain stable. We term this phenomenon the Edge of Convexity (EoC). After removing the effect of this LR-independent instability, we then propose a generalized definition of Edge of Stability (EoS) for stochastic training with adaptive optimizers, based on whether increasing the LR induces significant divergence in the quadratic dynamics while leaving the real dynamics largely unchanged. Across different LLM pretraining runs with various model sizes up to 787M, we find that (1) EoC is always observed across LLM pretraining, and (2) EoS is also prevalent but not universal: it disappears when the LR becomes sufficiently small (e.g., after decay) or when the batch size falls below a critical threshold that has a positive correlation with the critical batch size. Together, these findings characterize when and how quadratic approximations fail and provide a foundation for future work on understanding training dynamics.

1 INTRODUCTION

The paper studies where quadratic approximations fail in stochastic, adaptive, and large-scale training, addressing limits in prior evidence on negative curvature and LLM-scale behavior. It names the LR-independent negative-curvature failure EoC, generalizes EoS for stochastic adaptive training, and reports both phenomena in LLM pretraining.

Motivation

  • Prior work leaves open whether quadratic-approximation failures persist in stochastic training, with negative curvature, and at large scale.Direct measurements were largely limited to smaller settings such as CIFAR-10 and ImageNet.Source: §1 · p. 2

Contributions

  • The study systematically examines quadratic-approximation failure modes in stochastic, adaptive, and large-scale training, focusing on LLM pretraining.Source: §1 · p. 2
  • Persistent negative curvature can make quadratic dynamics diverge at every LR while real training remains stable; the paper names this Edge of Convexity (EoC).Source: §1 · p. 2
  • The paper defines stochastic adaptive EoS by LR-induced divergence in quadratic dynamics while real dynamics remain largely unchanged, and reports EoC consistently but EoS non-universally in LLM pretraining.EoS disappears with sufficiently small LR or below a batch-size threshold that scales linearly with the critical batch size.Source: §1 · p. 2

2 RELATED WORK

Prior work connected stability and sharpness to neural-network training and developed explanations of EoS, but the fidelity of quadratic approximations at large scale remained unclear. The paper distinguishes its study through its EoC analysis, scalable detection methods, and LLM-pretraining experiments.

Edge of Stability

  • Cohen et al. (2021) identified full-batch EoS, where the largest Hessian eigenvalue stays near 2/η while loss decreases non-monotonically; later work extended the phenomenon to adaptive methods.Source: §2 · p. 3

Training Stability of Neural Networks

  • Despite theoretical accounts of EoS, quadratic-approximation effectiveness in large-scale neural-network training had not been carefully studied.Source: §2 · p. 3

Hessian Spectrum of Neural Networks

  • Earlier Hessian studies noted negative eigenvalues, while this work measures population-Hessian negative curvature during large-scale training.Source: §2 · pp. 3–4

Comparison with Andreyev & Beneventano (2024)

  • Compared with Andreyev & Beneventano (2024), this work adds EoC, scalable detection with theoretical guarantees, and LLM pretraining with AdamW.It also says its full-batch EoS definition reduces to Cohen et al. (2021; 2022), unlike the comparison work's definition.Source: §2 · p. 4

3 PRELIMINARY

The preliminary section frames EoS as a mismatch between the stability boundary predicted by a local quadratic model and observed neural-network training. It recalls the full-batch threshold and uses fixed-checkpoint quadratic dynamics to make the discrepancy explicit.

Full-Batch Edge of Stability

  • Classical local analysis predicts full-batch GD stability on a positive-semidefinite quadratic only when ηλmax(Ht) ≤ 2.Source: §3 · p. 4
  • In observed full-batch EoS, λmax often approaches 2/η and remains near it while loss continues decreasing over longer time scales.Source: §3 · p. 4

Quadratic approximation and EoS

  • At a fixed checkpoint, the quadratic GD dynamics become unstable when ηλmax exceeds 2, whereas observed EoS places training at that predicted stability boundary despite continued loss decrease.Source: §3 · p. 4

4 STABILITY OF STOCHASTIC GRADIENT DESCENT ON QUADRATICS

The section separates stochastic-quadratic instability caused by negative curvature from LR-threshold instability on the positive-curvature subspace. Its mean-square analysis shows why stochastic EoS cannot be characterized by λmax alone.

4 STABILITY OF STOCHASTIC GRADIENT DESCENT ON QUADRATICS

  • Stochastic quadratic training has two instability modes: negative population curvature causes LR-independent divergence, while positive-curvature instability is governed by a second-moment operator.Mean-square stability tracks expected squared error because stochastic deviations can cancel in the mean while their squared norm grows.Source: §4 · p. 5

4.1 NEGATIVE CURVATURE DIVERGENCE

  • When the population quadratic has a negative eigenvalue, its quadratic dynamics diverge for every positive LR; lowering LR only slows the escape.Source: §4.1 · p. 5

4.2 POSITIVE CURVATURE DIVERGENCE AND CRITICAL LR

  • Projecting out nonpositive-curvature directions isolates positive-curvature stability; the resulting stochastic critical LR depends on more than λmax, while the full-batch definition recovers the 2/λmax threshold.Batch-to-batch Hessian covariance can destabilize second moments even when the mean transition is stable.Source: §4.2 · p. 6

5 STABILITY OF STOCHASTIC GRADIENT DESCENT ON NEURAL NETWORKS

The section distinguishes LR-independent instability from negative curvature (EoC) and LR-dependent instability from positive curvature (EoS). It extends these analyses from stochastic gradient descent to adaptive methods and examines when training approaches or leaves these boundaries.

5 STABILITY OF STOCHASTIC GRADIENT DESCENT ON NEURAL NETWORKS

  • Negative population curvature makes the quadratic model have an expanding direction at every positive LR, even when the real training dynamics remain stable.Reducing the LR slows this divergence but cannot prevent it; EoC describes small saddleness with negative curvature.Source: §5.1 · p. 7, §5 · p. 7

5.1 EDGE OF CONVEXITY

  • After warm-up, LLM pretraining has saddleness ν ≤ 0.1, and EoC occurs when negative curvature persists despite this small relative magnitude.The authors estimate the population minimum Hessian eigenvalue with stochastic power iteration because mini-batch estimates can be biased downward.Source: §5.1 · p. 7

5.2 EDGE OF STABILITY UNDER STOCHASTIC TRAINING

  • Even after removing nonpositive-curvature directions, quadratic dynamics can diverge above a critical LR, while neural-network training often lies near that threshold.A tiny LR increase can trigger divergence in the projected quadratic model.Source: §5.2 · p. 8
  • At neural-network scale, separating the curvature subspaces is infeasible, and stochastic EoS cannot be detected by checking ηλmax ≈ 2 alone.The stochastic threshold also depends on training noise.Source: §5.2 · p. 8
  • In CIFAR-10 ResNet runs, the bounded-to-exploding boundary lies at LR factor 1.0 across most middle-phase checkpoints, placing training near EoS.The detector compares local simulations across LR factors using final displacement norms.Source: §5.2 · p. 8

5.3 ADAPTIVE GRADIENT METHODS AT THE EDGE OF CONVEXITY AND EDGE OF STABILITY

  • After a larger LR drop, local simulations can remain bounded at the training LR as training enters a stable regime and does not return to EoS during the observed run.The bounded-to-exploding transition shifts to larger LR factors during cosine decay and after an abrupt multistep drop.Source: Figure 3 · p. 9
  • For adaptive gradient methods, EoC means negative effective curvature with small saddleness, while EoS means the training LR is near the projected dynamics’ critical LR.The framework freezes the preconditioner at a checkpoint; momentum SGD and SGD are special cases of the adaptive update.Source: §5.3 · pp. 8–9

6 EXPERIMENTS

Across GPT-style pretraining runs, EoC and EoS are widely observed, but EoS depends on batch size and the learning-rate schedule. Batch sizes above an EoS threshold show EoS more often, while sufficient LR decay can remove it and excessively large peak LRs can produce instability beyond EoS.

6.1 EXPERIMENTAL SETUP

  • The experiments use GPT-style models at four scales from 43M to 787M, generally trained on FineWeb with AdamW and a Warm–Stable–Decay learning-rate schedule.The default setup uses sequence length 1024, 4% warmup steps, and 20% decay steps.Source: §6.1 · p. 10

6.2 EOS AND EOC ARE WIDELY OBSERVED

  • Across more than 100 runs and over 3,000 checkpoints, EoS and EoC are widely observed, while EoC remains widespread as training configurations vary.At many checkpoints, estimated critical LR is close to training LR, and the minimum eigenvalue is negative with saddleness below 0.1, respectively.Source: §6.2 · p. 10

6.3 BATCH SIZE DEPENDENCE OF EOS

  • EoS emerges above the EoS Batch Size threshold: B = 512 and B = 1024 show it over a larger fraction of training than B = 256.Below the threshold, compensated local simulations after the early training phase remain bounded at every tested LR factor.Source: §6.3 · p. 10
  • Across datasets, EoS Batch Size and critical batch size are positively correlated; datasets with larger EoS thresholds also tend to sustain near-optimal loss at larger batch sizes.Critical batch size is the largest batch size whose final evaluation loss is within 0.3% of the best loss in that dataset’s sweep.Source: §6.3 · p. 10

6.4 LR DEPENDENCE OF EOS

  • After an LR drop, EoS can disappear; with a sufficiently large reduction, the estimated critical LR stays above the training LR and EoS does not return during the observed run.With smaller drops, the critical LR can fall toward the current training LR and EoS can return.Source: §6.4 · p. 11
  • Peak LRs substantially above the empirically optimal range can push training beyond EoS, where the estimated critical LR is far below the training LR and final loss is higher.At peak LR 0.04, the estimated critical LR is about one tenth of the training LR on average, and compensated local simulations explode at the training LR.Source: Figure 5 · p. 10, §6.4 · p. 11

7 CONCLUSION

The paper frames EoC and EoS as two failure modes of quadratic approximations and extends EoS to LLM training. It leaves optimizer generality and theoretical understanding of realistic neural-network dynamics as open directions.

7 CONCLUSION

  • The paper extends EoS to LLM training and clarifies EoC as another failure mode of the quadratic approximation.Experiments mainly focus on SGD and AdamW; the authors identify other optimizers and theoretical analysis of realistic training dynamics as open directions.Source: §7 · p. 11

A DETECTING EOC AND EOS FOR ADAPTIVE GRADIENT METHODS

  • The detection procedures for EoC and EoS in adaptive gradient methods use Hessian-vector products without forming the Hessian.The paper places procedure details and guarantees in Appendices A.2 and A.3, with computational cost in A.4.Source: A DETECTING EOC AND EOS FOR ADAPTIVE… · p. 17

A.1 ADAPTIVE GRADIENT METHODS

This appendix defines local adaptive-gradient dynamics and critical learning rates, presents procedures for detecting EoC and EoS while separating learning-rate-driven instability from negative effective curvature, and estimates Hessian-vector product costs. It reports that an HVP costs about three gradient evaluations, or nine forward passes when GEMMs dominate.

A.1 ADAPTIVE GRADIENT METHODS

  • Below the AGM critical learning rate η⋆, projected local dynamics are mean-square stable; above it, some initialization has an unbounded second moment, while boundary boundedness can depend on forcing.AGM training is at EoS when its training learning rate is close to η⋆.Source: §A.1 · p. 18

A.2 EOC DETECTION

  • Stochastic power iteration with bisection estimates the minimum effective curvature using fresh mini-batches; its accuracy guarantee separates curvature noise, finite-horizon error, and bisection resolution.The guarantee assumes a negative population minimum eigenvalue and requires the search interval, stepsize conditions, and initial eigenvector overlap as inputs.Source: §A.2 · pp. 19–21
  • EoC detection declares EoC when the estimated minimum effective curvature is negative and its magnitude-to-maximum-curvature ratio is at most the tolerance, set to 0.1 in experiments.The estimate uses stochastic power iteration because batchwise minimum-curvature eigenvectors are nearly orthogonal and their eigenvalues vary substantially, unlike the stable top eigenpair.Source: §A.2 · pp. 19–20, Figure 6 · p. 19

A.3 EOS DETECTION

  • The EoS detector compensates negative curvature, simulates local AGM trajectories across increasing learning-rate factors, and reports the first exploding rate as the critical-rate estimate.It declares EoS when q = η⋆/η is within 0.1 of 1; if no tested factor explodes, the result is censored and not counted as EoS.Source: §A.3 · pp. 22–23
  • For the expected detector, the finite-horizon threshold bound separates compensation, Hessian-noise, and finite-horizon effects; at fixed cutoff, the finite-horizon term decreases as 1/T.The analysis concerns expected squared displacement from unit initial displacements, not the single zero-initial-displacement trajectory used in experiments; estimated compensation does not ensure positive-definite effective curvature.Source: §A.3 · p. 23

A.4 COST OF A HESSIAN-VECTOR PRODUCT

  • For linear layers dominated by matrix multiplication, computing an HVP requires a second reverse-mode pass, with overhead comparable to another forward-and-backward pass plus extra GEMMs.The supplied passage gives no completed GEMM count.Source: §A.4 · p. 24
  • An HVP costs about 3× one gradient evaluation (9 versus 3 GEMMs per linear layer), or 9× a forward pass when GEMMs dominate.Transformer blocks contain multiple relevant linear layers: QKV projections, the attention output projection, and two MLP projections.Source: §A.4 · p. 24

B.1 MODELS, TRAINING, AND DETECTION SETTINGS

The experiments use several model architectures and training settings, with tokenization and data filtering designed to stabilize curvature detection. Detection uses specified iteration counts and a displacement cutoff.

B.1 MODELS, TRAINING, AND DETECTION SETTINGS

  • Table 1 lists the model architectures used in the experiments.Source: Table 1 · p. 24, §B.1 · p. 24
  • The experiments include ResNet20 and ResNet32 trained with MSE and SGD for 50 epochs at batch size 1024 and weight decay 5 × 10−4.The figure uses the raw Hessian and plots epoch on the horizontal axis.Source: Figure 7 · p. 25
  • The detection shard retains 92,060 of 97,656 sequences (94.3%) after filtering rare-token sequences and bursts, and is shared across model sizes, learning rates, and checkpoints.Rare tokens are removed to reduce distortions in Adam-preconditioned curvature estimates.Source: §B.1 · pp. 24–25
  • GPT-extra-small uses 100 stochastic power iterations and 50 compensated local-simulation steps; larger models use 50 and 25, respectively, with explosion cutoff M = 105.All models use sequence length 1024; QK-normalization and attention gating are applied for stability.Source: §B.1 · p. 25

B.2 ADDITIONAL RESULTS

Compensated local simulations show EoS across AdamW pretraining runs, while robustness testing indicates that the sharp transition depends on compensation.

B.2 ADDITIONAL RESULTS

  • Across AdamW pretraining runs from 43M to 787M parameters, compensated local simulations commonly transition from bounded trajectories to explosion near LR factor 1, but not at the smallest batch size or extreme peak LRs.The transition's location and duration depend on the LR schedule, peak LR, and batch size.Source: §B.2 · p. 25
  • At one checkpoint, the sharp transition is absent without compensation and remains visible at compensation coefficient 600 as well as at the estimated value.Source: §B.2 · p. 25

B.3 BATCH SIZE THRESHOLD EXPERIMENTS

Across five datasets, the estimated EoS batch threshold is strongly correlated with the critical batch threshold, and LR drops shift EoS behavior according to the post-drop LR.

B.3 BATCH SIZE THRESHOLD EXPERIMENTS

  • The five-dataset batch-size experiments use GPT-extra-small, AdamW, a WSD schedule, and square-root scaling of LR and weight decay while holding other hyperparameters fixed.β2 is also adjusted according to Marek et al. (2025).Source: §B.3 · p. 26
  • CBZ is the largest swept batch size whose final evaluation loss is within 0.3% of the best, while EBZ is the smallest batch size where fitted average critical-LR ratio reaches 1.1.EBZ uses checkpoint-averaged q = bη⋆/η; runs bounded across the LR grid are censored at its upper end.Source: §B.3 · p. 26
  • Pearson correlation between EBZ and CBZ is 0.95 across five datasets; EBZ is lower for C4 and FineWeb, but higher for OpenWebMath, DCLM, and OpenWebText.Neither threshold is uniformly smaller.Source: §B.3 · p. 26

B.4 MULTISTEP LR DROP ABLATIONS

  • Before multistep LR drops, all six runs are at EoS at most checkpoints; afterward, EoS returns only when the post-drop LR remains large.With peak LR 0.002, no post-drop checkpoint is at EoS for drop ratios 0.6, 0.45, or 0.3; at drop ratio 0.7, EoS returns at 26, 7, and 0 checkpoints for peak LRs 0.008, 0.004, and 0.0005, respectively.Source: §B.4 · p. 26

B.5 LINEAR LR DECAY ABLATIONS

Linear-decay settings affect when EoS re-emerges: longer schedules slow the transition, while EoC weakens as the LR becomes small.

B.5 LINEAR LR DECAY ABLATIONS

  • Longer linear-decay schedules produce a slower EoS transition, with transition time roughly proportional to the number of decay steps.The decay ratio also affects whether and when EoS re-emerges.Source: §B.5 · p. 26
  • During LR decay, −bλmin falls from about 850 to 56–540 and saddleness falls from about 0.09 to 0.03 or less, weakening EoC.The smaller values occur once the LR is below half its peak.Source: §B.5 · p. 26

B.6 EXPERIMENTS ACROSS MODEL SCALES

Across model scales, the experiments find EoS and EoC at lower peak LRs, while smaller batches can remove the observed EoS transition. The appendix also formalizes local stochastic dynamics and their stability conditions.

B.6 EXPERIMENTS ACROSS MODEL SCALES

  • The scale sweep trains GPT-small, medium, and large across eight peak LRs, measuring EoS and EoC at ten stable-phase checkpoints per run.Compensated local simulations test LR factors from 0.6 to 1.4, adding lower factors for runs with no bounded simulation at 0.6.Source: §B.6 · pp. 26, 30
  • At batch sizes 512 and 1,024, the EoS transition stays near LR factor 1.0 through much of training; at 256, no tested factor explodes after step 3,600.The batch-size comparison uses GPT-extra-small, FineWeb, and WSD schedules.Source: Figure 9 · p. 27
  • At the three smallest peak LRs, results show 17/18/19 EoS and 26/27/26 EoC checkpoints; from 0.01 on, only 12/90 are EoS and 38/90 lack bounded factors.The counts compare the three GPT sizes; the tested peak LRs begin at 0.01 for the latter pattern.Source: §B.6 · p. 30

C THEORETICAL ANALYSIS

  • The theoretical appendix analyzes quadratic-loss dynamics at fixed checkpoints with the preconditioner held fixed, deriving moment evolution, stability thresholds, and detector error bounds.Proofs are given in Appendices D to F.Source: C THEORETICAL ANALYSIS · p. 35

C.1 SETTING AND THE LOCAL LINEAR SYSTEM

  • The local stochastic model assumes independent fresh mini-batches and freezes the preconditioner over a short checkpoint simulation; this is an idealization requiring small preconditioner changes.A β2 near one motivates the approximation but does not remove that requirement.Source: §C.1 · p. 35
  • Compensation adds λP I to the preconditioned batch Hessian in the momentum update, while leaving the decoupled weight-decay coefficient unchanged.Setting λP = 0 recovers the uncompensated surrogate.Source: §C.1 · pp. 35–36
  • For the stochastic local system, ρ(E[Mt ⊗ Mt]) < 1 ensures finite limiting moments; strict homogeneous mean-square stability holds exactly under this condition.Strict homogeneous stability also implies bounded moments for the affine system on the same state space.Source: §C.1 · pp. 36–37

C.2 STABILITY OF THE LOCAL DYNAMICS

This section analyzes stability in the local dynamics, separating instability from negative curvature at every positive learning rate from learning-rate-driven instability on the positive-curvature subspace. It also examines stochastic curvature-estimation errors and EoS detector accuracy, with bounds and proofs for the compensated local system and its moment recursion.

C.2 STABILITY OF THE LOCAL DYNAMICS

  • A negative effective-curvature eigenvalue makes homogeneous AGM dynamics mean-square unstable for every η > 0, while a smaller learning rate only slows the divergence.For SGD, the mean displacement along a negative-Hessian eigenvector grows by (1 − ηλ)^t when its initial mean has a component in that direction.Source: §C.2 · p. 37
  • On the strictly positive-curvature subspace, strict mean-square stability holds exactly for 0 < η < η⋆; above η⋆, some deterministic initial state has an unbounded second moment.Zero-curvature directions are dropped because their mean transition has an eigenvalue of one, which precludes strict stability at any learning rate.Source: §C.2 · p. 38
  • Mean-square stability is governed by whether the second-moment map has spectral radius below one, and its threshold need not match the single-trajectory experimental diagnostic.Mean-square decay implies almost-sure exponential stability, but the converse can fail when rare expanding trajectories dominate the second moment.Source: §C.2 · p. 39

C.3 STOCHASTIC POWER ITERATION FOR EOC

  • The stochastic power-iteration estimator separates curvature-noise, finite-horizon, and bisection-resolution errors when estimating the minimum effective curvature.Its bound assumes the search interval contains the target, the initialization overlaps the growing direction, and the population minimum is negative; with other settings fixed, longer horizons reduce only the finite-horizon term.Source: §C.3 · pp. 39, 41

C.4 COMPENSATED LOCAL SIMULATION FOR EOS

  • The compensated local-simulation detector’s error relative to the projected critical learning rate has finite-horizon and compensation components.The compensation component remains even if the compensated threshold is known exactly.Source: §C.4 · pp. 42–43
  • For either mean system C and its centered noise, strict mean-square stability holds precisely when (1 − b)/α > r_C(b).The compensated and projected systems use different means and noise blocks, so the criterion applies to each system’s corresponding quantities.Source: §C.4 · pp. 42–43
  • Compensation shifts curvature on the positive eigenspace while retaining projected-out directions, with curvature and noise effects bounded separately.The increasing function Φ satisfies Φ(L) ≤ c⋆λP/2; the noise contribution vanishes when discarded and cross noise blocks vanish, while noise in lower-mean-curvature directions can limit stability.Source: §C.4 · pp. 43–44
  • Theorem C.16 combines compensation and finite-simulation bounds to characterize the expected EoS detector’s accuracy.The finite-simulation error lies in [−HT, 0] and decreases as T grows under the stated condition, but increasing T does not remove compensation error.Source: §C.4 · p. 44

D.1 PROOF OF PROPOSITION C.2

  • The compensated local system replaces the mini-batch Hessian with HBt + λPP, while leaving the weight-decay coefficient unchanged.In preconditioned coordinates, the Hessian replacement is DBt + λPI.Source: §D.1 · p. 45

D.2 PROOF OF PROPOSITION C.3

  • A strict second-moment spectral-radius bound makes the affine moment recursion converge and remain bounded for finite-second-moment initial states.The proof allows correlation between the current transition and forcing.Source: §D.2 · p. 46

D.3 PROOF OF COROLLARY C.4

The proof centers the SGD recurrence at the population stationary point to derive its moment dynamics. It also shows that negative effective curvature yields an expanding mean mode, while the projected positive-curvature system has a finite stability threshold.

D.3 PROOF OF COROLLARY C.4

  • Centering the SGD recurrence at the population stationary point removes mean forcing and yields moment equations for the centered displacement.The centered displacement is a stochastic linear system; the generic affine quadratic follows by applying the same proposition.Source: §D.3 · p. 47
  • Negative effective curvature gives the mean AGM dynamics an eigenvalue above one for every η > 0, so the homogeneous system is not mean-square stable.The proof decouples the mean dynamics along eigenvectors of D and finds an expanding mode when the effective curvature is negative.Source: §D.3 · p. 48
  • For SGD, a negative population-Hessian eigenvalue makes a nonzero initial component grow geometrically, causing the second moment to diverge.Along an eigenvector with λ < 0, the mean multiplier is 1 − ηλ > 1.Source: §D.3 · p. 48

D.5 PROOF OF COROLLARY C.8

  • Above the projected stability threshold, the second-moment spectral radius exceeds one and some deterministic initialization has diverging second moment.The proof uses a PSD eigenmatrix of the adjoint second-moment map to identify such an initialization.Source: §D.5 · pp. 48–49

D.6 PROOF OF THE AGM CRITICAL LR CHARACTERIZATION

The proof identifies the AGM critical learning rate with a spectral threshold for the projected affine system. Below that threshold, affine moments are bounded; above it, some initialization has a diverging second moment.

D.6 PROOF OF THE AGM CRITICAL LR CHARACTERIZATION

  • For learning rates below ηcrit, the affine moment recursion is bounded, so those rates belong to the stable set.The projected transition has spectral radius below one in this range.Source: §D.6 · p. 50
  • For η > ηcrit, a deterministic initialization has a homogeneous second moment that diverges, so the rate is not stable for every initialization.The proof constructs two affine trajectories with the same batches and forcing; their difference follows the unstable homogeneous system.Source: §D.6 · p. 50
  • The supremum of the stable learning-rate set equals ηcrit, although forcing can change boundedness at the endpoint.The proof gives examples with the same second-moment spectral radius but different boundedness behavior.Source: §D.6 · p. 50

E PROOF OF THE POWER-ITERATION BOUND

The appendices develop stability bounds and criteria, showing how bisection localizes a target shift and establishing a unique mean-square stability threshold along the AGM learning-rate curve. They extend the SGD criterion to momentum and indefinite-curvature samples and bound the gap between a detector’s accepted rate and the projected system’s critical rate, separating compensation error from finite-horizon error.

E.1 THE BISECTION BOUND

  • With probability at least 1 − δ, queries below λ⋆ − max{g−, ep} explode and those above λ⋆ + g+ do not; after N halvings, the search interval brackets λ⋆ within these margins.The finite-horizon term ep qualifies the lower margin; queries between the two regions may return either outcome.Source: §E.1 · pp. 51–52

E.2 PROOF OF THE EXPLOSION BOUND

  • For shifts below λ⋆, the proof normalizes the bottom-eigenvector component by its mean growth factor, stops the resulting martingale at the cutoff, and applies Azuma–Hoeffding.The finite horizon enters through Bρ^-Tp; the definition of ep ensures this term is at most ε/2.Source: §E.2 · pp. 52, 54, §E.4 · p. 56

E.3 PROOF OF THE NO-EXPLOSION BOUND

  • For shifts above λ⋆, a positive power of the norm forms a nonnegative supermartingale, enabling Hoeffding’s lemma and Ville’s maximal inequality to bound cutoff crossing.When λ − λ⋆ ≥ g+, the crossing-probability bound is at most δ/N.Source: §E.3 · pp. 54, 56, §E.4 · p. 56

E.4 AUXILIARY LEMMAS

  • The auxiliary results used in these proofs are conditional Hoeffding’s lemma, Ville’s maximal inequality, and the Azuma–Hoeffding inequality.The appendix attributes them to Hoeffding (1963), Ville (1939), Azuma (1967), and Hoeffding (1963), respectively.Source: §E.4 · p. 56

F PROOF OF THE COMPENSATED FINITE-HORIZON THRESHOLD BOUND

  • The appendix proof of Theorem C.16 compares projected and compensated critical learning rates with the finite-horizon detector output, using compensation and horizon bounds.The supporting stability results and auxiliary matrix facts are developed separately in the appendix.Source: F PROOF OF THE COMPENSATED… · p. 57

F.1 STABILITY CRITERION

  • Strict mean-square stability holds exactly under Theorem F.1’s criterion, and for fixed b the stable coefficients satisfy 0 < α < (1 − b)/rC(b).Here rC(b) is the maximum of the stability quotient introduced in the appendix.Source: §F.1 · p. 57
  • For fixed mean, noise, and test matrix, lowering momentum increases the stability quotient; a spectral-radius bound also controls how much the quotient can exceed its R = 1 bound.The shift identity for the quotient is used in the quantitative stability and compensation comparisons.Source: §F.1 · pp. 57–58
  • The AGM criterion has a unique critical LR ηcrit: strict mean-square stability holds exactly for 0 < η < ηcrit, and the second-moment spectral radius exceeds one above it.At ηcrit the spectral radius equals one.Source: §F.1 · p. 59
  • Under Mulayoff & Michaeli (2024, Theorem 5)’s sampling and positive-semidefinite sample-Hessian assumptions, maximizing the quotient recovers their exact SGD threshold on the positive-curvature subspace and extends to momentum and indefinite-curvature samples with a positive-definite mean.Source: §F.1 · p. 59
  • Momentum changes the deterministic-curvature coefficient and adds a commutator term, so it does not universally rescale the SGD threshold; noise cannot increase the critical LR.When β1λWD > 0, the threshold is the unique root of GC.Source: §F.1 · p. 59

F.2 PROOF OF THE THRESHOLD BOUND

  • The detector-to-projected-threshold gap separates finite-horizon error from compensation error, and the theorem bounds both without assuming small Hessian noise.The compensation error is bounded by Φ(L) + δcov; in the stated asymptotic regime, the finite-horizon term is O((1 + log M)/T).Source: §F.2 · pp. 60–61

F.3 COMPENSATION ERROR

This section develops quantitative bounds relating projected and compensated critical learning rates, including how mean shifts, discarded directions, noise, and rescaling affect compensation error. It also establishes finite-horizon detector cutoffs and stability criteria for stochastic recurrences, including momentum and AGM cases, with supporting moment, Lyapunov, Perron–Frobenius, and Jury results.

F.3 COMPENSATION ERROR

  • Compensation shifts the projected threshold, and restoring discarded directions and their noise increases the bound by at most δcov.The threshold after the shift is rA+(b⋆) + λPI(b⋆); the full-space upper bound adds δcov.Source: §F.3 · pp. 63–64
  • The projected and compensated critical learning rates are compared through their critical momenta and quotient thresholds, yielding the proposition’s clipped bounds.The proof handles both cases for the critical momenta and concludes that clipping is unnecessary in one case.Source: §F.3 · pp. 65–66
  • In the scalar SGD example, the compensation endpoints exactly equal the inverse-rate difference because both systems share the same noise.For homogeneous scalar SGD, the second-moment stability condition is η < 2a/(a^2 + σ^2).Source: §F.3 · p. 66

Then the projected and compensated critical LRs satisfy

  • The projected and compensated critical-rate comparison separates the curvature shift from noise contributions, with the latter controlled by a finite second-moment bound v.The bound need not be small, and v may depend on the compensation coefficient; small saddleness alone does not control it.Source: Then the projected and compensated… · pp. 68–69
  • When saddleness is positive, compensation λP = cνL with c > 1 makes the compensated mean positive definite; an estimated eigenvalue also requires accounting for estimation error.Saddleness controls the curvature shift, while v accounts for Hessian noise.Source: Then the projected and compensated… · pp. 68–69
  • Noise can only reduce the critical learning rate, and the projected-versus-compensated inverse-rate discrepancy is bounded using curvature shift and noise tolerance.The noise contribution lies in [0, v/2], and the expected-detector bound additionally includes H_T.Source: Then the projected and compensated… · pp. 69–71
  • Small saddleness alone is insufficient: discarded noise can keep compensation error nonvanishing, with the inverse-rate difference tending to 1/2 as s approaches zero.In the example, saddleness is s while the covariance mismatch δcov equals 1.Source: Then the projected and compensated… · p. 71

F.4.1 BOUNDING THE MOMENT RADIUS

  • Detector acceptance bounds the moment radius of the homogeneous compensated recurrence, connecting observed displacement cutoffs to second-moment growth.The proof couples affine runs from opposite unit displacements to cancel shared forcing, then uses a quadratic form whose expected value grows by the moment-map spectral radius.Source: §F.4.1 · pp. 71–72

F.4.2 BOUNDING THE INVERSE LR

  • Acceptance yields a uniform bound on the inverse learning-rate gap; in scalar systems without momentum, the quantitative stability bound can be tight.With momentum, the momentum comparisons can introduce a first-order loss when noise is small.Source: §F.4.2 · p. 76

F.4.3 CHOOSING THE CUTOFF

  • A polynomially growing cutoff in horizon T is sufficient to ensure acceptance at the projected critical learning rate.The critical rate’s moment radius being 1 does not by itself guarantee finite-cutoff acceptance; at β1 = 0.9, the example’s limiting amplitude is 19.Source: §F.4.3 · p. 76
  • At the compensated critical learning rate, acceptance holds when the finite-horizon bound keeps E ∥ϵt∥2 ≤ M2 for every unit initial displacement and all t ≤ T.The condition uses the actual values of KT, R0, and σf rather than requiring small noise; KT is finite for every T. The cutoff permits MT > 1 at every horizon and gives HT = O(log T/T) for a fixed model, while Tσf does not imply linear growth of the actual run.Source: §F.4.3 · pp. 77–78

F.5.1 PROOF OF THE EXACT CRITERION

  • The feedback criterion makes the recurrence strictly mean-square stable exactly when ρ(F) < 1 and ρ(KN) < 1.Here K accumulates the response of the mean dynamics to a noise injection, and N forms the noise second moment.Source: §F.5.1 · p. 79
  • For momentum, mean stability requires αλmax(C) < 2(1 + b), and the exact stochastic criterion reduces to positivity of the stability quadratic form on every nonzero PSD matrix.Equivalently, stability is determined by the maximum quotient condition; for fixed b, its right-hand side is independent of α.Source: §F.5.1 · pp. 80, 82, 83

F.5.2 PROOF OF THE QUANTITATIVE BOUND

  • In the quantitative-bound proof, scalar shifts cancel from the commutator term, while the noise term is unchanged because both systems use the same Δt.The passage begins the comparison by expanding the shifted mean terms.Source: §F.5.2 · p. 83
  • The quantitative bound is proved by geometrically rescaling trajectories to strict stability, applying the exact criterion at a fixed test matrix, and bounding the quotient difference.The rescaled recurrence is compared with the original quotient before the difference is bounded.Source: §F.5.2 · pp. 84–86

F.5.3 PROOF OF THE AGM CRITICAL LR

  • The AGM critical LR is the unique zero of GC: strict stability holds below it, the second-moment spectral radius equals one there, and exceeds one above.The proof establishes monotonicity and continuity of GC, then rules out stability above the critical rate.Source: §F.5.3 · pp. 87–88

F.6 AUXILIARY RESULTS

  • For i.i.d. linear dynamics with finite second moments, mean-square decay for every initial state holds exactly when the second-moment map has spectral radius below one.The second-moment map and the full Kronecker second-moment matrix have equal spectral radii; strict mean-square stability also bounds the squared spectral radius of the mean matrix.Source: §F.6 · p. 89
  • A PSD-preserving map U has spectral radius below one if and only if some W ≻ 0 satisfies W − U(W) ≻ 0.The proof constructs such a W from a convergent Neumann series when ρ(U) < 1.Source: §F.6 · p. 90
  • A PSD-preserving linear map with spectral radius R > 0 has a nonzero positive-semidefinite eigenmatrix W with eigenvalue R.This is the Perron–Frobenius theorem for a linear map preserving the cone of positive-semidefinite matrices (Vandergraft, 1968; Berman & Plemmons, 1994, Chapter 1).Source: §F.6 · p. 90
  • Both roots of x^2 − σx + π lie strictly inside the unit disk iff |π| < 1, 1 − σ + π > 0, and 1 + σ + π > 0.These are the Jury stability conditions for a real quadratic.Source: §F.6 · p. 90
  • Equivalently, the Jury test requires |π| < 1, p(1) > 0, and p(−1) > 0 for p(x) = x^2 − σx + π.These conditions are equivalent to both roots lying in the open unit disk.Source: §F.6 · p. 90
Loading 2610.36778v1…