[論文レビュー] Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability

Jeremy M. Cohen, Simran Kaur+3 more|arXiv (Cornell University)|Feb 26, 2021

Stochastic Gradient Optimization Techniques参考文献 49被引用数 22

ひとこと要約

この論文は、ニューラルネットワークにおけるフルバッチ勾配降下法が一貫して「安定性の縁」に位置することを示しており、最大ヘッセ固有値（鋭さ）が 2/η 程度のわずかに上回る領域で動作し、訓練損失は非単調的ではあるが最終的には減少する。この行動は標準的な最適化仮定に反しており、深層学習の訓練における収束性と滑らかさに関する現在の理論に挑戦する。

ABSTRACT

We empirically demonstrate that full-batch gradient descent on neural network training objectives typically operates in a regime we call the Edge of Stability. In this regime, the maximum eigenvalue of the training loss Hessian hovers just above the numerical value $2 / ext{(step size)}$, and the training loss behaves non-monotonically over short timescales, yet consistently decreases over long timescales. Since this behavior is inconsistent with several widespread presumptions in the field of optimization, our findings raise questions as to whether these presumptions are relevant to neural network training. We hope that our findings will inspire future efforts aimed at rigorously understanding optimization at the Edge of Stability. Code is available at https://github.com/locuslab/edge-of-stability.

研究の動機と目的

• フルバッチ勾配降下法がニューラルネットワークの訓練目的関数に与えるダイナミクスを調査すること。
• ニューラルネットワークの訓練において単調な低下とL-スムーズ性を仮定する最適化理論の現在の仮定に挑戦すること。
• 訓練中にヘッセ行列の最大固有値（鋭さ）の経験的特性を同定すること。
• 鋭さが 2/η よりわずかに上回る一貫した領域、すなわち『安定性の縁』を特定し、非単調であるが時間経過とともに訓練損失が減少することを示すこと。
• 安定性の縁における最適化ダイナミクスを説明できる新たな理論枠組みの創出を促すこと。

提案手法

• 標準データセット（CIFAR-10を含む）におけるさまざまなニューラルネットワークアーキテクチャにフルバッチ勾配降下法を適用する。
• 異なるステップサイズ η に対して、訓練損失ヘッセ行列の最大固有値（鋭さ）を訓練全体を通して計算・追跡する。
• 2次関数における勾配降下法の理論的安定性解析を基準として、臨界閾値 2/η を定義する。
• 経験的分析により、交差エントロピーとMSEの異なる損失関数およびステップサイズにおける鋭さの変化を比較する。
• 複数のアーキテクチャとデータセットを用いて一般化可能性を評価する。
• フルバッチGDと確率的勾配降下法（SGD）を対比し、鋭さダイナミクスの違いとSGDでは鋭さがフラットラインに近づかないことを見出す。

実験結果

リサーチクエスチョン

• RQ1 フルバッチ勾配降下法は、アーキテクチャやタスクにかかわらず一貫した動的領域を示すか？
• RQ2 訓練中にステップサイズ η とヘッセ行列の最大固有値（鋭さ）との関係は何か？
• RQ3 なぜ非単調的ではあるが、『安定性の縁』領域では訓練損失が減少し収束するのか？
• RQ4 『安定性の縁』領域は、最適化理論における単調な低下やL-スムーズ性といった標準的仮定とどのように矛盾するのか？
• RQ5 Jastrzębskiら（2020）のSGDモデルは、フルバッチ勾配降下法へどの程度一般化可能か？

主な発見

• フルバッチ勾配降下法は、多様なアーキテクチャとタスクにおいて一貫して『安定性の縁』に位置し、最大ヘッセ固有値が 2/η 程度のわずかに上回る領域で動作する。
• 鋭さが 2/η 未満の場合、徐々に鋭さが増し、臨界閾値に近づく。
• 鋭さが 2/η を超えると、勾配降下法は発散せず、むしろこの値の周辺で安定化し、『安定性の縁』領域に入る。
• 短期スケールでは非単調であるが、長期スケールでは訓練損失が減少しており、効果的な最適化が行われていることが示唆される。
• CIFAR-10における標準的なアーキテクチャでは、妥当なステップサイズを採用すれば常に『安定性の縁』領域に到達するため、これは例外的ではなく標準的である。
• 『安定性の縁』領域は、単調な低下、L-スムーズ性、2次テイラー近似といった最適化理論の主要仮定と矛盾する。