contrib: add MARS optimizer (arXiv:2411.10438)

[irhyl]

Summary

Adds MARS (Unleashing the Power of Variance Reduction for Training Large Models, Yuan et al. 2024 / ICML 2025) as a new optimizer in optax.contrib.

MARS augments Adam with a scaled stochastic recursive momentum correction. The key idea: instead of updating Adam's moments with the raw gradient $g_t$, MARS first computes a variance-reduced estimate

$$c_t = g_t + \gamma \frac{\beta_1}{1 - \beta_1}(g_t - g_{t-1})$$

which penalizes changes in gradient direction between consecutive steps. This acts as a low-cost variance reduction step without requiring an extra gradient evaluation. Adam's standard moment updates and bias correction then run on $c_t$ rather than $g_t$.

Setting $\gamma = 0$ recovers Adam exactly, making MARS a strict generalization.

Reported gains (from the paper, on GPT-2 training):

• 25–30% fewer steps to the same validation loss vs AdamW
• Equivalent wall-clock time (one gradient evaluation per step, same as Adam)

Files changed

• optax/contrib/_mars.py — scale_by_mars (primitive) and mars (full optimizer with weight decay and learning rate). Includes complete docstring with math, memory note, and usage example.
• optax/contrib/_mars_test.py — 13 unit tests covering: state shapes, last_grad tracking, γ=0/Adam equivalence, clipping (on/off), variance-reduction directional effect, quadratic convergence (2D and mixed params), JIT recompilation, weight decay, mu_dtype, count increment.
• optax/contrib/__init__.py — exports mars, scale_by_mars, ScaleByMARSState.
• optax/contrib/_common_test.py — adds mars to the shared optimizer test matrix (all 13 shared tests pass, including inject_hyperparams and dtype stability).

Implementation notes

Approximate vs exact MARS. The paper describes an exact variant that evaluates $\nabla f(x_{t-1}, \xi_t)$ — the previous-parameter gradient on the current mini-batch — requiring two gradient evaluations per step. This implementation uses the approximate variant: $g_{t-1}$ is simply the gradient stored from the previous step (different mini-batch). This is standard in practice (the official PyTorch reference implementation does the same), costs one gradient evaluation per step, and retains most of the variance-reduction benefit.

Clipping. The clip_threshold parameter (default 1.0) rescales $c_t$ to at most unit L2 norm before the moment updates. On the first step, last_grad is zero so $c_t = (1 + \gamma\beta_1/(1-\beta_1)) \cdot g_t$, which can be ~1.5× larger than the raw gradient for typical $\gamma = 0.025$, $\beta_1 = 0.95$. Clipping keeps the first step well-behaved. Setting clip_threshold=None disables it entirely.

inject_hyperparams compatibility. b1, b2, eps, and gamma are normalised to float32 at the top of scale_by_mars. This is necessary because inject_hyperparams reconstructs the inner optimizer from scratch each step with strongly-typed jnp.float32 hyperparameters; without normalisation, (1 - b2) evaluates differently for a Python float vs a JAX float32, producing numerically different moment updates and breaking the inject test.

Dtype stability. Moment buffers are cast back to their stored dtype after each float32 computation, following the same pattern as scale_by_soap. This ensures the state dtype is invariant across steps when params are float16 or bfloat16.

Test plan

• python -m pytest optax/contrib/_mars_test.py — 13/13 pass
• python -m pytest optax/contrib/_common_test.py — 562 pass, 4 skip
• python -m pytest optax/contrib/ — 697 pass, 7 skip

🤖 Generated with Claude Code

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.