Struct Muon 

pub struct Muon<B>
where
    B: Backend,
{ /* private fields */ }

Muon optimizer.

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.

Important Notes

1. Only for 2D+ parameters: Muon is designed for weight matrices. Use AdamW or SGD for biases, embeddings, and layer norms.

2. Learning rate adjustment: Muon automatically adjusts the learning rate based on parameter shape. See AdjustLrFn for details.

3. Weight decay timing: Unlike typical optimizers, Muon applies weight decay AFTER orthogonalization but uses the original (unadjusted) learning rate for it.

Trait Implementations

impl<B> Clone for Muon<B>

where B: Clone + Backend,

fn clone(&self) -> Muon<B>

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<B> SimpleOptimizer<B> for Muon<B>

where B: Backend,

fn step<const D: usize>( &self, lr: f64, tensor: Tensor<B, D>, grad: Tensor<B, D>, state: Option<<Muon<B> as SimpleOptimizer<B>>::State<D>>, ) -> (Tensor<B, D>, Option<<Muon<B> as SimpleOptimizer<B>>::State<D>>)

Perform a single Muon optimization step.

Algorithm

1. Apply momentum to gradient
2. Orthogonalize update via Newton-Schulz
3. Adjust learning rate based on parameter shape
4. Apply weight decay (using original lr)
5. Update parameter (using adjusted lr)

Notes

Unlike typical optimizers, the weight decay and parameter update use different learning rates:

• Weight decay uses the original lr
• Parameter update uses the shape-adjusted lr

Panics

This function will panic if the input tensors are not 2D.

type State<const D: usize> = MuonState<B, D>

The state of the optimizer. It also implements record, so that it can be saved.

fn to_device<const D: usize>( state: <Muon<B> as SimpleOptimizer<B>>::State<D>, device: &<B as Backend>::Device, ) -> <Muon<B> as SimpleOptimizer<B>>::State<D>

Change the device of the state.