The Loss Does Not See the Basis, but Adam Does

技術要約：「損失は基底を見ないが、Adamは見ている」

問題提起
因子分解モデル（$W = UV^\top$）における勾配降下法（GD）は、低ランク解への暗黙的なバイアスを示すことが知られており、劣決定の行列センシング・タスクにおいて最小核ノルム補間解に収束することが多い。しかし、Adamのような適応型オプティマイザは、同様の訓練損失を達成するにもかかわらず、しばしばこれらの低ランク構造の回復に失敗し、実効ランクが高く汎化性能の低い解をもたらす。近年の研究では、適応型手法が基底依存的（座標回転に敏感）であることが特定されているが、この依存性が因子分解モデルにおける低ランクの帰納的バイアスをどのように破壊するのかという具体的なメカニズムは、これまで明らかにされていなかった。本論文は、以下の問いに取り組む：オプティマイザが低ランクのバイアスを保持するか破壊するかを決定する基準はあるのか、そしてその基準によって、展開された手法をその回復性能によって分類できるのか？

手法および理論的枠組み
著者は、オプティマイザを分類するための中心的な基準として**ゲージ等変性（gauge equivariance）**を導入している。損失関数 $L(U,V) = f(UV^\top)$ は、ゲージ対称性を有している。すなわち、任意の直交行列 $Q \in O(k)$ に対して、変換 $(U,V) \mapsto (UQ, VQ)$ に対して不変である。この対称性は、潜在的な因子の基底を回転させてもモデルの予測は変化しないことを意味する。

1. 更新ルールの分類：

   • 等変オプティマイザ（Equivariant Optimizers）： オプティマイザがゲージ等変であるとは、その更新ルールがゲージ変換と可換であることを指す。本論文では、メモリレスな更新ルールがゲージ等変であるための必要十分条件は、グラム決定型の左前置行列 $\Delta = H(GG^\top)G$ の形式をとることであることを証明している（ここで $H$ はグラム行列 $GG^\top$ にのみ依存する）。
   • 非等変オプティマイザ（Non-Equivariant Optimizers）： 初期化時の勾配やオプティマイザの状態に対して固定の要素ごとの非線形性を適用する手法（標準的なAdam、RMSProp、Lion、SignSGDなど）は、この対称性を破壊する。具体的には、Adamの座標ごとの二次のモーメントへの適応は、直交回転と可換ではなく、これにより「第一ステップの欠陥（first-step defect）」が生じる。つまり、初期化を回転させるだけで、積行列 $W$ が即座に変化してしまうのである。

2. 転移定理（Transfer Theorem）：
   著者は、共有スカラー前置行列（すべての座標に対して単一のスカラー分母を持つもの）を持つオプティマイザは、再スケールされた時間パラメータ化の下での勾配流に対応することを示す転移定理を確立した。その結果、これら等変な手法に関する勾配流の低ランク・バイアスに関する理論的結果（Gunasekar et al., 2017; Li et al., 2021など）が直接適用可能となる。

3. スペクトル・スケジュール（Spectral Schedule）：
   等変性は低ランク多様体へアクセスするために必要であるが、低ランク回復には十分ではない。等変クラス内での具体的な挙動は、スペクトル・スケジュール（特異値がいかに成長するか）によって決定される。アグレッシブな等レート・スケジュール（Muonなど）は、純粋な低ランクターゲットに対してはほぼ正確な回復を実現できるが、ターゲットに顕著な「スペクトルの裾（spectral tail）」（低ランク部分以外のエネルギー）が存在する場合、失敗することがある。一方で、貪欲なスケジュール（GDなど）はこれらの裾を抑制する。

主要な貢献と実験

• オプティマイザ・ズー（Optimizer Zoo）の比較： 著者は、劣決定の行列センシング・タスク（$40 \times 40$、ランク3、$m=2 \times \text{dof}$）において、9つのオプティマイザを完全な補間まで評価した。

  • 結果： 明快な分離が見られた。等変な手法（GD、Scalar-Adam、Muon、Shampoo）は、回復誤差 $\le 0.286$ を達成した。一方、座標ごとの手法（Adam、RMSProp、Lion、SignSGD、Adafactor）は、低ランクのアトラクタに到達できず、誤差 $\ge 0.42$ となった。
  • Muonの性能： Muonは、純粋な低ランク・ターゲットに対しては極めて正確な回復（誤差 $6.8 \times 10^{-6}$）を達成したが、ターゲットのスペクトルの裾が増大すると急速に劣化し、裾のエネルギーが約4%に達した時点でGDの性能へとクロスオーバーした。

• 「ダイヤル（Dial）」実験： Adamの失敗の原因を分離するため、標準的なAdam（$p=1$、座標ごと）とScalar-Adam（$p=0$、共有スカラー）を補間する一パラメータ・ファミリーのオプティマイザ（Adam-$p$）を構築した。

  • 結果： $p$ が1から0へと減少するにつれて、回復誤差が減少し、実効ランクが向上するという単調な改善が見られた。これにより、前置行列における**座標ごとの異方性（coordinate-wise anisotropy）**こそが、低ランク・バイアスを破壊する具体的なメカニズムであり、それがモーメンタムや適応性とは独立したものであることが確認された。

• アテンション・ヘッド解析： 本論文は、トランスフォーマー内の個々のアテンション・ヘッドも、同一のゲージ対称性（$W_Q^\top W_K$）を共有していることを示している。

  • 結果： 2つのゲージ等価なトランスフォーマー・コピー（同一の関数だが異なる基底を持つもの）をAdamで訓練したところ、わずか1ステップで構造的に乖離し、ヘッドの不変量（$W_Q^\top W_K$）において56%の相対フロベニウス距離が生じた。等変なオプティマイザ（SGD、Scalar-Adam）は、マシン精度で区別がつかない状態を維持した。著者は、アライメント技術（permutation re-basinなど）はゲージを一致させることはできるが、不変量そのものが異なってしまった場合にはそれらを一致させることはできないと指摘している。つまり、基底依存の訓練は、モデル・マージングに必要な合意を損なう（どの回転を用いても不変量の差を埋めることができないため）。

• 実世界での検証： ハイパースペクトル画像補完タスク（Indian Pines, Pavia University）において、訓練損失を一致させた条件下で検証を行った。

  • サンプリング密度が低い場合、GDはAdamよりも保持誤差が43〜44%低かった。これは、GDが有意に低い実効ランクを持っていることに対応している。
  • Adamのテスト誤差は、より深く適合する（スペクトルの裾をフィッティングする）につれて増加したが、GDの誤差は減少した。

• アルゴリズムの修復： 著者は、自身の以前の研究であるFlowAdamに対してこの基準を適用した。座標ごとの速度クリッピング（velocity clipping）が等変性を壊していることを特定し、これをグローバル・ノルム・クリップに置き換えることで、低ランク・バイスを回復させ、回復誤差を37%改善し、GDの性能に近づけることに成功した。

意義と主張
本論文は、基底の選択は単なるチューニングの詳細ではなく、どの補間解を選択するかに関する根本的な決定であると主張している。

• メカニズムの説明： どのようなオプティマイザが因子分解モデルの低ランク帰納的バイアスを保持するかを予測するための、テスト可能かつ構成的な基準（ゲージ等変性）を提供している。
• 限界： 著者はその適用範囲について謙虚である。このバイアスは緩やかなものであり、最適な正則化が低く、かつ低ランク構造と整合している場合に主に有益である。強い正則化が必要な場合や、座標ごとの適応が有益な場合（PINNsにおける硬い問題など）には、座標ごとの手法が等変な手法よりも優れている可能性がある。
• ベンチマークの主張ではない： 本研究は新しい最先端のオプティマイザを提唱するものではなく、特定の状況下でなぜ特定のメソッド（GDやMuonなど）が成功し、他のメソッド（Adamなど）が失敗するのかという観察された現象に対するメカニズム的な説明を提供することを目的としている。

要約すれば、本論文は「損失は基底を見ない（ゲージ不変である）が、Adamは見ている（更新は基序依存的である）」と論じており、この依存性が、因子分解モデルにおける低ランク補間を選択することに失敗する根本原因であると結論付けている。