The Origin of Edge of Stability

1. Problem Statement

In deep learning, full-batch gradient descent (GD) often exhibits a phenomenon known as the Edge of Stability (EoS). Classical optimization theory dictates that for convergence, the learning rate $\eta$ must satisfy $\eta < 2/\lambda_{\max}$, where $\lambda_{\max}$ is the largest eigenvalue of the Hessian (sharpness). However, empirical observations (e.g., Cohen et al., 2021) show that during training, sharpness rises until it reaches the threshold $2/\eta$, at which point it saturates. The training loss then begins to oscillate on short timescales while continuing to decrease on longer timescales.

The Gap: While previous works have established that the system self-regulates near this edge (e.g., via cubic Taylor terms or gradient-norm penalties), they fail to explain why the trajectory is globally forced toward the $2/\eta$ threshold from arbitrary initializations. Existing explanations are inherently local and do not account for the global convergence to this specific value.

2. Methodology: The Edge Coupling

The author introduces a novel variational framework centered on a functional called the Edge Coupling, denoted as $\mathcal{A}_\eta(x, y)$.

• Definition: For consecutive iterates $(x, y)$, the edge coupling is defined as:
  $$\mathcal{A}_\eta(x, y) \triangleq L(x) + L(y) - \frac{1}{2\eta} \|x - y\|^2$$
• Variational Characterization: The criticality condition of this functional with respect to $x$ (i.e., $\nabla_x \mathcal{A}_\eta = 0$) recovers the standard gradient descent update $y = x - \eta \nabla L(x)$ if and only if the coefficient is fixed to $1/\eta$.
• Hamiltonian Interpretation: The functional is viewed as a discrete generating function in Hamiltonian mechanics, where the criticality conditions encode the gradient descent dynamics in a boundary-value form.

3. Key Contributions and Theoretical Results

A. Derivation of the Stability Boundary

By differencing the criticality condition $\nabla_x \mathcal{A}_\eta = 0$ between consecutive steps, the author derives a step-increment propagator:
$$d_{k+1} = (I - \eta \bar{H}_k) d_k$$
where $d_k$ is the step vector and $\bar{H}_k$ is a uniform average of the Hessian along the step segment.

• Result: Step reversal ($d_{k+1} \approx -d_k$) and two-step return ($w_{k+2} \approx w_k$) occur if and only if the eigenvalue of $\bar{H}_k$ approaches $2/\eta$. This establishes $2/\eta$ as the stability boundary for the recurrence.

B. Global Forcing Mechanism (The Origin)

The paper derives a loss-change formula via a second-order expansion of the edge coupling:
$$L(w_{k+1}) - L(w_k) = -\frac{\|d_k\|^2}{2\eta} \left( 2 - \eta \tilde{r}_k \right)$$
where $\tilde{r}_k$ is a triangularly weighted average of the Hessian curvature.

• Telescoping Sum: Summing this equation over $K$ steps yields a conservation law:
  $$\sum_{k=0}^{K-1} \|d_k\|^2 \left( \frac{2}{\eta} - \tilde{r}_k \right) = 2(L(w_0) - L(w_K))$$
• Forcing Theorem: Since the total loss drop is finite (bounded below), the weighted average of the curvature deviations $(2/\eta - \tilde{r}_k)$ must vanish as training progresses. This mathematically proves that the trajectory is globally forced to visit steps where the effective curvature is arbitrarily close to $2/\eta$.

C. Localization to True Hessian

A critical technical contribution is the use of the Mean Value Theorem. While the forcing theorems initially apply to averaged Hessians ($\bar{H}_k$ and $\tilde{H}_k$), the author proves that these averages are realized as exact pointwise values of the directional curvature at specific interior points of the step segment.

• Significance: This eliminates any "gap" between the averaged theory and the actual Hessian eigenvalues observed in practice. The forcing applies directly to the true sharpness $\lambda_{\max}(\nabla^2 L)$ at specific points along the trajectory.

D. Period-Two Orbits and Bifurcation Analysis

The paper analyzes the fixed points and period-two orbits of the edge coupling.

• Center Reduction: By reparametrizing the orbit as $(m-a, m+a)$, the problem reduces to a function of the half-amplitude $a$ alone.
• Bifurcation: The onset of period-two orbits (oscillations) is characterized as a bifurcation problem. The author derives a quartic expansion showing that period-doubling branches emerge continuously from zero amplitude when the learning rate crosses the critical threshold $\eta_c = 1/\sigma_1$ (where $\sigma_1$ is the largest singular value of the Hessian).
• Width Invariance: For two-layer linear networks, this bifurcation analysis is shown to be width-invariant, confirming that overparameterization adds only flat directions without altering the transverse edge dynamics.

E. Stability Mechanisms at the Edge

The paper explains why the system does not diverge once it hits the edge:

1. Growth Above Threshold: If curvature exceeds $2/\eta$, the step norm grows geometrically, forcing the trajectory back into the stable region.
2. Oscillatory Cancellation: Inside the stability window, the step direction alternates (multiplier $\approx -1$). The author proves that small asymmetries in this alternation cancel out rather than accumulating into secular drift, keeping the system bounded near the edge.

4. Experimental Validation

The theory is validated on:

• 3-layer MLPs: Demonstrating the saturation of sharpness at $2/\eta$ and the onset of loss oscillations.
• Two-layer Linear Networks: Confirming the continuous emergence of period-doubling branches and the $\sqrt{\eta - \eta_c}$ scaling of oscillation amplitude predicted by the bifurcation theory.
• Mini-batch SGD: The forcing theorems are extended to stochastic settings, showing that the telescoping structure holds in expectation (with an additional noise term).

5. Significance

This paper provides the first unified, global explanation for the Edge of Stability phenomenon.

• From Local to Global: It moves beyond local Taylor-series approximations to a global conservation law that forces curvature to the threshold.
• Exactness: By using the Mean Value Theorem, it bridges the gap between averaged curvature models and the true, instantaneous Hessian eigenvalues.
• Mechanistic Insight: It clarifies the interplay between step propagation, loss descent, and period-two bifurcations, offering a rigorous variational foundation for understanding why neural networks naturally operate at the edge of stability during training.

In summary, Litman demonstrates that the Edge of Stability is not an accidental artifact but a necessary consequence of the gradient descent update rule acting as a critical point of a specific variational functional, which inherently drives the system's curvature toward the $2/\eta$ threshold.