Gradient-Direction Sensitivity Reveals Linear-Centroid Coupling Hidden by Optimizer Trajectories

論文「Gradient-Direction Sensitivity Reveals Linear-Centroid Coupling Hidden by Optimizer Trajectories（勾配方向の感度が、オプティマイザの軌跡に隠された線形重心の結合を明らかにする）」の詳細な技術的サマリーを以下に示す。

1. 問題定義

機械的解釈可能性において、パラメータ空間内のどの方向が特定の機能（回路）の形成に対応するかを特定することが中心的な問いである。一般的な診断アプローチには、ローリングウィンドウを用いた特異値分解（SVD）によるオプティマイザの軌跡（パラメータ更新、$\Delta \theta$）の低ランク構造の分析が含まれる。この手法はUpdate SED（Spectral Edge Diagnostics）として知られており、オプティマイザが移動する方向が、機能形成に関連する方向を忠実に反映していると仮定している。

著者らは、この仮定が現代のトレーニングパイプライン（モメンタム、適応的スケーリング、重み減衰を備えた AdamW を使用）において、特にマルチタスク設定において欠陥があると主張する。彼らは、オプティマイザの更新が基礎となる勾配信号を集約・歪曲し、軌跡に基づく診断と実際の機能形成との間の整合性を隠蔽または破壊する可能性があると指摘している。

2. 手法

本研究は、2 層トランスフォーマーアーキテクチャ内のモジュラー演算タスク（加算、減算、乗算、および $a^2+b^2$）におけるgrokkingを用いてこの仮説を検証する。著者らは、タスク固有の重心の感度を最大化する方向に沿って機能が形成されると仮定する**線形重心仮説（LCH）**に対して、2 種類のスペクトル推定量を比較する。

比較された主要な推定量：

1. Update SED: 実際の AdamW パラメータ更新（$\Delta \theta_t$）のローリングウィンドウに対して SVD を実行する。これはオプティマイザがたどる「経路」を捉える。
2. Gradient SED: 現在のパラメータで評価された生勾配（$g_t = \nabla_\theta L$）のローリングウィンドウに対して SVD を実行する。これは、オプティマイザのアーティファクト（モメンタム、適応的スケーリング、重み減衰）を取り除いた損失関数の「最急降下」方向を捉える。
   • マルチタスク拡張: マルチタスク設定では、著者らは集約された総損失勾配を使用するのではなく、各タスクの損失を個別に逆伝播させることでタスク別勾配（$g_{op}$）を計算する。

診断指標：

スペクトル方向 $v_k$ と機能形成との結合は、摂動比 $R_k$ によって測定される：
$$R_k(t) = \frac{A(v_k, t)}{\text{median}_j A(r_j, t)}$$
ここで、$A(v, t)$ は注意パラメータを方向 $v$ に沿って摂動させた際のタスク重心の変化を測定する。高い $R_k$ は、方向 $v$ が機能形成に対して非常に敏感であることを示す。

因果的介入：

著者らは、SED 基底またはランダム基底によって定義された低ランク部分空間（ランク 3）に注意更新を制限することで因果的介入を行う。彼らは以下の 2 つの投影法をテストする：

• Update Projection: 最終的な AdamW 更新を投影する。
• Gradient Projection: AdamW オプティマイザがそのモメンタムと適応的スケーリングのモーメントを計算する前に勾配を投影する。

3. 主要な結果

A. 単一タスク結果：勾配対更新

• Update SED の失敗: 単一タスクモデルにおいて、Update SED は中程度の結合比（$\bar{R}_k \approx 3\text{--}9\times$）をもたらす。加算は異常に結合しているように見えるが、乗算は弱い。
• Gradient SED の成功: 更新を生の勾配に置き換えることで、結合比が30–100$\times$ 増加する（$\bar{R}_k \approx 100\text{--}330\times$）。
  • 乗算が最も強く結合されたタスクとなり、Update SED のランキングが逆転する。
  • 「op-spread」（タスク間の分散）は著しく縮小する。
• ローリングウィンドウの役割: $W=1$（瞬間的勾配）における高い結合は、勾配方向自体が主要な駆動力であることを示唆する。ローリングウィンドウ（$W=20$）は主にノイズ除去メカニズムとして機能し、比率を単調に減少させるが、ランダムな場合よりも桁違いに高いシグナルを維持する。

B. マルチタスク結果：集約の決定的な失敗

• Update SED の崩壊: マルチタスクトレーニングにおいて、Update SED は完全に失敗し、$\bar{R}_k \le 1$（ランダム以下）となる。4 つの競合する目的をバランスさせる集約されたオプティマイザ更新は、単一のタスクの機能関連部分空間に対して直交してしまう。
• Gradient SED の回復: タスク別勾配 SEDを使用することで、強い結合（$\bar{R}_k \approx 20\text{--}45\times$）が回復する。
  • 決定的なことに、マルチタスク勾配 SED シグナルは、単一タスクの Update SED シグナルよりも強い。
  • これは、機能形成がマルチタスク設定において頑健であるが、オプティマイザの更新段階における競合勾配の集約によって隠蔽されていることを証明する。

C. 因果的介入：ランク対方向

• 加速: ランク 3 部分空間に更新を制限することで、grokking が約2.3$\times$ 加速する。
• 方向の独立性: この加速は、部分空間が機能整合的なSED 基底であるか、ランダム基底であるかに関わらず発生する。
  • Gradient Projection: 最適化前に勾配を投影する場合、SED 部分空間を除去し（ランダムのみを保持する）ことは悪影響を与えず、SED 基底を保持する場合と同様に、ランダムなランク 3 のみを保持することがトレーニングを同様に加速する。
  • Update Projection: 初期の実験では SED を除去することが有益であるように見えたが、これは AdamW の内部モメンタム/適応的スケーリングが軌跡に情報を漏洩させるアーティファクトであった。勾配投影はこのアーティファクトを排除する。
• 結論: 加速の因果的要因は特定の方向ではなく、ランク制約（冗長性の削減）である。注意に対する自然な AdamW 更新は「ランク冗長」であり、grokking 遷移に対しては主に非機能的な運動を含んでいる。

4. 主要な貢献

1. 方法論的転換: 機能形成のスペクトル分析の正しい対象は、オプティマイザの軌跡ではなく、勾配構造であることを実証する。オプティマイザの更新は、機能関連方向を隠蔽する質的な歪曲をもたらす。
2. マルチタスク解釈可能性: 競合するタスク間の勾配集約を、回路レベルの解釈可能性に対する一般的な障壁として特定する。マルチタスクトレーニングにおける機能構造を回復するには、タスク解決型の勾配分析が必要である。
3. 因果的明確化: 勾配由来の方向は機能形成の強力な診断であるが、それらが唯一の因果的要因ではないことを示す。重要な要因は勾配基底との特定の整合性ではなく、更新の低ランク性である。
4. アーティファクトの特定: grokking における以前の「ランク 1 多様体」の観測は、AdamW 更新に固有のノイズと歪曲によって部分的に隠蔽されていたことを明らかにする。

5. 意義

• 解釈可能性にとって: この研究は、回路を見つけるためのパラメータ更新軌跡の分析（例えば、重み差の PCA）という標準的な実践に挑戦する。そのような手法は、特に複雑なマルチタスクや大規模設定（エキスパート混合など）において、機能形成を体系的に過小評価すると主張する。
• トレーニングダイナミクスにとって: AdamW などのオプティマイザがたどる「自然な」経路は、しばしば非効率的で冗長であることを示唆する。最適化をより低いランクに制限することで、性能を犠牲にすることなく学習を大幅に加速できる。
• 将来の方向性: この知見は、大規模モデル（LLM、指示微調整など）において、集約された軌跡分析に依存するのではなく、特定の回路がどのように形成されるかを理解するために、タスク別勾配（または経験的フィッシャーに対する Lanczos 法などの近似）を分析することが不可欠であることを示唆している。

要約すると、この論文は勾配がシグナルであり、オプティマイザがノイズであると主張する。ニューラルネットワークが機能をどのように学習するかを理解するためには、オプティマイザがパラメータ空間を通過する歪んだ経路ではなく、特にタスクごとに分解された生の勾配の地形を分析しなければならない。