Non-Euclidean Gradient Descent Operates at the Edge of Stability

本論文は、方向性滑らかさの概念を非ユークリッドノルムに拡張して一般化された鋭敏性を定義し、これが勾配降下法だけでなく$\ell_{\infty}$-降下やブロック座標降下など多様な最適化手法においても「安定性のエッジ」現象を統一的に説明できることを理論的・実験的に示しています。

要約

🏔️ 論文の核心：AI は「安定の限界」で踊っている

AI を学習させるには、誤り（損失）を減らすためにパラメータを少しずつ調整する必要があります。これを「勾配降下法（Gradient Descent）」と呼びます。

これまでの常識では、「ステップ（一歩の大きさ）が大きすぎると、山を越えすぎて転落してしまう（発散する）」と考えられていました。しかし、最近の研究で**「AI は、転落しそうなギリギリの崖っぷち（Edge of Stability）で、奇妙なダンスをしながら学習を進めている」**ことがわかってきました。

この論文は、その「崖っぷちのダンス」が、「普通の足（ユークリッド空間）」だけでなく、「特殊な靴（非ユークリッド空間）」を履いた場合でも同じように起こることを発見し、それを説明する新しい理論を提案しました。

---

🧐 3 つの重要な発見

1. 「滑らかさ」の新しい定義：道が曲がりくねっている度合い

AI が学習する際、地形（損失関数）は非常に複雑です。

• 従来の考え方: 「この道は全体的に滑らかだから、ゆっくり歩けば大丈夫」と考えます。
• この論文の視点: 「道は全体的に滑らかじゃなくても、今、私が歩いているこの一歩の区間だけを見れば、どれくらい急勾配か（曲がり具合）が重要だ」と考えます。

これを**「方向性のある滑らかさ（Directional Smoothness）」**と呼びます。

• 例え話: 雪原をスキーで滑るとき、全体が急斜面じゃなくても、今滑っているこのスロープだけが急すぎると転びます。AI はこの「今滑っているスロープの急峻さ」を常にチェックしています。

2. 「鋭さ（Sharpness）」の再定義：どんな靴でも通用するもの

AI が「崖っぷち」にいるかどうかを判断する指標として、「鋭さ（Sharpness）」という値が使われます。

• 従来の鋭さ: 地形の「一番急な傾き」を見る（ユークリッド距離を使う）。
• この論文の新しい鋭さ: **「どんな靴（ノルム）を履いていても通用する鋭さ」**を定義しました。

例えば：

• 普通の靴（ℓ2 ノルム）: 標準的な AI 学習。
• スパイク付きの靴（ℓ∞ノルム）: 特定の方向にだけ強く進む方法（SignGD など）。
• スケート靴（スペクトルノルム）: 行列の構造を考慮する方法（Muon など）。

この論文は、**「どんな靴を履いていても、AI はその靴に合わせた『鋭さ』を 2/η（ステップの大きさの逆数）という限界値に近づけ、そこで安定しようとする」**と示しました。

3. 驚きの現象：崖っぷちで「振動」しながら進む

AI は、この限界値（2/η）に近づくと、以下のような現象を起こします。

1. 鋭さが上がる: 地形が急になってくる。
2. 限界に到達: 急すぎて、一歩踏み出すと少し戻ってしまう（損失が増減する）。
3. 安定した振動: しかし、AI は転落せず、**「急斜面と緩斜面を行き来するリズム」**を見つけ出し、そのリズムの中で学習を続けます。

まるで、**「崖の縁を歩く達人」が、バランスを崩しそうになったら微調整して、また元の位置に戻るような、「安定した不安定さ」**の中で進んでいるのです。

---

🎒 具体的な実験結果

研究者たちは、従来の方法だけでなく、以下のような「特殊な靴」を履いた AI にも同じ現象があるか実験しました。

• ℓ∞-descent（スパイク靴）: 重みの符号（プラス/マイナス）だけを見て進む方法。
• Block CD（ブロック靴）: 層ごとにまとめて更新する方法。
• Spectral GD（スケート靴）: 行列の性質を考慮して更新する方法（Muon など）。

結果：
どの靴を履いていても、AI は**「その靴に合わせた『鋭さ』の指標」が、「2/η という限界値」の周りで振動しながら学習を続けることが確認されました。
つまり、「AI が崖っぷちで踊っている現象は、学習のアルゴリズムが何であれ、普遍的に起こっている」**ことがわかりました。

---

💡 なぜこれが重要なのか？

1. 新しい AI 開発の指針: これまで「なぜこの新しい最適化アルゴリズム（Muon など）がうまくいくのか」理論的に説明が難しかった部分に、**「崖っぷちの安定性」**という共通の理由が見つかりました。
2. 安全な学習: 「転落（発散）しないギリギリのライン」で学習させることで、より効率的に AI を鍛えることができます。
3. 統一された視点: 様々な異なる学習アルゴリズムを、一つの「幾何学的な視点」で理解できるようになりました。

📝 まとめ

この論文は、**「AI の学習は、滑らかな平地を歩くことではなく、急峻な崖の縁を、バランスを取りながら（振動しながら）進む高度な技術である」**と教えてくれました。

そして、**「どんな歩き方（アルゴリズム）を選んでも、AI はその歩き方に合わせた『バランスの限界値』を見つけ出し、そこで安定して学習を続ける」**という驚くべき共通法則を発見しました。

これは、AI のブラックボックスだった「なぜうまく学習するのか」という謎に、**「崖っぷちのダンス」**という美しいメタファーで光を当てた画期的な研究と言えます。

技術概要

論文要約：非ユークリッド勾配降下法は安定性のエッジで動作する

(Non-Euclidean Gradient Descent Operates at the Edge of Stability)

1. 概要と背景

この論文は、深層学習における最適化手法の重要な現象である「安定性のエッジ（Edge of Stability: EoS）」を、従来のユークリッド空間（$\ell_2$ノルム）に限定されずに、**任意のノルムに対する非ユークリッド勾配降下法（Non-Euclidean GD）**へ一般化する研究です。

EoS とは、勾配降下法（GD）の学習中に、損失関数の減少が単調でなくなり、損失が振動しながらも長期的には減少し続ける現象を指します。この際、ヘッセ行列の最大固有値（鋭度：Sharpness）が、ステップサイズ $\eta$ を用いて $2/\eta$ という閾値付近に収束・振動することが知られています。従来の研究は主に標準的な GD やアダプティブ手法（Adam, Adagrad など）に焦点を当てていましたが、Muon や SignGD などの非ユークリッド的な手法における EoS の振る舞いは未解明でした。

2. 問題定義と手法

2.1 非ユークリッド勾配降下法の定式化

著者らは、任意のノルム $\|\cdot\|$ に対する非ユークリッド GD を以下のように定義します。
現在の点 $w_t$ 周りの正則化された線形化を最小化する更新則です：
$$w_{t+1} = \arg\min_y \left( \langle \nabla L(w_t), y - w_t \rangle + \frac{1}{2\eta} \|y - w_t\|^2 \right)$$
これにより、双対ノルム $\|\cdot\|_*$ と双対勾配を用いて、更新方向 $d_t$ が決定されます。

• $\ell_2$ ノルム: 標準的な勾配降下法（Vanilla GD）に帰着。
• $\ell_\infty$ ノルム: $\ell_\infty$-descent（SignGD と関連）に帰着。
• スペクトルノルム ($\|\cdot\|_{2\to2}$): Muon アルゴリズム（モメンタムなし）に帰着。
• ブロック $\ell_{1,2}$ ノルム: ブロック座標降下法（Block CD）に帰着。

2.2 方向性滑らかさ（Directional Smoothness）の解釈

EoS のメカニズムを理解するために、Mishkin et al. [2024] が提案した方向性滑らかさ（Directional Smoothness） $D_{\|\cdot\|}(w, y)$ を導入します。これは、2 つの連続する反復点間の損失関数の平均曲率を表します。
$$D_{\|\cdot\|}(w, y) := \frac{L(y) - L(w) - \langle \nabla L(w), y - w \rangle}{\frac{1}{2}\|y - w\|^2}$$
この量を用いると、損失が減少するための条件は $D_{\|\cdot\|} \leq 2/\eta$ となります。損失が振動し始める（EoS に達する）と、この方向性滑らかさは $2/\eta$ 付近で振動することが示されます。

2.3 一般化された鋭度（Generalized Sharpness）の定義

方向性滑らかさを拡張し、任意のノルム $\|\cdot\|$ に対する一般化された鋭度 $S_{\|\cdot\|}(w)$ を以下のように定義します。
$$S_{\|\cdot\|}(w) := \max_{d \neq 0} \frac{d^\top \nabla^2 L(w) d}{\|d\|^2} = \max_{\|d\| \leq 1} d^\top \nabla^2 L(w) d$$

• ユークリッドノルムの場合、これはヘッセ行列の最大固有値 $\lambda_{\max}(\nabla^2 L(w))$ に一致します。
• 非ユークリッドノルムの場合、この最適化問題（凸集合上の二次関数の最大化）は一般的に NP 困難ですが、Frank-Wolfe アルゴリズムを用いて近似計算可能です。

3. 主要な貢献

1. EoS の一般化: 方向性滑らかさの概念を非ユークリッドノルムに拡張し、任意のノルムにおける GD のダイナミクスを統一的に記述する枠組みを提案しました。
2. 理論的洞察: 損失が減少する（または振動する）条件が、方向性滑らかさが $2/\eta$以下（または付近）であることと等価であることを示しました。これにより、EoS 現象が「鋭度が閾値$2/\eta$ に達するまで鋭くなり（Progressive Sharpening）、その付近で振動する」というメカニズムが、非ユークリッド手法でも同様に成立することを理論的に裏付けました。
3. 新しい鋭度の定義: 従来の $\ell_2$ 鋭度では捉えられなかった手法（$\ell_\infty$-descent, Block CD, Spectral GD/Muon など）においても、一般化された鋭度が $2/\eta$ 付近で安定することを定義しました。
4. 理論的保証の拡張: 二次関数に対する非ユークリッド GD の収束・発散条件を証明しました。ステップサイズ $\eta < 2/S$ なら収束し、$\eta > 2/S$ かつ特定の初期化からは発散することを示しました。

4. 実験結果

MLP、CNN、Transformer などの多様なアーキテクチャにおいて、以下の手法で実験を行いました：

• $\ell_\infty$-descent (SignGD)
• Block Coordinate Descent (Block CD)
• Spectral GD (Muon)
• 正規化された非ユークリッド GD

主な発見:

• Progressive Sharpening と EoS: 全ての手法において、損失が減少する初期段階で一般化された鋭度が上昇し（Progressive Sharpening）、その後 $2/\eta$ 付近で振動する EoS 現象が観測されました。
• 閾値の一致: 損失の振動と鋭度の振動は、常に $2/\eta$ という閾値を基準に行われています。
• 標準的鋭度との対比: 非ユークリッド手法（特に $\ell_\infty$ や Spectral ノルム）において、従来の $\ell_2$ 鋭度（最大固有値）は閾値 $2/\eta$ よりもはるかに低い値で推移し、EoS 現象を説明できないことが確認されました。一方、提案する一般化された鋭度は明確に閾値に収束します。
• 予備的振動領域: ユークリッド GD とは異なり、$\ell_\infty$ や Spectral GD においては、鋭度が閾値に達する前（$S < 2/\eta$）でも、方向性滑らかさが上昇し始め、重み空間で振動が始まる「中間的な振動領域」が存在することが観測されました。

5. 意義と結論

この研究は、深層学習における最適化の「安定性のエッジ」現象が、特定の最適化手法（Adam や SGD など）に限定されたものではなく、幾何学的な構造（ノルム）に依存した普遍的な現象であることを示しました。

• 理論的意義: 非ユークリッド空間における最適化の収束性と安定性の境界を、一般化された鋭度という単一の指標で統一的に説明する枠組みを提供しました。
• 実用的意義: Muon や SignGD などの最近の注目手法が、なぜ大規模モデルの学習で有効に機能するのか、その背後にある「安定性のエッジ」での動作メカニズムを解明しました。
• 今後の課題: 非ユークリッド GD における EoS の安定化メカニズム（なぜ鋭度が $2/\eta$ 付近で止まるのか）の詳細な理論的証明や、中間的な振動領域の性質の解明が今後の課題として挙げられています。

要約すると、本論文は「最適化手法の幾何学的性質（ノルム）を変えても、学習ダイナミクスは $2/\eta$ という普遍的な閾値を中心に振る舞う」という重要な知見を提供し、深層学習の最適化理論の基盤を強化するものです。