Optimizers
Muon(params, lr=0.02, weight_decay=0, momentum=0.95)
Bases: Optimizer
Muon - MomentUm Orthogonalized by Newton-schulz
https://kellerjordan.github.io/posts/muon/
Muon internally runs standard SGD-momentum, and then performs an orthogonalization post- processing step, in which each 2D parameter's update is replaced with the nearest orthogonal matrix. For efficient orthogonalization we use a Newton-Schulz iteration, which has the advantage that it can be stably run in bfloat16 on the GPU.
Muon should only be used for hidden weight layers. The input embedding, final output layer, and any internal gains or biases should be optimized using a standard method such as AdamW. Hidden convolutional weights can be trained using Muon by viewing them as 2D and then collapsing their last 3 dimensions.
Muon has been shown to improve the accuracy of neural implicit representations.
References
_Optimizing Rank for High-Fidelity Implicit Neural Representations, J. McGinnis et al., 2026
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
lr
|
The learning rate, in units of spectral norm per update. |
0.02
|
|
weight_decay
|
The AdamW-style weight decay. |
0
|
|
momentum
|
The momentum. A value of 0.95 here is usually fine. |
0.95
|
MuonWithAuxAdam(param_groups)
Bases: Optimizer
Distributed Muon variant that can be used for all parameters in the network, since it runs an
internal AdamW for the parameters that are not compatible with Muon. The user must manually
specify which parameters shall be optimized with Muon and which with Adam by passing in a
list of param_groups with the use_muon flag set.
The point of this class is to allow the user to have a single optimizer in their code, rather than having both a Muon and an Adam which each need to be stepped.
You can see an example usage below:
Example
hidden_matrix_params = [p for n, p in model.blocks.named_parameters() if p.ndim >= 2 and "embed" not in n]
embed_params = [p for n, p in model.named_parameters() if "embed" in n]
scalar_params = [p for p in model.parameters() if p.ndim < 2]
head_params = [model.lm_head.weight]
from muon import MuonWithAuxAdam
adam_groups = [dict(params=head_params, lr=0.22), dict(params=embed_params, lr=0.6), dict(params=scalar_params, lr=0.04)]
adam_groups = [dict(**g, betas=(0.8, 0.95), eps=1e-10, use_muon=False) for g in adam_groups]
muon_group = dict(params=hidden_matrix_params, lr=0.05, momentum=0.95, use_muon=True)
param_groups = [*adam_groups, muon_group]
optimizer = MuonWithAuxAdam(param_groups)
SingleDeviceMuonWithAuxAdam(param_groups)
Bases: Optimizer
Non-distributed variant of MuonWithAuxAdam.
Parameters are updated with torch._foreach_* operations, and the Newton-Schulz orthogonalization is batched over all parameters sharing the same shape. This gives the same updates as applying muon_update and adam_update to each parameter separately, with far fewer kernel launches.