PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Quadratic Loss: The Stability Phase Diagram of Adam

Gaoxiang Tang, Huanran Chen, Ziming Liu

arXiv:2609.18314v2cs.LGmath.OCstat.ML

TL;DR

The paper investigates how Adam’s two momentum timescales govern macroscopic loss spikes, a relationship that remains unclear. It maps stability boundaries across models and toy losses, finding near-linear scaling in neural-network settings and connecting it to finite-scale superquadratic geometry in confident cross-entropy landscapes.

  • Problem

    How Adam’s first- and second-moment timescales govern macroscopic loss spikes remains unclear despite reported links between spikes and optimizer dynamics.

    Source: §1 · pp. 1–2

  • Method

    The authors map Adam stability across β1 and β2, analyze quadratic and superquadratic one-dimensional losses, and inspect finite-scale loss landscapes along the preconditioned gradient.

    Source: §1 · p. 2

  • Results

    Neural-network settings show near-linear boundaries versus approximately cubic quadratic-loss scaling; superquadratic toy-loss boundaries depend on n, and confident cross-entropy has a core–wall landscape.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The effective exponent at the optimizer’s update scale predicts the empirical boundary coefficient, linking finite-scale landscape geometry to Adam’s beta phase diagram.

    Source: §5 · p. 8

    A complete analytical explanation of the empirical left boundary remains open.

    Source: §3 · p. 5

Abstract

from arXiv · show

Loss spikes are recurrent instabilities in neural-network training and can arise from multiple mechanisms. For Adam in particular, macroscopic loss spikes have been linked to optimizer dynamics, yet how its two momentum timescales govern them remains unclear. We investigate this dependence by mapping training dynamics across the $(\beta_1,\beta_2)$ plane. Across a range of model--task settings, an approximately linear boundary, $1-\beta_2=C(1-\beta_1)$, separates spiky from non-spiky dynamics, whereas a one-dimensional quadratic loss produces approximately cubic slope. A one-dimensional superquadratic loss $L(x)\propto|x|^n$ recovers the near-linear scaling and links the boundary coefficient to the effective loss exponent $n$. We further show that confident cross-entropy losses develop a core--wall landscape comprising a narrow quadratic core followed by a steep wall, which produces effective superquadratic behavior at the scale of an optimizer update. Together, these results connect Adam loss spikes to both the mismatch between momentum timescales and finite-scale superquadratic loss geometry beyond the Hessian.

1 Introduction

The introduction asks how Adam’s two momentum timescales relate to macroscopic loss spikes, finding near-linear boundaries in neural-network settings but approximately cubic scaling for a quadratic toy loss. It connects this difference to superquadratic loss geometry at finite optimizer-update scales.

1 Introduction

  • Adam loss spikes have multiple reported mechanisms, including effective-stability-threshold violations and mechanisms specific to low-precision or scale-invariant-layer settings.The introduction distinguishes Adaptive Edge of Stability dynamics from Numerical Feature Inflation and weight-norm criticality.Source: §1 · pp. 1–2
  • Across diverse models, the spiky/non-spiky boundary is near-linear in the memory gaps, unlike the approximately cubic boundary for a one-dimensional quadratic loss.At fixed β1, the spiky phase lies on the larger-β2 side of the boundary.Source: Figure 1 · p. 1, §1 · p. 2
  • For L(x) = |x|^n with n > 2, the spiky phase is wedge-shaped; its right boundary has C_R = 2(n −2)/n and its empirical left boundary has C_L ≃12(n −2)/(3n −4).The boundary observed in real models corresponds to the toy model’s left boundary.Source: §1 · p. 2
  • At finite scales, directional loss landscapes have a quadratic core, steep superquadratic wall, and outer rollover, with update-scale effective exponents capturing observed boundary coefficients.As the core contracts during training, a typical update can reach the wall.Source: §1 · p. 2

2 Stability Phase Boundaries Across Six Models

The section maps Adam’s phase boundaries across six model–task settings while varying β1 and β2. The fitted boundaries are near-linear, and the region above them generally contains no detected macroscopic spikes.

2 Stability Phase Boundaries Across Six Models

  • The six settings use Transformers, a character-level Transformer, an MLP, a CNN, and an autoencoder; within each setting, only β1 and β2 vary.Architecture, data, random seed, and other optimizer hyperparameters remain fixed.Source: §2 · p. 2
  • The study scans logarithmically spaced memory gaps and identifies macroscopic oscillations using qualifying low-frequency spectral peaks after extending trajectories with unresolved periods.The phase diagrams encode detected oscillation periods in optimizer steps.Source: §2 · pp. 2–3, Figure 2 · p. 3
  • p = 0.92, 0.84, 0.75, 0.95, 1.05, 0.95 across panels (a–f), respectively, show near-linear boundaries; larger γ/δ generally lies in the no-detected-spike region.The panels cover modular division, modular addition, Tiny Shakespeare, a CIFAR-10 MLP, a CIFAR-10 CNN, and an MNIST autoencoder.Source: §2 · p. 3, Figure 2 · p. 3

3 One-dimensional toy model of Adam stability

The one-dimensional model separates quadratic and superquadratic stability: the quadratic loss has an approximately cubic boundary, while superquadratic losses exhibit a wedge-shaped spiky region with near-linear boundaries. The two boundaries arise from different dynamics, and only the empirical left boundary matches those observed in neural networks.

Boundary mechanisms

  • For one-dimensional losses L(x)=k|x|^n/n with k=1, the empirical left boundary is γ=12(n−2)δ/(3n−4), while the predicted right boundary is γ=2(n−2)δ/n; both coefficients vanish at n=2, where the boundary is approximately cubic.Source: Figure 3 · p. 4

One-dimensional stability phases

  • For n = 2, the boundary is approximately γ ≃15.6δ^3.07; for n > 2, the spiky region is wedge-shaped between approximately linear boundaries.The quadratic boundary remains approximately cubic across tested k and η, and across ϵ when ϵ ≪√ˆvt.Source: §3 · p. 4, Figure 3 · p. 4

Boundary mechanisms

  • For neural networks, the observed boundary corresponds to the left boundary; their quadratic core prevents the asymptotically superquadratic right-boundary mechanism from applying.The left-boundary relation captures the measured trend with n, but its full analytical explanation remains open.Source: §3 · pp. 4–5

Trajectory regimes

  • Different trajectories yield suppression, repeated macroscopic spikes, or short-period microscopic oscillations, depending on how ηeff evolves relative to ηcrit.Repeated threshold crossings trigger spikes, whose large gradients replenish vt and lower ηeff.Source: §3 · pp. 4–5, Figure 4 · p. 5

Right-boundary analysis

  • For superquadratic losses, the right-boundary theorem says trajectories eventually cease exhibiting EoS-triggered spikes when its parameter condition holds.The limiting recurrence’s nonzero fixed points are stable exactly when b < 1.Source: §3 · p. 5

4 Core–wall landscape and its exponential-sum mechanism

Finite-scale directional slices reveal a quadratic core, superquadratic wall, and outer rollover, rather than a landscape fully described by its infinitesimal curvature. Confident cross-entropy naturally produces this structure, and the wall exponent at optimizer-update scale links the landscape to the empirical Adam phase boundary.

Finite-scale landscape

  • At pre-spike points across all six model–task settings, directional slices show a core–wall landscape whose quadratic core contracts toward the optimizer’s update scale.This finite-scale sharpening accompanies increasing preconditioned sharpness.Source: §4 · p. 6
  • The core–wall landscape is defined by a quadratic core, a superquadratic wall beginning at neff = 2.5, and an outer rollover where neff decreases.The slices are evaluated along the Adam-preconditioned gradient direction.Source: §4 · p. 6

Exponential-sum mechanism

  • Confident cross-entropy’s large margins suppress baseline loss, while large directional slopes make its quadratic approximation accurate only within O(1/σa).Under a Gaussian model for directional slopes, the quadratic-core scale is O(1/σa).Source: §4 · p. 7

Measured and synthetic slices

  • For the measured slice, neff is approximately 2.18 in the core and 10.90 along the wall; the transition near 10^-6 is comparable to the 6.3 × 10^-6 update scale.The rollover begins near 10^-5; the synthetic confident cross-entropy slice shows the same three regimes.Source: §4 · p. 7

Exponential-sum mechanism

  • Exponential reweighting favors increasingly extreme directional slopes beyond the core, producing a superquadratic wall that admits a finite-interval power-law approximation.At larger displacements, exact cross-entropy’s logarithm reduces the effective exponent and produces the outer rollover.Source: §4 · p. 7

Connecting geometry to the phase diagram

  • The measured wall exponent is compared with the empirical left-boundary relation because the neural-network quadratic core cuts off the pure-superquadratic right-boundary mechanism.The pure power-law right-boundary mechanism requires the loss to remain superquadratic as s → 0.Source: §4 · p. 8

5 Conclusion and outlook

The paper connects Adam’s near-linear spike boundary to first- and second-moment timescales and finite-scale superquadratic loss geometry. It concludes that confident cross-entropy’s core–wall structure links the effective exponent at update scale to the empirical phase boundary.

Conclusion

  • Across diverse models and tasks, Adam’s spiky and non-spiky dynamics are separated by a near-linear phase boundary.The conclusion identifies this as a paper-level finding.Source: §5 · p. 8
  • The first- and second-moment memory timescales set phase boundaries, while superquadratic geometry accounts for the shift from cubic to linear scaling.The effective exponent measured at optimizer-update scale predicts the empirical boundary coefficient.Source: §5 · p. 8

Conclusion and outlook

  • Confident cross-entropy produces a narrow quadratic core, superquadratic wall, and outer rollover, connecting finite-scale landscape geometry to the beta phase diagram.The outlook asks whether optimizers can estimate the evolving exponent at update scale and adapt learning rates or moment timescales.Source: §5 · p. 8

6 Related work

Prior work links stability thresholds, optimizer dynamics, and several setting-specific mechanisms to loss spikes. This paper’s related-work discussion distinguishes its phase-boundary and core–wall focus, while the appendix derives Adam’s EoS instability condition.

6 Related work

  • Prior studies connect quadratic and adaptive Edge-of-Stability dynamics to loss spikes, while other accounts target low-precision training, weight decay, or moment-timescale effects.The paper contrasts these mechanisms with its one-dimensional examples linking beta-boundary scaling to competing moment timescales and finite-scale superquadratic geometry.Source: §6 · pp. 8–9
  • The paper distinguishes its core–wall landscape, measured along a dynamically selected direction, from basin concepts used to study fine-tuning robustness.It also notes that parameter-space sharpness is not invariant to function-preserving reparameterizations.Source: §6 · p. 9
  • Prior non-quadratic-curvature work studies unstable convergence, subquadratic landscapes, and cubic self-stabilization, but not the directional O(1/σa) scale derived here for confident cross-entropy.The related-work discussion also distinguishes generalized Gauss–Newton and generalized self-concordance analyses from this finite-step scale.Source: §6 · p. 9

A.1 Proof of Lemma 1

  • The appendix derives Adam’s EoS instability condition by applying a stability condition to the largest eigenvalue of the preconditioned Hessian, including first-moment momentum.The derivation freezes the local preconditioner and Hessian and analyzes perturbations along a Hessian eigen-direction.Source: §A.1 · p. 13

A.2 Proof of Theorem 3

The proof reduces Adam’s dynamics on a superquadratic loss to a limiting recurrence and establishes when its nonzero fixed points are stable. Summable time-dependent perturbations preserve that local stability boundary, after which attracted trajectories cease crossing the EoS threshold.

A.2 Proof of Theorem 3

  • For bounded rescaled trajectories, V_t converges to a finite positive limit, with deviations and recurrence-coefficient errors decaying at a summable geometric rate.The proof uses V_t → V∞ > 0 and O(q^(2t)) convergence to control the time-dependent recurrence.Source: §A.2 · pp. 13–14
  • The nonzero fixed points exist when bq < 1, and their local stability requires b < 1.At b = 1 the roots lie on the unit circle; immediately beyond it, at least one root has modulus greater than one.Source: §A.2 · p. 14
  • The summable perturbation from gradient replenishment preserves the limiting recurrence’s local stability boundary for trajectories entering its basin sufficiently late.The resulting trajectories eventually remain below S_t = 1 and have no further threshold crossings.Source: §A.2 · pp. 13–14
  • For b < 1, trajectories attracted to either nonzero fixed point satisfy S_t → S* < 1 and eventually cease crossing the EoS threshold.The strict linear-stability boundary is b = 1, corresponding to β2 = β1^(2(n−2)/n).Source: §A.2 · p. 14

A.3 Proof of Lemma 5

The proof uses a generalized Gauss–Newton decomposition to characterize cross-entropy curvature. In the confident regime, binary cross-entropy has an exponential-sum form, while imperfect reconstruction can leave a nonzero-curvature core.

A.3 Proof of Lemma 5

  • For confident binary cross-entropy, locally linearized signed margins yield an exponential-sum loss form, and the Hessian-collapse argument extends to sample–coordinate pairs.This extension requires sigmoid predictions to approach their binary targets and the corresponding logit derivatives to remain bounded.Source: §A.3 · p. 15
  • If reconstruction remains imperfect so some pixel probabilities do not approach their targets, the vanishing-Hessian conclusion need not hold and a nonzero-curvature core may remain.Source: §A.3 · p. 15

A.4 Proof of Theorem 7

The proof uses a Gaussian model of the CE-weighted slope distribution to characterize confident-loss behavior. It identifies a quadratic core whose scale is set by the slope spread and extends the derivation to coordinate-wise BCE.

A.4 Proof of Theorem 7

  • The ratio of confident loss to its value at zero is treated as the moment-generating function of the CE-weighted slope distribution, then approximated with a Gaussian model.The supplied proof passage ends before stating the result of that replacement.Source: §A.4 · p. 15
  • The proof introduces an effective log–log exponent.Source: §A.4 · p. 15
  • The loss is quadratic for |u| ≪ 1 and departs from that form at |u| = O(1), giving a quadratic-core scale of score = O(1/σa).With neff = 2.5, the core boundary is |ucore| ≃ 0.964.Source: §A.4 · p. 15
  • The same derivation applies to confident coordinate-wise BCE because its signed-margin approximation has the same exponential-sum form.Source: §A.4 · p. 15

B.1 Experimental setup for the six models

The six AdamW beta-plane scans use seed 42 and adaptively extend unresolved runs; period detection analyzes the post-transient loss trace. The appendix also tests how quadratic-loss boundary fits respond to changes in k, η, and ϵ.

B.1.1 Modular division

  • Modular division uses a one-layer causal Transformer with 211,456 parameters, trained on 2,756 examples split evenly into training and test sets.It uses cross-entropy and AdamW; unresolved runs are extended from 8,000 steps up to 30,000.Source: §B.1.1 · p. 16

B.1.2 Modular addition

  • Modular addition uses the same one-layer Transformer architecture, with 2,809 examples split into 1,404 training and 1,405 test examples.Its AdamW runs start at 8,000 steps and may be extended to 30,000.Source: §B.1.2 · p. 16

B.1.3 Tiny Shakespeare

  • Tiny Shakespeare uses a 65,472-parameter character Transformer trained on fixed windows from a 1,003,854-character corpus.The base run is 12,000 steps, with selected unresolved runs extended as far as 100,000.Source: §B.1.3 · p. 16

B.1.4 CIFAR-10 MLP

  • The CIFAR-10 MLP has five 512-unit hidden layers and is trained full-batch on 200 randomly selected training images.Its runs start at 4,000 updates and are adaptively extended to at most 16,000.Source: §B.1.4 · pp. 16–17

B.1.5 CIFAR-10 CNN

  • The CIFAR-10 CNN is a VGG11-style model with 9,750,922 trainable parameters, evaluated on the same 200-image training and test samples as the MLP.Its scan begins at 8,000 updates and may be extended to 30,000.Source: §B.1.5 · p. 17

B.1.6 MNIST autoencoder

  • The MNIST autoencoder has 54,425 trainable parameters and reconstructs binarized images from 64 randomly selected training examples.Unresolved runs may be extended to 200,000 steps.Source: §B.1.6 · p. 17

B.2 Period-detection algorithm

  • Period detection discards the first 20% of the trace, smooths the log-transformed loss, and selects the accepted spectral peak with greatest power; runs without an accepted peak have no detected period.The CIFAR-10 MLP is exempt from the final-loss threshold used for the other models because dead ReLU units can keep its training loss above the threshold.Source: §B.2 · pp. 17–18

C.1 Sensitivity of the Adam boundary for L(x) = kx2/2

  • Changing k by one order of magnitude leaves the quadratic-loss boundary nearly unchanged; η preserves approximately cubic scaling with a moderate prefactor shift, while ϵ robustness requires ϵ to remain small relative to √ˆvt.If ϵ dominates the denominator, the expected boundary scaling is not obtained.Source: §C.1 · p. 18, Figure 6 · p. 18, Table 2 · p. 18

C.2 Effect of weight decay in the one-dimensional toy model

The shifted one-dimensional loss separates the task-loss minimum from the parameter origin, allowing weak decoupled weight decay to perturb rather than dominate Adam dynamics. The near-linear left boundary persists, but the closing boundary of the no-decay wedge does not.

C.2 Effect of weight decay in the one-dimensional toy model

  • The weight-decay update is −ηλwdθ; in high dimensions, its direction tends to have small overlap with the Adam-preconditioned gradient direction.The cited passage reports this as a tendency, not a universal property.Source: §C.2 · p. 18
  • The model shifts the minimum so weight decay toward the origin no longer points toward the task-loss minimum, unlike in the centered one-dimensional loss.The decay strength is chosen to perturb rather than dominate Adam dynamics.Source: §C.2 · pp. 18–19
  • Weak shifted weight decay preserves the approximately unit-slope left boundary and modestly enlarges the periodic region, but the lower/right boundary closing the no-decay wedge disappears over the scanned range.Source: §C.2 · p. 19, Figure 7 · p. 19

D.1 Locating the pre-spike points and progressive sharpening at finite scales

Pre-spike states are selected from completed low-loss valleys, then frozen directional slices are used to examine landscape shape at the optimizer’s update scale. The fitting procedure estimates a shared wall exponent from event-side branches and averages across beta settings.

D.1 Locating the pre-spike points and progressive sharpening at finite scales

  • Each completed temporal valley contributes at most one pre-spike checkpoint, selected from the loss trace alone; an unclosed final valley is right-censored and excluded from the exponent fit.Deterministic replay reconstructs the optimizer state when the selected checkpoint was not retained.Source: §D.1 · p. 19
  • Across examined models and beta settings, selected points occur immediately before a spike; the modular-division example marks loss-curve minima as pre-spike points.Source: §D.1 · p. 20, Figure 8 · p. 20
  • At the modular-division pre-spike point, a typical Adam update of 6.3 × 10−6 exceeds the core transition distances of 2.16 × 10−6 and 3.55 × 10−6, reaching the superquadratic wall.In the quiet phase, transition distances exceed the typical update by more than three orders of magnitude.Source: Figure 9 · p. 20, §D.1 · p. 21

D.2 Fitting and aggregating the wall exponent

  • For each event-side branch, the fit separates a quadratic core, steep wall, and outer rollover; joint refitting shares core and wall exponents within each beta setting.The middle regime must be steeper than both the core and rollover, preventing the outer rollover from being counted as wall behavior.Source: §D.2 · p. 21
  • The main-text exponent is the unweighted arithmetic mean of six beta-level estimates, so each beta pair receives equal weight despite differing numbers of usable valleys.The boundary coefficient is computed after this averaging step.Source: §D.2 · p. 21
Loading 2609.18314v2…