Improving Neural Network Training by Decoupling the Magnitude and Direction of Weight Vectors
Abstractの概要
本論文は、標準的なオプティマイザが各重み行列の大きさと方向という2つの異なる側面を結合させてしまっていると指摘している。これにより、方向の更新が現在のノルムに依存し、スケール不変の設定においても大きさの増大を引き起こすとする。著者らは、Magnitude-Direction（MD）Decouplingを提案し、各重みをハイパースフィア上の固定ノルムの方向成分と、個別の学習率で更新される学習可能な行ごと・列ごとのスケール（ゲイン）に分解する手法を導入した。この手法はオプティマイザ内部で実装されるため、計算グラフ上では結合済み重みテンソルとして機能し、AdamやMuonといった様々なベースオプティマイザと互換性を持つ。実証研究では、密結合モデルでのアブレーションと、より大規模なMixture-of-Expertsによる実験を通じ、正規化の選択肢、ゲインのパラメータ化、学習率の転移、スケジュールの挙動、スケーリングを検証している。
新規性
本研究の際立った貢献は、固定ノルムによる方向学習の適用と、行および列のゲインによるきめ細かいスケール学習の復元を同時に実現する、オプティマイザレベルの再パラメータ化である。明示的に大きさを学習せずノルムのみを固定する従来の球ベースのアプローチとは異なり、オプティマイザに依存しない形で方向とスケールを非結合化（Decoupling）する点に新規性がある。
成果
密な言語モデル全体において、MD Decouplingは十分に最適化されたAdamWやMuonのベースラインを上回る性能を示し、最適な行列の学習率が再調整なしでモデルの幅を超えて転移可能であることが報告されている。アブレーション解析により、最大の恩恵は方向制御（球面制約）から得られること、また、そこに行列の行と列のゲインを組み合わせることでさらに改善し、様々なゲイン設定の中で最高性能を示すことが確認された。大規模なMixture-of-Expertsの実験においても、全ての計算量範囲でベースラインを上回り、約半分の計算量でAdamWと同等の損失に到達したと報告されている。
論文の注目点
- MD Decouplingは、各重み行列を固定ノルムの方向成分と学習可能な大きさのゲインに分離し、それぞれに対して個別の更新ルールと学習率を用いる手法である。
- フロベニウス・ノルムの球面制約が汎用性の高いデフォルト値であり、これを用いることで単位ノルム埋め込みが制約なしのものと同等以上に機能し、行と列を組み合わせたゲインがその他のゲイン（スカラーや片側のみ）を上回ることが確認された。
- 本手法は重み減衰（Weight decay）やウォームアップへの依存を軽減し、モデル幅における学習率の転移性を維持しつつ、大規模なMixture-of-Expertsモデルへの拡張においても継続的に恩恵をもたらす。