Recent advances in spectral optimization, notably Muon, have demonstrated that constraining update steps to the Stiefel manifold can significantly accelerate training and improve generalization. However, Muon implicitly assumes an isotropic optimization landscape, enforcing a uniform spectral update norm across all eigen-directions. We argue that this "egalitarian" constraint is suboptimal for Deep Neural Networks, where the curvature spectrum is known to be highly heavy-tailed and ill-conditioned. In such landscapes, Muon risks amplifying instabilities in high-curvature directions while limiting necessary progress in flat directions. In this work, we propose Mousse (Muon Optimization Utilizing Shampoo's Structural Estimation), a novel optimizer that reconciles the structural stability of spectral methods with the geometric adaptivity of second-order preconditioning. Instead of applying Newton-Schulz orthogonalization directly to the momentum matrix, Mousse operates in a whitened coordinate system induced by Kronecker-factored statistics (derived from Shampoo). Mathematically, we formulate Mousse as the solution to a spectral steepest descent problem constrained by an anisotropic trust region, where the optimal update is derived via the polar decomposition of the whitened gradient. Empirical results across language models ranging from 160M to 800M parameters demonstrate that Mousse consistently outperforms Muon, achieving around ∼12% reduction in training steps with negligible computational overhead.
Figures & tables
Figure 1: The optimal result of Muon and Mousse optimizers on 800M models. Mousse achieves a ∼ 12% reduction in training steps to reach comparable loss levels against Muon .
Figure 2: Overview of the Mousse framework compared to baselines.
Figure 3: Validation Loss comparison on FineWeb (20B tokens). We report the final validation loss of Mousse against AdamW, Muon and SOAP across varying peak learning rates for model sizes ranging from 160M to 800M. Mousse consistently achieves the lowest validation loss across all model scales, demonstrating superior performance and scalability.
Figure 4: Mousse’s performance gains over Muon across 160M, 240M, 480M, and 800M models.
Figure 5: Scalability analysis of validation loss and training efficiency. Solid lines represent validation loss (left y-axis, lower is better), while dashed lines indicate total training time (right y-axis, lower is better) across model sizes ranging from 160M to 800M. Mousse consistently achieves the lowest validation loss across all scales. Crucially, Mousse maintains a training speed nearly identical to the efficient Muon optimizer.
Metric
AdamW
Muon
SOAP
Mousse (Ours)
Training Speed
++
++
−
++
Memory Efficiency
++
++
−
+
Final Performance
−
+
+
++
Table 1: Optimizer Comparison. We evaluate AdamW, Muon, SOAP, and our proposed Mousse across training speed, memory efficiency, and final performance. “ ++ ” indicates the best results.
Figure 6: Impact of Gradient Grafting on Update Stability. (Left) RMS Norm: We visualize the RMS norm of parameter updates throughout the training trajectory across varying learning rates. Solid lines represent settings without grafting, exhibiting a continuous downward drift independent of the learning rate schedule. The dashed line represents the constant update norm constrained by spectral normalization when grafting is enabled. (Right)Validation Loss trajectories for both settings around their respective optimal learning rates: By decoupling the magnitude from the optimization direction, grafting ensures superior convergence rates and lower final validation loss.
Figure 7: Ablation studies on 480M model. (a) We compare the validation loss trajectories under different preconditioning settings. The results demonstrate that Trace Normalization is essential for stability. Furthermore, Spectral Tempering ( α=0.125 ) consistently outperforms the aggressive curvature correction correction ( α=0.25 ), yielding the lowest final loss. (b) Comparison of single-sided and double-sided preconditioners.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Validation Loss comparison on FineWeb (20B tokens) under the WSD scheduler. We report the final validation across varying peak learning rates for model sizes ranging from 160M to 800M.
Figure 9: Quantification of Excessive Loss for WSD scheduler. We plot the magnitude of the excessive loss ( ΔL=Lstable−Ldecay ) during the learning rate decay phase against the peak learning rate for AdamW, Muon, SOAP and Mousse across various model scales. A consistent positive correlation is observed: higher learning rates sustain a higher "noise floor" (excessive loss) during the stable phase, which is subsequently released upon decay.
Figure 10: Evolution of Preconditioner Stability on 480M model trained by Mousse with cosine LR scheduler. We visualize the condition numbers and RMS norms of L and R across different network layers during training. Under the trace normalization, most curves showed a relatively stable downward trend, preventing numerical divergence. The preconditioner statistics are accumulated via an Exponential Moving Average (EMA) with β=0.95 .
Model
Embedding
dmodel
nlayers
nheads
162M
38M
768
12
6
247M
38M
768
24
6
494M
77M
1536
12
6
834M
77M
1536
24
6
Appendix
Table 2: Precise Model Architecture Configurations. All models use the architecture described above.
Hanoi University of Science and Technology Hanoi, Vietnam · Indiana University 107 S. Indiana Ave, Bloomington, IN 47405, USA · Monash University Clayton, VIC 3800, Australia
Department of Mathematics, Purdue University, West Lafayette, IN 47907, USA · Department of Mathematics & School of Mechanical Engineering, Purdue University, West Lafayette, IN 47907, USA