sMuon、MuonをLoRAで表現可能な部分空間に射影し、ReLoRAでRiemannionと同等の損失を達成しつつ約30％高速化

Muonは各層の重み行列の更新を直交化し、非ゼロの特異値を均一化する。しかしLoRAでは、低ランク行列 `BA` による重みの変更しか許されないため、完全なMuon方向は通常、この表現可能な空間の外にある。新たに提案されたsMuonは、まずMuonのスペクトルノルムに関する最急降下目標を線形化し、次に最小二乗問題を解くことで、直交化された勾配をLoRAの行空間と列空間へ射影する。さらに著者らは、低ランク構造を利用して演算を書き換え、完全な重み行列に対するQR分解やSVDを回避した。主要な処理に必要なのは、行列乗算、逆平方根、および小規模な行列演算のみである。

1億6,200万パラメータのTransformer、FineWebデータ、500ステップごとにadapterをマージするReLoRAを用いた実験では、sMuonとRiemannionはいずれも最終検証損失3.66を記録し、LoRA-Muonの3.68、因子ごとにMuonを適用する手法の3.70、AdamWの3.83を上回った。論文では同時に、sMuonが約30％高速であるとも報告されている。また、各パラメータについて保持するモメンタム状態は1組だけであり、理論上はAdamWの一次・二次モーメント推定より少ない。

ただし、すべての条件で優位だったわけではない。著者らはQwen2.5-3B、Llama-3.2-3B、DeepSeek-V2-Lite、およびMuonで事前学習されたMoonlight-16B-A3Bを対象に、11種類の常識推論・プログラミングタスクで評価した。sMuonはMoonlightでは6タスクで首位となったが、AdamWで事前学習されたモデルでは、AdamWやほかの低ランクMuon変種が上回る場合が多かった。例えばQwen2.5-3BのHumanEval pass@1は、AdamWが54.9％だったのに対し、sMuonは51.8％だった。これは、事前学習とファインチューニングで使用するオプティマイザの暗黙的バイアスが、引き続き相互作用することを示している。エンジニアリングチームは今後、公式実装、分散学習と混合精度における安定性、さらにLoRA rankを高めた場合の実際のVRAM使用量と実時間コストに注目すべきだ。現時点で、論文にはsMuon専用の公開ライブラリは付属していない。