PaperlayerView exact revision on arXiv

Source-linked AI summary

A Theory on Adam Instability in Large-Scale Machine Learning

Igor Molybog, Peter Albert, Moya Chen, Zachary DeVito, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, Binh Tang, Diana Liskovich, Puxin Xu, Yuchen Zhang, Melanie Kambadur, Stephen Roller, Susan Zhang

arXiv:2304.09871v2cs.LGcs.AImath.OC

TL;DR

Large-scale language-model training exhibits instability whose causes and remedies are not well established. The paper studies the phenomenon across model scales and develops a theory linking it to time-domain correlations in gradient estimates; it concludes that remedies depend on the training setup.

  • Problem

    Loss spikes in large-scale language-model training threaten further scaling, and their cause is not established by the prior observations described here.

    Source: §1 · p. 1

  • Method

    The paper reproduces instability across large-model training runs and develops a theory connecting it to time-domain correlation in gradient estimates of earlier layers.

    Source: §1 · p. 2, §9 · p. 16

  • Results

    Instability was reproduced at 546 billion parameters, was moderate at 65 billion, and did not appear as explosive divergence at 30 billion or 7 billion; the theory associates it with gradient-estimate correlation in earlier layers.

    Source: §1 · p. 2, §9 · p. 16

  • Takeaways & Limitations

    The authors propose several mitigation approaches but conclude that no single remedy fits every large-scale training setup.

    Source: §9 · p. 16

    Batch skipping works well mainly earlier in training, is operationally costly, and becomes less practical as spike frequency increases later.

    Source: §7.1 · p. 15

Abstract

from arXiv · show

We present a theory for the previously unexplained divergent behavior noticed in the training of large language models. We argue that the phenomenon is an artifact of the dominant optimization algorithm used for training, called Adam. We observe that Adam can enter a state in which the parameter update vector has a relatively large norm and is essentially uncorrelated with the direction of descent on the training loss landscape, leading to divergence. This artifact is more likely to be observed in the training of a deep model with a large batch size, which is the typical setting of large-scale language model training. To argue the theory, we present observations from the training runs of the language models of different scales: 7 billion, 30 billion, 65 billion, and 546 billion parameters.

1 Introduction

The paper reproduces large-scale training instability under a different hardware and software setup and investigates its causes. Instability was strongest at 546 billion parameters, moderate at 65 billion, and absent as explosive divergence at 30 billion and 7 billion.

1 Introduction

  • Earlier reports described dozens of loss spikes during training on roughly a trillion tokens, with spikes not recurring after recent batches were skipped.The authors note that replaying skipped data did not reproduce the spikes, suggesting an interaction between the batch and the training-run state rather than a data-only cause.Source: §1 · p. 1
  • The study analyzes four models from 7 billion to 546 billion parameters because smaller-scale reproductions may not share the causes or remedies of large-scale spikes.The runs used AdamW with a linear learning-rate schedule and modified architecture and optimizer settings relative to the earlier work.Source: §1 · p. 1
  • At 546 billion parameters, the loss curve again showed prominent spikes under a different hardware, codebase, and dataset setup.Most spikes recovered without batch skipping within about a dozen iterations; 65-billion-parameter training showed moderate instability, while the 30-billion- and 7-billion-parameter runs did not show explosive divergence.Source: §1 · p. 2, Figure 1 · p. 2
  • The authors attribute the nature of instability to Adam, while noting that severity depends on data and architecture and that tests with another learning algorithm are still needed.At 65 billion parameters, a cross-check using a different codebase and dataset showed no significant differences in results.Source: §1 · p. 2

2 Prerequisites

This section establishes the notation and describes the Adam training procedure analyzed in the paper. Adam updates model parameters using stochastic gradients estimated from batches, with optimizer-state vectors and hyperparameters governing the update.

2.1 Notations

  • For a vector x, x[G] selects entries indexed by G; for a matrix, x[G,t] selects those entries at iteration t.Vector and matrix operations such as square root are coordinate-wise, and E and D denote mean and variance.Source: §2.1 · p. 3

2.2 Adam algorithm

  • The training setup minimizes empirical risk for a model using a transformer architecture and an Adam variant.In this language-modeling context, cross-entropy loss is tied to perplexity.Source: §2.2 · p. 3
  • Adam updates parameters using a batch-based stochastic gradient estimate, whose variance is controlled in part by batch size.The paper focuses on the original Adam algorithm; its first- and second-moment vectors form the optimizer state, ε is a stability hyperparameter, and η_t is the learning rate.Source: §2.2 · p. 3

3 Assumptions underlying Adam efficacy

The section examines assumptions that help explain Adam's efficacy, focusing on time-domain independence of gradient estimates. Under additional local approximations, Adam's update can resemble a Newton-style update, while the analysis relates its second-moment estimate to expected gradient outer products.

3 Assumptions underlying Adam efficacy

  • Adam's efficacy is analyzed under the assumption that gradient estimates are independent over time, an assumption that may be subtly violated during optimization.The section motivates this by examining the expected gradient outer product and the covariance of model weights near a reference point.Source: §3 · p. 4
  • The second-moment vector v_t can approximate the diagonal of the expected gradient outer-product matrix over recent model weights.The time window depends on β2; the argument considers a covariance decomposition into a scaled identity and a residual, with the scaled identity dominating.Source: §3 · p. 4
  • If m_t approximates the gradient and v_t is much larger than ε, Adam's update approximates a Newton-style update up to a constant multiplier.This approximation is presented as potentially relevant to Adam's convergence and empirical success; the argument also assumes a nearly constant Hessian and considers a primarily diagonal Hessian.Source: §3 · pp. 4–5

4 The assumption of time-domain independence between gradient estimations

The section links Adam’s update distributions to whether gradient-estimation components are independent across parameters and time. Independent estimates yield bell-shaped updates, while strong time-domain correlation can produce bimodal updates associated with loss stalls or divergence.

4 The assumption of time-domain independence between gradient estimations

  • Time-domain independence requires gradient-estimation components to be approximately independent across both parameters and iterations, supporting nearly independent parameter dynamics.This corresponds to small off-diagonal covariance components over recent model weights.Source: §4 · p. 5

4.1 Manifestation of independent gradient estimates

  • With independent gradient estimates, update components approach a bell-shaped Gaussian distribution, consistent with the similar layerwise distributions observed during healthy convergence.The argument extends beyond Bernoulli gradients to distributions close to independent components.Source: §4.1 · pp. 5–6, Figure 2 · p. 6

4.2 Manifestation of time-domain correlated gradient estimates

  • Strong time-domain correlation produces a symmetric bimodal update distribution with modes centered near ±β1√β2, rather than the bell-shaped distribution associated with low correlation.The predicted distribution does not depend on the scale of gradient values.Source: §4.2 · p. 7
  • In observed snapshots, most layers remained bell-shaped, but bimodal layers became more numerous as larger models trained and were not detected in the 7b and 30b models without perplexity spikes.The authors associate severe time-domain correlation in even one layer with convergence properties of the entire model.Source: §4.2 · p. 7, Figure 3 · p. 7

5 Time-domain correlation leads to divergence of Adam

The section argues that time-domain correlation can make Adam diverge, even with exact gradients, and that the learning rate required for convergence may become impractically small as model dimension grows. It concludes that time-domain incoherence is necessary for Adam convergence.

5 Time-domain correlation leads to divergence of Adam

  • With exact, fully correlated gradients, the Adam update becomes zero and the stated condition for a loss-decreasing step is not satisfied.Source: §5 · p. 8
  • The analysis treats the update direction as unrelated to the Hessian spectrum and invokes near-orthogonality between a random vector and its sign.Source: §5 · pp. 7–8
  • When gradient and Hessian entries are assumed to have the same order, convergence requires η_t to scale as 1/n, an unrealistic learning-rate scale for large-scale machine learning.The section derives the condition for exact but correlated gradient estimates.Source: §5 · pp. 7, 9
  • Time-domain incoherence between gradient estimates is necessary for convergence of Adam, as well as for treating 1/√v_t as an efficient estimate of the Hessian inverse diagonal.Source: §5 · p. 9

6 Theory of the loss instability

The proposed account links vanishing updates in parameter groups to time-correlated gradient estimates and a transition in Adam’s update distributions. A rare large gradient event can then trigger an update-driven loss explosion, while recovery is not guaranteed.

6 Theory of the loss instability

  • In deeper and larger models, updates in low-gradient layers concentrate near zero, with the concentration becoming more severe as model size grows.The ratio distribution can remain bell-shaped early in training, even after updates become spiked, before shifting toward a bimodal form.Source: §6 · pp. 9, 11
  • Vanishing updates and large batches are associated with strongly time-correlated gradient estimates, while the ratio distribution shifts from unimodal toward bimodal.The proposed process begins when a group of parameters, often in earlier layers, has gradients much smaller than ε.Source: §6 · pp. 11–12
  • A rare gradient increase in a group with vanishing updates can make its updates approach the bimodal ratio distribution, triggering a divergence-like chain reaction and a gradient-norm rise before perplexity spikes.The account predicts the update distribution moves toward the ratio distribution, while gradient norms in earlier layers begin rising one step before perplexity.Source: §6 · pp. 9, 12
  • As gradient estimates rise above ε, updates become increasingly like the ratio distribution and unimodal again; training may recover quickly for a small affected group, but can also diverge.The account attributes the update’s sensitivity near zero to the large derivative 1/ε, so small gradient-estimate changes can disproportionately change update magnitude.Source: §6 · p. 12

6.1 Snapshots of the large model training run in the unstable regime

  • The 546-billion-parameter snapshots mark three distinct spike stages: t1 is near the pre-explosion loss, t2 marks onset, and t3 follows the explosion when loss stays flat.The snapshots examine the embedding layer’s update and ratio distributions alongside perplexity; other layers also show multimodal update distributions.Source: §6.1 · pp. 12–13
  • The theory does not explain why some spikes recover and others do not; the authors hypothesize that earlier-layer feature representations change more when loss fails to drop.Similar update and ratio distribution shapes appear at different phases of a perplexity spike, leaving the recovery distinction outside the proposed theory’s scope.Source: §6.1 · p. 13

7 Discussion

The discussion links instability to changes in Adam update distributions and the gradient estimates of earlier layers, while emphasizing that no single mitigation fits every training setup. It reviews several remedies whose effectiveness or costs depend on training stage, precision, batch size, and optimizer settings.

7 Discussion

  • A shift in r[:, t] from bimodal to unimodal and back can occur without a prominent perplexity spike when gradient-estimation components remain below ε.A norm increase alone does not imply a loss spike; the discussion says a spike would be expected if r[G, t] were mostly unimodal across all layers at that time.Source: §7 · p. 13
  • A bell-shaped distribution of r[:, t] can produce a gradient-norm spike without triggering the chain reaction, whereas a rare jump above ε can initiate it when the optimizer state is severely bimodal.Skipping the batch that caused the large gradient can avoid that rare event.Source: §7 · p. 15
  • Smaller or shallower models have near-normal update-component distributions, while deeper and larger models quickly develop layers with distributions spiked at zero.Small gradient estimates are more common in earlier layers, but the vanishing effect is not monotone by layer; model size, depth, and batch size correlate with instability proneness.Source: §7 · p. 15

7.1 How can the loss spikes be mitigated?

  • Batch skipping can recover the loss curve mainly earlier in training, but requires manual monitoring and rollback resources as spikes become more frequent later.The later increase occurs when gradient-estimation entries for most layers fall below ε.Source: §7.1 · p. 15
  • Lowering the learning rate helps theoretically and in practice but lengthens training; lowering ε conflicts with popular low-precision arithmetic, while reducing batch size is expected to reduce spike frequency but can underuse GPU acceleration.The authors also suggest setting ε to zero and mapping updates to zero when v_t[i] is zero as a possible way to avoid numerical issues and bimodal ratios.Source: §7.1 · p. 15
  • Reducing β1 and β2 would require gradient-estimate correlation to persist longer before severe bimodality appears, but makes updates less current during normal training.The change increases averaging across time steps.Source: §7.1 · p. 15
  • In smaller-scale experiments, high-quality, less varied text data reduced the number of layers with diminishing gradient values to zero.Source: §7.1 · p. 15
  • Monitoring ratio-distribution modality and lowering ε or reinitializing optimizer state when its shape changes prevented update bimodality in initial high-precision experiments.The dip statistic is offered as one possible modality measure.Source: §7.1 · p. 16

8 Related Work

The related work distinguishes rapid, spiky instability in large-scale training from previously studied forms of Adam divergence. Prior analyses offer convergence conditions or alternative algorithms, but address different dynamics or do not directly explain the phenomenon considered here.

8 Related Work

  • Reddi et al. [2019] and Wang and Klabjan [2022] study drifting divergence, rather than the rapid spiky divergence addressed here.Reddi et al. give divergent examples and propose AMSGrad and AdamNC; Wang and Klabjan identify gradient-estimate variance and propose variance-reduced Adam that converges under certain assumptions.Source: §8 · p. 16
  • Chen et al. [2018] identify sufficient Adam convergence conditions on dynamic quantities, while Zou et al. [2019] give hyperparameter-based convergence-rate guarantees that fail for commonly used Adam settings and discuss drifting divergence.Zaheer et al. [2018] proposed Yogi because Adam updates’ dependence on past gradient estimates decays exponentially with the number of steps.Source: §8 · p. 16
  • Zhou et al. [2018] attribute Adam divergence to negative correlation between gradient-estimate scale and learning rates, but the authors observe update-distribution shape independent of gradient scale when estimates greatly exceed ε.Source: §8 · p. 16

9 Conclusion

The paper associates large-scale training-loss instability with time-domain correlation between gradient estimates in earlier layers, proposes several mitigation approaches, and concludes that the appropriate remedy depends on the training setup.

9 Conclusion

  • The paper associates large-scale training-loss instability with time-domain correlation between gradient estimates in earlier layers, proposes several mitigations, and concludes that no single remedy suits every setup.Source: §9 · p. 16
Loading 2304.09871v2…