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 $\tilde{\mathcal{O}}(T^{-1/3})$, which improves upon $\tilde{\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/tree/main/MARS_M.

翻译：基于矩阵的预条件优化器（如Muon）近期被证明在训练大规模神经网络（包括大语言模型）时比基于标量的优化器更高效。另一方面，近期关于大语言模型预训练优化器的基准测试表明，采用方差缩减技术（如MARS）的优化器相比未采用该技术的标准优化器能实现显著加速。本文为兼取二者之长，提出了MARS-M——一种将MARS中的方差缩减技术与Muon相结合的新型优化器。在标准正则性条件下，我们证明了Muon-M以$\tilde{\mathcal{O}}(T^{-1/3})$的速率收敛至一阶稳定点，该速率优于Muon所达到的$\tilde{\mathcal{O}}(T^{-1/4})$。我们在语言建模与计算机视觉任务上的实验结果表明，MARS-M在多种下游基准测试中始终能获得更低的损失值与更优的性能。MARS-M的实现代码已发布于https://github.com/AGI-Arena/MARS/tree/main/MARS_M。