技術要約：安定性の境界（Edge of Stability）はデータ分布全体にわたって学習を選択的に形成する

1. 問題提起

既存の文献における「安定性の境界（Edge of Stability, EoS）」は、主に最適化のグローバルな特性として扱われてきた。そこでは、最大ヘッセ行列固有値が閾値 $2/\eta$（$\eta$ は学習率）付近で自己安定化する現象が中心である。このレジームにおけるダイナミクス（符号が交互に反転する振動と、緩やかな修正ドリフトによって特徴付けられる）は十分に文書化されているが、この閾値で動作することの機能的な帰結については不明なままである。具体的には、安定性の制約が均一な利益をもたらすのか、あるいは異なるデータサブセットに対して差別的な影響を与えるのかが分かっていない。中心となる問いは、**「EoSレジームは選択的な学習バイアスを誘発するのか、もしそうであれば、どのような幾何学的および動的な特性がどのデータサブセットに利益をもたらすのか？」**である。

2. 手法

2.1 実験設定とプロトタイプ分類学

著者らは、CIFAR-10（自動車 vs トラック）を用いた二値分類タスクと、フルバッチ勾配降下法で訓練された2層隠れ層MLPを用いた、制御された実験フレームワークを構築している。データの幾何学的影響を分離するため、彼らはクラス重心（$\mu_c$）に対する位置とラベルの一貫性に基づき、訓練分布を以下の4つのグループに分割するプロトタイプ分類学を定義した（これらは訓練モデルとは独立している）：

1. インライア（Inliers）: クラス重心付近の高密度点であり、正しいラベルを持つ。
2. 境界（Boundary）: クラス間の境界付近にある分布内（in-distribution）の点（ラベルの曖昧性が高い）。
3. 入力アウトライア（Input-outliers）: 幾何学的に非典型的な入力（$\mu_c$ から遠い）であり、正しいラベルを持つ。
4. 出力アウトライア（Output-outliers）: 誤ったラベルが付与された高密度な入力（合成的に構築されたもの）。

2.2 分岐介入（因果的特定）

EoSレジームに留まることと、そこから脱出することの効果を因果的に特定するために、著者らは**分岐介入（branching intervention）**を採用している：

• ベースライン・ブランチ: 元の学習率 $\eta$ で訓練を継続し、EoSレジームに留まる。
• 脱出ブランチ: EoSの発生時（最大ヘッセ固有値 $\lambda_1$ が $2/\eta$ に達した時点）に、学習率を半分にする（$\eta \to \eta/2$）。これにより、安定性閾値が $4/\eta$ に上昇し、オプティマイザがEoSレジームから脱出することを可能にする。
• 比較: 介入後の各グループの損失（$\ell_k$）の乖離を比較することで、著者らはEoSの安定性制約による因果的な影響を分離する。

2.3 メカニズムの分離

著者らは、あるサブセットがEoSから利益を得るための必要条件として、以下の2つの仮説を検証している：

1. 方向的一致（Directional Alignment）: サブセットの集約された勾配が、最大ヘッセ固有ベクトル（$v_1$）と整列している必要がある。これは、変位距離は維持しつつも方向をランダム化することで一貫した整列を破壊した「ランダム方向」の入力アウトライアを構築することでテストされる。
2. 勾配の持続性（Gradient Persistence）: サブセットが消失しない勾配の大きさを維持する必要がある。これは、MSE損失（勾配が持続する）と、Cross-Entropy（CE）損失（自信を持って分類された点で勾配が消失する）の比較によってテストされる。

3. 主な結果

3.1 選択的学習のトレードオフ

分岐介入により、EoSはグローバルではなく選択的であることが明らかになった。安定性の制約は、最適化の努力を不均等に再分配する：

• 受益者: 勾配が $v_1$ と整列し、かつ持続するサブセット（例：MSEにおける入力アウトライア、CEにおける出力アウトライア）は、脱出ブランチと比較して損失の減少が加速する。
• 抑制対象: 勾配が直交している、あるいは勾配の大きさが消失しているサブセット（例：特定の条件下でのインライアや境界点）は、脱出ブランチと比較してEoSブランチにおいて進捗が抑制される。
• このトレードオフは、アーキテクチャ（CNN, ResNet）、オプティマイザ（SGD, Momentum）、およびクラスペアを横断して再現される。

3.2 メカニズム：整列 $\times$ 持続

あるグループがEoSの利点を享受するためには、以下の2つの条件が共同で必要となる：

1. 整列: グループの集約された勾配が最大ヘッセ固有ベクトル（$v_1$）と整列していること。
   • 証拠: 入力アウトライアを（距離を保ちつつ方向をランダムに変えて）変位させたとき、その曲率の影響 $(\nabla \ell_k \cdot v_1)^2$ は1桁低下し、EoSの利点は消失する。
2. 持続性: グループが消失しない勾配の大きさを維持していること。
   • 証拠: Cross-Entropy損失の下では、自信を持って分類された入力アウトライアの勾配は飽和（消失）する。その結果、EoSの利点は、ラベルの不一致により勾配がアクティブなままである出力アウトライアへとシフトする。

3.3 幾何学が受益者を決定する

どのデータサブセットがEoSから利益を得るかは、データ分布の幾何学的構成によって決定される：

• 標準的なCIFAR-10の設定では、入力アウトライア（重心から最も遠い点）が $v_1$ を支配し、利益を得る。
• より困難なクラスペア（猫 vs 犬）や、異なるアウトライア構築パラメータを用いる場合、境界点が最も遠いグループとなり、$v_1$ を支配して主要な受益者となる可能性がある。
• 自然なデータ: 自然なCIFAR-10分布において、EoSの開始時に重心からの距離と $v_1$ への整列との間に正の相関が見られ、幾何学的なアウトライアが自然に不安定モードを駆動していることが確認された。

3.4 一般化とロバスト性の含意

本論文は、訓練時の選択性がテスト時の挙動に転移するという予備的な証拠を提示している：

• 敵対的ロバスト性: 境界点が $v_1$ を支配する場合にのみ、EoSは境界点に対するロバスト性を向上させる。もし入力アウトライアが $v_1$ を支配している場合、最適化の予算がそちらに転用されるため、EoSブランチは境界のロバスト性に関して性能が悪化する。
• 分布外（OOD）一般化: 外挿された入力が訓練中に $v_1$ を支配している場合にのみ、EoSはOOD性能を向上させる。

4. 理論的枠組み

著者らは、自己安定化の枠組み（Damian et al., 2023）を、グローバルな損失からサブセットごとの損失へと拡張している。彼らは、EoSレジームにおけるサイクル平均ドリフトには、$-\nabla S$（ここで $S$ はシャープネス）に比例する追加項が含まれることを導出した。サブセット $k$ への微分的な影響は、セレクターによって制御される：
$$Q_k = \langle \nabla \ell_k, \nabla S \rangle$$
単一モード近似の下では、このセレクターは最大固有ベクトルへの勾配の投影の二乗に比例する：
$$C_k = (\nabla \ell_k \cdot v_1)^2 = \|\nabla \ell_k\|^2 \cos^2 \theta_k$$
この分解は、あるサブセットがEoSのドリフトを経験するためには、方向的一致（$\cos^2 \theta_k$）と勾配の大きさの持続性（$\|\nabla \ell_k\|^2$）の両方が必要であることを裏付けている。

5. 意義と主張

本論文は、EoSの理解を、スカラー的な安定性の境界から、データ分布全体にわたる学習の割り当てを制御するメカニズムへと転換することを主張している。

• 平坦性（Flatness）の再定義: 著者らは、「平坦さ」は解の単なるスカラー特性ではなく、データ幾何学によって決定される方向的な特性であると主張している。低曲率の機能的利益は、どのサブセットが最大ヘッセ固有方向を支配するかに完全に依存する。
• 帰納バイアス: EoSは、同等の極小値の間で選択を行うのではなく、幾何学的な関係に基づいて特定のサブセットを最適化の対象として選択する帰納バイアスとして機能する。
• 視点の統合: 本研究は、パラメータ空間における暗黙的な正則化（曲率制約）と、データ分布に対する帰納バイアス（どの例が学習されるか）を結びつけている。
• 限定的な範囲: 著者らは、ヘッセ行列とベクトルの積の計算コストのため、実験が小規模なモデルとフルバッチ訓練に限定されているという限界を認めている。これらの選択的効果が、より大規模なデータセット、ミニバッチオプティマイザ、または表現空間のプロトタイプにおいても持続するかどうかは、未解決の課題であるとしている。

結論として、本論文は、安定性の境界が、安定性制約が最大ヘッセ固有方向に対して整列し、かつ持続するサブセットの学習を増幅させ、他のサブセットを抑制するという、選択的なレジームであることを示しており、それによってモデルの機能的特性を根本的に形成していることを実証している。