安定性の縁におけるSGD：確率的シャープネス・ギャップ

要旨: ニューラルネットワークをフルバッチ勾配降下法（GD）で、ステップサイズ eta を用いて学習するとき、ヘッセ行列の最大固有値――シャープネス S(boldsymbol{theta})――は 2/eta にまで上昇し、そこに留まる。この現象は安定性の端（Edge of Stability: EoS）と呼ばれる。
\citet{damian2023selfstab} は、この挙動が損失の3次の構造によって駆動される自己安定化メカニズムによって説明できること、さらにGDが暗黙のうちに制約 S(boldsymbol{theta})\leq 2/eta 上での射影勾配降下法（Projected Gradient Descent: PGD）に従っていることを示した。
ミニバッチ確率的勾配降下法（SGD）では、シャープネスは 2/eta 未満で安定化し、そのギャップはバッチサイズが小さくなるほど拡大するが、この抑制に対する理論的説明は存在しない。

我々は、自己安定化の枠組みをSGDへ拡張する形で、確率的自己安定化を導入する。我々の主要な洞察は、勾配ノイズが最上位のヘッセ固有ベクトル方向の振動的ダイナミクスに分散を注入し、それによって（3次の）シャープネスを低減する力が強められ、平衡点が 2/eta より下へ押し下げられるという点にある。
\citet{damian2023selfstab} のアプローチに従い、移動する射影勾配降下法の軌道に対する確率的な予測ダイナミクスを定義し、SGDがこれらの予測からどれだけ逸脱するかを抑える確率的結合（stochastic coupling）定理を証明する。
我々は閉形式の平衡シャープネス・ギャップを導出する：
\Delta S = \u007feta \beta \u007fsigma_{\u007fboldsymbol{u}}^{2}/(4\alpha)、ここで alpha は漸進的シャープ化率、beta は自己安定化の強さ、そして  sigma_{ \u007fboldsymbol{u}}^{2} は最上位固有ベクトルへ射影した勾配ノイズ分散である。
この式は、より小さなバッチサイズではより平坦な解が得られ、バッチが全データセットに等しいときGDが回復されることを予測する。