[論文レビュー] 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）における勾配降下法の数学的分析を提供し、非滑らかさのある損失関数と適応的学習率がゼロ損失多様体上での決定的流れを引き起こすことを示している。非滑らか性と適応的学習率の両方が、EoS状態に到達することを厳密に証明している。正規化勾配降下法と√L上での勾配降下法の両方が、明示的な正則化により平坦な極小値へと誘導される決定的流れに従って進化し、損失は全体的に減少するが、非単調な振動を伴う。

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.

研究の動機と目的

• 有限学習率を用いた勾配降下法が、古典的滑らかさ仮定を越えて、損失の減少を続ける理由を説明すること。
• 有限学習率下での損失減少を保証する条件としての「安定性」を定式化し、従来の滑らかさ仮定を超えること。
• 非滑らか損失関数と適応的学習率という2つのメカニズムが、EoS行動を明示的に引き起こすことを同定し、厳密に分析すること。
• √L上での勾配降下法と正規化勾配降下法が、λ₁(∇²L) を最小化する決定的流れに従って進化し、平坦な極小値へと向かうことを示すこと。

提案手法

• 非滑らか設定に一般化された標準的勾配降下補題を拡張する、S_L(x,η) = η·sup₀≤s≤η λ₁(∇²L(x−s∇L(x))) ≤ 2 という安定性の形式的定義を導入する。
• min L(x) = 0 である滑らかな損失関数 L を用い、√L 上での勾配降下法を分析し、極小値近傍で ∇²√L が発散することにより EoS 時相が誘発されることを示す。
• 正規化勾配降下法（学習率が η_t = η / ||∇L(x(t))|| に適応的となる）を研究し、極小値付近で学習率が大きくなり、EoS に明示的に到達することを証明する。
• ゼロ損失多様体 Γ = {x | L(x) = 0} 上での極限リーマン幾何的流れを導出し、更新が最大ヘッセ固有値 λ₁(∇²L) を最小化するように進化することを示す。
• トップヘッセ固有ベクトル v₁(x) をパワー反復法で計算し、勾配を解集合の接空間に射影することで、極限流れを模擬する。
• 合成的および分類タスクにおけるシミュレーションを通じて理論的結果を実証的に検証し、損失、鋭さ、流れのダイナミクスを追跡する。

実験結果

リサーチクエスチョン

• RQ1 有限学習率を用いた勾配降下法が、古典的条件 η < 2/λ を満たさない場合でも、なぜ損失を減少させ続けるのか？
• RQ2 安定性の縁（EoS）状態で鋭さが 2/η 付近に安定化する背後にある数学的メカニズムは何か？
• RQ3 損失関数の非滑らか性や適応的学習率が、決定的流れを誘発し、EoS 時相を説明可能か？
• RQ4 √L 上での勾配降下法や正規化勾配降下法が、EoS 時相に明示的に到達し、平坦な極小値へと進化するか？
• RQ5 EoS 動的挙動はゼロ損失多様体上での極限リーマン幾何的流れとして記述可能か？

主な発見

• √L 上での勾配降下法は、極小値近傍でのヘッセ曲率の発散により、EoS 時相に明示的に到達し、鋭さが 2/η 付近に安定化する。
• 勾配ノルムに応じて学習率が適応的に変化する正規化勾配降下法は、EoS に明示的に到達し、λ₁(∇²L) を最小化する決定的流れに従って進化する。
• 正規化勾配降下法の極限流れが、ゼロ損失多様体上でのリーマン勾配流れであり、更新が解集合の接空間に射影されることを示した。
• 実験により、√L および正規化勾配降下法の軌道がゼロ損失多様体の周囲を振動しながら、ゆっくりと平坦な極小値へと向かうことが確認され、鋭さは 2/η 付近を維持する。
• 理論的分析により、EoS 時相は、有限で非無限小のステップでさえも、λ₁(∇²L) が小さい極小値への暗黙的正則化に起因することが示された。
• 特に非滑らか領域においては、古典的指標 ηλ₁(∇²L(x)) < 2 よりも、S_L(x,η) ≤ 2 が損失減少の指標として優れていることが示された。