PaperlayerView exact revision on arXiv

Source-linked AI summary

MARS-M: When Variance Reduction Meets Matrices

Yifeng Liu, Angela Yuan, Quanquan Gu

arXiv:2510.21800v3cs.LGmath.OCstat.ML

TL;DR

High stochastic-gradient variance can slow and destabilize training, while prior work has separately reported benefits from matrix-based optimizers and variance reduction. This paper introduces MARS-M, combining MARS-style variance reduction with Muon/Moonlight, and reports improved empirical performance alongside a faster convergence rate than Muon.

  • Problem

    Stochastic gradients can have high variance that slows convergence and degrades stability, and prior benchmarks motivate combining matrix-based optimization with variance reduction.

    Source: §1 · pp. 1–2

  • Method

    MARS-M integrates MARS-style variance reduction with Moonlight’s matrix-based optimizer, with an approximate variant also developed.

    Source: §1 · p. 2, §4.1 · p. 6

  • Results

    MARS-M achieves Õ(T^-1/3) convergence versus Muon’s O(T^-1/4), and reports improved losses and benchmark performance across language-modeling and computer-vision tasks.

    Source: §1 · p. 2, §6 · p. 12

  • Takeaways & Limitations

    The results demonstrate an effective pathway for improving matrix-based optimizers by using variance-reduction techniques.

    Source: §6 · p. 12

    The analysis assumes a possibly nonconvex matrix-variate objective and access to an unbiased first-order gradient oracle; the introduction also assumes independent per-step training data.

    Source: §1 · p. 3, §3 · p. 4

Abstract

from arXiv · show

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). Recent benchmark studies of LLM pretraining optimizers have demonstrated that variance-reduction techniques such as MARS can substantially speed up training compared with standard optimizers that do not employ variance reduction. In this paper, we introduce MARS-M, a new optimizer that integrates MARS-style variance reduction with Muon. Under standard regularity conditions, we prove that MARS-M converges to a first-order stationary point at a rate of $\tilde{\mathcal{O}}(T^{-1/3})$, improving upon the $\tilde{\mathcal{O}}(T^{-1/4})$ rate attained by Muon. 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.

1 Introduction

The introduction motivates combining matrix-based optimization with variance reduction and presents MARS-M, a MARS instantiation tailored to Moonlight, alongside an approximate variant. It reports improved empirical performance and a convergence rate better than Muon’s established rate.

1 Introduction

  • Stochastic optimization can be slowed and destabilized by high gradient variance, while earlier variance-reduction methods had limited success at large scale because of incompatibilities with deep-learning practice and architectures.MARS addresses this compatibility challenge by scaling STORM’s correction; benchmark studies also report additional speedups from variance reduction alongside matrix-based optimizers.Source: §1 · pp. 1–2
  • Matrix-based optimizers preserve parameter matrices’ two-dimensional structure, and recent LLM-pretraining benchmarks report that they outperform scalar-based optimizers.Those benchmarks also report discernible additional speedups from variance-reduction approaches.Source: §1 · pp. 1–2
  • MARS-M integrates MARS-style variance reduction with Moonlight’s matrix-based optimizer, addressing an open question after MARS-Shampoo performed worse than vector-based MARS-AdamW.The introduction also proposes an approximate MARS-M as a practical variant and bridge to traditional preconditioned optimizers.Source: §1 · p. 2
  • The approximate MARS-M can be viewed as a Moonlight variant with adjusted momentum parameters.The paper also discusses its connections and differences from MARS-Shampoo.Source: §1 · p. 2
  • MARS-M converges at a rate of Õ(T^-1/3), improving on Muon’s O(T^-1/4) rate.The comparison is attributed to prior analyses of Muon by Li and Hong, Shen et al., and Pethick et al.Source: §1 · p. 2
  • MARS-M improves training and validation losses on GPT-2 language-model training and performs better on downstream language benchmarks and computer-vision test accuracy than baseline optimizers.The reported language-model evaluations use OpenWebText and FineWeb-Edu 100B.Source: §1 · p. 2

2 Related Work

Related work traces matrix-based preconditioning and variance-reduction methods, including their applications to adaptive optimizers. It situates MARS-M among Muon convergence analyses and notes that prior variance-reduced adaptive methods were evaluated on relatively modest benchmarks.

2 Related Work

  • Matrix-based optimization developed from practical second-order approximations such as Shampoo and SOAP to Muon and Moonlight, with later methods extending or unifying these approaches.SOAP applies AdamW in Shampoo’s preconditioner eigenbasis; Muon uses a Newton–Schulz approximation to the gradient-momentum SVD, while Moonlight adds weight decay and couples step sizes across components.Source: §2 · p. 3
  • Prior analyses establish Muon’s O(T^-1/4) convergence rate and study variants without momentum or under additional smoothness and large-batch assumptions.The cited work includes Li and Hong, Shen et al., Pethick et al., An et al., Kovalev, Sfyraki and Wang, and Sato et al.Source: §2 · p. 3
  • Variance-reduction research spans methods such as SAG, SVRG, SAGA, and SARAH, with later non-convex approaches including SPIDER, SNVRG, and STORM.SARAH uses a biased recursive gradient estimator, while STORM uses stochastic recursive momentum.Source: §2 · pp. 3–4
  • Adaptive variance-reduction methods have been evaluated on relatively simple vision and language-modeling benchmarks with modest model sizes; MARS uses scaled gradient correction, and MARS-M pairs it with Muon/Moonlight.Examples include Adam+, SuperAdam, VRAdam, AdaSPIDER, and MARS.Source: §2 · p. 4

3 Preliminaries

The paper frames stochastic matrix optimization as minimizing an expected possibly nonconvex loss with an unbiased first-order gradient oracle. It reviews Muon and MARS as background for the paper’s tailored MARS approach.

3 Preliminaries

  • The objective is to minimize a possibly nonconvex expected loss over a matrix variable, using an oracle that returns an unbiased gradient estimator.The random data vector is drawn from an unknown distribution, and the paper assumes m ≥ n.Source: §3 · p. 4

3.1 Muon

  • Muon uses matrix-parameter geometry, with Newton–Schulz iteration approximating U_tV_t from the momentum matrix’s singular value decomposition.The paper focuses on the original Muon algorithm.Source: §3.1 · p. 4, §3.2 · p. 5
  • Moonlight addresses mismatched update magnitudes across matrix-like and vector-like parameters; AdamW update RMS norms are approximately 0.2 to 0.4 during LLM training.Its described update uses decoupled weight decay.Source: §3.1 · p. 5

3.2 MARS

  • STORM’s gradient-difference correction is described as canceling data randomness and transferring the previous-step true-gradient estimate to the current parameter.The correction compares stochastic gradients at the current and previous parameters using the current step’s data.Source: §3.2 · p. 5
  • MARS scales the gradient-correction term, and clips the corrected gradient to improve training stability.The paper focuses on the original algorithm rather than the newest implementation described in the footnote.Source: §3.2 · pp. 5–6
  • MARS’s update framework covers adaptive methods including AdamW, Lion, and Shampoo with an approximated Hessian, and the paper views its approach as tailored to Muon or Moonlight.The paper also describes an approximate MARS version elsewhere; this passage establishes the general adaptive-gradient formulation.Source: §3.2 · p. 6

4 Method

Section 4 presents MARS-M, a MARS instantiation tailored to Moonlight, and establishes its convergence guarantees. It also develops an approximate version and relates it to Moonlight with adjusted momentum parameters.

4.1 MARS-M

  • MARS-M instantiates MARS for Moonlight by replacing MARS update rule (3.4) with update rule (3.8).Compared with MARS-Shampoo, it scales O_t approximately as U_tV_t^T/max(m, n) and incorporates weight decay λX_t.Source: §4.1 · p. 6

4.2 Approximated MARS-M

  • To avoid computing stochastic gradients twice, approximate MARS-M reuses the previous step's gradient and is Moonlight with adjusted momentum parameters.The paper reports that this difference in momentum parameters noticeably impacts loss.Source: §4.2 · pp. 6–7

4.3 Convergence Analysis

  • Under bounded-variance and L-smoothness assumptions, MARS-M achieves O(T^-1/3 log(T)) convergence, compared with Muon's O(T^-1/4).The theorem uses λ = 0, η_t = (s + t)^-2/3, and a time-varying β_t; constant β can be used in practice.Source: §4.3 · pp. 7–8
  • The convergence analysis adds MARS-M to the MARS family, extending the framework to more matrix-based optimizers with theoretical guarantees.The convergence theorem is deferred to Appendix C.Source: §4.3 · p. 8

5 Experiments

The experiments compare MARS-M with Moonlight and AdamW across GPT-2 language modeling, downstream evaluation, and CIFAR-10 vision, and examine the gradient-correction scale γ.

5.1 LLM Experiments

  • The language-model experiments use nanoGPT GPT-2 models on OpenWebText and FineWeb-Edu 100B, training for 100,000 steps with 2,000 warm-up steps.Models range from 125M to 1.5B parameters; the optimizer comparison uses AdamW for vector-like parameters and embeddings.Source: §5.1 · p. 9

5.2 Experiment Results

  • MARS-M steadily improves training and validation losses over Moonlight, and its losses fall below AdamW's in the final phase.AdamW has lower loss in the middle of training, which the authors attribute to its smaller maximum learning rate.Source: §5.2 · p. 9, Figure 1 · p. 10, Figure 2 · p. 10
  • MARS-M outperforms Moonlight on most reported downstream benchmarks for medium, large, and XL FineWeb-Edu 100B models.The reported 2-shot evaluations cover ARC, HellaSwag, OBQA, PIQA, SciQ, WinoGrande, and MMLU.Source: §5.2 · pp. 10–11

5.3 Computer Vision Experiments

  • On CIFAR-10 with ResNet-18, MARS-M achieves better test accuracy and lower test loss than both Moonlight and MARS-M-approx.The experiments use γ = 0.025 and run for 200 epochs.Source: Figure 5 · p. 11, §5.3 · pp. 11–12

5.4 Ablation Study

  • The γ ablation finds that the gradient-correction scale should be much smaller than 1, while losses are close for γ values from 0.005 to 0.025.The authors describe performance as relatively robust to variation within that range.Source: Figure 6 · p. 11, §5.4 · p. 12

6 Conclusion

The conclusion presents MARS-M as a combination of MARS-based variance reduction and Muon/Moonlight, with empirical gains over Moonlight on language-model and computer-vision tasks.

6 Conclusion

  • MARS-M integrates MARS-based variance reduction with Muon/Moonlight to target faster large-language-model training.The paper frames this as combining variance reduction with matrix-based optimization.Source: §6 · p. 12
  • Across language-model and computer-vision experiments, the authors report steady advantages over Moonlight in loss and downstream evaluation.They describe variance reduction as an effective pathway for improving matrix-based optimizers.Source: §6 · p. 12

A Hyper-parameter Settings

The appendix points to architecture settings and learning-rate settings for the GPT-2 experiments.

A Hyper-parameter Settings

B Additional Experiments

Additional experiments report comparisons across training curves and language-model evaluations for GPT-2 models trained on OpenWebText and FineWeb-Edu 100B.

B Additional Experiments

  • MARS-M performs better than AdamW and Moonlight in most reported additional experiments across the two datasets.The appendix includes training and validation curves and 0-shot and 2-shot evaluations.Source: B Additional Experiments · p. 18
  • The additional evaluation tables cover multiple model sizes and 0-shot or 2-shot results on OpenWebText and FineWeb-Edu 100B.The displayed captions identify small, medium, large, and XL models across the supplied tables.Source: B Additional Experiments · p. 18, Table 6 · p. 18, Table 7 · p. 19

C Proof of Theorem 4.3

The proof bounds the Lyapunov function by controlling the estimator-error term and telescoping its stepwise change. Under the stated assumptions, this yields a final bound containing a T^-1/3 term.

C Proof of Theorem 4.3

  • In Algorithm 2, variance-reduction Lemma C.1 assumes Assumptions 4.1 and 4.2, 1 ≥ βt+1 ≥ 0 for all t, and an approximate choice of γt+1; Lemma C.2 gives an inequality under Assumption 4.2.The proof states an upper bound for the function value according to Assumption 4.2.Source: C Proof of Theorem 4.3 · pp. 21–22
  • The analysis treats O_t as the exact polar factor and ignores the approximation error from its truncated Newton–Schulz computation.The passage notes that the factor 0.2 · p_max(m, n) is absorbed into η_t.Source: C Proof of Theorem 4.3 · pp. 22–23
  • For η_t = (s + t)^-2/3, Lemma C.3 bounds the change in inverse step sizes using η_t^1/2.The proof cites concavity of h(x) = x^2/3 for the inequality.Source: C Proof of Theorem 4.3 · pp. 23–24
  • The Lyapunov function combines the objective value with a weighted squared estimator-error term, whose change is bounded across consecutive iterations.The proof calculates the difference between consecutive Lyapunov values before bounding its terms.Source: C Proof of Theorem 4.3 · p. 23
  • The telescoping sum over t = 1, …, T, using ρt = 4, and the triangle inequality yield a final upper bound containing a T^-1/3 term and an additive 2G.The proof states ΦT+1 ≥ F(XT+1) ≥ minX F(X).Source: C Proof of Theorem 4.3 · pp. 24–26

Related papers

Loading 2510.21800v3…