MARS-M: 分散縮小法と行列の出会い

要旨

Muonに代表される行列ベースの前処理付きオプティマイザは、大規模言語モデル（LLM）を含む大規模ニューラルネットワークの学習において、スカラーベースのオプティマイザよりも効率的であることが最近示されている。一方、LLM事前学習におけるオプティマイザの最近のベンチマークでは、MARSのような分散低減技術が、分散低減を採用しない標準的なオプティマイザよりも大幅な高速化を達成できることが実証されている。本論文では、両者の利点を組み合わせるため、MARSの分散低減技術とMuonを統合した新しいオプティマイザであるMARS-Mを提案する。標準的な正則条件の下で、MARS-Mが一次停留点に \(\mathcal{O}(T^{-1/3})\) の速度で収束することを証明する。これはMuonが達成する \(\mathcal{O}(T^{-1/4})\) の収束率を改善するものである。言語モデリングとコンピュータビジョンタスクにおける実証実験の結果、MARS-Mが一貫して損失を低減し、様々なダウンストリームベンチマークで性能向上をもたらすことを示す。MARS-Mの実装はhttps://github.com/AGI-Arena/MARS/MARS_M で公開されている。

English

Matrix-based preconditioned optimizers, such as Muon, have recently been shown to be more efficient than scalar-based optimizers for training large-scale neural networks, including large language models (LLMs). On the other hand, recent benchmarks on optimizers for LLM pre-training have demonstrated that variance-reduction techniques such as MARS can achieve substantial speedups over standard optimizers that do not employ variance reduction. In this paper, to achieve the best of both worlds, we introduce MARS-M, a new optimizer that integrates the variance reduction technique in MARS with Muon. Under standard regularity conditions, we prove that Muon-M converges to a first-order stationary point at a rate of mathcal{O}(T^{-1/3}), which improves upon mathcal{O}(T^{-1/4}) rate attained by Muon. Our empirical results on language modeling and computer vision tasks demonstrate that MARS-M consistently yields lower losses and improved performance across various downstream benchmarks. The implementation of MARS-M is available at https://github.com/AGI-Arena/MARS/MARS_M.

MARS-M: 分散縮小法と行列の出会い

MARS-M: When Variance Reduction Meets Matrices

要旨

Support