On the Convergence of Muon and Beyond

Authors: Da Chang, Yongxiang Liu, Ganzhao Yuan

Paper: https://arxiv.org/abs/2509.15816

We introduce MuonMVR, a variance-reduced Muon optimizer. The implementation provides the one-batch MVR1 and same-sample MVR2 estimators in Algorithm 1 of the paper. Matrix parameters use Muon and other parameters can use the AdamW backup.

Gradient estimators

Let $g_t=\nabla f(X_t;\xi_t)$. The default momentum_init="zero" applies the recurrence from $M_0=g_0=0$. Without correction clipping this gives $M_1=[(1-\beta_1)+\gamma_1\beta_1]g_1$; for $\gamma_1=0$ it is the standard EMA initialization $M_1=(1-\beta_1)g_1$. These are valid initializations. momentum_init="gradient" explicitly sets $M_1=g_1$ as in Algorithm 1. The paper's theoretical schedules have $\beta_1=0$ or $\gamma_1=1$, so the unclipped zero-history recurrence already gives $M_1=g_1$ for those schedules.

For subsequent steps:

MVR1 (is_approx=True) reuses the previous step's gradient:

$$M_t=\beta_t M_{t-1}+(1-\beta_t)g_t +\gamma_t\beta_t\bigl(g_t-\nabla f(X_{t-1};\xi_{t-1})\bigr).$$

MVR2 (is_approx=False) evaluates the current batch at both parameter states:

$$M_t=\beta_t M_{t-1}+(1-\beta_t)g_t +\gamma_t\beta_t\bigl(g_t-\nabla f(X_{t-1};\xi_t)\bigr).$$

Setting gamma=0 gives the EMA recurrence. For fixed momentum $\beta=\mu$, MVR1 with $\gamma=1-\mu$ is equivalent to the Nesterov recurrence after rescaling $M_t/(1-\mu)$ under matched initialization, with clipping disabled. The default zero history matches $C_0=0$ in $C_t=\mu C_{t-1}+g_t$, $N_t=\mu C_t+g_t$. The explicit $M_1=g_1$ option instead matches $C_0=g_1/(1-\mu)$. Neither initialization is an implementation error. In general changing the initial momentum changes the initial tracking-error term and finite-time constants, rather than the convergence rate order under the analyzed schedules.

The former README's "MVR3" uses the previous batch at the current parameters. That estimator is different from MVR2 and is not part of the paper's Algorithm 1 or the supported training command.

Usage

A runnable small example with plain PyTorch:

import torch
from muon_mvr import (
    MuonMVR, MVR2GradientAccumulator, accumulation_windows, backward_microbatches,
)

model = torch.nn.Linear(3, 2)
muon_params = [p for p in model.parameters() if p.ndim == 2]
adamw_params = [p for p in model.parameters() if p.ndim != 2]
optimizer = MuonMVR(
    muon_params=muon_params,
    adamw_params=adamw_params,
    lr=0.01, momentum=0.95, gamma=1.0, is_approx=False,
)
accumulator = MVR2GradientAccumulator()
batches = [(torch.randn(4, 3), torch.randn(4, 2)) for _ in range(8)]

def loss_fn(current_model, batch):
    inputs, targets = batch
    return torch.nn.functional.mse_loss(current_model(inputs), targets)

for window in accumulation_windows(batches, steps=2):
    # First window: g_1 only. Later windows: old and current parameters on
    # exactly the same full window. Both gradients are averaged identically.
    loss = accumulator.backward(model, optimizer, window, loss_fn)
    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

For MVR1, initialize with is_approx=True and replace the accumulator.backward call with backward_microbatches(model, window, loss_fn). It needs one pass per window and stores the current gradient automatically at step().

Each microbatch loss should be a scalar mean, with equal microbatch sizes within a window. A final partial window uses its actual number of microbatches. The helper works with plain models and DDP. DDP performs one gradient synchronization per pass; MVR2 uses two after the first update. Its auxiliary old pass restores Torch RNG state and model buffers, so the current pass sees the same dropout randomness and only one pass updates running statistics.

For a custom MVR2 loop, call optimizer.update_last_grad() once after all old-parameter gradients have accumulated, clear .grad, then accumulate all current-parameter gradients and call step(). Never mix the two gradients in .grad or clear only part of an accumulation window. The next MVR2 correction must receive a fresh old-parameter gradient; stale history raises an error. The accumulator must persist across optimizer steps. For an exact checkpoint resume, preserve its previous-parameter snapshot as well as optimizer state; loading only model weights starts a new estimator history.

Practical clipping and initialization

The default max_correction_norm=None implements the unclipped MVR recurrence. Select momentum_init="gradient" to match Algorithm 1's explicit initialization for arbitrary coefficients. The former hard-coded MARS correction clipping is available as max_correction_norm=1.0. It clips $g_t+\gamma_t\beta_t(g_t-g_{\mathrm{old}})/(1-\beta_t)$ whenever the recurrence is applied, including the first step under zero history. It is an additional practical modification; the paper's estimator identities and Nesterov equivalence apply with clipping disabled. The first step always sets $M_1=g_1$ when explicitly using momentum_init="gradient".

Zero-history initialization remains the default; to enable the historical clipped recipe, set max_correction_norm=1.0. The optional initialization choice and fixes to gradient accumulation, normalization and the AdamW backup can change training results; previous CIFAR-10/C4 numbers should not be described as reruns of this revision. Newton--Schulz uses the original coefficients (3.4445, -4.7750, 2.0315), float32 Frobenius normalization, then bfloat16 matrix iterations. Decoupled weight decay uses the base learning rate, as in the previous implementation.

C4 training

torchrun_main.py retains the original Ascend NPU/torch_npu and HCCL setup. For MVR2, its training loop uses the same tested accumulation helper as above:

torchrun --standalone --nproc_per_node=8 torchrun_main.py \
  --model_config configs/llama_130m.json --optimizer muon-mvr2 \
  --batch_size 4 --total_batch_size 128 --max_length 4096 \
  --lr 0.002 --beta1 0.95 --betas 0.9 0.98 --mars_gamma 0.05 \
  --weight_decay 0.1 --num_training_steps 22000 --warmup_steps 2200

--mars_gamma is the correction coefficient for MARS/MVR. MVR1 with --beta1 0.95 --mars_gamma 0.05 selects its Nesterov-equivalent recurrence under matched initialization. MVR2's theoretical diminishing schedule instead uses gamma=1 and step-dependent learning rates/momentum; the command above is the practical constant-momentum/cosine recipe. --mvr_correction_clip 1.0 enables the historical clipping and --mvr_momentum_init gradient explicitly selects $M_1=g_1$; zero history remains the default. --grad_clipping independently clips complete gradients and defaults to zero.

Verification

Run the numerical and CPU distributed regression tests with PyTorch and pytest:

python -m pytest -q tests

These tests cover closed-form estimator recurrences, full-batch versus microbatch accumulation, matched Nesterov initialization, dropout pairing, buffer restoration and DDP gradient averaging. They do not reproduce the large-scale C4 or CIFAR-10 experiments.

We build on Muon, Moonlight-Muon and the optional clipping strategy in MARS.

Citation

@article{Chang2025OnTC,
  title={On the Convergence of Muon and Beyond},
  author={Da Chang and Yongxiang Liu and Ganzhao Yuan},
  journal={ArXiv},
  year={2025},
  volume={abs/2509.15816}
}