---
title: MARS-AdamW Optimizer for Scalable LLM Training
url: https://www.emergentmind.com/topics/mars-adamw
type: topic
---

# MARS-AdamW Optimizer for Scalable LLM Training

MARS-AdamW is an optimizer instance derived from the MARS (Make vAriance Reduction Shine) framework, which integrates preconditioned adaptive gradient methods with scalable stochastic variance reduction to improve the efficiency and convergence of large model training. Specifically, MARS-AdamW unifies AdamW-style momentum preconditioning with STORM-inspired variance-reduced gradient estimation and implements explicit stabilization via gradient clipping. Developed to address the observed gap between theoretical advances in variance reduction and their practical adoption in large-scale neural network training, MARS-AdamW demonstrates significant improvements in both token and time efficiency over vanilla AdamW for tasks such as large language model pretraining [2411.10438].

## 1. Mathematical Derivation and Update Steps

MARS-AdamW builds on a variance-reduced preconditioned gradient framework. The core mechanism involves mixing the standard stochastic gradient with a STORM-style recursive momentum correction and then applying AdamW's adaptive moment preconditioning:

- **Variance-Reduced Gradient Estimator:**
  $$
  c_t := \nabla f(x_t, \xi_t) + \gamma_t \frac{\beta_1}{1-\beta_1}\bigl[\nabla f(x_t,\xi_t) - \nabla f(x_{t-1},\xi_t)\bigr]
  $$
  Here, $x_t \in \mathbb{R}^d$ is the parameter vector at iteration $t$, $\xi_t$ denotes the stochastic mini-batch, $\gamma_t \in [0,1]$ modulates the variance reduction intensity (recovering vanilla AdamW at $\gamma_t = 0$ and full STORM at $\gamma_t = 1$).

- **Gradient Clipping:**
  $$
  \tilde{c}_t =
    \begin{cases}
    c_t / \|c_t\|_2, & \text{if } \|c_t\|_2 > 1\\
    c_t, & \text{else}
    \end{cases}
  $$

- **Moment Estimation:**
  $$
  m_t = \beta_1 m_{t-1} + (1-\beta_1) \tilde{c}_t,\;\;
  v_t = \beta_2 v_{t-1} + (1-\beta_2) \tilde{c}_t^2
  $$
  Bias corrections are then applied:
  $$
  \hat{m}_t = m_t / (1-\beta_1^t),\;\;
  \hat{v}_t = v_t / (1-\beta_2^t)
  $$

- **Preconditioned Parameter Update (with decoupled weight decay):**
  $$
  x_{t+1} = x_t - \eta_t\left(\frac{\hat{m}_t}{\sqrt{\hat{v}_t}+\epsilon} + \lambda x_t\right)
  $$

For $\gamma_t \equiv 0$, the procedure exactly recovers AdamW. When $\gamma_t \equiv 1$ and $\lambda \rightarrow 0$, it recovers Adam combined with full STORM recursion.

## 2. Algorithmic Workflow and Pseudocode

MARS-AdamW's step-by-step execution is summarized below:

```plaintext
Input: x₀ ∈ ℝᵈ, η_t, λ ≥ 0, Adam β₁, β₂, VR-scale γ_t, clip threshold =1, ε>0
Initialize: m₀ ← 0, v₀ ← 0, x₁ ← x₀
For t = 1, ..., T:
    1. Draw mini-batch ξ_t
    2. Compute gₜ = ∇f(xₜ, ξₜ)
    3. If t>1: δₜ = (β₁/(1−β₁)) · [gₜ − ∇f(xₜ₋₁, ξₜ)]
       Else: δₜ ← 0
    4. cₜ = gₜ + γ_t · δₜ
    5. If ∥cₜ∥₂ > 1: ṡcₜ = cₜ / ∥cₜ∥₂
       Else: ṡcₜ = cₜ
    6. mₜ = β₁·mₜ₋₁ + (1−β₁)·ṡcₜ
    7. vₜ = β₂·vₜ₋₁ + (1−β₂)·(ṡcₜ ⊙ ṡcₜ)
    8. ẑmₜ = mₜ / (1−β₁ᵗ), ẑvₜ = vₜ / (1−β₂ᵗ)
    9. xₜ₊₁ = xₜ − η_t·[ ẑmₜ / (√ẑvₜ + ε) + λ·xₜ ]
EndFor
```

Key distinguishing feature: lines 3–5 integrate STORM variance reduction into AdamW’s preconditioner.

## 3. Hyperparameters and Selection Properties

MARS-AdamW introduces several hyperparameters, some inherited from AdamW and others governing the VR component. The major hyperparameters and empirical guidelines are:

| Hyperparameter | Role in Optimization           | Recommended Setting / Range                   |
|----------------|-------------------------------|-----------------------------------------------|
| $\eta_t$       | Learning rate (schedule)      | Cosine decay with linear warmup; peak $\eta$: 6e-4 (GPT-2 small), 3e-4 (medium), 2e-4 (large) |
| $\beta_1$      | Moment-1 EMA                  | $[0.9,\,0.99]$, best $\approx0.95$            |
| $\beta_2$      | Moment-2 EMA                  | $[0.95,\,0.999]$, default $0.99$              |
| $\gamma_t$     | VR scale                      | Constant $\gamma \in [0.01,\,0.1]$; $0.025$ robust |
| $\epsilon$     | Numerical stability           | $10^{-8}$                                     |
| $\lambda$      | Weight decay                  | $0.1$–$0.5$ (model-dependent)                 |
| Clipping       | Gradient stabilization        | $\ell_2$-norm threshold $=1.0$                |

Batch size ($\approx$480) and warm-up steps ($\approx$2k) follow large-model conventions. A constant VR scale is favored over schedule-based variants.

## 4. Theoretical Properties and Convergence

Under standard assumptions that $F$ is $L$-smooth, stochastic gradients are unbiased with bounded variance, and preconditioners are positive definite, MARS-AdamW inherits and extends convergence properties of its constituent algorithms:

- **Incremental oracle complexity** to find $\|\nabla F(x)\| \leq \epsilon$ is $\mathcal{O}(\epsilon^{-3})$.
- For $\gamma \rightarrow 1$, the method recovers the nearly-optimal STORM convergence rate for non-convex smooth optimization, as formalized in Arjevani et al. (2023).
- Full formal proof of convergence for the preconditioned variant is indicated as a direction for future work; empirically, no divergence or instability was observed in large-scale runs [2411.10438].

## 5. Comparative Empirical Evaluation on GPT-2 Pretraining

Performance was assessed on GPT-2 models of varying scales (small: 125M, medium: 355M, large: 770M) using the OpenWebText dataset.

- **Token Efficiency:** For GPT-2 large, 27B tokens were required by MARS-AdamW to reach validation loss 2.58, while AdamW required 50B tokens. Final validation losses: 2.53 (MARS-AdamW) vs. 2.56 (AdamW).
- **Wall-Clock Efficiency:** The per-iteration cost of MARS-AdamW is approximately 5–10% higher than AdamW (due to the VR correction), but overall wall-clock time to achieve a given loss is reduced by 50–60%.
- **Ablation Studies:** Little difference was observed between the exact and approximate VR correction (MARS vs. MARS-AP), suggesting MARS-AP is preferable when computational cost is a concern. A constant $\gamma\approx0.025$ outperformed linear scheduling schemes in terms of final validation loss. MARS-AdamW also consistently exceeded MARS-Lion in performance under matched conditions.

| Metric                                   | AdamW          | MARS-AdamW      |
|-------------------------------------------|----------------|-----------------|
| Tokens to val loss 2.58 (GPT-2 large)     | $50$B          | $27$B           |
| Final validation loss (GPT-2 large)       | $2.56$         | $2.53$          |
| Wall-clock speed to given loss            | Baseline       | $1.5$–$2 \times$ faster |
| Hellaswag 0-shot acc. (GPT-2 medium/large)| $39.5\%/42.3\%$| $41.8\%/44.2\%$  |

- **Downstream Transfer:** For zero-shot transfer on Hellaswag (after 50B tokens), MARS-AdamW outperformed AdamW for both medium and large GPT-2 models.

## 6. Practical Considerations and Implementation

MARS-AdamW requires two gradient evaluations per iteration in the default setting (for the VR correction), though the approximate variant (MARS-AP) mitigates this overhead with minimal loss in quality. The optimizer introduces a modest (5–10%) per-iteration computational overhead relative to AdamW, but the net speedup from reduced training steps is substantial for large models. Stability is ensured by $\ell_2$-norm gradient clipping and careful hyperparameter tuning. Empirically, the method demonstrated robust convergence and efficiency across all tested LLM-scale configurations [2411.10438].

Source: https://www.emergentmind.com/topics/mars-adamw