[论文解读] Understanding Gradient Descent on Edge of Stability in Deep Learning

Sanjeev Arora, Zhiyuan Li+1 more|arXiv (Cornell University)|May 19, 2022

Stochastic Gradient Optimization Techniques被引用 4

一句话总结

本文对深度学习中边缘稳定性（EoS）处的梯度下降进行了数学分析，表明非光滑损失曲面与自适应学习率会诱导零损失流形上的确定性流。证明了归一化梯度下降和在√L上的梯度下降均可保证进入EoS阶段，此时曲率稳定在约2/η附近，尽管存在非单调振荡，损失仍整体下降，这是由于对更平坦极小值的隐式正则化所致。

ABSTRACT

Deep learning experiments by Cohen et al. [2021] using deterministic Gradient Descent (GD) revealed an Edge of Stability (EoS) phase when learning rate (LR) and sharpness (i.e., the largest eigenvalue of Hessian) no longer behave as in traditional optimization. Sharpness stabilizes around $2/$LR and loss goes up and down across iterations, yet still with an overall downward trend. The current paper mathematically analyzes a new mechanism of implicit regularization in the EoS phase, whereby GD updates due to non-smooth loss landscape turn out to evolve along some deterministic flow on the manifold of minimum loss. This is in contrast to many previous results about implicit bias either relying on infinitesimal updates or noise in gradient. Formally, for any smooth function $L$ with certain regularity condition, this effect is demonstrated for (1) Normalized GD, i.e., GD with a varying LR $η_t =\fracη{\| abla L(x(t)) \|}$ and loss $L$; (2) GD with constant LR and loss $\sqrt{L- \min_x L(x)}$. Both provably enter the Edge of Stability, with the associated flow on the manifold minimizing $λ_{1}( abla^2 L)$. The above theoretical results have been corroborated by an experimental study.

研究动机与目标

• 解释深度学习训练中使用确定性梯度下降时观察到的边缘稳定性（EoS）现象。
• 将稳定性形式化为一种条件，确保在有限学习率下损失减少，超越传统光滑性假设。
• 识别并严格分析两种可保证导致EoS行为的机制——非光滑损失函数与自适应学习率。
• 证明在√L上和归一化GD下的梯度下降均沿一条确定性流演化，以最小化λ₁(∇²L)，从而促进更平坦的极小值。

提出的方法

• 提出稳定性 S_L(x,η) = η·sup₀≤s≤η λ₁(∇²L(x−s∇L(x))) ≤ 2 的形式化定义，将标准下降引理推广至非光滑情形。
• 分析在√L上的梯度下降，其中L为光滑损失且min L(x) = 0，表明∇²√L在极小值附近发散，从而诱导EoS行为。
• 研究归一化梯度下降作为自适应学习率η_t = η / ||∇L(x(t))||的梯度下降，发现在极小值附近学习率变大，可保证进入EoS。
• 推导出在零损失流形Γ = {x | L(x) = 0}上的极限黎曼流，表明更新过程演化为最小化最大Hessian特征值λ₁(∇²L)。
• 使用幂迭代法计算最大Hessian特征向量v₁(x)，并将梯度投影到流形的切空间以模拟极限流。
• 通过在合成任务和分类任务上的模拟实验，实证验证了理论结果，追踪损失、曲率与流动力学。

实验结果

研究问题

• RQ1 在有限学习率下，梯度下降为何能在违反经典条件η < 2/λ的情况下仍持续减少损失？
• RQ2 边缘稳定性阶段中曲率稳定在2/η附近的数学机制是什么？
• RQ3 损失曲面的非光滑性或自适应学习率是否可诱导一种确定性流，从而解释EoS行为？
• RQ4 在√L上或归一化GD下的梯度下降是否可保证进入EoS阶段并演化至更平坦的极小值？
• RQ5 EoS动力学能否由零损失流形上的极限黎曼流描述？

主要发现

• 在√L上的梯度下降可保证进入边缘稳定性阶段，由于极小值附近Hessian曲率发散，曲率稳定在约2/η。
• 归一化梯度下降（学习率自适应于梯度范数）可保证进入EoS，并沿一条最小化λ₁(∇²L)的确定性流演化。
• 证明了归一化GD的极限流是零损失流形上的黎曼梯度流，其更新被投影到解集的切空间上。
• 实验结果证实，√L与归一化GD的轨迹在零损失流形附近振荡，同时缓慢移向更平坦的极小值，曲率始终稳定在2/η附近。
• 理论分析表明，即使在有限、非无穷小的步长下，EoS行为也源于对λ₁(∇²L)较小的极小值的隐式正则化。
• 稳定性条件S_L(x,η) ≤ 2被证明是损失减少的更优指标，尤其在非光滑情形下优于经典条件ηλ₁(∇²L(x)) < 2。