Adamがミューオンを改善：直交化運動量を用いた適応モーメント推定

Adam Improves Muon: Adaptive Moment Estimation with Orthogonalized Momentum

February 19, 2026

著者: Minxin Zhang, Yuxuan Liu, Hayden Scheaffer

cs.AI

要旨

効率的な確率的最適化手法は通常、決定論的領域で良好な性能を示す更新方向と、確率的摂動に適応する機構を統合する。Adamが適応的なモーメント推定を用いて安定性を促進する一方、Muonは直交化モーメントを通じて重み層の行列構造を利用し、大規模言語モデル学習において優れた性能を示す。本研究では、新たな最適化手法NAMOとその対角拡張版NAMO-Dを提案する。これらは、直交化モーメントとノルムベースのAdam型ノイズ適応を原理的に統合した初めての手法である。NAMOは単一の適応的ステップサイズを用いて直交化モーメントをスケーリングし、直交性を保持しながらMuonを上回る性能を無視できる追加コストで実現する。一方、NAMO-Dは直交化モーメントにクランプされた要素を持つ対角行列を右から乗算する。この設計はニューロン単位のノイズ適応を可能にし、一般的なほぼブロック対角なヘッシアン構造と整合する。標準的な仮定の下で、両アルゴリズムの決定論的設定における最適収束率を理論的に示し、確率的設定ではそれらの収束保証が確率的勾配のノイズレベルに適応することを示す。GPT-2モデルの事前学習実験により、NAMOとNAMO-DはAdamW及びMuonのベースラインと比較して性能が向上し、NAMO-Dは、条件の良い更新方向の維持と細粒度ノイズ適応の活用という競合する目標のバランスを取る追加のクランプハイパーパラメータを通じて、NAMOをさらに上回る改善を達成することが実証された。

English

Efficient stochastic optimization typically integrates an update direction that performs well in the deterministic regime with a mechanism adapting to stochastic perturbations. While Adam uses adaptive moment estimates to promote stability, Muon utilizes the weight layers' matrix structure via orthogonalized momentum, showing superior performance in large language model training. We propose a new optimizer and a diagonal extension, NAMO and NAMO-D, providing the first principled integration of orthogonalized momentum with norm-based Adam-type noise adaptation. NAMO scales orthogonalized momentum using a single adaptive stepsize, preserving orthogonality while improving upon Muon at negligible additional cost. NAMO-D instead right-multiplies orthogonalized momentum by a diagonal matrix with clamped entries. This design enables neuron-wise noise adaptation and aligns with the common near block-diagonal Hessian structure. Under standard assumptions, we establish optimal convergence rates for both algorithms in the deterministic setting and show that, in the stochastic setting, their convergence guarantees adapt to the noise level of stochastic gradients. Experiments on pretraining GPT-2 models demonstrate improved performance of both NAMO and NAMO-D compared to the AdamW and Muon baselines, with NAMO-D achieving further gains over NAMO via an additional clamping hyperparameter that balances the competing goals of maintaining a well-conditioned update direction and leveraging fine-grained noise adaptation.