[Paper Review] Analyzing Sharpness along GD Trajectory: Progressive Sharpening and Edge of Stability

Zhouzi Li, Zixuan Wang+1 more|arXiv (Cornell University)|Jul 26, 2022

Neural Networks and Applications4 citations

TL;DR

This paper analyzes gradient descent dynamics in deep neural networks, identifying a four-phase trajectory based on sharpness evolution—progressive sharpening followed by oscillation around the Edge of Stability (EOS) regime. It proposes that the output layer weight norm acts as a key indicator of sharpness, providing theoretical and empirical explanations for non-monotonic loss and sharpness behavior, with a rigorous proof in two-layer linear networks under simplifying assumptions.

ABSTRACT

Recent findings (e.g., arXiv:2103.00065) demonstrate that modern neural networks trained by full-batch gradient descent typically enter a regime called Edge of Stability (EOS). In this regime, the sharpness, i.e., the maximum Hessian eigenvalue, first increases to the value 2/(step size) (the progressive sharpening phase) and then oscillates around this value (the EOS phase). This paper aims to analyze the GD dynamics and the sharpness along the optimization trajectory. Our analysis naturally divides the GD trajectory into four phases depending on the change of the sharpness. We empirically identify the norm of output layer weight as an interesting indicator of sharpness dynamics. Based on this empirical observation, we attempt to theoretically and empirically explain the dynamics of various key quantities that lead to the change of sharpness in each phase of EOS. Moreover, based on certain assumptions, we provide a theoretical proof of the sharpness behavior in EOS regime in two-layer fully-connected linear neural networks. We also discuss some other empirical findings and the limitation of our theoretical results.

Motivation & Objective

• To understand the non-monotonic behavior of sharpness and loss during gradient descent in deep networks, particularly the emergence of the Edge of Stability (EOS) regime.
• To identify a practical indicator—specifically, the norm of the output layer weights—that correlates with sharpness dynamics during training.
• To provide a theoretical explanation for progressive sharpening and EOS behavior in a simplified two-layer linear neural network setting.
• To explain the non-monotonic loss decrease observed in training, which contradicts classical optimization theory.
• To clarify the mechanisms behind sharpness stabilization at 2/η despite learning rate violations of classical convergence rules.

Proposed method

• Proposes a four-phase division of the GD trajectory based on sharpness and loss dynamics: initial decrease, progressive sharpening, EOS regime, and final stabilization.
• Empirically identifies the output layer weight norm as a strong proxy for sharpness evolution, enabling analysis of sharpness dynamics through this surrogate.
• Derives theoretical dynamics of sharpness, loss, and weight norm under assumptions linking the output weight norm to sharpness behavior.
• Uses the Gauss-Newton approximation of the Hessian, focusing on the Fisher Information Matrix (FIM) and Neural Tangent Kernel (NTK) to approximate sharpness via the largest eigenvalue of M(t).
• Applies gradient flow dynamics to model the residual vector D(t), deriving the update rule D(t+1) - D(t) ≈ -ηM(t)D(t) to analyze sharpness evolution.
• Provides a rigorous theoretical proof for sharpness behavior in two-layer fully-connected linear networks under weak assumptions, including the convergence of sharpness to 2/η.

Experimental results

Research questions

• RQ1 What are the distinct dynamical phases of sharpness and loss along the gradient descent trajectory in modern deep networks?
• RQ2 Why does sharpness increase to 2/η and then stabilize around this value, despite violating classical convergence conditions?
• RQ3 How can the norm of the output layer weights serve as a reliable proxy for sharpness dynamics during training?
• RQ4 What drives the non-monotonic decrease in training loss observed in the EOS regime?
• RQ5 Under what conditions can the sharpness behavior in the EOS regime be theoretically justified in simplified network architectures?

Key findings

• The GD trajectory naturally divides into four phases based on sharpness and loss dynamics: initial loss decrease, progressive sharpening, Edge of Stability (EOS) regime with oscillating sharpness, and final stabilization.
• The norm of the output layer weights is empirically found to be a strong and reliable indicator of sharpness evolution, enabling indirect analysis of sharpness dynamics.
• In the EOS regime, sharpness oscillates around 2/η, the theoretical threshold that would normally cause divergence, yet training loss continues to decrease non-monotonically.
• Theoretical analysis in two-layer fully-connected linear networks shows that sharpness converges to 2/η under mild assumptions, providing a formal justification for the EOS phenomenon.
• The non-monotonic loss decrease is explained by the oscillatory behavior of sharpness and the interplay between weight norm growth and residual vector dynamics.
• The largest eigenvalue of the Neural Tangent Kernel (NTK) matrix M(t) closely approximates the Hessian’s largest eigenvalue (sharpness), validating its use as a proxy in analysis.