技術要約：Edge Flow

問題提起

ディープラーニングにおける勾配降下法（GD）は、しばしば**安定性の境界（Edge of Stability: EoS）**として知られる領域で動作します。この領域では、損失ヘッシアンの最大固有値（シャープネス、$S(w_t)$）が、学習率 $\eta$ による安定性の閾値 $2/\eta$ の付近を彷徨います。この間、訓練損失は一般に減少しますが、反復計算は準周期的な振動を示し、シャープネスは、発散するのではなく自己調節メカニズムを通じて安定化します。

勾配流（$\dot{w} = -\nabla L(w)$）や降下補題（descent lemma）といった古典的な連続時間解析ツールは、この領域では機能しません。なぜなら、それらは学習率を無視しており、$\eta \cdot S(w)$ の比率に依存する現象を捉えることができないからです。既存のEoSに関する連続時間モデルには以下の限界があります：

• Central Flow (Cohen et al., 2025): 振動の中心を投影勾配流によってモデル化しています。これは安定化しますが、投影制約があるため、システムのシミュレーションや理論的解析が困難です。
• Rod Flow (Regis and Chewi, 2026): 反復計算を剛体ロッドの端点としてモデル化しています。概念的には魅力的ですが、その振動振幅の更新規則は、EoSの開始時における自己安定化ダイナミクスを捉えることができません。

本論文が取り組む中心的な未解決問題は、**「EoSにおけるGDダイナミクスの、扱いやすく、忠実で、予測可能な連続時間モデルとは何か？」**という問いです。

手法：Edge Flow

著者らは、GDのダイナミクスを3つの緩やかに変化する量へと分解する、3つの結合された常微分方程式（ODE）の系であるEdge Flowを提案しています。

1. 中心 ($\bar{w}_t$): 振動のゆっくりと漂う中間点。
2. 振動方向 ($u_t$): ヘッシアンのトップ固有ベクトルを追跡する単位ベクトル。
3. 振動振幅 ($x_t$): 振動の半長。

GDの反復計算は $w_t \approx \bar{w}_t + (-1)^t x_t u_t$ とモデル化されます。Edge Flowの系は以下のように定義されます：

1. 中心のダイナミクス:
   $$\frac{d\bar{w}_t}{dt} = -\frac{1}{2} \left( \nabla L(\bar{w}_t + x_t u_t) + \nabla L(\bar{w}_t - x_t u_t) \right)$$
   これは、対称化された損失（2つの端点における勾配の平均）上の勾配流を表します。これは標準的な勾配流とは異なり、$x_t$ に関する3次微分補正項を含んでおり、中心をより曲率の低い領域へとバイアスさせます。

2. 方向のダイナミクス:
   $$\frac{du_t}{dt} = 2 \left( \nabla^2 L(\bar{w}_t) - (u_t^\top \nabla^2 L(\bar{w}_t) u_t) I \right) u_t$$
   これは、単位球面上のレイリー商勾配流です。これは、パワーメソッドの連続時間極限と同様に、$u_t$ を $\bar{w}_t$ におけるヘッシアンのトップ固有ベクトルへと駆動します。

3. 振幅のダイナミクス:
   $$\frac{dx_t}{dt} = \left( S(\bar{w}_t) - \frac{2}{\eta} \right) x_t$$
   このスカラー線形ODEは、振動振幅の増大または減衰を制御します。もしシャープネス $S(\bar{w}_t) > 2/\eta$ であれば、$x_t$ は指数関数的に増大し、$S(\bar{w}_t) < 2/\eta$ であれば、$x_t$ は減衰します。

離散化：Edge Gradient Descent (EGD)

著者らは、Edge Gradient Descent (EGD) と呼ばれる自然な離散化を提案しています。各反復には以下が必要です：

• 2回の勾配評価: $\nabla L(\bar{w}_k \pm x_k u_k)$
• 1回のヘッシアン・ベクトル積: $\nabla^2 L(\bar{w}_k) u_k$
• 固有値分解や最適化サブルーチンは不要。

このアルゴリズムは、2つの異なるステップサイズを使用します：$\eta$（元の学習率、シャープネスの閾値を定義）および $\rho$（ODEの積分ステップであり、$\rho \le \eta$ を満たす）。

主な貢献

1. 単純かつ扱いやすいモデル: Edge Flowは、第一原理とEoSにおけるGD構造の初等的な観察から導出されています。これは、シミュレーションや解析が容易な3つの結合されたODEの系です。
2. 創発的な自己安定化: 安定性を強制するために制約を課す従来のモデルとは異なり、Edge Flowは、結合されたダイナミクスの自然な帰結として自己安定化メカニズム（Damian et al., 2023によって特定されたもの）を再現します。シャープネスが閾値を超えると、振動振幅が増大し、対称化された勾配を介して中心を低曲率領域へと押し出し、それが結果としてシャープネスを減少させます。
3. 忠実な追跡と予測: 実験的に、Edge Flowは様々なアーキテクチャ（MLP, CNN, ResNet, ViT）および損失関数において、フルバッチGDの軌跡、損失、およびシャープネスを追跡します。決定的なのは、Central Flowの瞬時的な投影では捉えられない、EoS開始時におけるシャープネスのオーバーシュートと減衰振動を解明している点です。
4. 不安定性の洞察: 本モデルは、EoS開始時の損失スパイクの原因が、硬い（stiff）中心ダイナミクスの離散化によるアーティファクトであることを特定しています。これに対し、2つの原理的な解決策を提示しています：
   • 離散化の精緻化: より小さな積分ステップ $\rho$（より大きな $K = \eta/\rho$）を使用することで、損失スパイクを抑制できます。
   • ベースとなる振動の増加: 小さな初期振幅 $\epsilon$ を注入することで、自己安定化フィードバックループを加速させ、オーバーシュートを軽減できます。

結果

• 追跡性能: CIFAR-10のサブセットを用いた実験において、Edge FlowはGDおよびCentral Flowのバルクダイナミクスと一致します。しかし、Edge Flowは、EoS開始時におけるシャープネスの過渡的な振動挙動や、GDで見られる特有の「リンギング」を唯一再現します（図4）。
• 損失予測: モデルは、特徴的なスパイクを含む訓練損失の軌跡を正確に予測します。
• 安定化メカニズム: 結合されたダイナミクスは、明示的な投影制約なしに、システムがいかに自然に $2/\eta$ 閾値付近に落ち着くかを実証的に示しています。
• 離散化の影響: 実験により、離散化の解像度（$K$）を高めることで、自己安定化メカニズムを変えることなく損失スパイクを排除できることが示されました。これにより、スパイクがEoS現象の本質的な性質ではなく、数値的なアーティファクトであることが確認されました。

意義と主張

本論文は、Edge Flowが大きなステップサイズ領域における不安定性を理解し、軽減するための原理的な枠組みを提供すると主張しています。ダイナミクスを中心、方向、振幅に分解することで、複雑な投影フローや剛体ロッドモデルに代わる、扱いやすい選択肢を提示しています。

著者らは、Edge Flowが忠実（過渡現象を含む個々の軌跡を追跡する）、予測的（ノイズ注入や適応的ステップサイズなどの介入を提案できる）、そして扱いやすい（シミュレーションが安価である）ことを強調しています。彼らは、このモデルが以下のための現実的な経路を開くと考えています：

• 大きな学習率に対する暗黙的な正則化結果の証明。
• 損失の曲率、アーキテクチャ、および汎化性能の相互作用の理解。
• EoSダイナミクスの硬い性質に特化した、より効率的な離散化スキームの設計。

本論文は、導出において高次の結合補正を無視していることや、厳密な誤差解析および確率的な設定（SGD）や適応的手法への拡張が今後の課題であることを述べ、その範囲について謙虚な姿勢を保っています。