Summary (Overview)

• This paper investigates how Adam's two momentum timescales ($\beta_1$ and $\beta_2$) govern macroscopic loss spikes during neural network training, mapping training dynamics across the $(\beta_1, \beta_2)$ plane.
• The authors identify an approximately linear boundary $1-\beta_2 = C(1-\beta_1)$ separating spiky from non-spiky dynamics across diverse models, contrasting with the approximately cubic scaling found for one-dimensional quadratic losses.
• A one-dimensional superquadratic loss $L(x) \propto |x|^n$ recovers the near-linear scaling, with boundary coefficients linked to the effective loss exponent $n$ via $C_L = 12(n-2)/(3n-4)$ and $C_R = 2(n-2)/n$.
• The paper shows that confident cross-entropy losses develop a core-wall landscape (narrow quadratic core, steep superquadratic wall, outer rollover), producing effective superquadratic behavior at the scale of an optimizer update.
• Key theoretical contribution: a rigorous derivation of the stability boundary for superquadratic losses via the quiet-phase approximation, connecting Adam loss spikes to both momentum timescale mismatch and finite-scale loss geometry.

---

Introduction and Theoretical Foundation

Loss spikes are abrupt, macroscopic loss excursions that recur across neural network training settings. Under a local quadratic approximation, gradient descent is stable when $\lambda_{\max}(H_t) < 2/\eta$, and crossing this threshold induces Edge-of-Stability (EoS) dynamics characterized by non-monotonic loss excursions and spikes. For adaptive optimizers, the analogous criterion uses a preconditioned Hessian, leading to the Adaptive Edge of Stability for Adam and RMSProp.

Adam's first and second gradient moments have memory timescales set by $\beta_1$ and $\beta_2$ respectively. This paper investigates how the mismatch between these timescales governs loss spikes, focusing on macroscopic spikes triggered by violations of Adam's EoS condition rather than microscopic oscillations.

The central puzzle is the observed discrepancy between:

• Quadratic loss model: approximately cubic boundary $1-\beta_2 \propto (1-\beta_1)^3$
• Neural network training: near-linear boundary $1-\beta_2 \propto (1-\beta_1)$

This motivates analyzing superquadratic losses $L(x) = |x|^n$ with $n > 2$, which recover the near-linear scaling.

Key stability condition (Lemma 1): For a network in a locally quadratic region with Hessian $H_t$ and preconditioner $D_t = \mathrm{diag}[(\sqrt{\bar{v}_t} + \epsilon)^{-1}]$, Adam is unstable when:

$$\frac{1-\beta_1}{1+\beta_1}\lambda_{\max}(D_t H_t) > \frac{2}{\eta}. \tag{2}$$

Corollary 2 (1D specialization): For one-dimensional curvature $\lambda_t = L''(x_t) > 0$, define:

$$\eta_{\mathrm{eff},t} = \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon}, \qquad \eta_{\mathrm{crit},t} = \frac{2(1+\beta_1)}{(1-\beta_1)\lambda_t}. \tag{3}$$

Adam is unstable when $\eta_{\mathrm{eff},t} > \eta_{\mathrm{crit},t}$.

---

Methodology

Models and Tasks

Six model-task settings are studied:

1. One-layer Transformers for modular division and addition (modulo 53)
2. Character-level Transformer for next-character prediction on Tiny Shakespeare
3. MLP and VGG11-style CNN for CIFAR-10 classification
4. Convolutional autoencoder for binarized MNIST reconstruction

All models use AdamW with cross-entropy loss (or binary cross-entropy for reconstruction). Only $\beta_1$ and $\beta_2$ vary while other hyperparameters remain fixed. The memory gaps are $\delta = 1-\beta_1$ and $\gamma = 1-\beta_2$, scanned logarithmically over $\delta \in [10^{-3}, 0.2]$ and $\gamma \in [5\times 10^{-4}, 0.3]$.

Phase Diagram Construction

• Macroscopic oscillations are identified from low-frequency spectral peaks meeting a prominence criterion
• For each fixed $\gamma$, the boundary is the immediate right neighbor of the largest sampled $\delta$ with no accepted period
• Boundaries are fit using $\log \gamma = \log C + p \log \delta$

One-Dimensional Toy Model

The paper studies Adam on $L(x) = k|x|^n/n$ with $k=1$, $\eta=0.1$, $\epsilon=10^{-30}$, $x_0=1$, and $m_0 = v_0 = 0$. Each beta pair is trained for 250,000 updates with the first half discarded.

Core-Wall Landscape Analysis

Directional loss slices are evaluated along the Adam-preconditioned gradient direction:

$$\phi_t(s) = L(\theta_t + s\hat{d}_t), \qquad \hat{d}_t = \frac{D_t g_t}{\|D_t^{1/2} g_t\|_2}. \tag{9}$$

The effective logarithmic exponent is defined as:

$$n_{\mathrm{eff},\pm}(r) = \frac{\mathrm{d}\log\Delta L_{\pm}(r)}{\mathrm{d}\log r}. \tag{8}$$

---

Empirical Validation / Results

Phase Boundaries Across Six Models

All six models show near-linear scaling $\gamma \propto \delta$ with fitted exponents $p = 0.93, 0.85, 0.75, 0.95, 1.06, 0.95$ for the six settings respectively. The region above the boundary (where $(1-\beta_2)/(1-\beta_1)$ is larger) generally contains no detected macroscopic spikes.

One-Dimensional Superquadratic Loss

For $n > 2$, a wedge-shaped spiky region emerges between two lines of approximately unit log-log slope:

$$C_L(1-\beta_1) \geq 1-\beta_2 \geq C_R(1-\beta_1). \tag{4}$$

Theorem 3 (Right boundary): For $L(x) = k|x|^n$ with $\frac{1}{2} < \beta_1 < 1$, $0 < \beta_2 < 1$, and $n > 2$, trajectories attracted to nonzero fixed points cease to exhibit EoS-triggered spikes when:

$$\beta_1 \beta_2^{-n/[2(n-2)]} < 1. \tag{5}$$

This yields the right boundary:

$$\gamma_R(\delta, n) = 1 - (1-\delta)^{2(n-2)/n} = \frac{2(n-2)}{n}\delta + O(\delta^2). \tag{7}$$

The left boundary is described empirically by $C_L \simeq 12(n-2)/(3n-4)$, capturing the trend across $n$ values from 2.1 to 8.

Core-Wall Landscape

Figure 5 results: For the modular-division Transformer at a pre-spike point:

• Core exponent: $n_{\mathrm{eff}} \simeq 2.18$
• Wall exponent: $n_{\mathrm{eff}} \simeq 11.0$
• Core-wall transition at $O(10^{-6})$, comparable to the typical update scale $|s_{\mathrm{upd}}| \simeq 6.3 \times 10^{-6}$
• Rollover begins at $O(10^{-5})$
• $1/\sigma_a \simeq 4.35 \times 10^{-6}$ for the measured slice

Lemma 5 (Hessian collapse): For confident cross-entropy with one-hot targets, $\|\nabla_{\theta}^2 L\|_2 \to 0$ as predicted probabilities approach their hard targets.

Theorem 7 (Quadratic-core scale): Under a Gaussian model for the weighted slope distribution, the normalized confident loss is:

$$\frac{\Delta L_{\mathrm{conf}}(s)}{L_{\mathrm{conf}}(0)} = e^{(\sigma_a s)^2/2} - 1, \qquad n_{\mathrm{eff}} = \frac{(\sigma_a s)^2 e^{(\sigma_a s)^2/2}}{e^{(\sigma_a s)^2/2} - 1}. \tag{12}$$

Thus $n_{\mathrm{eff}} = 2 + (\sigma_a s)^2/2 + O((\sigma_a s)^4)$, with departure from quadratic behavior at $|s| = O(1/\sigma_a)$.

Effective Wall Exponents and Boundary Coefficients

---

Theoretical and Practical Implications

Theoretical Significance

1. Beyond quadratic approximation: The paper demonstrates that the conventional quadratic loss approximation fails to capture the stability phase diagram of Adam. The transition from cubic to linear boundary scaling is attributed to superquadratic loss geometry.

2. Momentum timescale competition: The first- and second-moment memory timescales ($\beta_1$, $\beta_2$) control the phase boundaries through their relative decay rates, formalized in the quiet-phase approximation where $v_{t+1} = \beta_2 v_t$.

3. Core-wall landscape mechanism: Confident cross-entropy losses naturally produce a core-wall landscape through the exponential-sum structure of the loss. The flat core arises from Hessian collapse, while the wall emerges at scale $O(1/\sigma_a)$ where $\sigma_a$ is the standard deviation of directional margin slopes.

4. Connection to existing results: The right-boundary condition matches the local-stability boundary of full Adam on even-degree degenerate polynomials (Bai et al., 2026a), recovered here for real $n>2$ under the quiet-phase approximation.

Practical Implications

1. Hyperparameter selection: The near-linear boundary provides practical guidance for selecting $(\beta_1, \beta_2)$ pairs to avoid loss spikes during training.
2. Loss landscape awareness: Practitioners should be aware that confident cross-entropy losses have a much smaller quadratic region than typically assumed, with the optimizer probing superquadratic geometry at update scales.
3. Adaptive optimization: The results suggest potential for optimizers that estimate $n_{\mathrm{eff}}$ at their update scale and adapt learning rates or moment timescales accordingly.

---

Conclusion

This work makes three principal contributions:

1. Near-linear phase boundary: Identification of an approximately linear boundary separating spiky and non-spiky Adam dynamics across diverse models and tasks, with the boundary coefficient linked to the effective loss exponent.

2. Superquadratic geometry: Attribution of the shift from cubic to linear boundary scaling to superquadratic loss geometry, showing how first- and second-moment memory timescales set the phase boundaries.

3. Core-wall landscape: Demonstration that confident cross-entropy forms a core-wall landscape (quadratic core, superquadratic wall, outer rollover), linking the effective exponent at the optimizer's update scale to the empirical boundary coefficient.

Future Directions

1. Can the core-wall exponential-sum mechanism explain basin-like landscapes in broader underparameterized neural network settings?
2. Can an optimizer estimate $n_{\mathrm{eff}}$ at its update scale and adapt learning rates or moment timescales to the evolving phase boundary?
3. What mechanism drives progressive sharpening at finite scales? How do margins $m_{ij}$ and their directional derivatives $a_{ij} = -dm_{ij}/ds$ evolve during training?