burn_optim/optim/
muon.rs

1use burn_core as burn;
2
3use crate::RecordState;
4
5use burn::config::Config;
6use burn::tensor::Device;
7use burn::tensor::Tensor;
8use serde::{Deserialize, Serialize};
9
10use super::{
11    Optimizer,
12    decay::WeightDecayConfig,
13    module_optimizer::ModuleOptimizer,
14    momentum::{Momentum, MomentumConfig, MomentumState},
15};
16use crate::LearningRate;
17
18#[cfg(not(feature = "std"))]
19#[allow(unused_imports)]
20use num_traits::Float as _;
21
22/// Learning rate adjustment method for Muon optimizer.
23///
24/// Muon adjusts the learning rate based on parameter shape to maintain consistent
25/// RMS across rectangular matrices.
26///
27/// # References
28///
29/// - Original: [Muon: An optimizer for hidden layers](https://kellerjordan.github.io/posts/muon/)
30/// - Moonshot: [Muon is Scalable for LLM Training](https://arxiv.org/pdf/2502.16982)
31#[derive(Clone, Default, Debug, Copy, PartialEq, Eq, Serialize, Deserialize)]
32pub enum AdjustLrFn {
33    /// Keller Jordan's original method: `lr * sqrt(max(1, A/B))`
34    ///
35    /// This scales the learning rate based on the aspect ratio of the weight matrix,
36    /// ensuring that tall matrices (more rows than columns) get proportionally larger
37    /// learning rates.
38    ///
39    /// # Example
40    ///
41    /// For a [1024, 512] matrix: `lr * sqrt(1024/512) = lr * 1.414`
42    #[default]
43    Original,
44
45    /// Moonshot's method: `lr * 0.2 * sqrt(max(A, B))`
46    ///
47    /// This method is designed to match AdamW's RMS, allowing Muon to directly reuse
48    /// learning rates and weight decay values tuned for AdamW without retuning.
49    ///
50    /// # Example
51    ///
52    /// For a [1024, 512] matrix: `lr * 0.2 * sqrt(1024) = lr * 6.4`
53    MatchRmsAdamW,
54}
55
56impl AdjustLrFn {
57    /// Calculate the learning rate adjustment ratio for a given parameter shape.
58    ///
59    /// # Arguments
60    ///
61    /// * `shape` - Parameter shape (uses first two dimensions)
62    ///
63    /// # Returns
64    ///
65    /// Adjustment ratio to multiply with the base learning rate
66    fn adjustment_ratio(&self, shape: &[usize]) -> f64 {
67        if shape.len() < 2 {
68            return 1.0;
69        }
70
71        let a = shape[0] as f64;
72        let b = shape[1] as f64;
73
74        match self {
75            Self::Original => {
76                // sqrt(max(1, A/B))
77                let ratio = a / b;
78                ratio.max(1.0).sqrt()
79            }
80            Self::MatchRmsAdamW => {
81                // 0.2 * sqrt(max(A, B))
82                0.2 * a.max(b).sqrt()
83            }
84        }
85    }
86}
87
88/// Muon configuration.
89///
90/// Muon is an optimizer specifically designed for 2D parameters of neural network
91/// hidden layers (weight matrices). Other parameters such as biases and embeddings
92/// should be optimized using a standard method such as AdamW.
93///
94/// # Learning Rate Adjustment
95///
96/// Muon adjusts the learning rate based on parameter shape to maintain consistent
97/// RMS across rectangular matrices. Two methods are available:
98///
99/// - **Original**: Uses `sqrt(max(1, A/B))` where A and B are the first two dimensions.
100///   This is Keller Jordan's method and is the default.
101///
102/// - **MatchRmsAdamW**: Uses `0.2 * sqrt(max(A, B))`. This is Moonshot's method
103///   designed to match AdamW's RMS, allowing direct reuse of AdamW hyperparameters.
104///
105/// # Example
106///
107/// ```ignore
108/// use burn_optim::{MuonConfig, AdjustLrFn};
109///
110/// // Using default (Original) method
111/// let optimizer = MuonConfig::new().init();
112///
113/// // Using MatchRmsAdamW for AdamW-compatible hyperparameters
114/// let optimizer = MuonConfig::new()
115///     .with_adjust_lr_fn(AdjustLrFn::MatchRmsAdamW)
116///     .init();
117/// ```
118///
119/// # References
120///
121/// - [Muon: An optimizer for hidden layers in neural networks](https://kellerjordan.github.io/posts/muon/)
122/// - [Muon is Scalable for LLM Training](https://arxiv.org/pdf/2502.16982)
123/// - [PyTorch Implementation](https://github.com/pytorch/pytorch/blob/main/torch/optim/muon.py)
124/// - [Original Implementation](https://github.com/KellerJordan/Muon)
125#[derive(Config, Debug)]
126pub struct MuonConfig {
127    /// [Weight decay](WeightDecayConfig) config.
128    weight_decay: Option<WeightDecayConfig>,
129
130    /// [Momentum](MomentumConfig) config.
131    ///
132    /// Muon always uses momentum. Default configuration:
133    /// - momentum: 0.95
134    /// - dampening: 0.0
135    /// - nesterov: true
136    #[config(default = "MomentumConfig { momentum: 0.95, dampening: 0.0, nesterov: true }")]
137    momentum: MomentumConfig,
138
139    /// Newton-Schulz iteration coefficients (a, b, c).
140    ///
141    /// These coefficients are selected to maximize the slope at zero for the
142    /// quintic iteration. Default values are from Keller Jordan's implementation.
143    #[config(default = "(3.4445, -4.775, 2.0315)")]
144    ns_coefficients: (f32, f32, f32),
145
146    /// Epsilon for numerical stability.
147    #[config(default = 1e-7)]
148    epsilon: f32,
149
150    /// Number of Newton-Schulz iteration steps.
151    #[config(default = 5)]
152    ns_steps: usize,
153
154    /// Learning rate adjustment method.
155    ///
156    /// Controls how the learning rate is adjusted based on parameter shape.
157    /// See [`AdjustLrFn`] for available methods.
158    #[config(default = "AdjustLrFn::Original")]
159    adjust_lr_fn: AdjustLrFn,
160}
161
162impl MuonConfig {
163    /// Build a [`Muon`] from the config.
164    ///
165    /// The bare optimizer, which
166    /// [`ModuleOptimizer::with_group`](crate::ModuleOptimizer::with_group) takes to
167    /// optimize one parameter group. [`init`](Self::init) is the whole-module
168    /// counterpart.
169    pub fn build(&self) -> Muon {
170        let momentum = Momentum::new(&self.momentum);
171        let weight_decay_penalty = self.weight_decay.as_ref().map(|wd| wd.penalty);
172
173        Muon {
174            momentum,
175            ns_params: NewtonSchulzParams::new(self.ns_coefficients, self.ns_steps),
176            weight_decay_penalty,
177            epsilon: self.epsilon,
178            adjust_lr_fn: self.adjust_lr_fn,
179        }
180    }
181
182    /// Initialize Muon optimizer.
183    ///
184    /// # Returns
185    ///
186    /// Returns an optimizer adaptor that can be used to optimize a module.
187    ///
188    /// # Example
189    ///
190    /// ```ignore
191    /// use burn_optim::{MuonConfig, AdjustLrFn, decay::WeightDecayConfig};
192    ///
193    /// // Basic configuration with default (Original) LR adjustment
194    /// let optimizer = MuonConfig::new()
195    ///     .with_weight_decay(Some(WeightDecayConfig::new(0.01)))
196    ///     .init();
197    ///
198    /// // With AdamW-compatible settings using MatchRmsAdamW
199    /// let optimizer = MuonConfig::new()
200    ///     .with_adjust_lr_fn(AdjustLrFn::MatchRmsAdamW)
201    ///     .with_weight_decay(Some(WeightDecayConfig::new(0.1)))
202    ///     .init();
203    ///
204    /// // Custom momentum and NS settings
205    /// let optimizer = MuonConfig::new()
206    ///     .with_momentum(MomentumConfig {
207    ///         momentum: 0.9,
208    ///         dampening: 0.1,
209    ///         nesterov: false,
210    ///     })
211    ///     .with_ns_steps(7)
212    ///     .init();
213    /// ```
214    pub fn init(&self) -> ModuleOptimizer {
215        ModuleOptimizer::from(self.build())
216    }
217}
218
219/// Parameters for Newton-Schulz orthogonalization.
220#[derive(Clone, Copy)]
221struct NewtonSchulzParams {
222    a: f32,
223    b: f32,
224    c: f32,
225    steps: usize,
226}
227
228impl NewtonSchulzParams {
229    fn new(coefficients: (f32, f32, f32), steps: usize) -> Self {
230        Self {
231            a: coefficients.0,
232            b: coefficients.1,
233            c: coefficients.2,
234            steps,
235        }
236    }
237}
238
239/// Muon optimizer.
240///
241/// Muon internally runs standard SGD-momentum, and then performs an orthogonalization
242/// post-processing step, in which each 2D parameter's update is replaced with the
243/// nearest orthogonal matrix. For efficient orthogonalization we use a Newton-Schulz
244/// iteration, which has the advantage that it can be stably run in bfloat16 on the GPU.
245///
246/// # Important Notes
247///
248/// 1. **Only for 2D+ parameters**: Muon is designed for weight matrices. Use AdamW
249///    or SGD for biases, embeddings, and layer norms.
250///
251/// 2. **Learning rate adjustment**: Muon automatically adjusts the learning rate based
252///    on parameter shape. See [`AdjustLrFn`] for details.
253///
254/// 3. **Weight decay timing**: Unlike typical optimizers, Muon applies weight decay
255///    AFTER orthogonalization but uses the original (unadjusted) learning rate for it.
256#[derive(Clone)]
257pub struct Muon {
258    momentum: Momentum,
259    ns_params: NewtonSchulzParams,
260    weight_decay_penalty: Option<f32>,
261    epsilon: f32,
262    adjust_lr_fn: AdjustLrFn,
263}
264
265impl Muon {
266    /// Adjust learning rate based on parameter shape.
267    ///
268    /// # Arguments
269    ///
270    /// * `lr` - Base learning rate
271    /// * `shape` - Parameter shape (uses first two dimensions)
272    ///
273    /// # Returns
274    ///
275    /// Adjusted learning rate
276    ///
277    /// ```ignore
278    /// // For a [1024, 512] weight matrix with lr=0.01:
279    /// // Original: 0.01 * sqrt(1024/512) = 0.01 * 1.414 = 0.01414
280    /// // MatchRmsAdamW: 0.01 * 0.2 * sqrt(1024) = 0.01 * 0.2 * 32 = 0.064
281    /// ```
282    fn adjust_lr(&self, lr: LearningRate, shape: &[usize]) -> LearningRate {
283        lr * self.adjust_lr_fn.adjustment_ratio(shape)
284    }
285
286    /// Perform Newton-Schulz orthogonalization on a gradient tensor.
287    ///
288    /// This computes the zeroth power (orthogonalization) of the input matrix G
289    /// using a quintic Newton-Schulz iteration.
290    ///
291    /// # Algorithm
292    ///
293    /// 1. Transpose if tall matrix (A > B)
294    /// 2. Normalize: X = X / ||X||
295    /// 3. For k steps:
296    ///    - A = X @ X^T
297    ///    - B = b*A + c*A^2
298    ///    - X = a*X + B@X
299    /// 4. Transpose back if needed
300    ///
301    /// # References
302    ///
303    /// - Original: https://github.com/KellerJordan/Muon/blob/master/muon.py
304    /// - PyTorch: https://github.com/pytorch/pytorch/blob/main/torch/optim/muon.py
305    fn zeropower_via_newtonschulz<const D: usize>(&self, g: Tensor<D>) -> Tensor<D> {
306        let shape = g.shape();
307        let dim_m2 = shape[D - 2];
308        let dim_m1 = shape[D - 1];
309
310        // Step 1: Transpose if tall matrix (more rows than columns)
311        let (mut x, needs_transpose) = if dim_m2 > dim_m1 {
312            (g.swap_dims(D - 2, D - 1), true)
313        } else {
314            (g, false)
315        };
316
317        // Step 2: Normalize by Frobenius norm
318        // X = X / (||X|| + epsilon)
319        let norm = x
320            .clone()
321            .powf_scalar(2.0)
322            .sum()
323            .sqrt()
324            .clamp_min(self.epsilon)
325            .unsqueeze();
326
327        x = x.div(norm);
328
329        // Step 3: Newton-Schulz iteration
330        // This is the quintic iteration with coefficients (a, b, c)
331        let NewtonSchulzParams { a, b, c, steps } = self.ns_params;
332
333        for _ in 0..steps {
334            // A = X @ X^T
335            let x_t = x.clone().swap_dims(D - 2, D - 1);
336            let a_matrix = x.clone().matmul(x_t);
337
338            // B = b*A + c*A@A
339            let a_squared = a_matrix.clone().matmul(a_matrix.clone());
340            let b_matrix = a_matrix.mul_scalar(b).add(a_squared.mul_scalar(c));
341
342            // X = a*X + B@X
343            x = x.clone().mul_scalar(a).add(b_matrix.matmul(x.clone()));
344        }
345
346        // Step 4: Restore transpose if it was a tall matrix
347        if needs_transpose {
348            x = x.swap_dims(D - 2, D - 1);
349        }
350
351        x
352    }
353}
354
355/// Muon state.
356#[derive(RecordState, Clone, new)]
357pub struct MuonState<const D: usize> {
358    /// Current momentum state
359    pub momentum: MomentumState<D>,
360}
361
362impl Optimizer for Muon {
363    type State<const D: usize> = MuonState<D>;
364
365    /// Perform a single Muon optimization step.
366    ///
367    /// # Algorithm
368    ///
369    /// 1. Apply momentum to gradient
370    /// 2. Orthogonalize update via Newton-Schulz
371    /// 3. Adjust learning rate based on parameter shape
372    /// 4. Apply weight decay (using original lr)
373    /// 5. Update parameter (using adjusted lr)
374    ///
375    /// # Notes
376    ///
377    /// Unlike typical optimizers, the weight decay and parameter update use
378    /// different learning rates:
379    /// - Weight decay uses the original `lr`
380    /// - Parameter update uses the shape-adjusted `lr`
381    ///
382    /// # Panics
383    /// This function will panic if the input tensors are not 2D.
384    fn step<const D: usize>(
385        &self,
386        lr: LearningRate,
387        tensor: Tensor<D>,
388        grad: Tensor<D>,
389        state: Option<Self::State<D>>,
390    ) -> (Tensor<D>, Option<Self::State<D>>) {
391        assert!(
392            D == 2,
393            "Newton-Schulz iteration requires 2D tensors, got {}D",
394            D
395        );
396
397        // Step 1: Apply momentum
398        let state_momentum = state.map(|s| s.momentum);
399        let (grad, new_momentum_state) = self.momentum.transform(grad, state_momentum);
400
401        // Step 2: Orthogonalize via Newton-Schulz
402        let update = self.zeropower_via_newtonschulz(grad);
403
404        // Step 3: Adjust learning rate based on parameter shape
405        let adjusted_lr = self.adjust_lr(lr, &tensor.shape());
406
407        // Step 4: Apply weight decay (using ORIGINAL lr, not adjusted)
408        // Muon applies weight decay AFTER orthogonalization
409        let tensor = if let Some(penalty) = self.weight_decay_penalty {
410            let decay_factor = 1.0 - lr * penalty as f64;
411            tensor.mul_scalar(decay_factor)
412        } else {
413            tensor
414        };
415
416        // Step 5: Update parameter (using ADJUSTED lr)
417        let delta = update.mul_scalar(adjusted_lr);
418        let new_state = MuonState::new(new_momentum_state);
419
420        (tensor - delta, Some(new_state))
421    }
422
423    fn to_device<const D: usize>(mut state: Self::State<D>, device: &Device) -> Self::State<D> {
424        state.momentum = state.momentum.to_device(device);
425        state
426    }
427}
428
429#[cfg(test)]
430mod tests {
431    use super::*;
432    use crate::{GradientsParams, Optimizer};
433    use burn::module::Param;
434    use burn::tensor::{Distribution, Tensor, TensorData};
435    use burn_nn::{Linear, LinearConfig};
436
437    const TOLERANCE: f64 = 1e-8;
438
439    fn given_linear_layer_no_bias(weight: TensorData) -> Linear {
440        let device = Device::default().autodiff();
441        Linear {
442            weight: Param::from_data(weight, &device),
443            bias: None, // No bias for Muon optimizer
444        }
445    }
446
447    #[test]
448    fn test_adjust_lr_fn_original() {
449        let method = AdjustLrFn::Original;
450
451        // Square matrix [512, 512] -> sqrt(1) = 1.0
452        let ratio = method.adjustment_ratio(&[512, 512]);
453        assert!((ratio - 1.0).abs() < TOLERANCE);
454
455        // Tall matrix [1024, 512] -> sqrt(2) ≈ 1.414
456        let ratio = method.adjustment_ratio(&[1024, 512]);
457        let expected = (2.0f64).sqrt();
458        assert!((ratio - expected).abs() < TOLERANCE);
459
460        // Wide matrix [512, 1024] -> max(1, 0.5) = 1.0
461        let ratio = method.adjustment_ratio(&[512, 1024]);
462        assert!((ratio - 1.0).abs() < TOLERANCE);
463    }
464
465    #[test]
466    fn test_adjust_lr_fn_match_rms_adamw() {
467        let method = AdjustLrFn::MatchRmsAdamW;
468
469        // [1024, 512] -> 0.2 * sqrt(1024) = 6.4
470        let ratio = method.adjustment_ratio(&[1024, 512]);
471        let expected = 0.2 * 1024.0f64.sqrt();
472        assert!((ratio - expected).abs() < TOLERANCE);
473
474        // [512, 512] -> 0.2 * sqrt(512) ≈ 4.525
475        let ratio = method.adjustment_ratio(&[512, 512]);
476        let expected = 0.2 * 512.0f64.sqrt();
477        assert!((ratio - expected).abs() < TOLERANCE);
478    }
479
480    #[test]
481    #[should_panic(expected = "Newton-Schulz iteration requires 2D tensors, got 1D")]
482    fn test_1d_tensor_panics() {
483        let device = Default::default();
484        let config = MuonConfig::new();
485        let optim = Muon {
486            momentum: Momentum::new(&config.momentum),
487            ns_params: NewtonSchulzParams::new(config.ns_coefficients, config.ns_steps),
488            weight_decay_penalty: None,
489            epsilon: config.epsilon,
490            adjust_lr_fn: config.adjust_lr_fn,
491        };
492
493        let tensor_1d = Tensor::<1>::zeros([512], &device);
494        let grad_1d = Tensor::<1>::ones([512], &device);
495
496        let _ = optim.step(0.01, tensor_1d, grad_1d, None);
497    }
498
499    #[test]
500    fn test_muon_optimizer_save_load_state() {
501        let device = Device::default().autodiff();
502        // Use Linear layer WITHOUT bias for Muon optimizer
503        let linear = LinearConfig::new(6, 6)
504            .with_bias(false) // No bias - only 2D weight matrix
505            .init(&device);
506
507        let x = Tensor::<2>::random([2, 6], Distribution::Default, &device);
508
509        let mut optimizer = MuonConfig::new().init();
510        let grads = linear.forward(x).backward();
511        let grads = GradientsParams::from_grads(grads, &linear);
512        let _linear = optimizer.step(0.01, linear, grads);
513
514        let state_before = optimizer.to_record();
515        let bytes = optimizer.into_bytes().unwrap();
516
517        let optimizer_loaded = MuonConfig::new().init().from_bytes(bytes).unwrap();
518        let state_after = optimizer_loaded.to_record();
519
520        assert_eq!(state_before.len(), state_after.len());
521    }
522
523    #[test]
524    fn test_muon_with_weight_decay() {
525        let device = Device::default().autodiff();
526        // Create Linear layer WITHOUT bias for Muon
527        let linear = given_linear_layer_no_bias(TensorData::from([
528            [1.0, 1.0, 1.0, 1.0],
529            [1.0, 1.0, 1.0, 1.0],
530            [1.0, 1.0, 1.0, 1.0],
531            [1.0, 1.0, 1.0, 1.0],
532        ]));
533
534        let x = Tensor::<2>::from_floats([[0.5, 0.5, 0.5, 0.5], [0.5, 0.5, 0.5, 0.5]], &device)
535            .require_grad();
536
537        let mut optimizer = MuonConfig::new()
538            .with_weight_decay(Some(WeightDecayConfig::new(0.01)))
539            .init();
540
541        let grads = linear.forward(x).backward();
542        let grads = GradientsParams::from_grads(grads, &linear);
543        let linear = optimizer.step(0.01, linear, grads);
544
545        let state = linear;
546        let weight = state.weight.to_data();
547
548        for val in weight.as_slice::<f32>().unwrap() {
549            assert!(
550                *val < 1.0,
551                "Weight should be reduced by weight decay, got {}",
552                val
553            );
554        }
555    }
556
557    #[test]
558    fn test_newton_schulz_orthogonalization() {
559        let device = Default::default();
560        let matrix = Tensor::<2>::from_floats([[1.0, 0.5], [0.5, 1.0]], &device);
561
562        let config = MuonConfig::new();
563        let muon = Muon {
564            momentum: Momentum::new(&config.momentum),
565            ns_params: NewtonSchulzParams::new(config.ns_coefficients, config.ns_steps),
566            weight_decay_penalty: None,
567            epsilon: config.epsilon,
568            adjust_lr_fn: config.adjust_lr_fn,
569        };
570
571        let orthogonalized = muon.zeropower_via_newtonschulz(matrix);
572        let o_t = orthogonalized.clone().transpose();
573        let product = orthogonalized.matmul(o_t);
574
575        let data = product.into_data();
576        let values = data.as_slice::<f32>().unwrap();
577
578        assert!(
579            (values[0] - 1.0).abs() < 0.1,
580            "Product[0,0] should be ~1.0, got {}",
581            values[0]
582        );
583        assert!(
584            (values[3] - 1.0).abs() < 0.1,
585            "Product[1,1] should be ~1.0, got {}",
586            values[3]
587        );
588    }
589
590    #[test]
591    fn test_tall_matrix_transpose() {
592        // Test that tall matrices (A > B) are transposed during Newton-Schulz iteration
593        // and then transposed back
594        let device = Default::default();
595
596        // Create a tall matrix: [8, 4] (more rows than columns)
597        let tall_matrix = Tensor::<2>::from_floats(
598            [
599                [1.0, 0.5, 0.3, 0.2],
600                [0.5, 1.0, 0.4, 0.1],
601                [0.3, 0.4, 1.0, 0.5],
602                [0.2, 0.1, 0.5, 1.0],
603                [0.1, 0.2, 0.3, 0.4],
604                [0.4, 0.3, 0.2, 0.1],
605                [0.2, 0.4, 0.1, 0.3],
606                [0.3, 0.1, 0.4, 0.2],
607            ],
608            &device,
609        );
610
611        let config = MuonConfig::new();
612        let muon = Muon {
613            momentum: Momentum::new(&config.momentum),
614            ns_params: NewtonSchulzParams::new(config.ns_coefficients, config.ns_steps),
615            weight_decay_penalty: None,
616            epsilon: config.epsilon,
617            adjust_lr_fn: config.adjust_lr_fn,
618        };
619
620        // Perform Newton-Schulz orthogonalization
621        let orthogonalized = muon.zeropower_via_newtonschulz(tall_matrix.clone());
622
623        // Verify shape is preserved (should be transposed internally but returned in original shape)
624        let original_shape = tall_matrix.shape();
625        let result_shape = orthogonalized.shape();
626        assert_eq!(
627            original_shape.dims::<2>(),
628            result_shape.dims::<2>(),
629            "Shape should be preserved: [8, 4]"
630        );
631
632        // Verify output is different from input (orthogonalization happened)
633        let original_data = tall_matrix.into_data();
634        let result_data = orthogonalized.into_data();
635        assert_ne!(
636            original_data.as_slice::<f32>().unwrap(),
637            result_data.as_slice::<f32>().unwrap(),
638            "Orthogonalized matrix should differ from input"
639        );
640
641        // For comparison, test a wide matrix [4, 8] should NOT be transposed
642        let wide_matrix = Tensor::<2>::from_floats(
643            [
644                [1.0, 0.5, 0.3, 0.2, 0.1, 0.4, 0.2, 0.3],
645                [0.5, 1.0, 0.4, 0.1, 0.2, 0.3, 0.4, 0.1],
646                [0.3, 0.4, 1.0, 0.5, 0.3, 0.2, 0.1, 0.4],
647                [0.2, 0.1, 0.5, 1.0, 0.4, 0.1, 0.3, 0.2],
648            ],
649            &device,
650        );
651
652        let orthogonalized_wide = muon.zeropower_via_newtonschulz(wide_matrix.clone());
653
654        // Verify wide matrix shape is also preserved
655        let wide_original_shape = wide_matrix.shape();
656        let wide_result_shape = orthogonalized_wide.shape();
657        assert_eq!(
658            wide_original_shape.dims::<2>(),
659            wide_result_shape.dims::<2>(),
660            "Wide matrix shape should be preserved: [4, 8]"
661        );
662    }
663
664    #[test]
665    fn test_zero_gradient() {
666        // Test that Muon handles zero gradients gracefully
667        let device = Default::default();
668
669        let tensor = Tensor::<2>::from_floats(
670            [
671                [1.0, 0.5, 0.3, 0.2],
672                [0.5, 1.0, 0.4, 0.1],
673                [0.3, 0.4, 1.0, 0.5],
674                [0.2, 0.1, 0.5, 1.0],
675            ],
676            &device,
677        );
678
679        // Zero gradient - all zeros
680        let zero_grad = Tensor::<2>::zeros([4, 4], &device);
681
682        let config = MuonConfig::new();
683        let muon = Muon {
684            momentum: Momentum::new(&config.momentum),
685            ns_params: NewtonSchulzParams::new(config.ns_coefficients, config.ns_steps),
686            weight_decay_penalty: None,
687            epsilon: config.epsilon,
688            adjust_lr_fn: config.adjust_lr_fn,
689        };
690
691        // Should not panic or produce NaN
692        let (updated_tensor, state) = muon.step(0.01, tensor.clone(), zero_grad, None);
693
694        // Verify state was created
695        assert!(state.is_some());
696
697        // With zero gradient and no weight decay, tensor should remain unchanged
698        let original_data = tensor.into_data();
699        let updated_data = updated_tensor.clone().into_data();
700
701        let original_vals = original_data.as_slice::<f32>().unwrap();
702        let updated_vals = updated_data.as_slice::<f32>().unwrap();
703
704        for (orig, upd) in original_vals.iter().zip(updated_vals.iter()) {
705            assert!(
706                (orig - upd).abs() < 1e-6,
707                "With zero gradient, tensor should remain unchanged (or very close)"
708            );
709        }
710
711        // Verify no NaN values
712        for val in updated_vals {
713            assert!(
714                !val.is_nan(),
715                "Result should not contain NaN values with zero gradient"
716            );
717        }
718
719        // Test with weight decay - should still work
720        let muon_with_decay = Muon {
721            momentum: Momentum::new(&config.momentum),
722            ns_params: NewtonSchulzParams::new(config.ns_coefficients, config.ns_steps),
723            weight_decay_penalty: Some(0.01),
724            epsilon: config.epsilon,
725            adjust_lr_fn: config.adjust_lr_fn,
726        };
727
728        let tensor2 = Tensor::<2>::from_floats(
729            [
730                [1.0, 0.5, 0.3, 0.2],
731                [0.5, 1.0, 0.4, 0.1],
732                [0.3, 0.4, 1.0, 0.5],
733                [0.2, 0.1, 0.5, 1.0],
734            ],
735            &device,
736        );
737        let zero_grad2 = Tensor::<2>::zeros([4, 4], &device);
738
739        let (updated_tensor_decay, _) =
740            muon_with_decay.step(0.01, tensor2.clone(), zero_grad2, None);
741
742        // With zero gradient but with weight decay, tensor should be slightly reduced
743        let updated_decay_data = updated_tensor_decay.into_data();
744        let updated_decay_vals = updated_decay_data.as_slice::<f32>().unwrap();
745
746        for val in updated_decay_vals {
747            assert!(
748                !val.is_nan(),
749                "Result should not contain NaN with zero gradient and weight decay"
750            );
751        }
752
753        // With weight decay, values should be slightly smaller than original
754        let original_vals2 = tensor2.into_data().as_slice::<f32>().unwrap().to_vec();
755        for (orig, upd) in original_vals2.iter().zip(updated_decay_vals.iter()) {
756            if orig.abs() > 1e-6 {
757                // Non-zero values should be reduced by weight decay
758                assert!(
759                    upd.abs() < orig.abs(),
760                    "Weight decay should reduce magnitude: original={}, updated={}",
761                    orig,
762                    upd
763                );
764            }
765        }
766    }
767}
burn_optim/optim/muon.rs

burn_optim/optim/
muon.rs
