MARS-M: When Variance Reduction Meets Matrices
Authors: name: Yifeng Liu, name: Angela Yuan, name: Quanquan Gu
Organization: University of California, Los Angeles

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.