# Muon

Source: https://www.tensorplay.cn/docs/generated/tensorplay.optim.Muon.html

# Muon

class tensorplay.optim.Muon(params, lr: [float](https://docs.python.org/3/library/functions.html#float) = 0.001, weight_decay: [float](https://docs.python.org/3/library/functions.html#float) = 0.1, momentum: [float](https://docs.python.org/3/library/functions.html#float) = 0.95, nesterov: [bool](https://docs.python.org/3/library/functions.html#bool) = True, ns_coefficients: [tuple](https://docs.python.org/3/library/stdtypes.html#tuple)[[float](https://docs.python.org/3/library/functions.html#float), [float](https://docs.python.org/3/library/functions.html#float), [float](https://docs.python.org/3/library/functions.html#float)] = (3.4445, -4.775, 2.0315), eps: [float](https://docs.python.org/3/library/functions.html#float) = 1e-07, ns_steps: [int](https://docs.python.org/3/library/functions.html#int) = 5, adjust_lr_fn: [str](https://docs.python.org/3/library/stdtypes.html#str) | [None](https://docs.python.org/3/library/constants.html#None) = None)

Implements Muon algorithm.

\[ \begin{align}\begin{aligned}\begin{split}\begin{aligned}
     &\rule{110mm}{0.4pt} \\
     &\textbf{input}      : \gamma \text{ (lr)},\ \lambda \text{ (weight decay)},\
        \mu \text{ (momentum)},\ \textit{nesterov}\in\{True,False\},\\
     &\hspace{13mm}(a,b,c)\ \text{ (NS coefficients)},\
        \varepsilon \text{ (epsilon)},\ k \text{ (NS steps)},\
        \theta_0 \text{ (params)},\ f(\theta) \text{ (objective)} \\
     &\textbf{initialize} : B_0 \leftarrow 0 \text{ (momentum buffer)} \\[-1.ex]
     &\rule{110mm}{0.4pt} \\
     &\textbf{for}\ t=1\ \textbf{to}\ \ldots\ \textbf{do} \\[0.25ex]
     &\hspace{5mm} g_t \leftarrow \nabla_{\theta} f_t(\theta_{t-1}) \\[0.25ex]
     &\hspace{5mm} B_t \leftarrow \mu B_{t-1} + g_t \\[0.25ex]
     &\hspace{5mm} \widetilde{B}_t \leftarrow
         \begin{cases}
            g_t + \mu B_t, & \text{if nesterov}=True \\
            B_t,           & \text{if nesterov}=False
         \end{cases} \\[1.0ex]
     &\hspace{5mm} O_t \leftarrow \mathrm{NS}^{(a,b,c)}_{k}\!\big(\widetilde{B}_t;\ \varepsilon\big) \\[0.5ex]
     &\hspace{5mm} \theta_t \leftarrow \theta_{t-1} - \gamma\,\lambda\,\theta_{t-1}
        \quad\text{(decoupled weight decay)} \\[0.25ex]\end{split}\\\begin{split}     &\hspace{5mm} \gamma \leftarrow \mathrm{AdjustLR}\!\big(\gamma;\ \mathrm{shape}\!\big(\theta_t \big) \big) \\[0.25ex]
     &\hspace{5mm} \theta_t \leftarrow \theta_t - \gamma\, O_t \\
     &\rule{110mm}{0.4pt} \\[-1.ex]
     &\mathbf{return}\ \theta_t \\[-1.ex]
     &\rule{110mm}{0.4pt}s
\end{aligned}\end{split}\end{aligned}\end{align} \]

Here, \(\mathrm{NS}^{(a,b,c)}_{k}(\cdot;\varepsilon)\) denotes \(k\) iterations of the
Newton–Schulz orthogonalization operator parameterized by coefficients \((a,b,c)\)
with numerical stabilization \(\varepsilon\).

The purpose for \(\mathrm{AdjustLR}\!\big(\gamma;\ \mathrm{shape}\!\big(\theta_t \big) \big)\)
is to make the orthogonalized update scale consistently across rectangular matrices.

Keller’s original implementation scales the update by \(\sqrt{\max\!\left(1, \frac{A}{B}\right)}\),
where \(A\) and \(B\) are dimensions of the matrix being optimized, which represent fan-out
and fan-in for a Linear weight matrix.

Moonshot’s implementation focuses on matching \(RMS\) of AdamW. The adjustment is computed as:
\(\gamma \leftarrow {0.2}\gamma\,\sqrt{\max\!\left({A}, {B}\right)}\)
The method is adopted from [Muon is Scalable for LLM Training](https://arxiv.org/pdf/2502.16982). Research
results show that with this adjustment Muon can directly reuse the learning rate
and weight decay tuned for AdamW.

Jeremy Bernstein in [Deriving Muon](https://jeremybernste.in/writing/deriving-muon) proposes a scaling condition on the spectral norm, which
scales the update by \(\sqrt{\frac{A}{B}}\). This is similar to the Keller’s “original”
implementation but removes clamping down to 1.

We provide these options for the learning rate adjustment: “original”, which follows Keller’s
implementation, “match_rms_adamw”, which refers to Moonshot’s implementation, and “spectral_unclamped”,
which matches Bernstein’s implementation. If adjust_lr_fn is not specified, the default is “original”.

For further details regarding the algorithm we refer to [Muon: An optimizer for hidden layers in neural networks](https://kellerjordan.github.io/posts/muon/),
[Muon is Scalable for LLM Training](https://arxiv.org/pdf/2502.16982), and [Deriving Muon](https://jeremybernste.in/writing/deriving-muon).

Parameters:

- Other ( params. Note that Muon is an optimizer for 2D parameters of neural network hidden layers. ) – parameters, such as bias, and embedding, should be optimized by a standard method such as AdamW.

- lr ([float](https://docs.python.org/3/library/functions.html#float) , Tensor , optional ) – learning rate (default: 1e-3).

- weight_decay ([float](https://docs.python.org/3/library/functions.html#float) , optional ) – weight decay (L2 penalty). (default: 0.1)

- momentum ([float](https://docs.python.org/3/library/functions.html#float) , optional ) – momentum factor (default: 0.95)

- nesterov ([bool](https://docs.python.org/3/library/functions.html#bool) , optional ) – enables Nesterov momentum. Only applicable when momentum is non-zero

- ns_coefficients ([tuple](https://docs.python.org/3/library/stdtypes.html#tuple) of three floats , optional ) – coefficients (a,b,c) for the Newton–Schulz orthogonalization polynomial (default: (3.4445, -4.775, 2.0315))

- eps ([float](https://docs.python.org/3/library/functions.html#float) , optional ) – term added to the denominator for numerical stability. (default: 1e-07)

- ns_steps ([int](https://docs.python.org/3/library/functions.html#int) , optional ) – number of Newton–Schulz iteration steps. (default: 5)

- adjust_lr_fn ([str](https://docs.python.org/3/library/stdtypes.html#str) , optional ) – function to adjust learning rate. One of “original”, “match_rms_adamw”, and “spectral_unclamped”. If not specified, we will default to use “original”. (default: None)

Example

```
>>> # xdoctest: +SKIP
>>> # Muon only supports 2D params; use a standard optimizer
>>> # such as AdamW for biases, embeddings, and other non-2D
>>> # parameters.
>>> muon_params = [
...     p for p in model.parameters() if p.ndim == 2
... ]
>>> other_params = [
...     p for p in model.parameters() if p.ndim != 2
... ]
>>> optim_muon = torch.optim.Muon(
...     muon_params, lr=0.02, momentum=0.95
... )
>>> optim_adamw = torch.optim.AdamW(
...     other_params, lr=3e-4, weight_decay=0.01
... )
>>> optim_muon.zero_grad()
>>> optim_adamw.zero_grad()
>>> loss_fn(model(input), target).backward()
>>> optim_muon.step()
>>> optim_adamw.step()
```

step(closure=None)[[source]](../_modules/tensorplay/optim/_muon.html#Muon.step)

Performs a single optimization step.