PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding Stability of Outer Nesterov Momentum on Stochastic Pseudo-Gradients

Nishant Sharma, Laura Gomezjurado Gonzalez, Ganesh Talluri, Sarath Chandar, Hiroki Naganuma

arXiv:2610.22623v1cs.LG

TL;DR

Distributed training often trades frequent communication for local optimization steps whose pseudo-gradients are passed to an outer optimizer, commonly with Nesterov momentum. This paper derives a mean-square stability criterion for random quadratics and probes predicted stability dynamics at neural-network checkpoints. It finds that batch size, local-step count, and worker averaging shape stability, while checkpoint measurements follow the theoretical trends.

  • Problem

    Stability is less understood when an outer optimizer acts on a stochastic K-step pseudo-gradient, and mean-only analysis can mispredict behavior because its variance can grow with the error.

    Source: §1 · p. 2

  • Method

    The paper derives an exact mean-square criterion for random quadratics and measures checkpoint stability by tracking perturbations between runs that share minibatches.

    Source: §1 · p. 2, §6.1 · p. 10

  • Results

    The inner batch size determines whether large K is eventually stable or unstable, worker averaging divides variance by W, and stability may require a number of workers growing exponentially in K.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Checkpoint perturbations can probe mean-square stability dynamics in neural networks, and measured amplification trends are consistent with the theory’s predictions.

    Source: §6.1 · p. 10, §6.3 · p. 11

    The analysis focuses on stability rather than convergence rate, does not capture possible longer-term recovery from local instability, and is limited to SGD inner steps.

    Source: §7 · p. 11

Abstract

from arXiv · show

Large neural networks are commonly trained across many workers, and communication between them is a major cost. Workers therefore take several local optimization steps before sending their net parameter displacements, or pseudo-gradients, to an outer optimizer, which often applies Nesterov momentum. We study how the outer momentum, the number of local steps and the averaging over workers control the training dynamics, and in particular their stability. For random quadratics, we derive an exact mean-square stability criterion in any dimension, which reveals that the inner batch size determines whether increasing the number of local steps eventually stabilizes or destabilizes the outer iteration. Moreover, while averaging the pseudo-gradients of independent workers can make larger outer momentum values admissible, the number of workers required for stability can grow exponentially with the number of local steps. Finally, we show that the mean-square dynamics governing stability can be probed directly at intermediate neural-network checkpoints using shared-minibatch perturbations. We test our theoretical findings across networks, and measure their consequences for training loss.

1 Introduction

Distributed training reduces communication by letting workers take local steps and send pseudo-gradients to an outer optimizer, often with Nesterov momentum. This paper analyzes how momentum, local-step count, and worker averaging shape stability, derives a random-quadratic criterion, and tests predicted trends in neural networks.

1 Introduction

  • Workers take K local minibatch-SGD steps and send their displacements for averaging before an outer Nesterov update.The setup uses inner batch size b and learning rate η, with outer learning rate α and momentum β.Source: §1 · p. 2
  • The two-phase template includes methods ranging from local SGD and federated averaging to Nesterov-based DiLoCo and MuLoCo.These methods differ in outer updates and inner optimizers.Source: §1 · p. 2
  • For random quadratics, the paper derives a necessary and sufficient mean-square stability condition for outer Nesterov momentum on pseudo-gradients from K stochastic inner SGD steps.The analysis targets stability beyond the mean dynamics, which can mispredict behavior when pseudo-gradient variance grows with the error.Source: §1 · p. 2
  • At fixed outer learning rate, increasing momentum can only shrink the noise budget; inner batch size determines whether large K is eventually stable or unstable.Worker averaging divides variance only by W, and the number of workers required for stability can grow exponentially with K.Source: §1 · p. 3
  • Shared-minibatch perturbations measure checkpoint stability without a known minimizer across CIFAR-10 MLPs and ResNet-20s and Pythia models from 70m to 1b parameters.The experiments test the theory’s predicted trends and measure consequences of instability for training loss.Source: §1 · p. 3

2 Related Work

Related work studies stability and convergence in two-phase optimizers, but the paper identifies a gap in exact mean-square criteria for outer momentum acting on stochastic pseudo-gradients. It also distinguishes its checkpoint analysis from work focused on minima or deterministic dynamics.

2 Related Work

  • Prior theory for Lookahead and LA-DiLoCo derives second-moment dynamics, while deterministic analyses study acceleration or stable outer learning-rate and momentum ranges.Other work proves local-SGD convergence guarantees using a uniform gradient-noise variance bound.Source: §2 · p. 3
  • Mulayoff and Michaeli (2024) derive an exact mean-square threshold for minibatch SGD, whereas this paper extends that perspective to outer momentum on a stochastic trajectory.Its criterion is a spectral-radius condition on a second-moment map, and shared-minibatch perturbations extend the investigation to intermediate network checkpoints.Source: §2 · p. 3
  • Whether two-phase training with outer momentum operates at the edge of stability remains open in the cited literature.Prior edge-of-stability observations include federated averaging with α = 1 and β = 0.Source: §2 · p. 3

3 Setup

The setup treats each block as K stochastic inner SGD steps whose displacement is passed to an outer Nesterov update. For random quadratics, this trajectory induces a random block operator, and stability depends on its mean and second moment.

3 Setup

  • Each block starts at θ, runs K minibatch-SGD steps, and produces a pseudo-gradient; independently drawn blocks feed an outer Nesterov update with learning rate α and momentum β.The momentum buffer starts at zero, and the initial stability setup assumes a shared minimizer; later analysis drops that assumption while retaining the criterion for boundedness.Source: §3 · pp. 3–4

3.1 Block operator and stability

  • For quadratic minibatch losses with a common minimizer, K inner SGD steps multiply the error by a product of one-step maps, making the pseudo-gradient linear in the error.The analysis assumes independent, identically distributed block operators with finite second moments; mean-square stability requires the expected squared state norm to vanish.Source: §3.1 · pp. 4–5
  • The random block operator maps a starting-point displacement to the corresponding pseudo-gradient change, replacing loss curvature for the outer Nesterov step.For K ≥ 2 it is generally nonsymmetric, so the pseudo-gradient need not be the gradient of a loss.Source: §3.1 · p. 4

3.2 Block moments and how K and b enter them

  • Stability depends on the block operator through its mean and second moment, which vary differently with K; analyzing only the mean misses how noise changes.For minibatches of b samples, batch size leaves the one-step mean unchanged and enters the second-moment map through a noise factor of 1/b.Source: §3.2 · p. 5

4 The Stability Boundary

The second-moment dynamics separate mean propagation from block noise, yielding an exact stability test: mean stability is necessary, but the noise must also remain within a budget set by the mean dynamics. The test applies to i.i.d. block operators under Assumption 1, with a simpler eigenvalue condition for the mean transition when its block-operator eigenvalues are real.

4.1 Second-moment map and the block noise

  • The second-moment map combines propagation through the mean transition with a positive block-noise term, so noise can only raise its spectral radius above ρ(A)^2.A mean-stability analysis keeps only the mean-transition term and can therefore predict stability for a setting that is actually unstable.Source: §4.1 · p. 6

4.2 The mean condition and the noise budget

  • The noise budget is the reciprocal of the loop gain: it measures how much block noise can grow before the stability criterion fails.The loop gain combines the noise injected by one block with the later error growth under the mean dynamics.Source: §4.2 · p. 6
  • Mean-square stability holds if and only if ρ(LK) < 1, equivalently the mean transition is stable and the block-noise loop gain stays below one.Mean stability alone is necessary but not sufficient; the block noise must fit within the noise budget set by the mean dynamics.Source: §4.2 · p. 6
  • Mean stability requires every eigenvalue λ of the mean block operator to give a stable 2 × 2 transition T(λ); for real eigenvalues, the criterion simplifies.The exact mean-square criterion and its split into mean stability and a noise budget use only Assumption 1 on the block operators.Source: §4.2 · p. 6
  • The leading-order small-α criterion has the form of the heavy-ball SGD noise scale, while K = 1 and (α, β) = (1, 0) recover the minibatch-SGD criterion.At (α, β) = (1, 0), any K corresponds to plain SGD observed every K steps.Source: §4.2 · p. 6
  • The figure links four mechanisms to stability: larger β shrinks the noise budget, block noise switches from growing to dying as K crosses bc, averaging divides loop gain by W, and splitting a fixed batch raises ρ(LK,W).Green and red denote stable and unstable settings, respectively.Source: Figure 3 · p. 7

5 Implications for β, K & W

Outer momentum tightens the stability constraint from block noise, while the inner batch size determines whether long blocks stabilize or destabilize. Worker averaging can reduce noise, but splitting a fixed total batch does not improve mean-square stability.

5.1 Outer momentum shrinks the noise budget

  • At fixed outer rate, increasing momentum cannot improve mean-square stability: it raises the noise burden, so lowering momentum or adding workers can help when block noise causes instability.The noise budget is monotone in momentum, although the spectral radius of the mean dynamics need not be.Source: §5.1 · p. 7

5.2 Local steps and the inner batch size

  • For sufficiently large K, blocks are mean-square stable when the inner batch b exceeds threshold bc, and unstable when b falls below it.The threshold grows linearly with the inner learning rate at small rates; raising the per-worker batch above bc or lowering the inner rate can permit longer blocks.Source: §5.2 · pp. 7–8

5.3 Averaging over workers

  • Averaging W independent workers divides block noise by W and makes stability possible when the mean condition holds and W exceeds ρ(VK ◦Gα,β); if it fails, no worker count suffices.Under the long-block unstable branch, the worker count required for stability can grow exponentially in K.Source: §5.3 · p. 8

5.4 Fixed total batch

  • At fixed total batch and fixed outer settings, splitting the batch across workers preserves the mean block operator but yields at least as much block noise as one worker.The noise penalty is equality at K = 1 and is nearly absent when per-worker batches remain large; parallel workers can make each outer step faster.Source: §5.4 · p. 8

6 Testing the Implications on Neural Networks

The section adapts the stability analysis to neural networks by linearizing block pseudo-gradients at checkpoints and tracking shared-minibatch perturbations. Experiments test how momentum, local steps, batch size, and worker averaging affect amplification and spread cost.

6.1 The block Jacobian

  • The frozen linearization fixes each block Jacobian at checkpoint θ0 across outer steps, while inner trajectories still run on the actual loss and use Hessians at visited points.Source: §6.1 · p. 9
  • Under the frozen linearization, shared-minibatch perturbations decay exactly when ρ(LK) < 1, while the spread of independent runs stays bounded at a noise floor under that condition.The spread generically grows without bound otherwise; at actual networks, shared-minibatch differences follow this recurrence only to first order.Source: §6.1 · pp. 9–10

6.2 Amplification and spread cost

  • Amplification measures expected shared-minibatch perturbation growth over T outer steps; values above 1 indicate amplification, and worker variation helps distinguish mean from mean-square instability.Under the frozen linearization, amplification is non-increasing in W; amplification with one worker that contracts after averaging enough workers is consistent with mean-square instability.Source: §6.2 · p. 10
  • Spread cost captures loss due to disagreement among independent runs; replacing those runs with their average recovers exactly this amount.Source: §6.2 · p. 10

6.3 Empirical results

  • As local steps K increase, settings contract more at inner batches 8 and 16 but amplify more at batches 1 and 2; more workers never raise amplification.At fixed outer rate, increasing momentum changes a setting from contracting to amplifying at most once; splitting a fixed total batch did not lower amplification in the tested settings.Source: Figure 4 · p. 9, §6.3 · p. 11
  • Spread cost rises with outer momentum, falls with worker count, and is not lowered by splitting a fixed total batch.As K grows, it first rises at every inner batch, then keeps rising for b ≤4 but turns down for b = 8 and 16.Source: Figure 5 · p. 9, §6.3 · p. 11

7 Discussion

The discussion frames stability as whether block-operator noise fits within the outer step’s noise budget once the mean condition holds. It notes practical constraints, scope limits, and the Nesterov update convention used by the optimizer.

7 Discussion

  • Stability depends jointly on the block operator, shaped by local steps, inner batch size, and worker count, and the outer-step noise budget, shaped by outer learning rate and momentum.Because training constraints can limit local steps, batch size, or workers, checkpoint measurements may help assess stability before committing to a planned change.Source: §7 · p. 11
  • The analysis focuses on stability rather than convergence rate and does not capture possible longer-term recovery from local instability.Its scope is inner-step SGD; adaptive optimizers such as AdamW or Muon may carry state across blocks, changing the block Jacobian and requiring Assumption 1 to be re-examined.Source: §7 · p. 11

A The Two-Phase Optimizer with Nesterov Momentum

B Mean and Mean-Square Stability

This section develops spectral criteria for mean convergence and mean-square stability, distinguishing stability of the mean from that of second moments. It examines how outer momentum, block length, and averaging independent workers affect these stability dynamics and second moments.

B Mean and Mean-Square Stability

  • Mean stability can coexist with divergent second moments: in Figure 6, ρ(A) ≈ 0.37 while ρ(LK) ≈ 3.9.The average step shrinks the error, but fluctuations proportional to the current error compound across outer steps; all 4,000 simulated runs diverged.Source: Figure 6 · p. 16, B Mean and Mean-Square Stability · pp. 16–17

E.1 Derivation of the block moments

  • Independence across inner steps gives E[MK] = Q̄^K and E[MK ⊗ MK] = Q^K, which determine the block moments.Source: §E.1 · p. 19

E.2 Proof of the moment recurrences, the decomposition and Theorem 1(i)

  • For the i.i.d. recurrence, mean-square stability holds exactly when ρ(LK) < 1; if ρ(LK) ≥ 1, an isotropic initial error with an empty buffer is already unstable.The second-moment recurrence contracts precisely when its spectral radius is below one.Source: §E.2 · pp. 19–20

E.3 Proof of Theorem 1(ii) and (iii)

  • When ρ(A) < 1, mean-square stability additionally requires ρ(N(I − A)^−1) < 1, where N is the positive noise map.The noise map is amplified by the stable mean dynamics before its spectral radius is tested.Source: §E.3 · p. 20
  • The mean transition is stable exactly when each two-dimensional transition T(λ_i), for an eigenvalue λ_i of the mean block operator, is stable.For real λ, the characteristic polynomial is μ^2 − (1 + β)(1 − αλ)μ + β(1 − αλ), yielding condition (13).Source: §E.3 · p. 20

E.4 The criterion in one dimension

  • At x = 1, the one-dimensional noise bound is rβ = 1/(1 + 2β + 2β^2), decreasing from 1 at β = 0 toward 1/5 as β → 1.Under the mean-stability condition, mean-square stability holds exactly when NK < r(x, β).Source: §E.4 · p. 21
  • Under the mean condition, the one-dimensional recurrence is mean-square stable exactly when NK < r(x, β).The critical value is the unique stability boundary; the bound can also be expressed as an impulse-response sum.Source: §E.4 · p. 22

E.5 The criterion at small outer steps

  • At small outer steps, the multidimensional criterion reduces to ν < 1 to leading order, with loop gain asymptotic to α κK/(1 − β).This criterion applies on the complement of common flat directions when the mean block operator is symmetric.Source: §E.5 · pp. 22–23
  • In one dimension, small-step stability is ν < F(x, β): F < 1 for β ≤ 0.60278, but F > 1 when β > 0.60278 and 0 < x < xc(β).Thus the familiar ν < 1 rule can be insufficient below the threshold and overly restrictive above it; at β = 0, F = 1 − x/2.Source: §E.5 · p. 23

E.6 Proof of Proposition 2

  • If (α, β′) is mean-square stable, then every lower momentum β ≤ β′ is also mean-square stable.The gain is monotone in momentum, and the proof does not require the mean operator to have real eigenvalues.Source: §E.6 · pp. 24–25

E.7 Proof of Proposition 3

  • Lengthening a stable block eventually stops improving contraction per inner gradient, so a finite block length is best.Once the mean block operator saturates, contraction per block is fixed while each block costs K gradients.Source: §E.7 · p. 26
  • As block length K grows, NK tends to 0, 1, or infinity according as sq is below, equal to, or above one.In the marginal case sq = 1, stability for sufficiently large K holds when r(α, β) > 1 and fails when r(α, β) < 1; with W workers the limiting boundary is 1/(W r(α, β)).Source: §E.7 · p. 26

E.8 Proof of Corollary 4 and Proposition 5

  • Averaging W independent workers leaves the mean block operator unchanged and divides the variance map by W, yielding stability exactly when ρ(VK ◦ Gα,β)/W < 1.In one dimension, this is equivalent to replacing NK by NK/W in the scalar criterion.Source: §E.8 · p. 27
  • The second-moment map for batch b is ordered below that for W workers using batch b/W, implying the corresponding spectral-radius ordering.The maps share A and have ordered noise terms; the comparison uses S ≠ 0. The isotropic zero-momentum case with outer relaxation corresponds to an Agarwala (2026) result for a high-dimensional linear-regression limit.Source: §E.8 · p. 27
  • The polynomial in 1/W has non-negative coefficients and is therefore non-increasing in W and convex in 1/W; with Ψ = I it is E||z_T||², and with Ψ = EᵀE it is E||e_T||².E||e_T||² is the numerator of (19).Source: §E.8 · p. 27
  • At finite horizon, ordering positive maps orders every positive-semidefinite readout of their iterates.If 0 ⪯ M ⪯ N, then M^T ⪯ N^T, so readouts with positive-semidefinite Ψ and Σ0 are ordered at every horizon T.Source: §E.8 · pp. 27–28

E.9 Proof of Proposition 7 (no common minimizer)

For i.i.d. affine outer updates with finite second moments, Proposition 7 characterizes mean-square stability through the block second-moment map, including when sample losses have different minimizers. It also gives covariance behavior in the stable and unstable cases.

Proposition 7

  • ρ(LK) < 1 if and only if every shared-minibatch perturbation is mean-square stable; under this condition, the affine process is mean-square bounded.The proposition assumes i.i.d. pairs (Dt, pt), independent of the initial state, with finite second moments.Source: §E.9 · p. 28
  • When ρ(LK) < 1, covariance across independent streams converges to (I − LK)^−1(F∞), with F∞ determined by the limiting mean and affine forcing.If E[p] = 0, the limit is (I − LK)^−1(BSpB⊤).Source: §E.9 · p. 28
  • If ρ(LK) ≥ 1 and persistent forcing has weight on its unstable mode, tr Ct diverges at least linearly at ρ(LK) = 1 and geometrically above it.When the mean is stable, the forcing condition is ⟨Φ, F∞⟩ > 0.Source: §E.9 · pp. 28–29

Moment recurrences

  • The mean and covariance follow separate recurrences: mt+1 = Amt + BE[p], while Ct+1 = LK(Ct) + Ft with positive-semidefinite forcing.The forcing Ft is a covariance conjugated by B.Source: §E.9 · p. 28

Offset and second moment

  • With i.i.d. inner steps, the offset has zero mean at the full-loss minimizer, and its second moment contributes to the stationary covariance floor.The offset covariance is determined by the minibatch-gradient second moment together with the block quantities.Source: §E.9 · p. 29

Scope

  • When common flat directions exist, the analysis is performed on their complement; with a common minimizer, the offset vanishes.Source: §E.9 · pp. 28–29

F The Exact Model Used in the Figures

The figures evaluate the exact criterion on finite-sum quadratic models with shared minimizers. They examine stability under momentum, block length and worker averaging, and note a limitation for correlated block draws.

F The Exact Model Used in the Figures

  • The figure models use finite sets of feature vectors whose sample losses share a minimizer, with minibatches drawn with replacement.The sample curvatures are rank-one; the figures use Gaussian features with power-law covariance or features of equal norm.Source: F The Exact Model Used in the Figures · p. 29
  • The second-moment map is formed from the block moments and evaluated by dense eigenvalues, so the plotted model quantities are computed exactly.The models use an eight-dimensional rank-one setup with 1,024 samples for Figures 1, 3, 6 and 8.Source: F The Exact Model Used in the Figures · pp. 29–30
  • In the 12-dimensional regression model, the loop gain rises with momentum at each outer rate, while the second-moment map’s spectral radius need not.The stable momentum range ends where the loop gain crosses one: β = 0.85 at α = 1.1 and β = 0.575 at α = 1.4.Source: Figure 7 · p. 30
  • With b = 1, longer blocks make both settings stable when ρ(Q) = 0.906 and unstable when ρ(Q) = 1.376, although the mean dynamics remain stable.The changes are attributed to block noise; per-gradient growth returns toward one in the stable case and approaches ρ(Q)^1/2 = 1.173 in the unstable case.Source: Figure 8 · p. 30
  • The figures omit blocks with correlated draws, for which the mean block operator can have complex eigenvalues, although Proposition 2’s proof covers them.The plotted i.i.d. inner steps make the mean block operator symmetric and its spectrum real.Source: F The Exact Model Used in the Figures · p. 30
  • Table 3 compares the predicted worker count W⋆ with the second-moment map, reporting ρ(LK,W) at W⋆−1 and W⋆ workers.W⋆ is computed from the single-worker loop gain.Source: F The Exact Model Used in the Figures · p. 30, Table 3 · p. 31

G The Measures at a Checkpoint

Checkpoint measures compare perturbed or independently initialized runs to quantify finite-horizon amplification and loss from run-to-run spread. Frozen predictions apply exactly to shared-minibatch perturbations under the stated i.i.d. block-Jacobian condition, while other results are tested empirically.

G The Measures at a Checkpoint

  • The frozen linearization uses the block Jacobian at checkpoint θ0 throughout, whereas an actual run uses the Jacobian at its current parameters.Predictions and actual runs agree while block Jacobians change little, at least over short horizons.Source: G The Measures at a Checkpoint · p. 31
  • Under frozen linearization, i.i.d. block Jacobians make the moment recurrences, Theorem 1, Proposition 2 and Corollary 4(a) exact for shared-minibatch perturbations.Proposition 3, Proposition 5 and Corollary 4(b) require products of i.i.d. one-step maps; because network Hessians depend on earlier minibatches, those results are tested rather than derived at networks.Source: G The Measures at a Checkpoint · p. 31
  • Amplification measures growth over T outer steps in the full state, and with momentum it can rise temporarily before decaying.The parameter ratio and full-state ratio tend to zero when ρ(LK) < 1; the finite horizon T is therefore reported.Source: G The Measures at a Checkpoint · p. 31
  • At T = 10 rather than T = 20, only 16 of 1,404 settings change contraction-versus-amplification status, and 241 of 244 predicted orderings hold at both horizons.The study uses Aθ(T) as a finite-horizon measure of mean-square dynamics, while Theorem 1’s stability criterion is asymptotic.Source: G The Measures at a Checkpoint · p. 31
  • The initial perturbation is drawn in the range of a block operator, using an isotropic vector and a block disjoint from the stream.The buffer is empty.Source: G The Measures at a Checkpoint · p. 31
  • For independent runs from a common start, the second-order spread cost is proportional to curvature-weighted parameter covariance.It is expressed as 1/2 tr(Ht Ĉt), with Ht the full-data Hessian at the runs’ average and Ĉt their covariance about that average, normalized by the run count.Source: G The Measures at a Checkpoint · pp. 31–32

H Setup of the Network Experiments

The experiments probe trained checkpoints from CIFAR-10 networks and Pythia language models using inner SGD trajectories, with settings chosen to measure stability-related quantities across model types and checkpoints.

H Setup of the Network Experiments

  • The study measures an MLP, ResNet-20 and Pythia models at several training checkpoints.The MLP and ResNet-20 are trained on CIFAR-10; the appendix summarizes the models and checkpoints.Source: H Setup of the Network Experiments · p. 32
  • The reported checkpoints are steps 60k, 90k, 120k and 150k for the MLP; 90k, 120k and 150k for ResNet-20; and 50k, 100k and 143k for Pythia-70m.The MLP and ResNet checkpoints follow 150,000-step training; Pythia-70m was trained for 143,000 steps.Source: H Setup of the Network Experiments · p. 32
  • The CIFAR MLP and ResNet-20 use plain SGD with batch 4 and squared loss on the same 8,192 CIFAR-10 images.The MLP has widths 3,072, 512, 512 and 10; the 269,722-parameter ResNet-20 uses eight-group group normalization.Source: H Setup of the Network Experiments · p. 32
  • Pythia-70m was pretrained with Adam, but the experiments use plain SGD on its transformer blocks and final layer norm, with other parameters fixed.The measured parameters total 18.9 million of 70.4 million; batches contain sequences of 2,048 tokens and loss is evaluated on 128 held-out sequences.Source: H Setup of the Network Experiments · p. 32
  • Unless otherwise specified, experiments use K = 14, b = 4 and W = 1, with the MLP at step 120,000 in Section 6.3.Pythia uses b = 8 for the momentum panel; its worker and split panels use α = 0.5 and β = 0.6.Source: H Setup of the Network Experiments · p. 32
  • Some unstable settings diverged within 20 outer steps, so reported averages sometimes include only surviving runs or bounded frozen trajectories.For Figure 5, divergence occurred only at settings where runs also left the checkpoint basin, except one split setting where 4 of 32 runs diverged.Source: H Setup of the Network Experiments · p. 32

I Experiments across Networks and Checkpoints

The experiments track amplification and spread cost across checkpoints and model scales, then map measured contraction regions at a CIFAR MLP checkpoint. Results broadly align with the predicted trends, with some differences in Pythia-70m spread cost.

I Experiments across Networks and Checkpoints

  • The CIFAR MLP and ResNet-20 checkpoint figures report amplification and spread cost at successive training steps.The figures cover the MLP at steps 60,000, 90,000 and 150,000, and ResNet-20 at 90,000, 120,000 and 150,000.Source: Figure 9 · p. 33, Figure 10 · p. 33, Figure 11 · p. 34

I.3 Pythia-70m

  • On Pythia-70m, shared-minibatch loss gaps follow amplification: they grow with K at b = 1 and shrink at b = 4, 8 and 16.The gaps also grow for the 8 × 1 split and shrink for 1 × 8, matching the amplification trend.Source: §I.3 · p. 36, Figure 18 · p. 38, Figure 19 · p. 38
  • The Pythia-70m checkpoint figures track amplification and spread cost at 50,000, 100,000 and 143,000 steps; larger-model amplification is shown at step 100,000.The Pythia-70m loss-gap figures use pairs of runs sharing minibatches; Figure 18 specifies a 10^-5 initial separation and 20 outer steps.Source: Figure 15 · p. 36, Figure 16 · p. 37, Figure 17 · p. 37

I.4 Larger language models

  • Across Pythia-160m, Pythia-410m and Pythia-1b, small inner batches amplify more as K grows, while large batches amplify less.At long blocks, the most split total batch amplifies most.Source: §I.4 · p. 39, Figure 21 · p. 39

J Phase Diagrams at a Checkpoint

  • At the CIFAR MLP checkpoint, the contracting region expands with worker count, and its boundary shifts from larger α at low β to smaller α at high β.The frozen linearization matches the sign of measured amplification at every setting; excluding runs that left the basin, spread cost rises with α and β and falls with W.Source: J Phase Diagrams at a Checkpoint · p. 40, Figure 22 · p. 40
  • At the CIFAR MLP checkpoint, increasing K shrinks the contracting region to nothing at b = 1, barely changes it at b = 4, and preserves its size at b = 16.At b = 16 and K = 56, the boundary follows the mean boundary for β up to 0.5; the frozen linearization matches measured amplification signs at almost every setting.Source: J Phase Diagrams at a Checkpoint · p. 40, Figure 23 · p. 41

K Continuations

K Continuations show that amplification can raise CIFAR training loss through runs spreading apart, while adding a worker can reverse that behavior. In L.1, two inner laws with the same reference-point observables have opposite mean-square stability, outside the exact model of Section 3.1.

K Continuations

  • At (0.25, 0.6), both CIFAR networks’ losses keep falling and spread cost stays small with either one or two workers.The continuation uses K = 14 and b = 4.Source: K Continuations · p. 41, Figure 24 · p. 42
  • At (0.5, 0.6), one worker’s amplification raises mean CIFAR-run loss above checkpoint—threefold on the MLP—while a second worker contracts both networks and loss falls.Averaging over workers lowers spread cost at every setting.Source: K Continuations · p. 41, Figure 24 · p. 42
  • On Pythia, loss falls at every setting over 100 outer steps, while spread cost remains far below training progress.The checkpoint is at step 100,000 of 143,000.Source: K Continuations · p. 41, Figure 24 · p. 42

L.1 Mean-stability analyses

  • At α = 18/5, two scalar inner laws with the same E[Dω], E[pω], and Var(pω) yield a stationary second moment for one and mean-square divergence for the other.Their shared observables are E[Dω] = 1/2, E[pω] = 0, and Var(pω) = σ².Source: §L.1 · pp. 42–43
  • The two laws share a mean block operator and pseudo-gradient covariance, but their sample losses do not share a minimizer, placing the example outside Section 3.1’s exact model.Thus those reference-point quantities alone cannot distinguish their stability.Source: §L.1 · p. 42

L.2 Bounded-noise guarantees

The bounded-noise examples show that matching worst-case guarantees or mean behavior can coexist with different local mean-square stability. The examples trace this difference to second-moment dynamics that vary with noise and block length.

L.2 Bounded-noise guarantees

  • Both c settings have the same convex expected objective, 1-Lipschitz gradient, minimizer, and uniformly bounded gradient variance, yet their local mean-square behavior differs.Each sample loss is convex and has minimizer 0.Source: §L.2 · p. 43
  • At η = 1/8, the one-step error multiplier’s second moment is 0.78125 for c = 1 and 1.25 for c = 31, despite matching worst-case guarantees.The guarantees take the same value at any common initialization.Source: §L.2 · p. 43
  • For K-step blocks, the two problems have the same mean block operator at every K but different second moments, so their guarantees do not resolve the local stability split.The block second moment includes q_c^K, which varies with c.Source: §L.2 · pp. 43–44
  • At K = 8, c = 31 has ρ(G) = 1.196 versus 0.990 at K = 4, while c = 1 falls from 0.637 to 0.545; only c = 31 becomes unstable.At K = 4, both problems are stable under α = 0.15, β = 0.6, and η = 1/8.Source: §L.2 · p. 44
  • The instability is local: noise saturates away from the minimizer, so it does not contradict a guarantee allowing an error level that depends on noise.Source: §L.2 · p. 44
Loading 2610.22623v1…