技術要約：重みベクトルの大きさと方向を分離することによるニューラルネットワーク学習の改善

1. 問題提起

現代のニューラルネットワーク学習は、重み行列を単一のモノリシックなオブジェクトとして扱うAdamやMuonといったオプティマイザに依存しています。しかし、あらゆる重み行列 $W$ は本質的に、大きさ ($\|W\|$) と 方向 ($\hat{W} = W/\|W\|$) という2つの異なる量を持っています。標準的な最適化では、これら2つの量が結合されているため、不安定で予測不可能な学習ダイナミクスを引き起こします。

• 方向変化の依存性: 更新によって誘発される角度の変化は、現在の大きさに反比例します。その結果、学習率（LR）は方向の変化率を直接制御できず、層間や時間経過とともに変化する可能性があります。
• 意図しない大きさのドリフト: 更新は（スケール不変性やノイズにより）現在の重みに対して垂直になる傾向があり、放射状の勾配がない場合でも大きさが膨張します。この「暴走的な成長」により、実務家は学習を安定させるために重み減衰（weight decay）などの間接的なメカニズムや、初期の不安定性を防ぐためのウォームアップに頼らざるを得なくなります。
• スケーリングの限界: モデルの幅や深さを変更する際にハイパーパラメータの再調整を必要とする既存の手法（例：$\mu$P）が多く存在します。さらに、重みを固定ノルムの球面上に制約する最近のアプローチは、学習可能な大きさのパラメータを欠いており、スケールの制御を正規化層のみに委ねている場合があります。

2. 手法：大きさ–方向（MD）デカップリング

著者らは、各重み行列を固定ノルムの方向と学習可能な大きさのゲインへと分解し、それらを個別の学習率で更新する、オプティマイザの修正案である大きさ–方向（MD）デカップリングを提案しています。

コアとなる分解

各重み行列 $W$ は、次のように再パラメータ化されます：
$$W = \text{diag}(\gamma_{\text{row}}) \, \hat{W} \, \text{diag}(\gamma_{\text{col}})$$
ここで：

• $\hat{W}$ は方向であり、固定された超球面（例：フロベニウスノルム）上に制約されています。
• $\gamma_{\text{row}}$ および $\gamma_{\text{col}}$ は、それぞれ行と列に作用するベクトルである学習可能な大きさのゲインです。

最適化手順

この手法は、モデルのフォワードパスを変更することなく、完全にオプティマイザのステップ内で動作します（モデルは単一の融合されたテンソル $W$ を見ています）：

1. 復元: 融合された重み $W$ から $\hat{W}$ とゲイン $\gamma$ を復元します。
2. 勾配の分割: 勾配 $G = \partial L / \partial W$ を、方向勾配 ($\partial L / \partial \hat{W}$) とゲイン勾配（それぞれの軸にわたって合計されたもの）に分割します。
3. 方向の更新: ベースとなるオプティマイザ（Adam、Muonなど）を用いて $\hat{W}$ を学習率 $\eta_W$ で更新し、その後、固定ノルムを維持するために球面上に投影します。
4. ゲインの更新: ゲイン $\gamma$ を（正の値を保証するためにsoftplusのような滑らかな写像をオプションで使用して）別の学習率 $\eta_\gamma$ で更新します。
5. 再構成: 次のフォワードパスのために $W$ を再構成します。

主な特性

• オプティマイザに依存しない: この手法は、任意のベース・オプティマイザ（Adam, Muon, Shampooなど）のラッパーとして機能します。
• 直接的な学習率制御: 方向を球面上に固定することで、学習率は相対的な角度更新を直接設定できるようになり、現在の重みの大きさに依存する関係を取り除きます。
• ヒューリスティックの排除: 重みはすでにノルムが制約されているため、重み減衰（weight decay）の必要性がなくなり、また、初期の不安定な大きな更新が最初から制御されるため、ウォームアップも不要になります。
• きめ細かなスケール制御: 正規化層のゲインのみに依存する方法とは異なり、MDデカップリングは個々の行や列のスケールに対して明示的かつ学習可能な制御を提供し、モデルが特定の機能を増幅または減衰させることを可能にします。

3. 主な貢献と知見

実証結果

著者らは、高密度モデル（181Mから1.29Bパラメータ）および大規模な混合エキスパート（MoE）モデル（総計最大6.7Bパラメータ）にわたってMDデカップリングを検証しています。

• 性能向上: MDデカップリングは、適切にチューニングされたベースライン（AdamWおよびMuon）を上回ります。大規模なMoEモデルにおいて、MuonMD（MuonにMDデカップリングを適用したもの）は、AdamWと比較して約2倍少ない計算量で同等の損失に到達します。
• 学習率の転送: この手法は、モデルの幅を超えたゼロショット学習率転送を可能にします。小さなモデルで見つかった最適な行列学習率は、再調整なしで大幅に大きなモデルに対しても最適であり続けます。これは、幅と深さの両方のスケーリングにおいて成立します。
• ウォームアップフリーの学習: ウォームアップを廃止しても安定性が維持されるだけでなく、初期の学習ステップが減少した学習率によって無駄にされないため、最終的な損失が改善されることが実験で示されています。
• アブレーション研究の洞察:
  • 正規化軸: 方向をフロベニウスノルム（行列全体）によって制約することが最も柔軟であり、行/列の制約と同等の性能を示すことが分かりました。
  • ゲインのパラメータ化: 行と列を組み合わせたゲイン（$\gamma_{\text{row}} + \gamma_{\text{col}}$）は、スカラーまたは単一軸のゲインよりも優れた性能を示しました。ゲインのパラメータ化においてsoftplusを使用すると、直接的または指数的な更新よりもわずかに一貫した優位性があります。
  • 埋め込み（Embeddings）: 埋め込みベクトルを単位L2ノルムに制約することは、制約なしの場合よりもわずかに高い性能を示しました。
• 関連研究との比較: 本論文は、MDデカップリングを、球面学習や再パラメータ化に関する一連の研究の中に位置づけています。
  • nGPTとは異なり（nGPTはL2正規化や補間残差などの特定のアーキテクチャ変更と球面制約をセットにしていますが）、MDデカップリングは標準的なアーキテクチャに適用可能なオプティマイザ側の修正です。nGPTのアーキテクチャに適用した場合、MDデカップリングはさらに性能を向上させます。
  • MuownやHyperPとは異なり（これらは学習可能な大きさを持たずにノルムを固定したり、特定の手法を特定のオプティマイザに結びつけたりしていますが）、MDデカップリングは明示的に大きさを学習し、オプティマイザに依存しません。
  • また、先行研究（Kosson et al., 2024a）で指摘された「大きさ–方向の干渉」に対し、2つの量のダイナミクスを分離することで対処しています。

4. 重要性と主張

本論文は、大きさと方向を別々に制御される量として扱うことで、より予測可能な学習ダイナミクスが得られ、現代的なオプティマイザに対するシンプルで広く適用可能な改善となることを主張しています。

• 簡潔さ: アーキテクチャの変更を必要とせず、オプティマイザのステップにおける計算オーバーヘッドも最小限です（要素ごとの演算）。
• 堅牢性: 重み減衰やウォームアップの必要性を排除することで、学習レシピを簡素化し、ハイパーパラメータの探索空間を縮小します。
• スケーラビリティ: 学習率の幅方向への転送を実現する原理的なメカニズムを提供し、大規模言語モデルのスケーリングにおける主要なボトルネックに対処します。

著者らは、オプティマイザのステップにおける追加のロジックがわずかなオーバーヘッドをもたらすものの、学習の安定性、計算効率、およびハイパーパラメータ転送におけるメリットが、MDデカップリングを大規模ニューラルネットワーク学習における重要な進歩にしていると結論付けています。