Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MARS-M: When Variance Reduction Meets Matrices

Yifeng Liu, Angela Yuan, Quanquan Gu

TL;DR

This work addresses the inefficiencies of stochastic variance and matrix-based optimization in large-scale neural networks by introducing MARS-M, a matrix-based instantiation of MARS that integrates variance reduction with the Muon/Moonlight framework. It presents both a concrete algorithm and an approximated variant, and proves a non-asymptotic convergence rate of $\tilde{\mathcal{O}}(T^{-1/3})$, improving upon the rate for Muon. Empirically, MARS-M yields lower training and validation losses and better downstream performance on large language models and computer vision tasks, with a public implementation provided. Overall, the paper demonstrates that combining variance reduction with matrix preconditioning offers meaningful speedups and accuracy gains in practical, large-scale learning settings.

Abstract

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.

MARS-M: When Variance Reduction Meets Matrices

TL;DR

This work addresses the inefficiencies of stochastic variance and matrix-based optimization in large-scale neural networks by introducing MARS-M, a matrix-based instantiation of MARS that integrates variance reduction with the Muon/Moonlight framework. It presents both a concrete algorithm and an approximated variant, and proves a non-asymptotic convergence rate of , improving upon the rate for Muon. Empirically, MARS-M yields lower training and validation losses and better downstream performance on large language models and computer vision tasks, with a public implementation provided. Overall, the paper demonstrates that combining variance reduction with matrix preconditioning offers meaningful speedups and accuracy gains in practical, large-scale learning settings.

Abstract

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 , which improves upon 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.
Paper Structure (17 sections, 4 theorems, 40 equations, 11 figures, 17 tables, 2 algorithms)

This paper contains 17 sections, 4 theorems, 40 equations, 11 figures, 17 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Muon
  5. MARS
  6. Method
  7. MARS-M
  8. Approximated MARS-M
  9. Convergence Analysis
  10. Experiments
  11. LLM Experiments
  12. Experiment Results
  13. Computer Vision Experiments
  14. Ablation Study
  15. Hyper-parameter Settings
  16. ...and 2 more sections

Key Result

Theorem 4.3

In Algorithm alg:MARS_moonlight, under Assumptions assum:variance and assum:L-smooth, when choosing $\lambda=0$, $\eta_t=(s+t)^{-2/3}, s\geq 2$, suppose $\beta_{t+1}= 1 - 2\eta_t$, then for $\forall T\ge s$, it holds that where $G=F(\mathbf{X}_1)-\min_\mathbf{X} F(\mathbf{X})+\frac{\sqrt{2}s^{1/3}\sigma^2}{4L}+\frac{3Ln}{2s^{1/3}}$, $B=(\frac{2\sqrt{2} \sigma^2}{L}+\frac{3\sqrt{2}Ln}{2})$ and $\w

Figures (11)

  • Figure 1: The training and validation loss of XL-size models (1.5B) trained with different optimizers on the FineWeb-Edu 100B dataset.
  • Figure 2: The training and validation loss of XL-size models (1.5B) trained with different optimizers on the FineWeb-Edu 100B dataset.
  • Figure 3: The training and validation loss of large-size models (770M) trained with different optimizers on the OpenWebText dataset.
  • Figure 4: The test loss and test accuracy for different optimizers on CIFAR-10 dataset.
  • Figure 5: The training and validation loss of small-size models (125M) trained with MARS-M with different $\gamma$s on the OpenWebText 100B dataset.
  • ...and 6 more figures

Theorems & Definitions (4)

  • Theorem 4.3
  • Lemma C.1: yuan2024mars
  • Lemma C.2
  • Lemma C.3