Muon-VR2: Optimal Variance-Reduced Optimizer

• Muon-VR2 is a variance-reduced variant of the Muon optimizer that uses a same-sample differencing mechanism to cancel stochastic noise.
• The update rule employs a momentum-like recursion followed by a Newton-Schulz transformation, ensuring optimal convergence in non-convex optimization.
• Empirical evaluations on vision and language tasks confirm faster convergence and improved accuracy compared to standard methods like SGD and Adam.

Searching arXiv for the Muon-VR2 optimizer paper and related Muon papers to ground the article. Muon-VR2 is a variance-reduced variant of the Muon framework introduced in “On the Convergence of Muon and Beyond” (Chang et al., 19 Sep 2025). It is designed for optimizing a possibly non-convex objective over matrix-valued parameters and was constructed to explore the theoretical limits of the Muon framework in stochastic non-convex settings. Its defining feature is a same-sample variance-reduction term inside the momentum-like recursion, and its main theoretical result is the first rigorous proof that such a mechanism enables a Muon-style optimizer to attain the optimal convergence rate $\tilde{\mathcal O}(T^{-1/3})$, matching the theoretical lower bound for this class of problems. The analysis also establishes convergence guarantees for Muon variants under the Polyak–Łojasiewicz condition.

1. Optimization problem and update rule

Muon-VR2 considers the optimization problem

$$f(X)\;=\;\mathbb E_{\xi\sim\mathcal D}[\,f(X;\xi)\,],\quad X\in\mathbb R^{m\times n},$$

and iterates according to

$$\begin{aligned} g_t &:= \nabla f(X_t;\xi_t),\quad g_{t-1}'\;:=\;\nabla f(X_{t-1};\xi_t)\quad(\text{same sample }\xi_t),\ M_t &= \beta_t\,M_{t-1}\;+\;(1-\beta_t)\,g_t \;+\;\gamma\,\beta_t\,(g_t - g_{t-1}'),\ O_t &=\text{Newton-Schulz}(M_t)\quad\bigl(\,O_t^\top O_t=I_n,\;\|O_t\|_F=\sqrt n\bigr),\ X_{t+1} &= X_t - \eta_t\,O_t\,. \end{aligned}$$

In practice one sets $\gamma=1$, $\beta_t=1-\eta_t$, and $g_{t-1}'=\nabla f(X_{t-1};\xi_t)$ so that the difference $g_t-g_{t-1}'$ is computed on the same mini-batch. Pseudocode for this variant is given as Option VR2 in lines 5–7 of Algorithm 1.

The update rule is notable because the recursion is not expressed directly in terms of the raw stochastic gradient step. Instead, Muon-VR2 first forms $M_t$, then applies Newton-Schulz to obtain $O_t$, and finally updates $X_t$ with stepsize $$f(X)\;=\;\mathbb E_{\xi\sim\mathcal D}[\,f(X;\xi)\,],\quad X\in\mathbb R^{m\times n},$$0. This places Muon-VR2 within the family of Muon-style methods while isolating the effect of variance reduction on convergence.

2. Variance-reduction mechanism

Standard Muon, denoted Option EMA, uses

$$f(X)\;=\;\mathbb E_{\xi\sim\mathcal D}[\,f(X;\xi)\,],\quad X\in\mathbb R^{m\times n},$$1

which accumulates an exponentially weighted average of stochastic gradients. Muon-VR1 adds a correction $$f(X)\;=\;\mathbb E_{\xi\sim\mathcal D}[\,f(X;\xi)\,],\quad X\in\mathbb R^{m\times n},$$2, but those two gradients come from different samples, so only an approximate variance reduction occurs.

Muon-VR2 replaces that approximation with the two-batch term $$f(X)\;=\;\mathbb E_{\xi\sim\mathcal D}[\,f(X;\xi)\,],\quad X\in\mathbb R^{m\times n},$$3 computed using the same sample $$f(X)\;=\;\mathbb E_{\xi\sim\mathcal D}[\,f(X;\xi)\,],\quad X\in\mathbb R^{m\times n},$$4 at steps $$f(X)\;=\;\mathbb E_{\xi\sim\mathcal D}[\,f(X;\xi)\,],\quad X\in\mathbb R^{m\times n},$$5 and $$f(X)\;=\;\mathbb E_{\xi\sim\mathcal D}[\,f(X;\xi)\,],\quad X\in\mathbb R^{m\times n},$$6. Because

$$f(X)\;=\;\mathbb E_{\xi\sim\mathcal D}[\,f(X;\xi)\,],\quad X\in\mathbb R^{m\times n},$$7

the difference $$f(X)\;=\;\mathbb E_{\xi\sim\mathcal D}[\,f(X;\xi)\,],\quad X\in\mathbb R^{m\times n},$$8 is an unbiased estimator of $$f(X)\;=\;\mathbb E_{\xi\sim\mathcal D}[\,f(X;\xi)\,],\quad X\in\mathbb R^{m\times n},$$9 and cancels a large part of the stochastic noise.

This mechanism is the central structural distinction of Muon-VR2. The paper’s theoretical claims are organized around the proposition that same-sample differencing, rather than merely adding a correction term, is what enables the optimal $$\begin{aligned} g_t &:= \nabla f(X_t;\xi_t),\quad g_{t-1}'\;:=\;\nabla f(X_{t-1};\xi_t)\quad(\text{same sample }\xi_t),\ M_t &= \beta_t\,M_{t-1}\;+\;(1-\beta_t)\,g_t \;+\;\gamma\,\beta_t\,(g_t - g_{t-1}'),\ O_t &=\text{Newton-Schulz}(M_t)\quad\bigl(\,O_t^\top O_t=I_n,\;\|O_t\|_F=\sqrt n\bigr),\ X_{t+1} &= X_t - \eta_t\,O_t\,. \end{aligned}$$0 rate. A plausible implication is that the variance-reduction component is the decisive ingredient in closing the gap between the practical performance of Muon-style optimizers and their earlier theoretical analyses.

3. Convergence guarantees in the general non-convex setting

The general non-convex analysis assumes that $$\begin{aligned} g_t &:= \nabla f(X_t;\xi_t),\quad g_{t-1}'\;:=\;\nabla f(X_{t-1};\xi_t)\quad(\text{same sample }\xi_t),\ M_t &= \beta_t\,M_{t-1}\;+\;(1-\beta_t)\,g_t \;+\;\gamma\,\beta_t\,(g_t - g_{t-1}'),\ O_t &=\text{Newton-Schulz}(M_t)\quad\bigl(\,O_t^\top O_t=I_n,\;\|O_t\|_F=\sqrt n\bigr),\ X_{t+1} &= X_t - \eta_t\,O_t\,. \end{aligned}$$1 is bounded below by $$\begin{aligned} g_t &:= \nabla f(X_t;\xi_t),\quad g_{t-1}'\;:=\;\nabla f(X_{t-1};\xi_t)\quad(\text{same sample }\xi_t),\ M_t &= \beta_t\,M_{t-1}\;+\;(1-\beta_t)\,g_t \;+\;\gamma\,\beta_t\,(g_t - g_{t-1}'),\ O_t &=\text{Newton-Schulz}(M_t)\quad\bigl(\,O_t^\top O_t=I_n,\;\|O_t\|_F=\sqrt n\bigr),\ X_{t+1} &= X_t - \eta_t\,O_t\,. \end{aligned}$$2, that $$\begin{aligned} g_t &:= \nabla f(X_t;\xi_t),\quad g_{t-1}'\;:=\;\nabla f(X_{t-1};\xi_t)\quad(\text{same sample }\xi_t),\ M_t &= \beta_t\,M_{t-1}\;+\;(1-\beta_t)\,g_t \;+\;\gamma\,\beta_t\,(g_t - g_{t-1}'),\ O_t &=\text{Newton-Schulz}(M_t)\quad\bigl(\,O_t^\top O_t=I_n,\;\|O_t\|_F=\sqrt n\bigr),\ X_{t+1} &= X_t - \eta_t\,O_t\,. \end{aligned}$$3 and each $$\begin{aligned} g_t &:= \nabla f(X_t;\xi_t),\quad g_{t-1}'\;:=\;\nabla f(X_{t-1};\xi_t)\quad(\text{same sample }\xi_t),\ M_t &= \beta_t\,M_{t-1}\;+\;(1-\beta_t)\,g_t \;+\;\gamma\,\beta_t\,(g_t - g_{t-1}'),\ O_t &=\text{Newton-Schulz}(M_t)\quad\bigl(\,O_t^\top O_t=I_n,\;\|O_t\|_F=\sqrt n\bigr),\ X_{t+1} &= X_t - \eta_t\,O_t\,. \end{aligned}$$4 are $$\begin{aligned} g_t &:= \nabla f(X_t;\xi_t),\quad g_{t-1}'\;:=\;\nabla f(X_{t-1};\xi_t)\quad(\text{same sample }\xi_t),\ M_t &= \beta_t\,M_{t-1}\;+\;(1-\beta_t)\,g_t \;+\;\gamma\,\beta_t\,(g_t - g_{t-1}'),\ O_t &=\text{Newton-Schulz}(M_t)\quad\bigl(\,O_t^\top O_t=I_n,\;\|O_t\|_F=\sqrt n\bigr),\ X_{t+1} &= X_t - \eta_t\,O_t\,. \end{aligned}$$5-smooth in Frobenius norm, and that stochastic gradients satisfy the bounded-variance condition

$$\begin{aligned} g_t &:= \nabla f(X_t;\xi_t),\quad g_{t-1}'\;:=\;\nabla f(X_{t-1};\xi_t)\quad(\text{same sample }\xi_t),\ M_t &= \beta_t\,M_{t-1}\;+\;(1-\beta_t)\,g_t \;+\;\gamma\,\beta_t\,(g_t - g_{t-1}'),\ O_t &=\text{Newton-Schulz}(M_t)\quad\bigl(\,O_t^\top O_t=I_n,\;\|O_t\|_F=\sqrt n\bigr),\ X_{t+1} &= X_t - \eta_t\,O_t\,. \end{aligned}$$6

Under these assumptions, Theorem 1.3 sets

$$\begin{aligned} g_t &:= \nabla f(X_t;\xi_t),\quad g_{t-1}'\;:=\;\nabla f(X_{t-1};\xi_t)\quad(\text{same sample }\xi_t),\ M_t &= \beta_t\,M_{t-1}\;+\;(1-\beta_t)\,g_t \;+\;\gamma\,\beta_t\,(g_t - g_{t-1}'),\ O_t &=\text{Newton-Schulz}(M_t)\quad\bigl(\,O_t^\top O_t=I_n,\;\|O_t\|_F=\sqrt n\bigr),\ X_{t+1} &= X_t - \eta_t\,O_t\,. \end{aligned}$$7

and establishes

$$\begin{aligned} g_t &:= \nabla f(X_t;\xi_t),\quad g_{t-1}'\;:=\;\nabla f(X_{t-1};\xi_t)\quad(\text{same sample }\xi_t),\ M_t &= \beta_t\,M_{t-1}\;+\;(1-\beta_t)\,g_t \;+\;\gamma\,\beta_t\,(g_t - g_{t-1}'),\ O_t &=\text{Newton-Schulz}(M_t)\quad\bigl(\,O_t^\top O_t=I_n,\;\|O_t\|_F=\sqrt n\bigr),\ X_{t+1} &= X_t - \eta_t\,O_t\,. \end{aligned}$$8

The significance of this result lies in two points. First, the paper states that this is the first rigorous proof that incorporating a variance-reduction mechanism enables Muon-VR2 to attain an optimal convergence rate in stochastic non-convex optimization. Second, the rate matches the theoretical lower bound for this class of problems. In the paper’s framing, this resolves the suboptimal $$\begin{aligned} g_t &:= \nabla f(X_t;\xi_t),\quad g_{t-1}'\;:=\;\nabla f(X_{t-1};\xi_t)\quad(\text{same sample }\xi_t),\ M_t &= \beta_t\,M_{t-1}\;+\;(1-\beta_t)\,g_t \;+\;\gamma\,\beta_t\,(g_t - g_{t-1}'),\ O_t &=\text{Newton-Schulz}(M_t)\quad\bigl(\,O_t^\top O_t=I_n,\;\|O_t\|_F=\sqrt n\bigr),\ X_{t+1} &= X_t - \eta_t\,O_t\,. \end{aligned}$$9 convergence behavior previously associated with the standard Muon variant in stochastic non-convex settings.

4. Proof structure and Polyak–Łojasiewicz analysis

The proof strategy proceeds through a sequence of estimates. The first step is a descent lemma for $\gamma=1$0, together with control of cross-terms by Young’s inequality. The second step introduces the momentum error

$\gamma=1$1

and derives the recursive bound

$\gamma=1$2

The third step chooses the coupling $\gamma=1$3 and $\gamma=1$4 so that $\gamma=1$5 and $\gamma=1$6 behave like $\gamma=1$7, and then applies the discrete stability lemma to show

$\gamma=1$8

The final step plugs this estimate back into the descent inequality, telescopes, and uses Cauchy–Schwarz and Jensen to convert $\gamma=1$9 control into $\beta_t=1-\eta_t$0 control, yielding the $\beta_t=1-\eta_t$1 bound.

The paper also analyzes Muon variants under the Polyak–Łojasiewicz condition

$\beta_t=1-\eta_t$2

In this regime one obtains a non-ergodic rate for suboptimality. Muon-EMA and Muon-VR1, with their stated parameter schedules, satisfy

$\beta_t=1-\eta_t$3

whereas Muon-VR2, with $\beta_t=1-\eta_t$4, $\beta_t=1-\eta_t$5, and $\beta_t=1-\eta_t$6, satisfies the faster rate

$\beta_t=1-\eta_t$7

Within the paper’s taxonomy, this makes Muon-VR2 the fastest sublinear Muon variant under the PL condition. This suggests that the same variance-reduction mechanism that improves the ergodic gradient bound in the general non-convex case also sharpens single-point suboptimality guarantees when PL structure is available.

5. Empirical evaluation

The empirical study covers both vision and language workloads. On CIFAR-10 with ResNet-18, the hardware setup is 4×Ascend 910C NPUs plus 4×RTX 4090 GPUs. Training uses 100 epochs, batch size 128, cosine decay to 10% of initial learning rate, weight decay 0.01, and a learning rate chosen by grid search. The reported metrics are train and test loss and accuracy versus epochs, together with test accuracy versus wall-clock time. The results state that all Muon variants converge faster than SGD and Adam in loss, and that Muon-VR2 achieves the highest final test accuracy and fastest per-iteration progress (Chang et al., 19 Sep 2025).

On C4 with LLaMA2-130M, the same multi-accelerator setup is used. Training is from scratch for 10 000 steps with global batch 512, constant learning rate $\beta_t=1-\eta_t$8, linear warmup plus cosine decay, weight decay 0.1, and bfloat16. The reported metrics are training and validation loss curves. The results state that Muon-VR2 reduces train loss more rapidly than AdamW, Muon-EMA, and Muon-VR1, and yields the lowest validation loss throughout.

The empirical claims are framed as corroboration of the theoretical findings on per-iteration convergence. In the paper’s presentation, the experiments are not merely illustrative benchmarks; they are used to support the claim that the improved theoretical rate corresponds to faster practical progress on both a vision benchmark and a large-scale language pretraining workload.

6. Position within the Muon literature and nomenclature

Muon-VR2 is defined relative to three variants discussed in the paper: standard Muon as Option EMA, Muon-VR1 as a one-batch variance-reduced alternative, and Muon-VR2 as the same-sample two-batch construction. The key distinction is that standard Muon uses only an exponentially weighted average of stochastic gradients, Muon-VR1 uses a correction with gradients from different samples, and Muon-VR2 uses $\beta_t=1-\eta_t$9 with the same sample $g_{t-1}'=\nabla f(X_{t-1};\xi_t)$0. The article’s overall conclusion is that Muon-VR2 provides the first proof of optimality for a Muon-style optimizer and clarifies the path toward developing more practically efficient, accelerated variants.

A recurring source of confusion is the name itself. In the optimization literature, “Muon-VR2” denotes the variance-reduced Muon optimizer analyzed in (Chang et al., 19 Sep 2025). Separately, the term “Muon-VR2” is also used in JUNO for a muon shower vertex reconstruction technique based on waveform information (Zhang, 2022). This suggests that the label is context-dependent rather than a unique cross-domain term.

The optimization usage is therefore best understood as a specific theoretical and algorithmic object: a Muon-style method with a same-sample variance-reduction term, Newton-Schulz transformation, and provable $g_{t-1}'=\nabla f(X_{t-1};\xi_t)$1 stochastic non-convex convergence. Within that scope, Muon-VR2 occupies a distinct place as the variant for which optimality is proved rather than only empirically suggested.