MARS: Unleashing the Power of Variance Reduction for Training Large Models

🎯 Pitch

Variance reduction, long dismissed as ineffective for deep learning, actually accelerates GPT-2 training by nearly 2× when properly scaled. MARS introduces a tunable recursive momentum correction that finally bridges adaptive preconditioners—like AdamW, Lion, and Shampoo—with provable variance reduction, slashing the tokens needed to reach AdamW’s loss from 50B to just 28B.

1. Executive Summary

This paper proposes MARS (Make vAriance Reduction Shine), a unified optimization framework that integrates variance reduction into preconditioned adaptive gradient methods for training large models, reconciling preconditioned gradient updates with a scaled stochastic recursive momentum technique. Evaluated on GPT-2 models (125M–770M parameters) trained on the OpenWebText dataset, MARS introduces a scaling parameter $\gamma_t$ that controls the strength of gradient correction in the variance-reduced estimator (operationalized as a STORM-style momentum with controllable recursive correction), combined with preconditioning matrices drawn from AdamW, Lion, and Shampoo to form three instantiated variants. On GPT-2 large, MARS achieves a final validation loss of 2.51 versus AdamW's 2.58 and reaches AdamW's 50-billion-token validation loss of 2.58 in only 28 billion tokens—a 1.8× token efficiency improvement—while improving Hellaswag 5-shot accuracy from 41.70% to 44.64%. The paper establishes that variance reduction can be effectively applied to large model training, but only when the recursive momentum estimator is properly scaled and its second-order momentum is redefined to match the variance-reduced gradient rather than the raw stochastic gradient, resolving the long-standing ineffectiveness of variance reduction in deep learning.

2. Context and Motivation

The Core Problem: Variance Reduction Has Failed in Deep Learning

This paper confronts a genuinely puzzling disconnect in optimization research. On one side, variance reduction techniques — methods that construct gradient estimators with lower variance than standard stochastic gradients — have been a dominant focus in the optimization theory community for over a decade. Algorithms like SVRG (Johnson & Zhang, 2013), SPIDER (Fang et al., 2018), and STORM (Cutkosky & Orabona, 2019) provably reduce the gradient complexity of finding stationary points from $\mathcal{O}(\varepsilon^{-4})$ (standard SGD) to $\mathcal{O}(\varepsilon^{-3})$ for nonconvex smooth optimization (Arjevani et al., 2023). This is not a minor improvement — it converts a quartic dependence on desired accuracy to a cubic one, which in theory should translate to dramatically faster convergence, especially in noisy regimes.

On the other side, variance reduction has not found widespread success in training deep neural networks or large language models. As the paper states bluntly in its abstract:

"Despite the development of numerous variance reduction algorithms in the past decade aimed at accelerating stochastic optimization in both convex and nonconvex settings, variance reduction has not found widespread success in training deep neural networks or large language models. Consequently, it has remained a less favored approach in modern AI."

The paper explicitly cites Defazio & Bottou (2019) to explain why this happened historically: data augmentation, batch normalization, and dropout disrupt the "finite-sum structure" that variance reduction methods rely upon. In the standard variance reduction setup (e.g., SVRG), the objective is assumed to be a finite sum over a fixed training set: $F(x) = \frac{1}{n}\sum_{i=1}^{n} f_i(x)$ . The correction term $\nabla f_i(x_t) - \nabla f_i(\tilde{x})$ in SVRG works because the same data point $i$ is evaluated at two different parameter values, enabling the noise from $i$ at the current and reference points to partially cancel. But when data augmentation randomizes each "data point" dynamically — so that calling $f_i(x)$ twice yields different outputs even at the same $i$ — or when batch normalization makes the loss depend on statistics of the current minibatch rather than just the parameters and data, the finite-sum assumption breaks. The correction term ceases to cancel noise and may even inject additional variance.

The door has now reopened. The paper's key motivating observation is that modern language model training has evolved in ways that remove these obstacles:

Data augmentation is rarely used in LLM pretraining — the data is tokenized text, and while there may be preprocessing, it is deterministic rather than on-the-fly randomization.
Batch normalization has been largely replaced by Layer Normalization and RMSNorm in transformer architectures, eliminating the minibatch-dependent loss issue.
Dropout is often set to 0.0 in modern LLM training (the paper's GPT-2 experiments explicitly use dropout 0.0).

This creates a new environment where the structural barriers to variance reduction are gone. The paper frames this as an explicit research question:

"Can variance reduction technique be applied to improve the performance of training large models?"

The question matters because LLM training is inherently a high-variance optimization problem. McCandlish et al. (2018) established that language model training gradients exhibit substantial variance, particularly at moderate batch sizes. This variance forces the use of small learning rates or large batch sizes, both of which slow convergence or increase computational cost. If variance reduction could be made to work, it would directly address this bottleneck.

Why This Matters: The Stalled Progress in LLM Optimizers

Beyond the variance reduction problem specifically, the paper identifies a broader stagnation in optimizer design for large models. It surveys the landscape of adaptive gradient methods and makes a pointed observation: despite a decade of innovation since Adam (Kingma & Ba, 2015) and AdamW (Loshchilov & Hutter, 2019), no method has convincingly and consistently outperformed AdamW in LLM pretraining at scale. The paper cites two recent comprehensive studies that support this claim:

"recent studies (Kaddour et al., 2024; Zhao et al., 2024) have shown that these optimizers perform on par with AdamW in LLM pretraining, yet do not outperform it."

This is striking when you list the candidates that have been tried:

LAMB (You et al., 2019): layerwise adaptation that boosted BERT training but hasn't displaced AdamW for GPT-style models.
Lion (Chen et al., 2023): discovered via symbolic program search, claims faster training and reduced memory, but the paper and the cited studies find it performs comparably to AdamW, not better.
Sophia (Liu et al., 2023): uses stochastic diagonal Hessian estimators (Hutchinson's method or Gauss-Newton-Bartlett) with clipping-by-value. Sophia showed improvements over AdamW in some settings, but Kaddour et al. (2024) and Zhao et al. (2024) found its gains inconsistent at larger scales.
Shampoo (Gupta et al., 2018; Anil et al., 2020): full-matrix preconditioning over tensor spaces, approximating the Gauss-Newton component of the Hessian via Kronecker factorization. More powerful in theory than diagonal preconditioners, but heavier computationally, and Zhao et al. (2024) found it "on par" with AdamW.
SOAP (Vyas et al., 2024): stabilizes Shampoo by combining it with Adam in the eigenbasis of Shampoo's preconditioner. Improves Shampoo's stability but still competes with rather than dominates AdamW.
Muon (Jordan et al., 2024): uses Newton-Schulz iteration to perform eigenspace-based preconditioning (essentially doing SVD on the gradient matrix and updating with $U V^\top$ ). Demonstrated faster convergence than AdamW in some settings, but again, the paper's experiments show Muon trailing MARS on GPT-2 large (validation loss of 2.606 vs. 2.511 for MARS-AdamW).

The implication is clear: first-order momentum (AdamW, Lion) and second-order preconditioning (Shampoo, Muon) represent two families of ideas that have been extensively explored, yet neither has yielded decisive improvements over the AdamW baseline. The paper's thesis is that variance reduction is a third, orthogonal dimension that has been neglected, and that combining it with existing preconditioning approaches could break through this performance ceiling.

Where Prior Variance Reduction + Adaptive Method Attempts Fell Short

The paper is careful to acknowledge that it is not the first to try combining variance reduction with adaptive gradient methods. It cites several predecessors, but argues that their approaches had critical design flaws that limited their effectiveness:

1. Adam+ (Liu et al., 2020): Attempts to reduce variance in Adam's first-order momentum by estimating gradients at extrapolated points rather than the current parameters. The paper mentions this but does not elaborate on its limitations — it is included as evidence that the combination has been tried but not successfully scaled to large models, having been validated only on simple tasks like MNIST and CIFAR-10.

2. SuperAdam (Huang et al., 2021): The closest predecessor to MARS, and the one the paper engages with most critically. SuperAdam incorporates STORM-style recursive momentum into an Adam-like framework. However, the paper identifies three specific design flaws that MARS corrects:

"SuperAdam's design focuses on diagonal precondition matrix and draws heavily from the design used in Adam, AdaGrad-Norm, and AdaBelief. Furthermore, their preconditioner matrix is designed following Adam's structure but does not account for the revised definition of variance-reduced momentum, resulting in a significant mismatch between the first-order and second-order momentum."

This mismatch is the central critique. In standard Adam, the second-order momentum $v_t$ is the EMA of squared gradients $[\nabla f(x_t, \xi_t)]^2$ , and the first-order momentum $m_t$ is the EMA of the same gradients. These two quantities are coherent: $v_t$ estimates the scale of the gradient components, and $m_t$ estimates the direction, both derived from the same underlying random variable. But in a variance-reduced method, the momentum $m_t$ is not a direct EMA of raw stochastic gradients — it includes a gradient correction term $\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t)$ that changes its statistical properties. If $v_t$ continues to be defined using raw stochastic gradients $[\nabla f(x_t, \xi_t)]^2$ , it no longer reflects the scale of the updates being applied. This mismatch can cause the preconditioner (which divides by $\sqrt{v_t}$ ) to improperly scale the coordinates of the variance-reduced gradient, potentially undoing the benefits of variance reduction.

Moreover, SuperAdam lacks:

A scaling parameter $\gamma_t$ for controlling the strength of variance reduction (MARS introduces this, allowing interpolation between pure AdamW at $\gamma_t = 0$ and full STORM variance reduction at $\gamma_t = 1$ ).
Gradient clipping on the intermediate gradient estimator $c_t$ (MARS clips $c_t$ by norm before feeding it into the EMA).
Bias correction and decoupled weight decay (MARS-AdamW uses both, following the AdamW recipe).

3. VRAdam (Li, 2024) and AdaSPIDER (Kavis et al., 2022): Mentioned briefly as additional instances of variance-reduced adaptive methods. VRAdam integrates variance reduction with AdamW for improved convergence rates; AdaSPIDER introduces adaptive step sizes into the SPIDER algorithm. Like Adam+ and SuperAdam, these were validated only on basic computer vision tasks (MNIST, CIFAR-10) and small-scale language modeling (SWB-300 with 2-layer LSTMs or 2-layer Transformers), never at the scale of GPT-2 or larger models.

The paper's diagnosis is unambiguous:

"a significant gap remains in the successful application of variance reduction techniques to adaptive gradient methods, particularly in the rapidly evolving domain of large language models."

How This Paper Positions Itself

The paper's positioning is a classic "revisit a failed idea when conditions have changed" narrative, combined with a "the reason it failed before was poor engineering, not fundamental limitations" argument. The components of this positioning:

1. Structural barriers have been removed. As discussed above, dropout, batch normalization, and data augmentation — the "factors disrupting the finite-sum structure" identified by Defazio & Bottou (2019) — are absent in modern LLM training. This is not a new discovery but rather a re-interpretation of the landscape: what was once a fundamental obstacle is now a non-issue, provided you are training transformers at scale.

2. Previous attempts had correctable design flaws. Rather than treating the failure of SuperAdam et al. as evidence that variance reduction cannot work with adaptive methods, the paper argues that those specific implementations had identifiable problems — the first/second-order momentum mismatch in particular — that can be fixed. MARS is positioned as the correction of these flaws.

3. Variance reduction is orthogonal to existing optimizer innovation. This is a crucial strategic claim. The paper compiles a taxonomy of optimizer components:

Momentum mechanisms: Nesterov's acceleration, STORM variance reduction, standard EMA momentum.
Preconditioning mechanisms: Diagonal approximations (Adam, Lion's sign operation), full-matrix approximations (Shampoo, K-FAC), eigenspace-based approximations (Muon).

The argument is that most prior work has focused on improving preconditioning (better Hessian approximations) while keeping the momentum component standard. MARS flips this by keeping preconditioning designs from AdamW, Lion, and Shampoo but replacing the standard momentum with a variance-reduced variant. The two axes are independent, so improvements along one should combine with any choice along the other. This is why MARS is a "framework" rather than a single algorithm — it plugs into existing preconditioner designs.

4. The unification provides theoretical guarantees. The paper includes convergence analysis (Theorems B.5 and B.6 in the Appendix) showing that MARS achieves $\mathcal{O}(T^{-1/3})$ convergence rate, improving over AdamW's $\mathcal{O}(T^{-1/4})$ rate under standard assumptions (bounded variance, $L$ -smoothness, preconditioner lower-bounded by $\rho I$ ). This is not the main contribution — the paper is primarily empirical — but it provides theoretical backing for the claim that variance reduction genuinely improves convergence when properly integrated.

5. The "scaled" variance reduction is key. The introduction of $\gamma_t$ is not just a hyperparameter for tuning — the paper provides a control variate argument in Section 3.1 showing that an optimal choice of $\gamma_t$ exists that minimizes the variance of the gradient estimator. When $\gamma_t = 1$ , the algorithm recovers full STORM variance reduction; when $\gamma_t = 0$ , it recovers standard SGD momentum. In practice, the paper finds $\gamma_t = 0.025$ works best (from the ablation in Appendix E.5, Figure 14), meaning the optimal setting is much closer to standard momentum than to full variance reduction — a finding that would not have been possible without the scaling parameter. This explains why prior attempts that implicitly used $\gamma_t = 1$ (full STORM) may have underperformed: the full variance reduction correction can over-correct or amplify noise in some settings.

The Practical Stakes

The paper's motivation is ultimately empirical and practical, not just theoretical. The introduction frames the problem in terms of training efficiency:

"Adaptive gradient algorithms like Adam, AdamW, and their variants have been central to this task [training large models]... This suggests the ongoing challenge in developing adaptive gradient methods superior to Adam and AdamW for large-scale model training."

The concrete claim in the abstract — that "AdamW requires 50 billion tokens to reach a validation loss of 2.58, whereas MARS only requires 28 billion tokens" — translates to a ~44% reduction in training tokens. For GPT-2 large (770M parameters) this is substantial but not transformative. But if the gains scale to larger models (GPT-3 scale at 175B parameters, or Llama-scale at 70B+), the cost savings would be enormous. The paper stops short of making this claim explicitly, but the trajectory of results across GPT-2 small → medium → large (where the gap between MARS and AdamW widens with scale) suggests that variance reduction becomes more beneficial as models grow, which is consistent with the intuition that larger models have higher-dimensional, noisier gradient spaces where variance reduction provides more value. This scaling trend is, implicitly, the strongest motivation the paper offers for why this work matters beyond the specific GPT-2 experiments.

3. Technical Approach

3.1 Reader Orientation

The paper develops a unified optimization framework — a family of algorithms, not a single optimizer — that modifies existing adaptive gradient methods (AdamW, Lion, Shampoo) by replacing their standard momentum with a variance-reduced gradient estimator while keeping their preconditioning mechanisms intact. The system solves the high-variance stochastic gradient problem in large-model training by injecting a scaled version of the STORM (stochastic recursive momentum) correction term into the gradient estimation pipeline, then redefining the second-order momentum to be consistent with this corrected gradient rather than the raw stochastic gradient. The "shape" of the solution is: take an existing optimizer, replace its gradient estimator with a variance-reduced one, adjust the adaptive preconditioner to match, and add a single scalar $\gamma_t$ that smoothly interpolates between no variance reduction and full variance reduction, enabling practitioners to tune the strength of correction without changing any other architectural choices.

3.2 Big-Picture Architecture (Diagram in Words)

The MARS framework has five major components, organized as a sequential pipeline that transforms raw data and parameters into updated model weights:

Gradient Oracle — the standard mechanism that computes the stochastic gradient $\nabla f(x_t, \xi_t)$ for current parameters $x_t$ on a minibatch $\xi_t$ . This is unchanged from standard training; MARS wraps around it, not inside it.
Variance-Reduced Gradient Estimator ( $c_t$ ) — a new intermediate quantity that combines the raw stochastic gradient with a scaled gradient correction term $\gamma_t \frac{\beta_1}{1-\beta_1}(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t))$ . This correction term subtracts the gradient at the previous parameters $x_{t-1}$ evaluated on the same data $\xi_t$ , which cancels the noise component shared across adjacent parameter states. The scaling parameter $\gamma_t$ controls how much correction is applied. When $\gamma_t = 0$ , $c_t$ reduces to the standard stochastic gradient; when $\gamma_t = 1$ , $c_t$ recovers the full STORM variance-reduced gradient.
Gradient Clipping Module — applies standard clipping-by-norm to $c_t$ (not to the preconditioned update, unlike Sophia), producing $\tilde{c}_t = \text{Clip}(c_t, 1)$ . If $\|c_t\|_2 > 1$ , the vector is rescaled to unit norm; otherwise it passes through unchanged. This prevents variance-reduced estimates from producing excessively large updates during early training or when the correction term amplifies gradient differences.
Preconditioned Momentum State — maintains two running exponential moving averages (EMAs). The first-order momentum $m_t = \beta_1 m_{t-1} + (1-\beta_1)\tilde{c}_t$ is the variance-reduced analog of Adam's momentum, but computed from $\tilde{c}_t$ rather than from the raw gradient. The second-order momentum $v_t = \beta_2 v_{t-1} + (1-\beta_2)\tilde{c}_t^2$ (for MARS-AdamW) or $H_t$ (for the general framework) is the critical redesign: it uses the variance-reduced estimate $\tilde{c}_t$ rather than the raw stochastic gradient to compute the adaptive learning rate scale, ensuring the preconditioner reflects the actual magnitude of the updates being applied.
Preconditioned Parameter Update — performs a mirror descent step $x_{t+1} = \arg\min_x \{\eta_t \langle m_t, x\rangle + \frac{1}{2}\|x - x_t\|^2_{H_t}\}$ , which reduces to the familiar Adam-style update $x_{t+1} = x_t - \eta_t \frac{m_t}{\sqrt{v_t} + \epsilon} - \lambda x_t$ when $H_t$ is diagonal and weight decay is included. The three MARS instances differ only in how $H_t$ is constructed: diagonal from squared $\tilde{c}_t$ (MARS-AdamW), sign-based from $m_t^2$ (MARS-Lion), or Kronecker-factored from matrix-shaped $m_t$ via SVD (MARS-Shampoo).

Information flow: At each iteration $t$ , the system samples a minibatch $\xi_t$ , computes $\nabla f(x_t, \xi_t)$ , retrieves the previously computed $\nabla f(x_{t-1}, \xi_t)$ (which requires either storing the previous gradient or recomputing it with the same data), forms $c_t$ by combining these with $\gamma_t$ , clips $c_t$ to unit norm to produce $\tilde{c}_t$ , updates the two EMA states $m_t$ and $v_t$ (or $H_t$ ), and applies the preconditioned update to $x_t$ . The approximate variant MARS-approx replaces $\nabla f(x_{t-1}, \xi_t)$ with $\nabla f(x_{t-1}, \xi_{t-1})$ — using the previous step's minibatch rather than requiring the same data twice — which is computationally cheaper at the cost of a small empirical performance degradation (shown in Appendix E.2).

3.3 Roadmap for the Deep Dive

First, the core variance-reduced gradient estimator $c_t$ — how it is constructed, why the subtraction $\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t)$ reduces variance, and the role of the scaling parameter $\gamma_t$ with the control variate derivation that shows an optimal $\gamma_t$ exists.
Second, the gradient clipping step and why it clips $c_t$ by norm rather than clipping the final preconditioned update — a distinction from Sophia that matters for stability with variance-reduced estimates.
Third, the preconditioned momentum definitions — how $m_t$ and $v_t$ (or $H_t$ ) are redefined to be coherent with the variance-reduced gradient, and why using raw gradients for $v_t$ (as SuperAdam does) creates a critical mismatch.
Fourth, the three instantiated algorithms (MARS-AdamW, MARS-Lion, MARS-Shampoo) — their specific preconditioner constructions, their connections to existing optimizers (Adan, Lion, Muon), and the exact hyperparameter configurations used.
Fifth, the MARS-approx practical variant — how replacing $\xi_t$ with $\xi_{t-1}$ in the correction term changes the algorithm and the empirical trade-off.
Sixth, the convergence theory — the main theorems (B.5 and B.6), the Lyapunov function construction, and why MARS achieves $\mathcal{O}(T^{-1/3})$ versus AdamW's $\mathcal{O}(T^{-1/4})$ .

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an empirical systems paper with theoretical backing whose core idea is that variance reduction can be made to work in large-scale training if and only if (a) the strength of variance reduction is tunable via a scaling parameter, (b) the adaptive preconditioner is redefined to be coherent with the variance-reduced gradient, and (c) gradient clipping is applied to the variance-reduced estimate before it enters the EMA pipeline. The paper demonstrates this by instantiating the framework with three existing preconditioner designs (AdamW, Lion, Shampoo) and evaluating on GPT-2 at scales up to 770M parameters.

The Core Gradient Estimator $c_t$ : Scaled Variance Reduction

The heart of MARS is the intermediate gradient estimator $c_t$ , which replaces the raw stochastic gradient $\nabla f(x_t, \xi_t)$ as the signal that enters the momentum pipeline. The estimator is defined as:

$c_t = \nabla f(x_t, \xi_t) + \gamma_t \frac{\beta_1}{1-\beta_1}\left(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t)\right)$

where $\nabla f(x_t, \xi_t)$ is the stochastic gradient at the current parameters on data $\xi_t$ , $\nabla f(x_{t-1}, \xi_t)$ is the stochastic gradient at the previous parameters evaluated on the same data $\xi_t$ , $\beta_1 \in (0, 1)$ is the first-order momentum decay parameter, and $\gamma_t \in [0, 1]$ is the newly introduced scaling parameter that controls the strength of variance reduction.

What it computes: $c_t$ takes the standard stochastic gradient and adds a correction term proportional to the difference between the current and previous gradients evaluated on the same data. This difference $\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t)$ has a specific statistical property: the noise component from $\xi_t$ is present in both gradient evaluations at approximately the same magnitude (since the parameters haven't moved far), so subtracting them partially cancels the noise. What remains is primarily the signal component — the change in the true gradient $\nabla F(x_t) - \nabla F(x_{t-1})$ plus a residual noise term that shrinks as $x_t$ approaches $x_{t-1}$ . The scaling factor $\frac{\beta_1}{1-\beta_1}$ adjusts the correction to be consistent with the EMA weighting scheme that will be applied subsequently; when $\gamma_t = 0$ , the correction term is disabled and $c_t$ reduces to the raw stochastic gradient; when $\gamma_t = 1$ , the magnitude matches the standard STORM formulation.

Why this form: There are three design choices here that warrant explanation:

Why the same data $\xi_t$ in both terms? This is the defining characteristic of the STORM/SPIDER family of variance reduction methods, as opposed to SVRG (which uses an anchor point with a different, independent batch). By evaluating both gradients on the same random variable $\xi_t$ , the noise $\nabla f(x_t, \xi_t) - \nabla F(x_t)$ is approximately cancelled by $\nabla f(x_{t-1}, \xi_t) - \nabla F(x_{t-1})$ , assuming $F$ is smooth (so $\nabla F(x_t) \approx \nabla F(x_{t-1})$ ). This is fundamentally different from Nesterov's momentum (discussed in Appendix B.1), which uses $\nabla f(x_{t-1}, \xi_{t-1})$ — a different random batch from the previous step — and therefore does not achieve noise cancellation, only acceleration along the smoothed gradient direction. The paper explicitly contrasts these in Appendix B.1:

"In Nesterov's acceleration, $\nabla f(x_t, \xi_t)$ is subtracted by $\nabla f(x_{t-1}, \xi_{t-1})$ to determine a direction of improvement. In contrast, STORM variance reduction subtracts $\nabla f(x_{t-1}, \xi_t)$ from $\nabla f(x_t, \xi_t)$ to cancel out the noise introduced by $\xi_t$ ."

Why the $\frac{\beta_1}{1-\beta_1}$ scaling? This factor ensures that when $c_t$ is subsequently combined into the EMA $m_t = \beta_1 m_{t-1} + (1-\beta_1)\tilde{c}_t$ , the effective weight on the correction term matches the STORM formulation in Equation (2.5), where the correction appears as $\frac{\beta_1}{1-\beta_1}(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t))$ inside the parentheses multiplied by $(1-\beta_1)$ outside. Without this scaling, the correction would have the wrong magnitude relative to the momentum carryover $\beta_1 m_{t-1}$ . The factor can be derived by substituting the unclipped $c_t$ into the EMA:

$m_t = \beta_1 m_{t-1} + (1-\beta_1)\left[\nabla f(x_t, \xi_t) + \gamma_t \frac{\beta_1}{1-\beta_1}(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t))\right]$ $= \beta_1 m_{t-1} + (1-\beta_1)\nabla f(x_t, \xi_t) + \gamma_t \beta_1(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t))$

which matches Equation (2.4) when $\gamma_t = 1$ , confirming that the STORM momentum is recovered exactly.

Why the scaling parameter $\gamma_t$ ? This is the paper's most important architectural innovation for practical deployment. The paper provides a control variate derivation (Section 3.1, "Why $\gamma_t$ improves convergence") showing that the variance reduction correction is not always beneficial at full strength. In the standard control variate framework, if we have an estimator $X$ and a correlated control variate $Y$ with known mean $\mathbb{E}[Y]$ , the variance-minimizing estimator is $X - \gamma^*(Y - \mathbb{E}[Y])$ where:

$\gamma^* = \frac{\mathbb{E}[(X - \mathbb{E}[X])(Y - \mathbb{E}[Y])]}{\text{Var}(Y)}$

The optimal $\gamma^*$ is not necessarily 1 — it depends on the correlation between the noise in $X$ and $Y$ . In the STORM context, the paper defines $X = (1-\beta)\nabla f(x_{t+1}, \xi_{t+1})$ , $Y = \beta[\nabla f(x_{t+1}, \xi_{t+1}) - \nabla f(x_t, \xi_{t+1})]$ , and shows that the optimal $\gamma^*$ in this setting is:

$\gamma^* = 1 - \frac{\mathbb{E}[UY] + \text{Var}(Y)}{\mathbb{E}[Y^2]}$

where $U = X - \mathbb{E}[X] + Z_t$ and $Z_t = m_t - \nabla F(x_t)$ . This optimal value is generally less than 1, and the paper's empirical finding that $\gamma = 0.025$ works best (Appendix E.5, Figure 14) — far below 1 — validates this analysis. Setting $\gamma = 1$ (as SuperAdam implicitly does by using full STORM) over-corrects, potentially amplifying noise in early training where parameter movement is large (so $\nabla f(x_t, \xi_t)$ and $\nabla f(x_{t-1}, \xi_t)$ are genuinely different, making the correction term a poor noise estimate). The scaling parameter allows MARS to interpolate between pure AdamW ( $\gamma = 0$ ) and full STORM ( $\gamma = 1$ ), and the optimal point in practice is much closer to the AdamW end.

Gradient Clipping on $c_t$

After computing $c_t$ , MARS applies gradient clipping-by-norm with a threshold of 1:

$\tilde{c}_t = \text{Clip}(c_t, 1) = \begin{cases} \frac{c_t}{\|c_t\|_2} & \text{if } \|c_t\|_2 > 1, \\ c_t & \text{otherwise.} \end{cases}$

where $\tilde{c}_t$ is the clipped gradient estimator that feeds into the momentum EMA.

What it computes: If the variance-reduced gradient estimate has $\ell_2$ norm exceeding 1, it is rescaled to have unit norm while preserving its direction. If its norm is already $\leq 1$ , it passes through unchanged.

Why this form and placement: This is a deliberate departure from existing optimizers that use clipping. Sophia (Liu et al., 2023) applies clipping-by-value to the preconditioned gradient (i.e., after dividing by the Hessian estimate), which controls the magnitude of the final parameter update. MARS clips before the preconditioning step, on the raw gradient estimator $c_t$ itself. The paper explains:

"their approach does clipping upon the preconditioned gradient with clipping-by-value, while our method applies clipping to the intermediate gradient estimate using the more standard technique of clipping-by-norm."

The rationale is that the variance-reduced estimator $c_t$ can occasionally produce very large values in early training or when the correction term amplifies gradient differences rather than cancelling noise. This is particularly important because the correction term involves a difference $\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t)$ which, when parameters are moving rapidly (high learning rate early in training, or in sharp loss landscapes), can be large even if both individual gradients are moderate — their difference can point in unexpected directions with significant magnitude. Clipping $c_t$ prevents these outliers from destabilizing the momentum states $m_t$ and $v_t$ , which have long memories due to the EMA. The threshold of 1 is not tuned per-model in the reported experiments; it is fixed across all GPT-2 scales, implying it is chosen based on the observation that well-conditioned gradient estimates should typically have norm $\leq 1$ after the scaling inherent in the $\frac{\beta_1}{1-\beta_1}$ factor. The paper also notes that this follows "standard gradient clipping technique performed in neural network training," making it compatible with existing training infrastructure.

The clipping-by-norm choice (rather than clipping-by-value) preserves the direction of $c_t$ when it exceeds the threshold, only scaling down its magnitude. This is important because the direction of the variance-reduced gradient carries the noise-cancellation information — scaling it uniformly preserves the relative contributions of different coordinates, whereas clipping-by-value would independently cap each coordinate, potentially distorting the noise-cancellation structure.

Preconditioned Momentum States: Redefining $m_t$ and $v_t$

With $\tilde{c}_t$ computed and clipped, MARS updates two EMA states that mirror Adam's structure but are semantically different because they operate on the variance-reduced estimate rather than the raw gradient.

First-order momentum ( $m_t$ ):

$m_t = \beta_1 m_{t-1} + (1-\beta_1)\tilde{c}_t$

where $\beta_1$ is the first-order momentum decay rate, $m_{t-1}$ is the previous momentum state (initialized to 0), and $\tilde{c}_t$ is the clipped variance-reduced estimator.

What it computes: $m_t$ is an exponentially weighted moving average of past variance-reduced gradient estimates. It combines the historical direction information in $m_{t-1}$ (weighted by $\beta_1$ ) with the new variance-reduced signal $\tilde{c}_t$ (weighted by $1-\beta_1$ ). The result is a smoothed direction estimate that benefits from both the noise-cancellation in $c_t$ and the temporal smoothing from the EMA.

Why this form: This is structurally identical to Adam's first-order momentum, but the input $\tilde{c}_t$ is fundamentally different from $\nabla f(x_t, \xi_t)$ . The effect is that $m_t$ now tracks a variance-reduced direction, meaning its variance is lower than standard momentum for the same $\beta_1$ , or equivalently, it can use a higher $\beta_1$ (more aggressive smoothing) without introducing excessive bias. The paper uses $\beta_1 = 0.95$ for MARS, compared to $\beta_1 = 0.9$ for AdamW in their GPT-2 experiments (Table 12 in Appendix F), consistent with this intuition: lower variance in the input allows heavier temporal smoothing. When $\gamma_t = 0$ , $\tilde{c}_t$ reduces to the clipped raw gradient, and $m_t$ becomes identical to Adam's momentum (modulo the clipping, which Adam typically does not apply to the gradient before EMA).

Second-order momentum ( $v_t$ ) — the critical redesign:

For MARS-AdamW:

$v_t = \beta_2 v_{t-1} + (1-\beta_2)\tilde{c}_t^2$

where $\beta_2$ is the second-order momentum decay rate, $v_{t-1}$ is the previous second-order state (initialized to 0), and $\tilde{c}_t^2$ denotes element-wise squaring of the clipped variance-reduced estimate.

What it computes: $v_t$ is an EMA of squared variance-reduced gradient estimates. Each coordinate $i$ of $v_t$ approximates the recent average squared magnitude of the $i$ -th component of $\tilde{c}_t$ . When used as a preconditioner (dividing $m_t$ by $\sqrt{v_t}$ ), it adaptively rescales each coordinate's learning rate inversely proportional to its typical magnitude, as in standard Adam.

Why this form — the paper's key design insight: This is where MARS diverges crucially from SuperAdam and other prior variance-reduced adaptive methods. In standard Adam, $v_t$ tracks $\mathbb{E}[(\nabla f)^2]$ , and $m_t$ tracks $\mathbb{E}[\nabla f]$ , so $\frac{m_t}{\sqrt{v_t}}$ approximates a signal-to-noise ratio: the mean gradient direction divided by its typical scale. Both quantities describe the same random variable. In SuperAdam, the first-order momentum uses a STORM variance-reduced estimator, but the second-order momentum continues to track raw squared gradients $[\nabla f(x_t, \xi_t)]^2$ . This creates a mismatch: $m_t$ has lower variance than $\nabla f(x_t, \xi_t)$ , but $v_t$ estimates the scale of the unreduced gradient, which has higher variance. When dividing $m_t$ by $\sqrt{v_t}$ , the denominator is inflated (because the raw gradient has higher variance), causing the update to be too small in coordinates where variance reduction was most effective. The paper states this explicitly:

"SuperAdam's preconditioner matrix is designed following Adam's structure but does not account for the revised definition of variance-reduced momentum, resulting in a significant mismatch between the first-order and second-order momentum."

MARS fixes this by making $v_t$ an EMA of $\tilde{c}_t^2$ — the squared variance-reduced estimate — so both $m_t$ and $v_t$ describe the same underlying signal. When $\gamma_t = 0$ , $\tilde{c}_t$ is the raw gradient, and $v_t$ reduces to Adam's second-order momentum; when $\gamma_t = 1$ , $v_t$ tracks the squared STORM estimator, whose magnitude is generally smaller than the raw squared gradient because the noise has been partially cancelled. The adaptive learning rate $\frac{1}{\sqrt{v_t}}$ is therefore appropriately scaled for the actual update magnitude.

Bias correction and weight decay:

Following the AdamW formulation, MARS-AdamW applies bias correction to both momentum states:

$\hat{m}_t = \frac{m_t}{1-\beta_1^t}, \quad \hat{v}_t = \frac{v_t}{1-\beta_2^t}$

These corrections compensate for the initialization at zero, which biases the early EMAs toward zero at small $t$ . The bias-corrected estimates are then used in the update with decoupled weight decay (Loshchilov & Hutter, 2019):

$x_{t+1} = x_t - \eta_t\left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} + \lambda x_t\right)$

where $\eta_t$ is the learning rate (scheduled via cosine decay in the GPT-2 experiments), $\epsilon$ is a small constant for numerical stability, and $\lambda$ is the weight decay coefficient. The decoupled form applies weight decay directly to the parameters $x_t$ rather than incorporating it into the gradient, avoiding interference with the adaptive learning rate mechanism. The paper notes that this completes the mirror descent update from Equation (3.4) with $H_t$ defined as:

$H_t := \sqrt{\text{diag}(v_t)} \cdot \frac{1-\beta_1^t}{\sqrt{1-\beta_2^t}}$

where the $\frac{1-\beta_1^t}{\sqrt{1-\beta_2^t}}$ factor accounts for the bias correction in the mirror descent formulation.

Hyperparameter values for MARS-AdamW in GPT-2 experiments: The paper uses $\beta_1 = 0.95$ , $\beta_2 = 0.99$ , $\gamma_t \equiv 0.025$ (constant across all steps), and learning rates of $6 \times 10^{-3}$ (small), $3 \times 10^{-3}$ (medium), and $2 \times 10^{-3}$ (large), with a minimum learning rate of $3 \times 10^{-5}$ (small), $6 \times 10^{-5}$ (medium), and $1 \times 10^{-5}$ (large) under cosine decay, weight decay searched over $\{10^{-1}, 10^{-2}, 10^{-3}\}$ , gradient clipping threshold 1.0, batch size 480, context length 1024, 100,000 training steps with 2,000 warmup steps, and dropout 0.0 (Tables 10–12 in Appendix F).

The General Preconditioned Variance-Reduced Update

Beyond the AdamW instantiation, the paper frames the entire MARS family within a unified mirror descent formulation. The general update is:

$m_t = \beta_1 m_{t-1} + (1-\beta_1)\left[\nabla f(x_t, \xi_t) + \gamma_t \frac{\beta_1}{1-\beta_1}(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t))\right]$

$x_{t+1} = \arg\min_{x \in \mathbb{R}^d} \left\{\eta_t \langle m_t, x\rangle + \frac{1}{2}\|x - x_t\|^2_{H_t}\right\}$

where $H_t$ is a positive-definite preconditioning matrix (general, not necessarily diagonal), and the argmin has the closed-form solution $x_{t+1} = x_t - \eta_t H_t^{-1} m_t$ (plus weight decay when included).

What it computes: The first equation is the variance-reduced momentum (without clipping, which is added in the practical implementations). The second equation performs a proximal step that minimizes a linear approximation of the loss ( $\langle m_t, x\rangle$ ) plus a Bregman divergence $\frac{1}{2\eta_t}\|x - x_t\|^2_{H_t}$ that penalizes movement in directions where $H_t$ is large (i.e., where the loss is sharply curved). The penalty structure is determined by $H_t$ : directions with large $H_t$ entries (high curvature or high gradient variance) permit little movement, while directions with small $H_t$ entries (flat or low-variance) permit larger steps.

Why this general form: The paper explicitly frames this as an online mirror descent (OMD) generalization, following Gupta et al. (2018), because it accommodates both diagonal preconditioners (AdamW, Lion) and full-matrix preconditioners (Shampoo, Muon) within the same mathematical framework. The different MARS instantiations (AdamW, Lion, Shampoo) differ only in their choice of $H_t$ , demonstrating that variance reduction is compatible with any preconditioner design. This modularity is the paper's architectural claim: variance reduction is an orthogonal axis of improvement that does not constrain preconditioner choice.

MARS-Lion: Sign-Based Preconditioning

MARS-Lion adapts the Lion optimizer's sign-based update rule, discovered through symbolic program search by Chen et al. (2023). The standard Lion update is:

$u_{t+1} = \beta_2 u_t + (1-\beta_2)\nabla f(x_t, \xi_t)$ $m_t = \beta_1 u_t + (1-\beta_1)\nabla f(x_t, \xi_t)$ $x_{t+1} = x_t - \eta_t(\text{sign}(m_t) + \lambda x_t)$

where $u_t$ is an auxiliary momentum state with decay $\beta_2$ , $m_t$ is a combination of $u_t$ and the current gradient, and the parameter update uses only the sign of $m_t$ (rather than its magnitude), creating a uniform-magnitude update for all coordinates. This can be interpreted as a mirror descent with $H_t = \sqrt{\text{diag}(m_t^2)}$ , which yields the sign operation when solving the argmin.

MARS-Lion replaces the raw gradient with the variance-reduced estimator $\tilde{c}_t$ and simplifies the Lion momentum structure to a single EMA (eliminating the auxiliary $u_t$ ):

$c_t = \nabla f(x_t, \xi_t) + \gamma_t \frac{\beta_1}{1-\beta_1}(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t))$ $\tilde{c}_t = \text{Clip}(c_t, 1)$ $m_t = \beta_1 m_{t-1} + (1-\beta_1)\tilde{c}_t$ $x_{t+1} = x_t - \eta_t(\text{sign}(m_t) + \lambda x_t)$

What it computes: MARS-Lion uses the same variance-reduced gradient estimator and clipping as MARS-AdamW, but replaces the adaptive per-coordinate scaling $\frac{1}{\sqrt{v_t}}$ with the sign function, which outputs $\pm 1$ for each coordinate based on the direction of $m_t$ . This means all coordinates receive updates of equal magnitude $\eta_t$ , regardless of their typical gradient scale.

Connection to standard Lion (Lemma 3.4): The paper proves that standard Lion's two-state momentum is equivalent to a single momentum update with a built-in gradient correction term. By Lemma 3.4, the Lion updates can be rewritten as:

$m_t = \beta_2 m_{t-1} + (1-\beta_2)\nabla f(x_t, \xi_t) + (\beta_2 - \beta_1)(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_{t-1}))$

where the last term $(\beta_2 - \beta_1)(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_{t-1}))$ is a gradient correction term that resembles variance reduction but uses the previous batch $\xi_{t-1}$ rather than the current batch $\xi_t$ . MARS-Lion with $\beta_1$ set to $\beta_2$ (Lion's $\beta_2$ ) and $\gamma_t = \frac{\beta_2 - \beta_1}{\beta_2}$ (where $\beta_1$ here means Lion's $\beta_1$ ) recovers a version using the same batch $\xi_t$ for both terms:

$m_t = \beta_2 m_{t-1} + (1-\beta_2)\nabla f(x_t, \xi_t) + (\beta_2 - \beta_1)(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t))$

The paper notes: "the only difference lies in the stochasticity used, specifically, $\xi_t$ versus $\xi_{t-1}$ when calculating $\nabla f(x_{t-1}, \cdot)$ ." This means Lion's built-in gradient correction, discovered inadvertently by symbolic program search, is actually an approximate form of variance reduction — the search algorithm independently found a mechanism that provides a similar benefit to STORM, albeit with a different noise structure. MARS-Lion makes this variance reduction explicit and controllable via $\gamma_t$ .

Why sign-based preconditioning: The sign operation provides a form of adaptive scaling that is coarser than Adam's diagonal preconditioner but computationally cheaper (no need to maintain $v_t$ or compute square roots) and memory-efficient (no second-order momentum state). The uniform magnitude can also be beneficial for training stability, as it prevents any single coordinate from dominating the update. The MARS framework shows that this preconditioner choice is independent of the variance reduction mechanism — both can be freely combined.

MARS-Shampoo: Eigenspace Preconditioning

MARS-Shampoo applies variance reduction to Shampoo's full-matrix preconditioning, which operates on the eigenspace of the gradient matrix rather than coordinate-wise scaling. For a weight matrix $x_t \in \mathbb{R}^{m \times n}$ and gradient matrix $m_t \in \mathbb{R}^{m \times n}$ (the variance-reduced momentum, reshaped), the update is:

$U_t, \Sigma_t, V_t = \text{SVD}(m_t)$ $x_{t+1} = x_t - \eta_t(U_t V_t^\top + \lambda x_t)$

where $U_t \in \mathbb{R}^{m \times r}$ and $V_t \in \mathbb{R}^{n \times r}$ are the left and right singular vectors of $m_t$ (with $r = \min(m, n)$ ), $\Sigma_t$ is the diagonal matrix of singular values, and $U_t V_t^\top$ is the eigenspace-preconditioned update direction.

What it computes: The SVD decomposes the momentum matrix $m_t$ into $U_t \Sigma_t V_t^\top$ . The Shampoo-inspired update ignores the singular values $\Sigma_t$ and uses only the singular vectors $U_t V_t^\top$ , which is a matrix with the same shape as $m_t$ but all singular values set to 1. Geometrically, this performs a rotation: $U_t V_t^\top$ is the orthogonal matrix closest to $m_t$ in Frobenius norm (by the Eckart-Young-Mirsky theorem), and applying it as an update rotates the weight matrix toward the directions emphasized by the momentum, with a uniform step size across all eigenmodes. The update can be seen as preconditioning on the eigenspace of $m_t$ with the preconditioner $H_t = (m_t m_t^\top)^{-1/4} \otimes (m_t^\top m_t)^{-1/4}$ using the Kronecker-product structure of Shampoo (Gupta et al., 2018).

Why SVD-based preconditioning: Standard Shampoo maintains running accumulators $L_t = L_{t-1} + m_t m_t^\top$ and $R_t = R_{t-1} + m_t^\top m_t$ and computes $L_t^{-1/4} m_t R_t^{-1/4}$ as the update, which requires matrix root computations. The SVD-based formulation $U_t V_t^\top$ is equivalent when using the instantaneous gradient $m_t$ rather than accumulated statistics (i.e., $L_t = m_t m_t^\top$ and $R_t = m_t^\top m_t$ , so $L_t^{-1/4} = U_t \Sigma_t^{-1/2} U_t^\top$ and $R_t^{-1/4} = V_t \Sigma_t^{-1/2} V_t^\top$ , yielding $L_t^{-1/4} m_t R_t^{-1/4} = U_t V_t^\top$ ). The paper notes that this formulation is equivalent to Muon (Jordan et al., 2024) in the SVD instantiation, and that practical implementations can use sketching, Newton iteration, or Newton-Schulz iteration instead of full SVD for computational efficiency.

Connection to Muon (Equation 3.24 vs. 3.25): Muon uses a Nesterov-style momentum:

$u_t = \mu u_{t-1} + \nabla f(x_t, \xi_t)$ $m_t = \mu u_t + \nabla f(x_t, \xi_t)$

which, by Lemma D.1, is equivalent to:

$m_t = \mu m_{t-1} + \nabla f(x_t, \xi_t) + \mu(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_{t-1}))$

This again contains a gradient correction term, but using $\xi_{t-1}$ rather than $\xi_t$ . MARS-Shampoo with $\beta_1 = \mu$ and $\gamma_t = 1-\mu$ yields (after dividing by $1-\mu$ ):

$\frac{m_t}{1-\mu} = \mu \frac{m_{t-1}}{1-\mu} + \nabla f(x_t, \xi_t) + \mu(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t))$

which is a rescaled version of Muon's momentum but with the variance-reducing property of using the same batch $\xi_t$ in the difference. The paper remarks that "in practice, we observe little difference" between using $\xi_t$ vs. $\xi_{t-1}$ , which is consistent with the approximate variant MARS-approx performing nearly as well as the exact MARS.

MARS-approx: The Practical Variant

The exact MARS formulation requires evaluating $\nabla f(x_{t-1}, \xi_t)$ — the gradient at the previous parameters on the current minibatch — which costs an additional forward-backward pass per iteration (or requires storing the previous gradient and replaying the same data). To avoid this cost, the paper proposes MARS-approx, which substitutes $\nabla f(x_{t-1}, \xi_{t-1})$ (the gradient computed at the previous step with its own minibatch) for $\nabla f(x_{t-1}, \xi_t)$ :

$c_t^{\text{approx}} = \nabla f(x_t, \xi_t) + \gamma_t \frac{\beta_1}{1-\beta_1}(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_{t-1}))$

What changes: The correction term now uses two different batches: the current batch $\xi_t$ for $\nabla f(x_t, \xi_t)$ and the previous batch $\xi_{t-1}$ for $\nabla f(x_{t-1}, \xi_{t-1})$ . This breaks the exact noise cancellation property, because the noise in $\nabla f(x_{t-1}, \xi_{t-1})$ is independent of the noise in $\nabla f(x_t, \xi_t)$ . However, if the parameters $x_{t-1}$ and $x_t$ are close (small learning rate or late in training), the signal components are similar, and the noise components, while independent, still have zero mean — so the correction term provides a biased but potentially still useful estimate. The paper notes that while "MARS and MARS-approx differ in their updates and may theoretically exhibit distinct convergence guarantees," the empirical difference is small: "MARS provides only marginal improvements over MARS-approx in practice" (Appendix E.2, Figures 4–5). All main GPT-2 experiments in Section 4 use MARS-approx as the default, with exact MARS compared only in the ablation.

Practical advantage: MARS-approx has the same computational cost per iteration as standard AdamW (one forward-backward pass, plus the stored previous gradient), making it directly comparable in wall-clock time. The paper's wall-clock time measurements in Figures 1(c), 2(c), and 3(c) show MARS-AdamW and MARS-Lion have "slightly higher per-iteration cost compared to AdamW but much faster than Muon," and they achieve lower validation loss within the same wall-clock time.

Convergence Theory (Theorems B.5 and B.6)

The paper provides a convergence analysis in Appendix B.2, proving two main theorems: Theorem B.5 for the general MARS framework (Algorithm 1) and Theorem B.6 for MARS-AdamW with weight decay (Algorithm 2). The analysis uses standard nonconvex optimization assumptions: Assumption B.1 (bounded gradient variance $\sigma^2$ ), Assumption B.2 ( $L$ -smoothness of $f$ ), and Assumption B.3 (preconditioner $H_t$ lower-bounded by $\rho I$ , satisfied by adding $\epsilon$ to the diagonal).

Key result (Theorem B.5): Under carefully chosen time-varying parameters $\eta_t = (s+t)^{-1/3}$ , $\beta_{1,t+1} = 1 - c\eta_t^2$ (with $c \geq 32L^2\rho^{-2} + 1$ ), and $\beta_{2,t+1} = 1 - \eta_t^6$ , MARS achieves:

$\frac{1}{T}\sum_{t=1}^T \mathbb{E}\|\nabla F(x_t) - m_t\|_2^2 \leq \mathcal{O}\left(\frac{\log T}{T^{2/3}}\right) - \mathcal{O}\left(\frac{\sum M_{t+1}}{T^{1/3}}\right)$

where $\nabla F(x_t) - m_t$ is the error between the true gradient and the variance-reduced momentum, and $M_{t+1} \geq 0$ is a term that captures the variance reduction benefit from choosing $\gamma_{t+1}$ optimally (defined in Lemma C.2, Equation C.2). When $\gamma_{t+1} = 1$ (full STORM), $M_{t+1} = 0$ ; for any $\gamma_{t+1} \neq 1$ , $M_{t+1} > 0$ , increasing the bound.

What this means: The dominant term decays as $\mathcal{O}(T^{-2/3})$ , which after converting to gradient norm (via the parameter movement bound) yields $\mathcal{O}(T^{-1/3})$ convergence to a stationary point. This is faster than AdamW's $\mathcal{O}(T^{-1/4})$ rate under similar assumptions. The improvement arises because the variance-reduced momentum $m_t$ tracks $\nabla F(x_t)$ more closely than standard EMA momentum — the error $\|\nabla F(x_t) - m_t\|$ shrinks faster. The $M_{t+1}$ term in the bound quantifies the benefit of optimal $\gamma_t$ selection: a flexible schedule that adapts $\gamma_t$ to minimize the gradient estimator variance yields a strictly smaller upper bound than any fixed $\gamma_t$ , including $\gamma_t = 1$ . This provides theoretical justification for the scaling parameter's existence.

Lyapunov function construction: The proofs use a Lyapunov (potential) function that combines the objective value and the momentum estimation error:

$\Phi_t = \mathbb{E}\left[F(x_t) + \frac{\rho}{16L^2\eta_{t-1}}\|\nabla F(x_t) - m_t\|_2^2\right]$

The proof shows $\Phi_{t+1} - \Phi_t \leq -\frac{\eta_t}{\rho}\mathbb{E}\|\nabla F(x_t) - m_t\|^2 - \frac{3\rho}{8\eta_t}\mathbb{E}\|x_{t+1} - x_t\|^2 + \mathcal{O}(\eta_t^3) - \mathcal{O}(M_{t+1})$ , establishing that the potential decreases at each step provided the learning rate and momentum parameters are properly scheduled. The $\mathcal{O}(\eta_t^3)$ term comes from the variance of the gradient estimator and decays as training progresses. The telescoping sum over $T$ iterations yields the final rate.

Why time-varying $\beta_1$ and $\beta_2$ : The theoretical analysis requires $\beta_{1,t} \to 1$ and $\beta_{2,t} \to 1$ as $t \to \infty$ to achieve the $\mathcal{O}(T^{-1/3})$ rate, with specific schedules $\beta_{1,t} = 1 - c\eta_{t-1}^2$ and $\beta_{2,t} = 1 - \eta_{t-1}^6$ . This is a standard technique in convergence proofs for adaptive methods (following STORM's original analysis) — it ensures the momentum states average over increasingly longer histories, reducing bias. In practice, the paper uses constant $\beta_1 = 0.95$ , $\beta_2 = 0.99$ , which the theory acknowledges as a special case where the asymptotic rate may degrade but practical performance remains strong.

4. Key Insights and Innovations

Innovation 1: Variance Reduction as a Tunable Dial, Not a Binary Switch

The paper's most conceptually distinctive move is reconceptualizing variance reduction from a fixed algorithmic choice — you either use it or you don't — into a continuous, tunable mechanism controlled by a single scalar $\gamma_t \in [0, 1]$ . This is a fundamental reframing of how variance reduction relates to standard momentum, and it represents the intellectual key that unlocks practical deployment.

The pre-MARS assumption. Prior work combining variance reduction with adaptive methods effectively treated variance reduction as a binary state. SuperAdam (Huang et al., 2021) uses the full STORM estimator, implicitly setting what would be MARS's $\gamma_t$ to 1. VRAdam and AdaSPIDER similarly adopt the variance-reduced estimator at full strength. The implicit assumption was: if variance reduction is theoretically beneficial (it reduces gradient complexity from $\mathcal{O}(\varepsilon^{-4})$ to $\mathcal{O}(\varepsilon^{-3})$ in the nonconvex smooth setting), then applying it at maximum strength should yield the best empirical results. This assumption is natural — it mirrors how momentum strength is typically treated as a fixed $\beta_1$ rather than something that should interpolate between no momentum and full momentum.

The MARS reframing. The paper introduces $\gamma_t$ with an explicit control variate derivation (Section 3.1, "Why $\gamma_t$ improves convergence") showing that the variance-minimizing value of $\gamma_t$ is not generally 1; it depends on the correlation between the noise in the stochastic gradient and the noise in the correction term. The optimal $\gamma^*$ is derived as:

$\gamma^* = 1 - \frac{\mathbb{E}[UY] + \text{Var}(Y)}{\mathbb{E}[Y^2]}$

where $U$ captures the current momentum error and $Y$ captures the gradient difference term. When the momentum error is large (early in training, or when parameters are changing rapidly), the optimal $\gamma^*$ is small, meaning the correction should be applied weakly. When the optimizer is near convergence and parameter movement is small (so $\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t)$ genuinely approximates noise cancellation rather than signal subtraction), $\gamma^*$ approaches 1. The theoretical analysis in Theorem B.5 embeds this insight: the term $M_{t+1}$ in the convergence bound is zero at $\gamma_t = 1$ (full STORM) and positive otherwise, but the bound is stated as a minimum over $\gamma_t$ schedules, acknowledging that the path to optimality may involve varying $\gamma_t$ over time.

Why this is more than a hyperparameter. If $\gamma_t$ were merely a hyperparameter to tune, the contribution would be incremental — "we found that turning down variance reduction helps." But the paper frames it as a diagnostic concept: the reason prior variance reduction attempts failed is not that variance reduction doesn't work in deep learning, but that it was being applied at the wrong strength for the optimization phase. This transforms the narrative from "variance reduction is ineffective" (Defazio & Bottou, 2019) to "variance reduction must be calibrated to the current noise regime." The fact that the optimal constant $\gamma = 0.025$ in the GPT-2 experiments (Appendix E.5, Figure 14) — two orders of magnitude below the full STORM value of 1 — quantifies just how far off the binary assumption was.

Evidence. The sensitivity analysis in Appendix E.5 (Figure 14) tests $\gamma \in \{0.0001, 0.001, 0.01, 0.025, 0.05, 0.1, 0.2\}$ and finds a clear optimum at 0.025, with performance degrading on both sides. Values above 0.1 (still far below 1) already show visible degradation. This empirical finding retroactively explains why SuperAdam (implicitly $\gamma = 1$ ) failed to scale: it was operating in a regime where the correction term was so strong that it amplified noise rather than cancelling it. The tunable-dial framing makes MARS robust in a way that binary variance reduction cannot be — practitioners can dial $\gamma$ down to near-zero and recover standard AdamW, or dial it up if the training regime permits stronger correction.

Significance beyond performance. This conceptual move opens a new axis for optimizer design that was previously invisible. Rather than asking "should we use variance reduction?" the question becomes "how much variance reduction should we use, and how should it vary over training?" The paper only scratches the surface of this — it uses constant $\gamma_t$ in all experiments — but the theoretical framework and the control variate derivation point toward dynamic $\gamma_t$ schedules as a natural next step. This is a genuine reframing, not an incremental refinement, because it changes the design space from a binary choice to a continuous one, with the potential for significantly more sophisticated optimization strategies.

Innovation 2: Diagnosing and Fixing the First-Order / Second-Order Momentum Mismatch

The paper identifies a specific, previously unarticulated design flaw in prior variance-reduced adaptive methods: the semantic mismatch between the first-order and second-order momentum states. This is simultaneously a sharp diagnostic insight (explaining why prior attempts underperformed) and a constructive fix (redesign $v_t$ to track the variance-reduced gradient rather than the raw gradient).

The mismatch, stated precisely. In standard Adam, $m_t$ and $v_t$ are both derived from the same random variable — the stochastic gradient $\nabla f(x_t, \xi_t)$ . The first-order momentum $m_t = \beta_1 m_{t-1} + (1-\beta_1)\nabla f(x_t, \xi_t)$ tracks the mean gradient direction, and the second-order momentum $v_t = \beta_2 v_{t-1} + (1-\beta_2)[\nabla f(x_t, \xi_t)]^2$ tracks the per-coordinate variance. When dividing $m_t$ by $\sqrt{v_t}$ , the ratio approximates a signal-to-noise ratio: coordinates with high variance (noisy) get suppressed, coordinates with low variance (reliable signal) get amplified. This works because $m_t$ and $v_t$ describe the same underlying signal — the scale that $v_t$ measures is the scale of the update that $m_t$ proposes.

When you replace the gradient in $m_t$ with a variance-reduced estimator $\tilde{c}_t$ (as SuperAdam does), but keep $v_t$ tracking $[\nabla f(x_t, \xi_t)]^2$ (the raw squared gradient), this coherence breaks. The variance-reduced $m_t$ has lower variance than the raw gradient — that's the whole point. But $v_t$ estimates the scale of the unreduced gradient, which has higher variance than $\tilde{c}_t$ . The denominator $\sqrt{v_t}$ is therefore inflated relative to the numerator $m_t$ , causing the update to be systematically too conservative in precisely those coordinates where variance reduction was most effective. The more variance reduction helps (the noisier the gradient), the more the denominator overestimates the needed scaling, and the more the benefit is attenuated.

Why this was missed. Prior work treated the first-order momentum modification as the only change needed — add variance reduction to $m_t$ , and leave the rest of Adam unchanged. This is a natural impulse: Adam's $v_t$ is a well-engineered adaptive learning rate mechanism, and changing it risks destabilizing the optimization. But the paper argues that this conservatism is precisely what doomed those attempts. The second-order momentum is not an independent module — it is semantically coupled to the first-order momentum through their shared input. Changing one without the other creates a statistical inconsistency that degrades performance.

The fix. MARS redefines $v_t$ as the EMA of $\tilde{c}_t^2$ — the squared variance-reduced estimate — restoring the coherence between the two momentum states. This is a small change in the code (one line: track $\tilde{c}_t^2$ instead of $[\nabla f(x_t, \xi_t)]^2$ ) but a large change in the statistical interpretation of the optimizer. The ratio $m_t / \sqrt{v_t}$ now properly represents the signal-to-noise ratio of the variance-reduced signal, not a mismatched hybrid.

Evidence. The paper does not provide a clean ablation that isolates this fix from the other MARS innovations (scaling $\gamma_t$ , clipping placement). This is a limitation — we cannot directly attribute what fraction of MARS's gains come from the $v_t$ fix versus the $\gamma_t$ scaling versus the clipping design. However, the conceptual argument is strong: if $v_t$ were tracking the wrong quantity, the entire adaptive learning rate mechanism would be misaligned, systematically under-utilizing the variance reduction benefit. The fact that MARS significantly outperforms both AdamW (no variance reduction) and, by implication from the literature, SuperAdam (variance reduction with the mismatch) is consistent with this diagnosis.

Significance. This insight generalizes beyond the specific algorithms in this paper. Any future attempt to integrate variance reduction with adaptive preconditioning must ensure that the preconditioner is computed from the same statistical estimator as the update direction. This is not an Adam-specific issue — it applies to any optimizer that separately tracks first- and second-order statistics (or, in Shampoo's case, full-matrix statistics). The MARS-Shampoo variant implicitly respects this by applying SVD to the variance-reduced $m_t$ rather than to a raw gradient accumulator. This is a design principle that the paper establishes by example rather than by explicit statement, but it is arguably the most transferable insight for future optimizer development.

Innovation 3: Unifying Momentum Correction Across Optimizer Families via Lemma 3.4

The paper proves a simple algebraic lemma (Lemma 3.4) that reveals a hidden structural commonality across seemingly unrelated optimizers: Lion's two-state momentum and Muon's Nesterov-style momentum both contain implicit gradient correction terms that approximate variance reduction. This unification is elegant because it demonstrates that the optimizer design space is more constrained than it appears — independent search procedures (symbolic program search for Lion, manual design for Muon) converged on mechanisms that approximate the same mathematical operation.

The lemma. Lemma 3.4 states that any two-state momentum system of the form $m_t = b_1 u_t + b_2 g_t$ , $u_{t+1} = a_1 u_t + a_2 g_t$ is equivalent to the single momentum update:

$m_t = a_1 m_{t-1} + (b_1 a_2 - a_1 b_2 + b_2) g_t + (a_1 b_2 - b_1 a_2)(g_t - g_{t-1})$

The critical term is the last one: $(a_1 b_2 - b_1 a_2)(g_t - g_{t-1})$ , which is a gradient correction term — the difference between current and previous gradients multiplied by a scalar coefficient. When $a_1, b_1, a_2, b_2$ are set to specific values, this recovers the correction terms of various optimizers.

What it reveals about Lion. Substituting Lion's parameters ( $a_1 = \beta_2, a_2 = 1-\beta_2, b_1 = \beta_1, b_2 = 1-\beta_1$ ) into Lemma 3.4 yields:

$m_t = \beta_2 m_{t-1} + (1-\beta_2)\nabla f(x_t, \xi_t) + (\beta_2 - \beta_1)(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_{t-1}))$

The correction term $(\beta_2 - \beta_1)(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_{t-1}))$ has the same structure as STORM's variance reduction, but with a crucial difference: it uses the previous batch $\xi_{t-1}$ rather than the current batch $\xi_t$ . Lion was discovered through symbolic program search — an automated procedure that tested thousands of update rules and selected this one — and yet it independently converged on a structure that approximates variance reduction. The paper interprets this as evidence that variance reduction is a natural and beneficial operation that even naive search procedures will rediscover. MARS-Lion makes this connection explicit, replacing the approximate correction with the exact STORM correction (same batch $\xi_t$ ) and adding the $\gamma_t$ scaling parameter.

What it reveals about Muon. Applying the related Lemma D.1 to Muon's Nesterov-style updates ( $a_1 = \mu, a_2 = 1, b_1 = \mu, b_2 = 1$ ) yields:

$m_t = \mu m_{t-1} + \nabla f(x_t, \xi_t) + \mu(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_{t-1}))$

Muon's momentum is Nesterov acceleration with an additional gradient correction term — again using the previous batch. MARS-Shampoo with $\beta_1 = \mu$ and $\gamma_t = 1-\mu$ recovers a rescaled version with the same batch correction, strengthening the variance reduction property.

Why this matters — the unification argument. This lemma reframes three independently developed optimizers (Lion, Muon, Adan) as approximations to a single underlying algorithm: a momentum update with a gradient correction term that partially cancels noise. The differences between them reduce to (a) whether the correction uses the same batch $\xi_t$ (variance reduction) or the previous batch $\xi_{t-1}$ (approximate), (b) the coefficient on the correction term, and (c) the choice of preconditioner (sign, eigenspace, or diagonal). MARS unifies them by making the correction exact (same batch), adding a tunable scaling parameter $\gamma_t$ , and allowing the preconditioner to be freely chosen among AdamW, Lion, and Shampoo styles.

This unification is significant beyond the specific algorithms because it collapses the apparent diversity of optimizer designs into a small set of fundamental choices: (1) how to precondition (diagonal, sign, eigenspace, Kronecker), and (2) how much variance reduction to apply ( $\gamma_t$ ). The paper doesn't state this taxonomy explicitly, but it is implicit in the framework — and it makes the optimizer design space more legible and navigable for future work.

Limitation. The unification is retrospective, not predictive. The paper doesn't use Lemma 3.4 to propose a new optimizer or to predict which correction coefficient would be optimal — it only shows that existing optimizers can be expressed in this common form. The value is in explaining why these diverse algorithms perform similarly (they share a correction mechanism) and in suggesting that further improvements should focus on refining the correction (e.g., through optimal $\gamma_t$ scheduling) rather than searching for entirely new update rules.

Innovation 4: The Negative Result That Defines the Design Space

The paper contains an important negative result that, while not highlighted as an "innovation," significantly shapes the intellectual contribution: MARS's optimal $\gamma_t = 0.025$ is so close to zero that the algorithm is effectively AdamW with a very small variance reduction perturbation. This is not a failure — it is a diagnostic finding that redefines what "variance reduction works" means in the context of large-scale training.

What the field expected. If you surveyed optimization researchers before this paper and asked "if variance reduction can be made to work for LLM training, what strength of correction would you expect?", the implicit answer would likely be closer to $\gamma = 1$ (full STORM) than to $\gamma = 0.025$ . The entire edifice of variance reduction theory — from SVRG through SPIDER to STORM — is built around the idea that the full correction term provides the optimal variance reduction, leading to the improved $\mathcal{O}(\varepsilon^{-3})$ complexity. The SuperAdam design implicitly encodes this assumption. The expectation was that if the integration challenges could be solved (mismatch, clipping, etc.), the full variance reduction would shine.

What MARS found. The optimal constant $\gamma$ in the GPT-2 experiments is 0.025 — meaning the variance reduction correction is weighted at only 2.5% of its STORM magnitude. At $\gamma = 0.025$ , the gradient estimator is:

$c_t = \nabla f(x_t, \xi_t) + 0.025 \cdot \frac{\beta_1}{1-\beta_1}(\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t))$

with $\beta_1 = 0.95$ , so $\frac{\beta_1}{1-\beta_1} = 19$ , making the correction term $0.025 \times 19 \times (\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t)) \approx 0.475 \times (\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t))$ . The correction is still non-trivial in absolute terms (about half the raw gradient difference), but far below the STORM value of $19 \times$ the raw gradient difference. MARS is operating in a regime where variance reduction is a subtle adjustment rather than a dominant mechanism.

What this means. This finding reframes what "making variance reduction work" means in practice. It is not about deploying the full theoretical apparatus of STORM. It is about adding a small, controlled amount of noise cancellation to standard momentum — enough to measurably improve the gradient signal quality without destabilizing the adaptive preconditioning. This is a qualitatively different claim than "variance reduction works after all." It's more precise: weak variance reduction works; strong variance reduction still fails. The failure mode at higher $\gamma$ values (degradation visible at $\gamma = 0.05$ in Figure 14) suggests that even with the mismatch fixed, there are limits to how much correction the adaptive preconditioner can tolerate. The variance-reduced estimates deviate from the raw gradient distribution in ways that the preconditioner (which adapts slowly via EMA) cannot fully track, and at some threshold the benefit of noise cancellation is outweighed by the distribution shift.

Theoretical nuance. Theorem B.5 sheds light on this: the term $M_{t+1}$ in the convergence bound captures the benefit of choosing $\gamma_t$ optimally versus $\gamma_t = 1$ . When $\gamma_t$ deviates from the optimal value, $M_{t+1} > 0$ and the bound loosens. The theory does not claim $\gamma_t = 1$ is optimal — it claims a schedule can be chosen to minimize the bound, and the optimal constant value in practice turns out to be 0.025. This is consistent with the theoretical framework, but the magnitude of the gap (40× below 1) is surprising and suggests that the theoretical assumptions (smoothness, bounded variance) are not tight enough to predict the empirical optimum.

Significance as a negative result. This finding is intellectually important because it sets the boundary conditions for future work on variance reduction in deep learning. The design implication is: don't try to replace momentum with variance reduction; try to augment momentum with a small amount of variance reduction. This is a more modest but more practical goal than what the variance reduction literature has historically pursued, and it suggests that the theoretical gains of $\mathcal{O}(\varepsilon^{-3})$ complexity may be unattainable in practice even as the $\mathcal{O}(\varepsilon^{-4})$ to $\mathcal{O}(\varepsilon^{-3})$ improvement remains a valid theoretical insight. The paper's contribution is not "we achieved the full theoretical benefit of variance reduction" but rather "we found the pragmatic middle ground where variance reduction provides empirically measurable gains without destabilizing training." This is a fundamentally empirical, engineering contribution rather than a theoretical one, and the paper's honesty about the small $\gamma$ value is a strength — it tells the field where to look next (dynamic $\gamma_t$ scheduling, better preconditioner adaptation to distribution shift) rather than claiming a complete solution.

5. Experimental Analysis

Evaluation Methodology

Dataset. All main experiments use the OpenWebText corpus (Gokaslan et al., 2019), an open-source recreation of OpenAI's WebText dataset. The training set contains approximately 9 billion tokens, and the validation set contains approximately 4.4 million tokens, both preprocessed using the GPT-2 tokenizer. For supplementary experiments in Appendix E.3, the paper also uses FineWeb-Edu 100B, a ~100B-token subset of the filtered educational web pages from FineWeb-Edu (Lozhkov et al., 2024), with approximately 99.9B tokens for training and ~0.1B tokens held out for validation.
Base model(s). The paper experiments on the GPT-2 model series (Radford et al., 2019) at three scales: small (125M parameters, 12 layers, 12 heads, embedding dim 768), medium (355M parameters, 24 layers, 16 heads, embedding dim 1024), and large (770M parameters, 36 layers, 20 heads, embedding dim 1280). An additional XL variant (1.5B parameters, 48 layers, 25 heads, embedding dim 1600) is used only in the FineWeb-Edu experiments (Appendix E.3). The choice of GPT-2 is justified by the paper's goal: demonstrating that variance reduction can work on transformer-based language models at non-trivial scale, where prior variance-reduced adaptive methods had never been validated. These sizes span roughly an order of magnitude in parameters, allowing the paper to observe scaling trends. All experiments use the nanoGPT implementation (Karpathy, 2022) with biases disabled, GeLU activations, and Dropout set to 0.0 (matching the structural conditions the paper argues are necessary for variance reduction to apply — no batch normalization, no data augmentation, no dropout).
Metrics. The primary metric throughout is training loss and validation loss (cross-entropy), plotted against both training tokens processed and wall-clock time (seconds). Token-based comparison is the standard fairness metric in LLM training, as it accounts for differences in per-iteration cost (some optimizers do more computation per step). Wall-clock time comparisons (Figures 1c, 2c, 3c) provide a deployment-realistic view that accounts for optimizer overhead. For downstream evaluation, the paper reports 0-shot and 5-shot accuracy on standard benchmarks using lm-evaluation-harness (Gao et al., 2024), including ARC-Easy, ARC-Challenge, BoolQ, HellaSwag, OpenBookQA, PIQA, WinoGrande, MMLU, and SciQ, with average accuracy across all tasks reported as a summary metric. All validation loss numbers are reported without confidence intervals or error bars in the main text (Appendix plots similarly lack uncertainty quantification).
Baselines. The paper compares against five optimizers:
1. AdamW (Loshchilov & Hutter, 2019): the dominant optimizer in LLM training, with hyperparameter grid search over learning rates $\{1\times 10^{-4}, 1.5\times 10^{-4}, 3\times 10^{-4}, 6\times 10^{-4}, 1\times 10^{-3}, 1.5\times 10^{-3}, 3\times 10^{-3}, 6\times 10^{-3}\}$ and weight decay over $\{10^{-1}, 10^{-2}, 10^{-3}\}$ . The paper uses the "golden standard" learning rates from the literature (validated by Liu et al., 2023), with final chosen values of $6\times 10^{-4}$ (small), $3\times 10^{-4}$ (medium), and $2\times 10^{-4}$ (large), all with $\beta_1 = 0.9$ , $\beta_2 = 0.95$ .
2. Lion (Chen et al., 2023): the symbolically-discovered optimizer using sign-based updates. The paper does not detail its Lion hyperparameter search procedure in the main text, but the results appear in Figures 1–3 and Tables 1–6.
3. Muon (Jordan et al., 2024): an eigenspace-preconditioning optimizer using Newton-Schulz iteration to compute $U V^\top$ from the momentum matrix. Learning rates are $2\times 10^{-2}$ (small), $1\times 10^{-2}$ (medium), and $6.67\times 10^{-3}$ (large) with momentum parameter $\mu = 0.95$ (Table 12).
4. MARS-AdamW: the AdamW instantiation of the MARS framework, using $\beta_1 = 0.95$ , $\beta_2 = 0.99$ , $\gamma_t \equiv 0.025$ , and learning rates of $6\times 10^{-3}$ (small), $3\times 10^{-3}$ (medium), and $2\times 10^{-3}$ (large). This is the paper's primary proposed method.
5. MARS-Lion: the Lion instantiation of the MARS framework, using the same $\beta_1 = 0.95$ and $\gamma_t = 0.025$ as MARS-AdamW, with learning rates of $6\times 10^{-3}$ (small), $3\times 10^{-3}$ (medium), and $2\times 10^{-3}$ (large).
Notably absent as baselines: Sophia (Liu et al., 2023), SuperAdam (Huang et al., 2021), and any non-MARS variance-reduced optimizer. The paper explicitly positions against SuperAdam (Section 3, Remark 3.2) but never benchmarks it, meaning the claim that the $v_t$ redesign is critical is inferred from MARS's performance relative to AdamW and from the $\gamma_t$ ablation, not from direct comparison to a mismatched-vt baseline. Sophia's absence is noted but not discussed — the paper cites Sophia's clipping-by-value design as a point of contrast (Section 3.1) but doesn't explain why it's excluded from experiments.
Generation budget / compute accounting. There is no "generation budget" in this paper — it studies training optimization, not test-time inference compute. The relevant compute metric for fair comparison is training tokens processed, which is standard in the LLM training literature. The paper reports all loss curves against this x-axis (Figures 1a-b, 2a-b, 3a-b). For wall-clock time comparisons, total seconds on specific hardware are reported: small models on 16 NVIDIA A100 GPUs, medium on 32 NVIDIA A100 GPUs, and large on 32 NVIDIA H100 GPUs. The paper states that MARS-AdamW and MARS-Lion have "slightly higher per-iteration cost compared to AdamW but [are] much faster than Muon" (Section 4.2). This cost difference comes from MARS-approx needing to store and access the previous step's gradient $\nabla f(x_{t-1}, \xi_{t-1})$ rather than recomputing it — a minor memory and data-movement overhead, not an additional forward-backward pass. The per-iteration FLOPs are not quantified precisely; the paper's wall-clock time curves (Figures 1c, 2c, 3c) implicitly account for this overhead.
Cross-validation / statistical protocol. There is no cross-validation or statistical significance testing reported. All experiments are single runs with a fixed base seed of 5000 (Table 11). The paper does not report error bars, confidence intervals, or standard deviations on any result. The training and validation curves in the main figures are single-run traces (training loss curves are EMA-smoothed for visual clarity, but this is a plotting convention, not a statistical procedure). The $\gamma_t$ sensitivity analysis (Appendix E.5, Figure 14) and the learning rate/batch size ablations (Appendices E.6, E.7) are similarly single runs. The learning rate grid search for AdamW and MARS is the closest the paper comes to a tuning protocol, but it reports only the best-found configuration, not the distribution of outcomes across seeds or hyperparameters. This makes the results potentially fragile — with training runs at this scale taking tens of thousands of GPU-hours, single-seed reporting is a genuine limitation, though one shared with much of the LLM training literature.

Main Quantitative Results

Training and Validation Loss: MARS Consistently Outperforms AdamW, Lion, and Muon

The paper's central result is that MARS-AdamW and MARS-Lion achieve lower training and validation losses than all baselines across all three GPT-2 model scales, with the performance gap widening as model size increases. The headline numbers for GPT-2 large (770M) are:

"AdamW requires 50 billion tokens to reach a validation loss of 2.58, whereas MARS only requires 28 billion tokens, and it achieves a final validation loss of 2.51."

This represents a 1.8× token efficiency improvement (28B vs. 50B tokens to reach the same loss) and a 0.07 absolute improvement in final validation loss (2.51 vs. 2.58) after 50B tokens of training. The final validation losses on GPT-2 large (Table 13, inferred from Figure 1b and Section 4.2 text) are:

Optimizer	Final Validation Loss (GPT-2 large, 50B tokens)
MARS-AdamW	2.511
MARS-Lion	2.534
AdamW	2.568
Lion	2.565
Muon	2.606

Figure 1 (GPT-2 large) shows the training and validation loss curves:

Figure 1a (Training Loss): All MARS variants lie visibly below all baselines throughout training. At 50B tokens, MARS-AdamW achieves approximately 2.40 training loss versus AdamW at approximately 2.44 — a gap of ~0.04. MARS-Lion tracks slightly above MARS-AdamW but below Lion and Muon throughout. AdamW and Lion are nearly overlapping; Muon trails both. The advantage is visible from early training (by ~5B tokens, the gap is already established) and is sustained rather than being a late-training effect.
Figure 1b (Validation Loss): MARS-AdamW reaches a validation loss of 2.51 at 50B tokens, while AdamW reaches 2.58 — a gap of 0.07. MARS-AdamW crosses AdamW's final validation loss of 2.58 at approximately 28B tokens, which is where the headline "28 billion tokens vs. 50 billion tokens" comes from. MARS-Lion reaches approximately 2.53–2.54, outperforming all non-MARS baselines. The curves show no sign of convergence — MARS continues to improve relative to AdamW at the end of training.
Figure 1c (Wall-clock time): Despite slightly higher per-iteration cost, MARS-AdamW and MARS-Lion achieve lower validation loss at equivalent wall-clock time. At 40,000 seconds (~11 hours on 32×H100), MARS-AdamW reaches approximately 2.55 while AdamW is at approximately 2.65. Muon is substantially slower per step (due to Newton-Schulz iteration) and reaches higher loss at all time points. The wall-clock advantage narrows slightly compared to the token advantage (the per-iteration overhead partially offsets the efficiency gain), but remains substantial.

Figures 2 and 3 (Appendix, GPT-2 small 125M and medium 355M) show the same qualitative pattern:

GPT-2 medium (Figure 3): The gap between MARS-AdamW and AdamW at 50B tokens is approximately 0.04–0.05 in validation loss (roughly 2.68 vs. 2.72–2.73). MARS-Lion performs similarly to MARS-AdamW on medium. The relative improvement appears smaller in absolute terms than on large, suggesting the benefit of variance reduction scales with model size.
GPT-2 small (Figure 2): The gap narrows further — MARS-AdamW achieves approximately 2.94 validation loss versus AdamW at approximately 2.96 at 50B tokens, a gap of ~0.02. MARS-Lion and AdamW are nearly overlapping on small. This scaling pattern (gap widens from small → medium → large) is consistent with the intuition that larger models have noisier gradients (due to higher-dimensional parameter spaces and more complex loss landscapes), making variance reduction more beneficial.

Training loss validation: The paper does not report perplexity (exponentiated cross-entropy loss), which would be more interpretable. Converting: validation loss 2.511 on GPT-2 large corresponds to perplexity ~12.3; validation loss 2.568 corresponds to perplexity ~13.0. The difference of 0.057 in loss translates to approximately 0.7 perplexity points — noticeable but not transformative.

Downstream Task Performance: MARS Improves Few-Shot Accuracy

The paper evaluates the pretrained checkpoints on downstream benchmarks to verify that the validation loss improvement translates to useful capability gains. Table 1 (5-shot, GPT-2 large) and Tables 2–6 (0-shot and 5-shot for small/medium/large) report accuracy across 9 benchmarks.

Key result (Table 1, GPT-2 large, 5-shot): MARS-AdamW achieves the highest average accuracy at 50.15%, compared to AdamW at 48.64%, Lion at 48.61%, and Muon at 47.31%. MARS-Lion reaches 49.80%. The improvement is not uniform across tasks:

"On the downstream task Hellaswag, MARS improved accuracy to 44.64%, outperforming AdamW's 41.70% after training on 50 billion tokens."

This is a +2.94 percentage point improvement on HellaSwag — the largest single-task gain. Other notable improvements from MARS-AdamW over AdamW (5-shot, Table 1): BoolQ +6.45 (62.78 vs. 56.33), OBQA +2.20 (31.60 vs. 29.40), SciQ +1.60 (84.50 vs. 82.90). MARS-AdamW slightly underperforms AdamW on ARC-C (26.28 vs. 28.67, a loss of -2.39) and is roughly tied on WinoGrande (52.49 vs. 52.01).

Scaling pattern: The advantage on downstream tasks grows with model size. On GPT-2 small (Table 2, 0-shot), MARS-AdamW averages 42.94% vs. AdamW's 42.63% — a gain of only 0.31 points, within plausible noise. On GPT-2 medium (Table 3, 0-shot), MARS-AdamW averages 45.41% vs. 44.33% — a gain of 1.08 points. On GPT-2 large (Table 4, 0-shot), MARS-AdamW averages 48.27% vs. 46.78% — a gain of 1.49 points. While these differences are not tested for statistical significance, the monotonic increase with scale is consistent with the loss curves.

Limitation of downstream results: The paper evaluates checkpoints at the end of training (50B tokens), not at matched validation loss. Since MARS reaches lower validation loss than AdamW, the downstream comparison confounds two effects: (1) MARS may produce better representations at the same loss, or (2) MARS simply reached a lower loss because it converged faster. The paper provides no downstream evaluation at matched loss (e.g., MARS at 28B tokens vs. AdamW at 50B tokens, where they have equal validation loss, or MARS and AdamW both at 2.60 validation loss). This makes it impossible to determine whether MARS's downstream gains come from better optimization (lower loss → better task performance, which any optimizer achieving lower loss would show) or from qualitatively better representations (same loss but better downstream performance). The Consistent improvement on HellaSwag (+2.94) suggests some task-specific benefit, but the lack of loss-controlled comparison is a significant gap.

MARS vs. MARS-approx: Exact Variance Reduction Provides Marginal Gains

The paper compares the exact MARS formulation (using $\xi_t$ for both gradient evaluations in the correction term, requiring either gradient storage and data replay or a second forward-backward pass) against MARS-approx (using $\xi_{t-1}$ for the previous gradient, which comes for free from the previous training step). Appendix E.2, Figures 4 and 5 show the comparison on GPT-2 small and medium.

Finding: The exact version consistently but marginally outperforms the approximate version. On GPT-2 small (Figure 4, left; Figure 5, left), MARS-exact training loss curves are slightly lower throughout training — the gap is on the order of 0.01 in training loss and perhaps 0.01–0.02 in validation loss at 50B tokens. On GPT-2 medium (Figure 4, right; Figure 5, right), the pattern is similar but the gap is even smaller. The paper characterizes this as:

"Models trained with MARS exhibit consistently better performance than those trained with MARS-approx. This suggests that: (a) The exact version, which employs the variance reduction formulation, is more fundamental than the approximate version. (b) The approximate version serves as a practical alternative in scenarios where computational efficiency is a priority, as it incurs only minimal performance loss."

Practical implication: Since the exact version requires either doubling the forward-backward cost per iteration (recomputing $\nabla f(x_{t-1}, \xi_t)$ ) or storing and replaying the previous batch, and the approximate version achieves nearly the same performance with zero additional cost, MARS-approx is the recommended practical variant. All main experiments in Section 4 use MARS-approx. This is a practically important finding: it means the computational overhead of MARS over AdamW is essentially just the cost of storing the previous gradient (a minor memory increase, likely a few percent) plus the arithmetic to compute $c_t$ (a few vector additions and scalar multiplications), with no additional forward-backward passes needed.

Theoretical caveat: The paper notes that "MARS and MARS-approx differ in their updates and may theoretically exhibit distinct convergence guarantees." The exact version's noise cancellation relies on $\mathbb{E}[\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_t)] = \nabla F(x_t) - \nabla F(x_{t-1})$ (the noise from $\xi_t$ cancels in expectation). MARS-approx's correction term involves $\nabla f(x_t, \xi_t) - \nabla f(x_{t-1}, \xi_{t-1})$ , where the noise terms are independent — the correction does not achieve the same variance reduction property. The fact that this matters so little empirically suggests that the benefit of MARS is not primarily from the STORM-style noise cancellation per se, but from the broader design changes (redefined $v_t$ , clipping placement, $\gamma_t$ scaling) that would apply even with approximate correction. This is consistent with the small optimal $\gamma_t = 0.025$ — at such weak correction strength, the exact vs. approximate distinction matters less.

Ablation Studies and Robustness Checks

Sensitivity to $\gamma_t$ (Appendix E.5, Figure 14): The paper sweeps $\gamma \in \{0.0001, 0.001, 0.01, 0.025, 0.05, 0.1, 0.2\}$ on GPT-2 small with MARS-AdamW-approx. The optimal constant $\gamma$ is 0.025, with visible degradation at both lower and higher values. At $\gamma = 0.0001$ (nearly pure AdamW), validation loss is roughly 0.01–0.02 higher than at $\gamma = 0.025$ — a small but consistent gap, suggesting that even the minimal variance reduction provides measurable benefit. At $\gamma = 0.05$ , performance visibly degrades; at $\gamma = 0.2$ , the degradation is substantial (validation loss roughly 0.03–0.04 higher than optimum). The paper does not test $\gamma = 1$ (full STORM), but the trend strongly suggests it would perform worse than all tested values. This is the paper's most important ablation: it establishes that (a) variance reduction helps but only at very weak strength, and (b) the optimal $\gamma$ is an order of magnitude closer to 0 (no variance reduction) than to 1 (full STORM).
Constant learning rate (Appendix E.6.1, Figures 15–16): To isolate the effect of MARS from the cosine learning rate scheduler, the paper runs GPT-2 small, medium, and large with constant learning rates, comparing AdamW and MARS-AdamW-approx at two different LR values each. Across all scales and both LR values, MARS maintains lower training and validation losses throughout training. The gap is similar in magnitude to the cosine schedule experiments. This rules out the hypothesis that MARS's gains come from interacting favorably with the LR schedule (e.g., MARS benefiting more from the high-LR early phase or the decay phase). The constant-LR results also suggest MARS has good "continuous training" properties — the loss curves continue improving without divergence even without LR decay, consistent with the WSD scheduler results below.
WSD scheduler (Appendix E.6.2, Figures 17–18): The paper tests the Warmup-Stable-Decay (WSD) scheduler (Hu et al., 2024) on GPT-2 small and medium, comparing MARS-AdamW-approx against AdamW at multiple maximum learning rates. WSD consists of linear warmup, a long constant-LR phase, and a final short decay phase. MARS consistently outperforms AdamW during the stable (constant LR) phase, and the gap persists or widens during the decay phase. On GPT-2 medium, MARS with a single LR outperforms AdamW at both tested LRs. The paper interprets this as evidence that "MARS has a better potential for continuous training and exhibits an explicit edge over baseline algorithm." However, no comparison is made against AdamW + WSD at matched total tokens beyond what is shown; the test is whether MARS works with WSD, not whether WSD + MARS beats cosine + MARS.
Batch size sensitivity (Appendix E.7, Figure 19): The paper tests batch sizes of 240, 480, and 960 on GPT-2 small with MARS-AdamW-approx and AdamW. MARS-AdamW outperforms AdamW at all three batch sizes, and the performance gap widens at smaller batch sizes. At batch size 240, the validation loss gap at 80,000 steps is approximately 0.05–0.06; at batch size 480, roughly 0.02–0.03; at batch size 960, the curves nearly overlap. This is consistent with the theoretical motivation: smaller batch sizes produce noisier gradient estimates, making variance reduction more beneficial. The result also suggests a practical guidance: if computational constraints force the use of small batch sizes (e.g., limited GPU memory per device), MARS becomes more valuable relative to AdamW than at larger batch sizes.
FineWeb-Edu 100B dataset (Appendix E.3, Figures 6–7, Tables 7–8): To test robustness to dataset quality, the paper trains GPT-2 small (125M) and XL (1.5B) on FineWeb-Edu 100B, a higher-quality filtered educational dataset. MARS-AdamW-approx outperforms both AdamW and Muon on both scales. On GPT-2 XL at 50B tokens, MARS-AdamW achieves validation loss of approximately 2.44 versus AdamW at approximately 2.48 (Figure 7, right) — a gap of ~0.04. Downstream evaluation (Tables 7–8) shows MARS-AdamW achieves the highest average accuracy on both small and XL models, though the margins are modest (e.g., GPT-2 XL 0-shot average: MARS-AdamW 56.41% vs. AdamW 55.43%, a gain of ~1 point). Notably, the paper includes the OpenAI released GPT-2 checkpoints ("OpenAI-Comm.") in the comparison, and all optimizers trained on FineWeb-Edu substantially outperform them, confirming dataset quality as a dominant factor. MARS's gains are additive to dataset quality improvements.
Computer vision experiments (Appendix E.4, Figures 8–13, Table 9): MARS is also tested on CIFAR-10 and CIFAR-100 using ResNet-18, with all three instantiations (MARS-AdamW, MARS-Lion, MARS-Shampoo) compared against their non-variance-reduced counterparts (AdamW, Lion, Shampoo) plus Muon. Key findings: (1) All MARS variants achieve lower test loss and higher test accuracy than their base counterparts, but the gains primarily manifest after the learning rate decay at epoch 100 — before decay, the curves are largely overlapping. (2) The exact version of MARS marginally outperforms MARS-approx across all three families (e.g., MARS-AdamW exact reaches 95.26% on CIFAR-10 vs. MARS-AdamW-approx at 95.29% — actually slightly worse in this case; on CIFAR-100, exact reaches 77.38% vs. approx 76.97%). (3) The best test accuracy on CIFAR-100 is MARS-AdamW exact at 77.38% vs. Muon at 74.64% and AdamW at 73.70% — a substantial 3.68% improvement over AdamW. On CIFAR-10, the gaps are smaller (MARS-AdamW-approx at 95.29% vs. AdamW at 94.81%). (4) MARS-Shampoo performs worse than MARS-AdamW on CIFAR-100 (75.83% exact) but better than baseline Shampoo (74.27%). This suggests the framework's benefits generalize beyond language modeling to vision tasks, but the gain is concentrated in the later stages of training (post-LR-decay), hinting that variance reduction helps fine-grained convergence rather than initial exploration.
MARS-Shampoo (no GPT-2 experiments): The paper introduces MARS-Shampoo (Algorithm 4) but does not evaluate it on any GPT-2 experiment. It is tested only on CIFAR-10/100 (Appendix E.4, Figures 8–13, Table 9). This is a significant gap: the paper's central claim is about training large models, and Shampoo is a second-order method designed for exactly this setting. Without GPT-2 experiments for MARS-Shampoo, the claim that the MARS framework "accommodates all existing full matrix or diagonal Hessian approximations" is only partially validated — it works for diagonal (AdamW, Lion) in language modeling and for full-matrix (Shampoo) in vision, but not for full-matrix in language modeling. The computational cost of SVD on large transformer weight matrices (GPT-2 large's largest layer would require SVD of roughly 1280×1280 matrices for attention projections) is likely the limiting factor, but the paper does not discuss this or propose a scalable approximation (beyond mentioning that sketching, Newton iteration, or Newton-Schulz iteration could be used).
Lion and Muon learning rate sweep (incomplete reporting): While the paper details AdamW and MARS hyperparameter tuning, it provides minimal information about the Lion and Muon baselines. For Muon, the paper states $\mu = 0.95$ and learning rates in Table 12, but does not mention whether a grid search was performed or whether these are taken from the Muon paper's recommendations. For Lion, neither hyperparameters nor search procedure are reported in the main text or Table 12. This asymmetry in tuning effort could bias the comparison — if MARS received more careful tuning than Lion and Muon, the performance gaps may be overstated. The paper acknowledges this implicitly by stating that Lion and Muon were run with recommended settings, but recommended settings may not be optimal for the specific GPT-2 / OpenWebText / training-horizon configuration.

Critical Assessment

The experimental section provides credible evidence that MARS-AdamW and MARS-Lion outperform AdamW, Lion, and Muon on GPT-2 models trained on OpenWebText, with the caveat that all comparisons are single-seed, single-configuration runs without statistical testing. The central empirical claims from the paper's abstract and introduction require careful scrutiny against what the experiments actually demonstrate.

Claim: "AdamW requires 50 billion tokens to reach a validation loss of 2.58, whereas MARS only requires 28 billion tokens"

What the experiments show: This claim is directly supported by Figure 1b. MARS-AdamW reaches validation loss 2.58 at approximately 28B tokens on the x-axis. AdamW reaches the same loss at approximately 50B tokens. The ratio is indeed about 1.8×. This is a straightforward reading of the learning curve.

What the experiments do NOT show: The claim applies specifically to GPT-2 large (770M) on OpenWebText with the specific hyperparameters reported in Table 12, using MARS-approx (not exact MARS) with $\gamma_t = 0.025$ , $\beta_1 = 0.95$ , $\beta_2 = 0.99$ , LR = $2\times 10^{-3}$ (cosine), and weight decay searched over $\{10^{-1}, 10^{-2}, 10^{-3}\}$ . The claim's generalizability to other model architectures, datasets, or hyperparameter configurations is not tested. Specifically:

The 1.8× factor may be hyperparameter-sensitive. If AdamW's learning rate ( $2\times 10^{-4}$ for large) were retuned specifically for this training budget or dataset, the gap might narrow. The paper did a grid search for AdamW but reports that search was based on "golden standard learning rates in literature" and validated by Liu et al. (2023) — it's unclear whether the search was exhaustive or simply confirmed that known values work well.
The 1.8× factor applies to a specific loss threshold (2.58). At lower thresholds (e.g., 2.70), the efficiency ratio may be different. At higher thresholds (e.g., MARS's final loss of 2.51), AdamW never reaches the threshold in the 50B-token training run, so no comparison is possible. The claim could be stated more precisely as "MARS reaches AdamW's 50B-token loss 1.8× faster for this specific loss value."
The single-seed nature of the experiment means we cannot assess whether the 28B crossing point is a reliable estimate. With only one run each, the crossing point could shift by several billion tokens under different random seeds. The seed is fixed (5000), which controls for one source of variation but doesn't provide a variance estimate.

Claim: "MARS achieves a final validation loss of 2.51 [vs. AdamW's 2.58]"

What the experiments show: Figure 1b shows MARS-AdamW at approximately 2.51 and AdamW at approximately 2.58 at the 50B-token mark. The gap of 0.07 in validation loss is clearly visible and sustained over the final ~10B tokens of training (the curves are not converging). This supports the claim.

What the experiments do NOT show: The claim leaves open whether the gap would continue to widen with more training, or whether AdamW would eventually catch up. The learning rate scheduler (cosine decay to a minimum of $1\times 10^{-5}$ for large models) means both optimizers are taking very small steps near the end of training — the fact that MARS continues to improve relative to AdamW under these conditions is somewhat surprising and suggests the variance reduction provides benefit even in the low-LR regime. However, with the cosine schedule decaying to near-zero, neither optimizer can make substantial further progress past 50B tokens without a schedule change, so the "final" loss is contingent on the chosen training horizon.

Claim: "MARS consistently outperforms AdamW by a large margin"

What the experiments show: On GPT-2 large, the margin is 0.07 in validation loss and 1.8× in token efficiency — a clear and practically meaningful gap. On GPT-2 medium, the gap narrows to approximately 0.04–0.05 in validation loss. On GPT-2 small, the gap is only approximately 0.02. Whether these qualify as "large" is subjective, but they are consistent in direction across all scales, all tested configurations (constant LR, WSD, different batch sizes, different dataset), and both MARS variants (AdamW and Lion). The consistency of the sign of the effect across all experiments is the strongest evidence.

What the experiments do NOT show: The claim of "large margin" is not tested against the null hypothesis that the gain comes from the specific hyperparameter changes rather than from variance reduction per se. MARS-AdamW differs from AdamW in at least five ways: (1) variance-reduced gradient estimator, (2) $\beta_1$ changed from 0.9 to 0.95, (3) $\beta_2$ changed from 0.95 to 0.99, (4) learning rate increased by 10× ( $2\times 10^{-3}$ vs. $2\times 10^{-4}$ for large), and (5) gradient clipping applied to the intermediate $c_t$ rather than the final update. The paper provides no ablation that isolates the variance reduction effect from these confounded changes. It is possible that AdamW with $\beta_1 = 0.95$ , $\beta_2 = 0.99$ , and the 10× higher learning rate would perform substantially better than the baseline AdamW configuration, even without the variance reduction term. This would not invalidate MARS — the point is that the MARS framework enables these hyperparameter choices, but it would change the interpretation from "variance reduction causes the improvement" to "variance reduction enables more aggressive hyperparameters which cause the improvement." The $\gamma_t$ sensitivity analysis (showing $\gamma = 0$ performs worse than $\gamma = 0.025$ ) partially addresses this — at $\gamma = 0$ , MARS reduces to a clipped, differently-configured AdamW, and it performs worse than with $\gamma = 0.025$ , so the correction term is doing something. But the ablation doesn't compare MARS $\gamma = 0$ against a properly-tuned AdamW baseline with matched $\beta_1, \beta_2$ , and LR, so the confound remains.

Specific weaknesses in the experimental design:

No SuperAdam baseline. The paper explicitly identifies SuperAdam's design flaws (first/second-order momentum mismatch, no $\gamma_t$ scaling, no clipping on $c_t$ , no bias correction/weight decay) and argues that MARS fixes them. But without benchmarking SuperAdam, there is no direct evidence that these specific fixes matter. If SuperAdam with appropriate tuning performed similarly to MARS, the paper's diagnostic narrative would be undermined. The exclusion may be practical (SuperAdam was designed for small-scale vision tasks and may not have a readily available GPT-2 implementation) but it leaves a clear gap.
No Sophia baseline. Sophia (Liu et al., 2023) is perhaps the most prominent recent optimizer claiming to outperform AdamW on language model training, using diagonal Hessian estimates and clipping. The paper cites Sophia's clipping design as a point of contrast but never compares against it. If Sophia's gains over AdamW are similar in magnitude to MARS's gains, the unique value of variance reduction would be less clear.
No validation of the "mismatch" hypothesis in isolation. The paper could have tested an intermediate optimizer: AdamW with variance-reduced $m_t$ but raw-gradient $v_t$ (the SuperAdam configuration, essentially MARS with the $v_t$ fix disabled). Comparing this against full MARS would isolate the contribution of the $v_t$ redesign. Without this ablation, the paper's claim that "the second-order momentum must match the first-order momentum" is supported by intuition and by the overall MARS performance, but not by controlled experiment.
Single dataset. The main experiments use only OpenWebText (9B tokens), a relatively small and potentially idiosyncratic dataset. The FineWeb-Edu 100B experiments partially address this by showing benefits persist on a different, higher-quality dataset, but FineWeb-Edu is also a filtered web text corpus — both may share structural properties (e.g., token distribution, sequence length distribution) that affect optimizer behavior. Testing on a radically different domain (code, scientific papers, multilingual text) would strengthen the generalizability claim.
Single model architecture family. All experiments use GPT-2 (decoder-only transformer). The CIFAR experiments use ResNet-18, which is a different architecture (CNN vs. transformer) and task (classification vs. language modeling), but MARS's claimed contribution is specifically about "large models" — the kind trained with AdamW at scale. No experiments cover encoder-only architectures (BERT), encoder-decoder architectures (T5), or state-space models (Mamba). Each of these might interact differently with variance reduction due to different gradient structures.
No comparison at matched computational budget. The wall-clock time curves (Figures 1c, 2c, 3c) partially address this, but they compare single training runs with fixed step counts, not runs where the budget is reallocated (e.g., run AdamW for more steps to match MARS's total FLOPs, or run MARS for fewer steps and compare loss). The token-efficiency metric (1.8×) is clear, but a FLOPs-matched comparison would account for MARS's per-step overhead more precisely. MARS-approx's overhead is described as "slight" but not quantified.
No experiments on models larger than GPT-2 large (770M). The paper's motivation section discusses Llama 3, PaLM, and DeepSeek-V2 — models orders of magnitude larger than 770M parameters — and implies that optimizer improvements at GPT-2 scale will transfer. But the literature on optimizer scaling is mixed: Liu et al. (2023) found Sophia's advantages narrowed at larger scales, and Kaddour et al. (2024) found most optimizer innovations fail to beat AdamW at scale. The 770M ceiling leaves open the question of whether MARS's 1.8× token efficiency holds at 7B, 70B, or 175B parameters. The scaling trend from small → medium → large is encouraging (gap widens with scale) but extrapolation from three points is unreliable.
Constant $\gamma_t = 0.025$ in all experiments. The theoretical analysis (Appendix B.2, Lemma C.2) suggests that the optimal $\gamma_t$ may depend on the current variance structure, which changes during training (higher variance early, lower variance late). The paper uses constant $\gamma_t$ , missing the opportunity to test dynamic schedules (e.g., start high and decay, or start low and increase). The sensitivity analysis (Figure 14) uses constant values, not dynamic schedules. This is an acknowledged limitation (Section 5 mentions future work on dynamic $\gamma_t$ ), but it means the reported results are a lower bound on what MARS could achieve with an optimal schedule — and conversely, the optimal constant value of 0.025 is the best among constant schedules, not necessarily the optimal overall.
No statistical significance testing. All experiments are single runs. Given the cost of GPT-2 large training (~32 H100 GPUs for the duration shown), running multiple seeds would be expensive, but even 2-3 seeds on the small or medium model would provide error bars and allow informal assessment of whether the observed gaps exceed run-to-run variation. The paper's claim that MARS "consistently outperforms" is based on direction consistency (MARS < AdamW in all comparisons), not statistical testing. The downstream evaluation tables show mixed results — MARS-AdamW wins on some tasks and loses on others (e.g., 5-shot ARC-C: AdamW 28.67 vs. MARS-AdamW 26.28, a loss of 2.39 points) — and without error bars, it's unclear whether these task-level variations are noise or signal.
The Shampoo and Muon baselines are underexplored. MARS-Shampoo is introduced but never tested on GPT-2, presumably due to computational cost (SVD on weight matrices at scale). The Muon baseline, which the paper positions as a comparator (Muon uses Newton-Schulz iteration for eigenspace preconditioning), is tested only with specific hyperparameters and may not be optimally tuned. Muon's learning rates ( $2\times 10^{-2}$ for small, etc.) are an order of magnitude higher than MARS's — this is appropriate for Muon's design but means the comparison involves very different optimization dynamics, making it hard to attribute differences specifically to variance reduction versus preconditioner choice.

Positive aspects of the experimental design that deserve recognition:

Multi-scale evaluation. The small/medium/large scaling across an order of magnitude in parameters is more thorough than many optimizer papers, which often test only one model size. The widening gap with scale is genuinely informative and increases confidence that the effect is not an artifact of a particular model size.
Multiple learning rate schedules and datasets. The constant LR (Figures 15-16), WSD (Figures 17-18), and FineWeb-Edu (Figures 6-7) experiments test MARS under distinct conditions and all show consistent improvement over AdamW. This reduces the concern that MARS's gains are specific to the cosine schedule or the OpenWebText dataset.
Wall-clock time reporting. Many optimizer papers report only per-iteration or per-token metrics, ignoring the computational cost of the optimizer itself (which is substantial for methods like Shampoo or K-FAC). The paper's wall-clock time curves (Figures 1c, 2c, 3c) show that MARS's per-step overhead does not negate its token-efficiency gains — it still wins on wall-clock time.
$\gamma_t$ sensitivity analysis. Despite the limitations noted above, the sweep across $\gamma_t$ values (Figure 14) is a genuinely informative ablation — it shows a clear optimum, quantifies the degradation from suboptimal $\gamma$ , and demonstrates that the benefit of variance reduction is real ( $\gamma = 0.025$ outperforms $\gamma = 0.0001$ , which is nearly pure AdamW with matched $\beta_1, \beta_2$ , LR).
Approximate vs. exact comparison. The paper honestly reports that the exact STORM correction provides only marginal gains over the approximation, and recommends MARS-approx for practical use. This is pragmatically useful and demonstrates that the authors are not overclaiming the importance of the exact variance reduction formulation.

Summary assessment: The experiments provide strong evidence that MARS-AdamW and MARS-Lion, as configured in this paper, outperform the specified baselines on GPT-2 models trained on OpenWebText. The 1.8× token efficiency gain on GPT-2 large is the paper's most compelling result. However, the experimental design has several confounds (simultaneous changes to $\beta_1$ , $\beta_2$ , learning rate, and gradient estimator), missing baselines (SuperAdam, Sophia), and limited scale/diversity (single architecture family, maximum 770M parameters, single-seed runs) that narrow the scope of the supported claims. The paper demonstrates that a particular combination of design choices — weak variance reduction ( $\gamma = 0.025$ ), redefined second-order momentum, and clipping on the intermediate gradient estimator — can improve GPT-2 training efficiency, but it does not isolate which of these choices are necessary or sufficient, nor does it establish that the improvement mechanism is specifically variance reduction rather than a beneficial interaction between the new hyperparameter regime and the transformer loss landscape.

6. Limitations and Trade-offs

The Optimal γt Is Found via Per-Model Tuning, Not a Transferable Principle

The assumption or constraint. The paper introduces γt as the central mechanism for controlling variance reduction strength, and the sensitivity analysis in Appendix E.5 (Figure 14) demonstrates that the optimal constant γ depends sensitively on the training configuration— testing γ ∈ {0.0001, 0.001, 0.01, 0.025, 0.05, 0.1, 0.2} on GPT-2 small finds a clear optimum at 0.025, with visible degradation on both sides. The paper uses this single value (γ = 0.025) across all model sizes (small, medium, large), all MARS instantations (AdamW, Lion), and both datasets (OpenWebText, FineWeb-Edu). There is no systematic study of whether the optimal γ transfers across model scale, architecture family, dataset, or training horizon — the sweeping is done only on GPT-2 small.

The consequence. A practitioner wanting to deploy MARS on a different model (e.g., Llama, BERT, a vision transformer) or at a different scale cannot rely on γ = 0.025 being optimal — they would need to perform their own sweep, which is expensive (each point in the sweep requires a full or partial training run). The sensitivity is sharp enough that a suboptimal choice visibly degrades performance (Figure 14: γ = 0.05 already shows loss increase of ~0.01–0.02 on GPT-2 small; γ = 0.2 shows ~0.03–0.04 degradation). At large model scales where a single training run costs tens of thousands of GPU-hours, the cost of tuning γt could exceed the efficiency gains from using MARS, undermining the practical value proposition. Furthermore, the theoretical analysis (Appendix B.2) indicates that the optimal γt should be time-varying — depending on the current gradient variance, momentum error, and parameter movement — yet all experiments use a constant value, suggesting the reported results are a lower bound that could be improved with dynamic scheduling, but providing no guidance on how to design such schedules without expensive per-problem tuning.

What evidence exists in the paper. The sensitivity analysis in Appendix E.5 (Figure 14) is the only direct evidence, performed on a single configuration (GPT-2 small, MARS-AdamW-approx, OpenWebText, cosine schedule). The paper does not report γ-sweeps for GPT-2 medium or large, nor for MARS-Lion, nor for the FineWeb-Edu dataset. The constant γ = 0.025 is carried through all other experiments without justification beyond the GPT-2 small sweep. The paper does not test dynamic γt schedules at all, despite the theoretical motivation in Section 3.1 and Lemma C.2.

Mitigation status. The paper does not address this limitation directly. It treats γ = 0.025 as a fixed hyperparameter (listed in Table 12 alongside learning rates and β values) rather than as a quantity requiring per-configuration tuning. The theoretical analysis (Theorem B.5, Lemma C.2) provides an expression for the optimal γt+1 (Equation D.14 in Appendix D.5), but this involves quantities (expectations of inner products between gradient differences and momentum errors) that are not computable in practice without additional estimation. The paper flags dynamic γt schedules as future work only implicitly through the theoretical framework — there is no explicit discussion of the tuning burden or transferability problem. A practitioner reading the paper would reasonably assume γ = 0.025 is a safe default, but the sensitivity results suggest it may not be.

All GPT-2 Experiments Are Single-Seed Runs with No Uncertainty Quantification

The assumption or constraint. The paper runs exactly one training run per (optimizer, model size, configuration) combination, with a fixed base seed of 5000 (Table 11). No error bars, confidence intervals, standard deviations, or reproducibility metrics are reported for any training or validation loss curve, any downstream evaluation, or any ablation study. The downstream evaluation tables (Tables 1–8) report accuracies to two decimal places but are based on single pretrained checkpoints, single evaluation runs — there is no assessment of whether the differences between optimizers exceed run-to-run variation from different random seeds.

The consequence. Without uncertainty quantification, the paper's quantitative claims are point estimates whose reliability is unknown. The headline claim — "MARS achieves final validation loss of 2.51 versus AdamW's 2.58" — could be sensitive to the particular random seed. A gap of 0.07 in validation loss, while visually clear in Figure 1b, may not be statistically significant if the standard deviation across seeds is, say, 0.02–0.03. This is not a hypothetical concern: LLM training is known to exhibit non-trivial seed sensitivity, particularly for downstream task evaluations where the paper's Table 1 shows task-level variation (MARS-AdamW beats AdamW on HellaSwag by +2.94 points but loses on ARC-C by -2.39 points). Without knowing the run-to-run variance, a practitioner cannot assess whether the HellaSwag improvement is a reliable effect of MARS or a chance fluctuation that would reverse under a different seed. Similarly, the downstream evaluation is done at the end of training (50B tokens), not at matched validation loss — a single data point that compounds the uncertainty of both the training dynamics and the evaluation sampling.

What evidence exists in the paper. Zero. The paper provides no uncertainty quantification anywhere. The seed is reported (5000) but only one value is used. The GPT-2 small results (Figure 2) show MARS-AdamW and AdamW curves that are separated by roughly 0.02 in final validation loss — a gap small enough that a second seed could plausibly reverse the ordering or produce overlapping curves. The downstream evaluation tables (Tables 1–8) show mixed results across tasks with no error estimates, making it impossible to distinguish signal from noise for individual task comparisons.

Mitigation status. The paper does not acknowledge this as a limitation. It presents all results as deterministic curves and precise numerical comparisons without caveats about statistical reliability. The consistent direction of the effect across model sizes and configurations (MARS outperforms AdamW in all 20+ comparisons) provides some informal reassurance — a systematic advantage across independent experiments is unlikely under a null hypothesis of no effect — but this is a qualitative argument, not a statistical one. The paper could have run 2–3 seeds on GPT-2 small (which trains on 16 A100 GPUs and is the cheapest configuration) to provide error bars and demonstrate that the observed gaps exceed seed-level variation, but it did not.

The 1.8× Token Efficiency Claim Conflates Variance Reduction with Simultaneous Hyperparameter Changes

The assumption or constraint. MARS-AdamW differs from the AdamW baseline in at least four hyperparameters beyond the variance reduction mechanism: (1) β1 is changed from 0.9 (AdamW) to 0.95 (MARS-AdamW); (2) β2 is changed from 0.95 (AdamW) to 0.99 (MARS-AdamW); (3) the learning rate is increased by a factor of 10 for GPT-2 large ( $2 × 10^{-3}$ vs. $2 × 10^{-4}$ ) and by similar factors for small and medium (Table 12); (4) gradient clipping is applied to the intermediate estimator ct rather than to the final update (or not at all, depending on the AdamW baseline's clipping configuration). The paper provides no ablation that isolates the effect of the variance reduction term (γ > 0) from the effects of these hyperparameter changes — the comparison is always MARS-AdamW with its full configuration versus AdamW with the standard AdamW configuration.

The consequence. The paper's central causal claim — "variance reduction improves training efficiency" — is confounded with "we found a better hyperparameter configuration." It is possible that AdamW with β1 = 0.95, β2 = 0.99, and the 10× higher learning rate (with appropriate clipping) would substantially outperform the baseline AdamW configuration, even without the variance reduction term. If this were the case, the MARS framework's contribution would shift from "variance reduction is the mechanism" to "the MARS formulation enables (or was discovered alongside) a higher-learning-rate, higher-momentum training regime." The paper's γ sensitivity analysis (Figure 14) partially addresses this: at γ = 0.0001 — which is essentially MARS with the variance reduction term effectively disabled but with all the other hyperparameter changes (β1 = 0.95, β2 = 0.99, higher LR, clipping on ct) — performance is worse than at γ = 0.025, suggesting the correction term itself contributes. However, the γ = 0.0001 configuration is not compared against a properly retuned AdamW baseline with matched β1, β2, and LR, so it remains unclear how much of the gain comes from the MARS hyperparameter regime versus the variance reduction specifically.

What evidence exists in the paper. The γ-sweep (Appendix E.5, Figure 14) is the closest the paper comes to an isolation experiment. At γ = 0.0001, the gradient estimator ct is approximately the raw gradient ∇f(xt, ξt) plus a negligible correction (0.0001 × 19 × gradient_difference ≈ 0.002 × gradient_difference), so the optimizer is effectively AdamW with MARS's hyperparameters (β1 = 0.95, β2 = 0.99, higher LR, clipping on ct). This configuration achieves validation loss roughly 0.01–0.02 higher than optimal (γ = 0.025), meaning the variance reduction term adds approximately 0.01–0.02 improvement over the already-improved hyperparameter regime. The baseline AdamW configuration (β1 = 0.9, β2 = 0.95, standard LR) is not shown on the same plot, so we cannot decompose the total MARS-vs-AdamW gap into "hyperparameter improvement" versus "variance reduction improvement."

Mitigation status. The paper does not acknowledge this confound explicitly. It presents MARS as a single integrated framework and compares it against standard AdamW, implicitly treating the hyperparameter changes as part of the MARS design rather than as separable variables. This is defensible to a degree — the argument is precisely that variance reduction enables these more aggressive hyperparameters (higher LR, higher β1) because the gradient estimates are less noisy — but the paper does not make this enabling argument clearly or provide evidence for it (e.g., showing that AdamW with β1 = 0.95, β2 = 0.99, and 10× LR diverges or underperforms without variance reduction). The constant learning rate experiments (Appendix E.6.1, Figures 15–16) test MARS at two LRs and AdamW at two LRs, but the LR values differ between the two optimizers, so the comparison still confounds optimizer choice with learning rate choice.

The Optimal γ = 0.025 Is So Small That MARS Is Effectively AdamW with a Minimal Perturbation — Challenging the "Variance Reduction" Narrative

The constraint. The paper's empirical finding that the optimal constant γ is 0.025 (Appendix E.5, Figure 14) means the variance reduction correction term is weighted at only 2.5% of its STORM magnitude. With β1 = 0.95, the scaling factor on the gradient difference is γ × β1/(1-β1) = 0.025 × 19 = 0.475, so the correction term magnitude is roughly half the raw gradient difference — substantially smaller than the STORM correction (which has coefficient 19). Performance degrades at γ = 0.05 (already visible) and severely at γ = 0.2, and the paper never tests γ = 1 (full STORM). The framework that MARS proposes — variance reduction as a tunable dial — is supported, but the dial's effective range is squeezed into a narrow band near zero.

The consequence. This finding fundamentally recasts what "variance reduction works" means. The paper's narrative — that variance reduction can be made to work for large model training where prior attempts failed — is technically correct, but the magnitude of the effect is far smaller than the variance reduction literature (and the paper's own theoretical analysis) would predict. A practitioner hoping for a breakthrough in training efficiency would find that MARS provides a ~0.07 validation loss improvement on GPT-2 large (1.8× token efficiency) — real and valuable, but not transformative in the way that the optimizer scaling literature sometimes promises. More importantly, the finding suggests that variance reduction at meaningful strength (γ approaching 1) still fails in deep learning, even with the paper's architectural fixes. The mismatch fix, clipping placement, and scaling parameter make variance reduction possible at very weak strength, but they do not solve the underlying limitation that strong variance reduction destabilizes adaptive preconditioning. This implies a fundamental tension between variance reduction and adaptive learning rates that the paper mitigates but does not resolve.

What evidence exists in the paper. The γ-sweep in Figure 14 is definitive: γ = 0.025 is optimal; γ = 0.1 is already worse; the trend lines clearly point downward as γ increases. The paper never tests γ > 0.2, but extrapolation suggests that γ = 1 (SuperAdam regime, full STORM) would substantially underperform even γ = 0.0001 (no variance reduction). The paper's own theoretical analysis (Theorem B.5) shows that the convergence bound has a term -M_{t+1} that is maximized (most negative, improving the bound) at the optimal γt+1, and that M_{t+1} = 0 at γt+1 = 1. The theory does not claim γ = 1 is optimal, but the empirical finding that the optimum is at γ = 0.025 — meaning the optimal γ is two orders of magnitude closer to 0 than to 1 — is not predicted or explained by the theory. The paper is honest about the γ = 0.025 finding (it reports it clearly in Figure 14 and discusses it), but does not grapple with the implication that "variance reduction" in the classical sense is not what is helping.

Mitigation status. The paper does not explicitly frame the small optimal γ as a limitation. It presents the γ-sweep as evidence that the scaling parameter is useful (which it is) without fully acknowledging that the optimal value being 0.025 substantially weakens the claim that variance reduction per se is the mechanism. The paper could have discussed dynamic γt schedules as a way to use stronger variance reduction at certain training phases (e.g., late in training when parameter movement is small, so the correction term is a better noise estimator), but it does not. The constant γ = 0.025 result is treated as a hyperparameter finding rather than as a diagnostic about the limits of integrating variance reduction with adaptive methods.

No Experiments on Models Larger Than GPT-2 Large (770M), Despite the Paper's Framing Around "Large Models"

The assumption or constraint. The paper's title, abstract, and introduction frame the contribution as addressing training efficiency for "large models" — citing Llama 3 (Dubey et al., 2024), PaLM (Chowdhery et al., 2023), DeepSeek-V2 (Liu et al., 2024), and models trained with AdamW at scale. The motivation section (Section 1) argues that recent studies (Kaddour et al., 2024; Zhao et al., 2024) show most optimizers fail to outperform AdamW in LLM pretraining, and positions MARS as the exception. However, all GPT-2 language modeling experiments are conducted at scales of 125M, 355M, and 770M parameters — an order of magnitude smaller than even modest "large" models by contemporary standards (Llama 3 8B, let alone 70B or 405B). The maximum training compute is 32 H100 GPUs — a non-trivial but modest budget compared to the thousands of GPUs used in production LLM training. The FineWeb-Edu experiments go up to GPT-2 XL (1.5B), but these are supplementary (Appendix E.3) and still far below the scale the paper's framing implies.

The consequence. The paper's most policy-relevant claim — that organizations should consider MARS as an alternative to AdamW for large-scale training — rests on an extrapolation from three data points (125M, 355M, 770M) with no validation at 7B, 13B, 70B, or beyond. The observed scaling trend (the gap between MARS and AdamW widens from small to medium to large) is encouraging but based on only three points — extrapolating to 100× larger models is unreliable. The literature contains cautionary examples: Sophia (Liu et al., 2023) showed promising gains at GPT-2 scale that narrowed at larger scales in Kaddour et al. (2024). Lion (Chen et al., 2023) was discovered via symbolic search on smaller models and performs on par with AdamW at large scale (per the cited studies). It is entirely possible that MARS's advantage stems from improved optimization in a regime where gradient noise is moderate, and that at truly large scales (where massive batch sizes or model parallelism change the noise structure), the benefit of weak variance reduction diminishes or vanishes. Conversely, it is also possible that variance reduction becomes more beneficial at larger scales — the paper argues this is plausible because larger models have higher-dimensional, noisier gradient spaces — but the argument is speculative without evidence.

What evidence exists in the paper. The scaling trend from Figures 1–3: on GPT-2 small, the validation loss gap between MARS-AdamW and AdamW at 50B tokens is approximately 0.02; on medium, approximately 0.04–0.05; on large, approximately 0.07. The trend is monotonic and the absolute gap approximately triples from small to large. The FineWeb-Edu GPT-2 XL experiments (Appendix E.3, Figures 6–7) show a gap of approximately 0.04 in validation loss at 50B tokens — slightly smaller than the GPT-2 large gap on OpenWebText but still substantial. This provides four data points (125M, 355M, 770M, 1.5B) all showing MARS outperforming AdamW, with the gap generally increasing with scale over the first three points and remaining significant at the fourth. This is the strongest evidence in the paper for scalability, but it remains within the same model family (GPT-2) and training paradigm.

Mitigation status. The paper does not explicitly acknowledge the scale limitation or discuss the risks of extrapolation. The abstract and introduction freely use language about "large models" and "LLM pretraining" without qualifying that the experiments stop at 770M parameters (1.5B in supplementary material). The conclusion (Section 5) states that MARS "contribut[es] to the advancement of optimizers in large model training" without noting that this claim is based on models three orders of magnitude smaller than current frontier models. The paper does call for future work including "training large-scale models such as LLMs" — but this is framed as validation of an already-established method rather than as a missing piece of evidence that would be needed to support the paper's claims. The gap between the framing and the experimental scale is one of the most significant limitations for a practitioner evaluating whether to invest in implementing and tuning MARS for a production-scale training run.

The MARS-Shampoo Instantiation Is Never Tested on Language Models, and No Second-Order MARS Variant Is Validated at Scale

The assumption or constraint. The paper presents MARS as a "unified framework for preconditioned variance reduction" that "accommodates all existing full matrix or diagonal Hessian approximations" (Section 1, contributions). It instantiates the framework with three preconditioners: AdamW (diagonal), Lion (sign-based), and Shampoo (full-matrix via SVD). The Shampoo instantiation (Algorithm 4) is described in Section 3.2.3 and connected theoretically to Muon via Lemma D.1, showing that Muon's Nesterov-style momentum is a rescaled, approximate version of MARS-Shampoo's variance-reduced momentum. However, MARS-Shampoo is evaluated only on CIFAR-10 and CIFAR-100 with ResNet-18 (Appendix E.4, Figures 8–13, Table 9). It is never tested on any GPT-2 model or any language modeling task. The paper provides no discussion of why, though the computational cost of SVD on transformer weight matrices at scale is the likely reason.

The consequence. The "unified framework" claim is only partially validated. The diagonal and sign-based preconditioners (MARS-AdamW, MARS-Lion) are validated on GPT-2 at scale; the full-matrix preconditioner (MARS-Shampoo) is validated only on small-scale vision tasks. A practitioner interested in second-order methods for large model training — the setting where Shampoo and its variants (SOAP, Muon) are most relevant — learns nothing from this paper about whether variance reduction helps in that regime. Furthermore, the theoretical connection to Muon (Equations 3.24–3.25) demonstrates that Muon's momentum is algebraically related to MARS-Shampoo's variance-reduced momentum, but this analysis is not accompanied by an empirical comparison between MARS-Shampoo and Muon on language modeling. The Muon baseline in the GPT-2 experiments uses Muon's standard formulation (Nesterov-style momentum), not a MARS-augmented version, so the paper cannot claim that adding variance reduction to eigenspace preconditioning improves language model training — it can only claim that MARS with diagonal/sign preconditioners works.

What evidence exists in the paper. The CIFAR experiments (Appendix E.4, Table 9) show that MARS-Shampoo exact achieves 75.83% test accuracy on CIFAR-100, compared to 74.27% for baseline Shampoo — a +1.56 percentage point improvement. MARS-AdamW exact achieves 77.38%, outperforming MARS-Shampoo. On CIFAR-10, the gaps are smaller. This suggests that variance reduction helps Shampoo on vision tasks, but does not address whether the benefit transfers to language modeling or whether the SVD overhead is justified by the improvement. The GPT-2 experiments use Muon (which approximates Shampoo-style eigenspace preconditioning via Newton-Schulz iteration for efficiency) rather than Shampoo directly, and Muon underperforms both MARS-AdamW and MARS-Lion on GPT-2 large (validation loss 2.606 vs. 2.511 and 2.534) — but this comparison is between MARS's diagonal preconditioning and Muon's non-variance-reduced eigenspace preconditioning, not a test of MARS-Shampoo.

Mitigation status. The paper acknowledges this gap implicitly by including MARS-Shampoo as an algorithmic contribution but relegating it to vision experiments only. Section 3.2.3 mentions that "our algorithm design accommodates any of these SVD solvers to best fit specific computational needs" (referring to SVD, sketching, Newton iteration, Newton-Schulz iteration), suggesting that efficient implementations exist, but none are tested on language models. The paper does not state that MARS-Shampoo GPT-2 experiments were attempted and failed, or were too computationally expensive to run, or were deprioritized — it simply omits them. The conclusion (Section 5) states that "we have developed three optimization algorithms based on the ideas of AdamW, Lion, and Shampoo" and that "MARS consistently outperforms baseline algorithms" — but the "consistently" claim for Shampoo is based only on CIFAR, while the "baseline algorithms" for GPT-2 do not include Shampoo. This is a significant gap for a paper whose central thesis is that variance reduction is a modular, orthogonal improvement compatible with any preconditioner.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper shifts the narrative around variance reduction in deep learning from "it doesn't work" to "it works, but only at surprisingly weak strength and only when architectural details are done right." It is best understood as a diagnostic reframing rather than a paradigm shift — the contribution is not that variance reduction is fundamentally more powerful than previously believed, but that it has been applied incorrectly, and that correcting a specific set of design flaws (the first/second-order momentum mismatch, the absence of a scaling parameter, the placement of gradient clipping) unlocks consistent, measurable improvements. The magnitude of the practical gain is real but modest: 1.8× token efficiency on GPT-2 large, ~0.07 validation loss improvement, ~1.5 percentage point average downstream accuracy gain. These are meaningful for production training cost reduction, but they do not represent a breakthrough that renders existing optimizers obsolete.

The key landscape change is a new axis for optimizer design. Prior work on improving AdamW fell into roughly two categories: better preconditioning (diagonal Hessian estimates in Sophia, full-matrix approximations in Shampoo and Muon, sign-based updates in Lion) and better momentum (Nesterov acceleration in Adan, symbolic discovery in Lion). The paper demonstrates that weak variance reduction is a third, orthogonal axis that can be combined with any preconditioner choice. This is conceptually significant because it organizes the optimizer design space more cleanly: choose a preconditioner (diagonal, sign, eigenspace, Kronecker), choose a momentum scheme (standard EMA, Nesterov, variance-reduced), and tune the strength. The paper's Lemma 3.4, which shows that Lion and Muon's momentum updates are algebraically equivalent to approximate variance reduction, collapses what appeared to be diverse algorithmic discoveries into a single underlying mechanism with different parameterizations. This unification makes the design space more legible and suggests that further innovation should focus on refining the correction mechanism (optimal γt scheduling, adaptive correction strength) rather than searching for entirely new update rules via symbolic search or manual intuition.

The paper resolves a specific contradiction in the optimization literature without fully resolving the broader tension. The contradiction: Defazio & Bottou (2019) argued that variance reduction is ineffective in deep learning due to data augmentation, batch normalization, and dropout disrupting the finite-sum structure. The paper correctly observes that these factors are absent in modern LLM training, reopening the door for variance reduction. The empirical results validate that variance reduction can help — at γt = 0.025. But the deeper tension — that the theoretical benefit of variance reduction (O(ε⁻³) gradient complexity vs. O(ε⁻⁴) for SGD) should translate to dramatically faster convergence in practice — remains unresolved. The fact that the optimal γt is 0.025, two orders of magnitude below the full STORM value, suggests that the theoretical framework overestimates how much variance reduction the adaptive preconditioner can tolerate. The paper's convergence analysis (Theorem B.5) provides an upper bound that improves with optimal γt, but the bound is loose — it doesn't predict that the optimal constant γt would be 0.025, nor does it explain why larger values degrade performance. This leaves a theoretical puzzle: why does strong variance reduction fail even when the preconditioner is properly matched to the variance-reduced gradient? The paper provides the empirical finding but not the explanatory mechanism.

Research directions that become more attractive: (1) Dynamic γt scheduling — the paper uses constant γt = 0.025, but the theoretical framework (Lemma C.2) strongly suggests that optimal γt should vary over training, potentially starting low (when parameter movement is large and the correction term is a poor noise estimator) and increasing as training converges (when the correction becomes more accurate). (2) Verifier/optimizer co-design — the paper's finding that γ = 0.025 is optimal implies that variance reduction interacts with the preconditioner in ways that may be specific to the optimizer architecture; designing preconditioners that are robust to stronger variance reduction could unlock the theoretical O(ε⁻³) benefits. (3) Principled γt transfer — understanding whether the optimal γt depends on model scale, architecture, dataset, or batch size in predictable ways, and whether it can be estimated cheaply without per-configuration sweeps.

Research directions that become less attractive: (1) Naïve integration of full-strength variance reduction (γ = 1) with adaptive methods — the paper's γ-sweep (Figure 14) strongly suggests this will underperform even the non-variance-reduced baseline. (2) Symbolic search for optimizer rules — the paper's unification (Lemmas 3.4, D.1) shows that Lion and Muon's momentum, discovered by very different methods, reduce to approximate variance reduction with different coefficients. This suggests that the design space is more constrained than symbolic search assumes, and that future search efforts should parameterize the correction strength and preconditioner choice rather than searching over arbitrary expression trees. (3) Optimizer comparisons that don't control for hyperparameter regime — the paper's confounded comparison (simultaneously changing β₁, β₂, LR, and gradient estimator) highlights that much of the apparent optimizer advantage in the literature may come from hyperparameter choices rather than algorithmic mechanisms. This raises the bar for future optimizer papers: isolating the mechanism requires controlled ablation, not just comparison against off-the-shelf AdamW.

Follow-Up Research This Work Enables

Dynamic γt scheduling via online variance estimation. The paper uses constant γt = 0.025 in all experiments, but the theoretical analysis (Lemma C.2, Equation D.14 in Appendix D.5) provides an expression for the variance-minimizing γt+1 that depends on the correlation between the gradient difference ∇f(xt, ξt) − ∇f(xt−1, ξt) and the current momentum error mt − ∇F(xt). A natural follow-up would estimate these quantities online (e.g., by maintaining a running estimate of the gradient variance and the momentum error, or by using small-batch statistics) to adapt γt dynamically without per-problem tuning. The concrete experiment: compare a dynamically scheduled γt against the constant γ = 0.025 baseline on GPT-2 medium and large, measuring whether dynamic scheduling can (a) match the constant γ performance without tuning, and (b) potentially exceed it by using stronger variance reduction late in training when parameter movement is small. The theoretical prediction is that γt should start near 0 (parameters move rapidly early, making the correction term a poor noise estimator) and increase toward 1 as the optimizer converges (when ∇f(xt, ξt) ≈ ∇f(xt−1, ξt) and the correction genuinely cancels noise). A negative result — dynamic scheduling underperforms constant γ — would suggest that the theoretical variance decomposition fails to capture practical noise structures, pointing toward better noise models.

Isolating the contribution of the second-order momentum redesign. The paper argues that SuperAdam's critical flaw is defining vt on raw squared gradients while mt uses variance-reduced gradients, creating a statistical mismatch. However, this claim is never tested directly — MARS is never compared against a "MARS-with-mismatched-vt" baseline. A clean ablation would run three configurations on GPT-2 medium: (a) standard AdamW, (b) MARS with variance-reduced mt but raw-gradient vt (the SuperAdam configuration, essentially MARS with the vt fix disabled), and (c) full MARS with matched mt and vt. All three would use the same β₁ = 0.95, β₂ = 0.99, γ = 0.025, and learning rate to isolate the vt redesign. The hypothesis is that (b) underperforms (c) because the denominator √vt overestimates the needed scaling, causing overly conservative updates. If (b) performs similarly to (c), the vt redesign is not the critical mechanism, and the paper's diagnostic narrative is undermined. If (b) performs worse than (a) — meaning variance reduction without the vt fix is harmful — that would provide strong evidence for the paper's central design principle and would be a practically important warning for anyone attempting to retrofit variance reduction into existing optimizers.

Scaling MARS to 7B+ parameter language models. The most urgent follow-up is a replication at a scale where the paper's framing ("large models," citing Llama 3, PaLM, DeepSeek-V2) is directly relevant. Train a Llama-2-style 7B model on a 100B+ token corpus (C4, FineWeb-Edu, or SlimPajama) with MARS-AdamW versus AdamW, using the same hyperparameter transfer approach the paper uses (scale the learning rate, keep γ = 0.025, β₁ = 0.95, β₂ = 0.99). Measure validation loss, downstream benchmark performance, and the token efficiency ratio at matched loss. The key question: does the 1.8× token efficiency observed at 770M parameters hold, grow, or shrink at 7B? Kaddour et al. (2024) and Zhao et al. (2024) found that many optimizer innovations that beat AdamW at GPT-2 scale fail to do so at larger scales. If MARS maintains gains at 7B, it becomes a credible production alternative. If the gains shrink, the paper's contribution narrows to a specific scale regime. This experiment is expensive (~hundreds of GPU-hours on 64+ GPUs) but is the minimal validation needed to support the paper's claims about "large model training." Additionally, testing MARS-Shampoo (using an efficient SVD approximation like Newton-Schulz iteration, as Muon does) at this scale would close the paper's most conspicuous experimental gap, validating whether variance reduction helps full-matrix preconditioning in language modeling.

Understanding why strong variance reduction (γ → 1) fails. The paper's γ-sweep (Figure 14) shows optimal performance at γ = 0.025 with monotonic degradation as γ increases, and the paper never tests γ > 0.2. A diagnostic follow-up would run GPT-2 small with γ from 0 to 1 in finer increments and analyze why larger γ values degrade — is it gradient estimate variance (the correction term adding noise rather than cancelling it), preconditioner mismatch (the variance-reduced distribution drifts too far from what √vt can track), or instability in the momentum EMA (large correction terms causing mt to oscillate)? This could be diagnosed by tracking gradient variance, momentum norm, and parameter update magnitude across γ values. If the degradation is primarily from gradient variance (the correction term amplifies noise when parameter movement is large), then dynamic scheduling that uses low γ early and higher γ late should help. If the degradation is from preconditioner mismatch (√vt can't adapt fast enough to the variance-reduced distribution), then increasing β₂ or using a different preconditioner adaptation rate might help. If the degradation is from momentum instability, then stronger clipping or lower β₁ might help. This is the kind of systematic failure analysis that the paper's findings motivate but don't conduct, and it's essential for turning the empirical observation (γ = 0.025 is best) into a design principle (here's how to choose γ as a function of training conditions).

Testing MARS on non-transformer architectures and non-language domains. The paper evaluates MARS on GPT-2 (decoder-only transformer) and ResNet-18 (CNN for image classification), but an entire class of architectures — encoder-only transformers (BERT, RoBERTa), encoder-decoder transformers (T5, BART), state-space models (Mamba, S4), and mixture-of-experts models — is unexplored. Each has different gradient structures: BERT's bidirectional attention may produce different noise patterns than GPT's causal attention; Mamba's recurrent state-space parameterization has fundamentally different gradient dynamics. Training BERT-base on a masked language modeling task (e.g., C4 or Wikipedia) with MARS-AdamW versus AdamW would test whether variance reduction benefits transfer to encoder-only architectures, where the training objective (masked token prediction) differs structurally from autoregressive language modeling. Similarly, training a vision transformer (ViT) on ImageNet would test whether the benefits persist when attention mechanisms are applied to continuous rather than discrete inputs. A negative result in any of these domains wouldn't invalidate MARS but would define the boundaries of its applicability — which is as valuable as positive results for practitioners deciding whether to adopt it.

Cheap γt estimation via gradient statistics. The paper's γt sensitivity analysis (Figure 14) requires running multiple full training runs to find the optimum — prohibitive at scale. A practical follow-up would develop a method to estimate the optimal γt from a short training run (e.g., 1,000 steps) by measuring gradient variance, momentum error, and the correlation between the correction term and the true gradient direction. The theoretical expression for optimal γt+1 in Lemma C.2 (Equation D.14) involves expectations that could be approximated from batch statistics during a warmup phase, potentially yielding a γt estimate without a full sweep. The concrete experiment: on GPT-2 medium, run 1,000 steps with γ = 0 (or a small exploratory value), compute the online estimates of the quantities in Equation D.14, predict an optimal γ, then run full training with that γ and compare against the oracle optimum from a full sweep. If the online estimate reliably predicts the optimum within ±0.01, it would make MARS practically deployable without expensive per-configuration tuning. If it fails, it would suggest that the control variate model underlying Equation D.14 doesn't capture the relevant noise structure in LLM training — a finding that would motivate richer noise models for optimizer analysis.

Practical Applications and Downstream Use Cases

Cost reduction for mid-scale LLM pretraining (100M–1B parameters). Organizations training models in the GPT-2 small to large range — for domain-specific language models, academic research, or on-device deployment — can directly adopt MARS-AdamW with the hyperparameters reported in Table 12 (γ = 0.025, β₁ = 0.95, β₂ = 0.99, learning rate 10× the AdamW baseline, cosine schedule). The 1.8× token efficiency on GPT-2 large translates to a ~44% reduction in training tokens to reach the same validation loss. For a 770M-parameter model trained on ~50B tokens, this saves approximately 22B tokens of training compute — at the paper's configuration (32 H100 GPUs for ~100,000 seconds), roughly 40,000 GPU-seconds or ~11 GPU-hours on H100 equivalents. For organizations training multiple such models (hyperparameter sweeps, ablation studies, model variants), the aggregate savings are substantial. The wall-clock time curves (Figure 1c) show that MARS's slight per-step overhead does not negate the token-efficiency gain, so the savings are realized in practice. The MARS-approx variant is recommended (Appendix E.2, Figures 4-5) since it achieves nearly identical performance to exact MARS without the computational cost of recomputing ∇f(xt−1, ξt).

Improving training efficiency at small batch sizes. The batch size sensitivity experiment (Appendix E.7, Figure 19) shows that MARS's advantage over AdamW widens at smaller batch sizes — the validation loss gap is approximately 0.05–0.06 at batch size 240 versus ~0.01–0.02 at batch size 960 on GPT-2 small. This makes MARS particularly valuable in settings where batch size is constrained by GPU memory: single-GPU fine-tuning, on-device training, or distributed training with limited communication bandwidth. In these regimes, the gradient estimates are inherently noisier, and even weak variance reduction (γ = 0.025) provides disproportionate benefit. A practitioner fine-tuning a 125M–355M parameter model on a single A100 or H100 GPU, where memory constraints force batch sizes of 8–32, can expect MARS to recover some of the efficiency that would otherwise be lost to gradient noise — potentially making single-GPU training viable for model scales that would otherwise require multi-GPU setups to achieve reasonable batch sizes.

Continuous training and self-improvement loops. The paper's experiments with constant learning rates (Appendix E.6.1, Figures 15–16) and the WSD scheduler (Appendix E.6.2, Figures 17–18) demonstrate that MARS outperforms AdamW under training regimes where the learning rate does not decay to near-zero — specifically, the WSD scheduler maintains a constant learning rate for most of training, and MARS shows a persistent advantage throughout this phase. This matters for "continuous training" scenarios where a model is periodically updated on new data without restarting from scratch — a common pattern in production systems serving constantly-evolving domains. The improved stability under constant learning rates suggests MARS may reduce the need for careful learning rate schedule tuning in these settings, simplifying the operational overhead of maintaining and updating deployed models. For self-improvement loops (where a model generates training data, is fine-tuned on it, and repeats), the improved gradient signal quality from variance reduction could make each iteration more sample-efficient — though this is speculative without direct experiments.

When a strong verifier/reward model exists for online optimization. While the paper focuses on supervised pretraining, MARS's variance-reduced gradient estimator could be applied to reinforcement learning from human feedback (RLHF) or other online optimization settings where the "gradient" comes from a learned reward model. The variance reduction correction ∇f(xt, ξt) − ∇f(xt−1, ξt) would cancel noise from the reward model's stochasticity (different prompts or response pairs at adjacent parameter states) while preserving the signal about how the policy is changing. Given the high variance typical of RLHF training (due to small batch sizes, diverse prompts, and noisy human preference labels), MARS's finding that variance reduction helps most at small batch sizes is directly relevant. The practical application: plug MARS into the PPO optimization step of a standard RLHF pipeline, using γ = 0.025 (or a small dynamically-tuned value) and observing whether policy improvement is faster or more stable. This is speculative — the paper provides no RLHF evidence — but the structural similarity between stochastic gradient noise in pretraining and reward gradient noise in RLHF makes this a low-risk, high-potential extension.