Abstract
ミューオン・オプティマイザは、大規模言語モデルの学習において強力な性能を示すことから、かなり注目を集めてきました。しかし、行列勾配の直交化に関する設計原理は、いまだにほとんど解明されていません。本論文では、ミューオンの設計に新たな光を当てるだけでなく、より重要な点として、新しいオプティマイザにつながる代理モデル（サロゲートモデル）を導入します。ニュートン法の導出と同じ精神に基づき、代理モデルは、重み行列 W への摂動に対する損失を、3つの行列だけを用いて（すなわち）勾配 G、出力空間における曲率行列 H、そして層への入力を縦に積み重ねたデータ行列 Z を用いて、二次関数として近似します。この代理モデルを1ステップで最小化し、重みに対してある等方的（アイソトロピック）仮定を採用することで、閉形式の更新則（運動量と weight decay を除く）W \leftarrow W - \eta \cdot \mathrm{msgn}(G(ZZ^\top)^{-1}) が得られます。ここで
eta は学習率であり、コンパクト特異値分解 X=USV^\top に対して
mathrm{msgn}(X)=UV^\top です。この新しい最適化手法を Newton-Muon と呼びます。これは、標準的なミューオンが、入力の2次モーメントによって誘起される右側の前処理を無視する、暗黙的なニュートン型手法として解釈できることを示します。実験的には、GPT-2 の事前学習において Muon を用いた、最も初期に公開された Modded-NanoGPT speedrun 構成の再現により、Newton-Muon は目標の検証損失に到達するまでに 6少ない反復ステップ数で済み、ウォールクロックの学習時間を約 4% 削減します。
