Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning

Yechen Zhang; Shuhao Xing; Junhao Huang; Kai Lv; Yunhua Zhou; Xipeng Qiu; Qipeng Guo; Kai Chen

Source-linked AI summary

Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning

Yechen Zhang, Shuhao Xing, Junhao Huang, Kai Lv, Yunhua Zhou, Xipeng Qiu, Qipeng Guo, Kai Chen

arXiv:2603.09697v2Mar 10, 2026cs.LGcs.AIcs.CL

TL;DR

Muon’s isotropic spectral constraint does not match the anisotropic, ill-conditioned curvature of neural-network landscapes. Mousse addresses this mismatch by whitening gradients with Shampoo-derived statistics before spectral orthogonalization, and reports consistently lower validation loss with about 12% fewer training steps and negligible overhead.

Problem
Muon’s isotropic trust region imposes uniform spectral norms across eigen-directions despite anisotropic and ill-conditioned neural-network curvature.
Method
Mousse applies Shampoo-based Kronecker-factored preconditioning in a whitened coordinate system before Newton-Schulz orthogonalization under anisotropic spectral steepest descent.
Results
Mousse consistently outperforms Muon, requiring approximately 12% fewer training steps to reach the Muon baseline’s final converged loss with negligible computational overhead.
Takeaways & Limitations
Mousse provides a curvature-aware framework for combining spectral optimization with second-order preconditioning in large-scale language-model training.
Takeaways & Limitations
The current implementation is a proof of concept focused on algorithmic correctness rather than hardware efficiency, leaving room for system-level optimization.

Abstract

from arXiv · show

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 \textbf{Mousse} (\textbf{M}uon \textbf{O}ptimization \textbf{U}tilizing \textbf{S}hampoo's \textbf{S}tructural \textbf{E}stimation), 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 $\sim$12\% reduction in training steps with negligible computational overhead.

1. Introduction

Mousse addresses Muon’s isotropic trust-region assumption by combining spectral optimization with Shampoo-based curvature adaptation. The proposed whitened-coordinate approach is theoretically motivated and reported to improve sample efficiency with minimal overhead.

Motivation: Muon’s isotropic trust region enforces uniform spectral norms despite substantial curvature disparities in neural-network loss landscapes.The authors describe this as an egalitarian constraint that treats dimensions as geometrically equivalent.
Approach: Mousse combines second-order preconditioning with spectral constraints in a unified optimizer.Its name expands to Muon Optimization Utilizing Shampoo’s Structural Estimation.
Approach: Mousse whitens gradients using Shampoo’s Kronecker-factored curvature statistics before applying Newton-Schulz orthogonalization.This change of basis is intended to align spectral constraints with the loss surface’s geometry.
Theory: Mousse is theoretically grounded as the optimal solution to a dual-norm maximization problem under anisotropic geometry.The formulation bridges spectral optimization and second-order preconditioning through a whitened coordinate system.

2. Related Works

Related work spans coordinate-wise adaptive optimization, structural second-order preconditioning, and spectral methods. Mousse is positioned at their intersection by combining curvature-aware preconditioning with spectral efficiency.

Adaptive and second-order optimization: AdamW uses coordinate-wise heuristics, whereas K-FAC and Shampoo capture parameter correlations through Kronecker-factored curvature statistics.SOAP extends this direction by combining Shampoo-style structural preconditioning with Adam-style momentum and adaptive step scaling.
Spectral optimization: Muon introduces strict spectral constraints on parameter updates as an efficient optimization paradigm.Mousse explores integrating advanced preconditioning with this spectral approach.
Positioning: Mousse combines structural preconditioning with spectral efficiency rather than treating these as separate optimization paradigms.The paper frames this integration as its position within the development of spectral methods.
Spectral optimization: Prior spectral methods address issues including curvature correction, gradient magnitude retention, and neuron-level scale imbalance.The cited approaches include Shampoo’s relation to Muon, PolarGrad’s nuclear-norm scaling, and NorMuon’s neuron-level treatment.

3. Methods

Mousse formulates spectral optimization as steepest descent under curvature-aware geometry. It whitens gradients using Shampoo-style Kronecker-factored statistics, applies Muon’s spectral constraint in that basis, and reports lower validation loss across model scales.

Geometric formulation: Steepest descent selects updates by minimizing the local linearized loss within a norm-defined trust region.The chosen norm determines the optimization geometry.
Muon baseline: Muon uses the spectral norm and approximates the polar decomposition of the negative gradient with Newton-Schulz iteration.Its update restricts the step to the Stiefel manifold while bounding the maximum singular-direction step.
Motivation: Applying Muon directly to raw parameters is geometrically valid only when local curvature is spherical.Mousse instead measures update size under a curvature-induced quadratic form.
Mousse framework: Mousse whitens the optimization geometry and imposes Muon’s spectral constraint in the transformed coordinate system.The whitening makes the curvature-induced quadratic form equivalent to the standard Euclidean norm.
Preconditioning: Mousse uses Shampoo’s layerwise Kronecker-factored curvature approximation, with factors estimated from exponential moving averages of GG^T and G^TG.The resulting preconditioned problem reduces to the standard Muon form in whitened coordinates.
Method outcome: Mousse applies spectral optimization in a rectified geometric frame, combining Muon’s spectral structure with Shampoo’s curvature awareness.Figure 3 reports the lowest final validation loss for Mousse across model scales from 160M to 800M parameters.

4. Experiments

Experiments evaluate Mousse against AdamW, Muon, and SOAP across GPT-2-style models from 160M to 800M parameters. Mousse improves convergence quality and sample efficiency while maintaining Muon-like training speed and low overhead.

Experimental Setup: Experiments use a modified distributed Muon library and standard decoder-only GPT-2 architectures across multiple model scales.Models are trained on FineWeb for 10,000 steps with a global batch size of 2M tokens, totaling 20 billion tokens.
Experimental Setup: Mousse is compared with AdamW, Muon, and SOAP using peak-learning-rate grid searches across model sizes and settings.Matrix-wise optimizers use Lion for embedding and lmhead modules.
Main Results: Mousse demonstrates a strict convergence advantage across all evaluated model scales.On the 800M model, it reduces final validation loss by approximately 0.012 compared with the best Muon baseline.
Main Results: Mousse requires approximately 12% fewer training steps to reach Muon’s final converged loss level.Its training speed remains nearly identical to Muon across model sizes, whereas SOAP suffers significant throughput degradation.
Efficiency: Mousse achieves lower validation loss without the computational penalty typical of curvature-corrected methods.Its memory footprint is approximately 88% of SOAP’s and about 1.05x Muon’s, while retaining superior convergence.

5. Ablation & Analysis

The ablations examine update-magnitude control, curvature correction, and preconditioner structure in Mousse. They identify grafting, trace normalization, milder spectral tempering, and single-sided whitening as stability or efficiency choices supported by the experiments.

5.1. Grafting: Gradient grafting stabilizes Mousse by combining a robust update magnitude with a sophisticated optimizer’s direction.Mousse can also use the intrinsic unit spectral norm of its orthogonalized update for step-size control.
5.1. Grafting: Without grafting, update RMS norms drift downward despite cosine scheduling, while grafting maintains stable magnitudes and improves later-training performance.The ungrafted baseline shows degradation in later training stages, whereas grafting yields superior performance and stability.
5.2. Condition Control: Negative fractional powers can amplify noisy, low-curvature directions, creating a stability challenge for Mousse’s Hessian approximations.Small eigenvalues produce large scaling factors under negative powers.
5.2. Condition Control: Trace normalization rescales L and R so their mean eigenvalues equal one, addressing layer-to-layer scale variation before decomposition.This normalization is intended to make a uniform damping value more applicable across layers.
5.2. Condition Control: α=0.125 consistently outperforms α=0.25, indicating that milder spectral tempering is preferable to aggressive curvature correction in Mousse.The stronger correction can distort update directions, while larger damping can moderate excessive steps in flat directions.
5.3. Single-Sided Preconditioner: Single-sided preconditioning halves eigen-decomposition cost and preconditioner memory while achieving comparable or slightly better performance than full Mousse.The experiments support single-sided whitening as a feasible approximation.

6. Future Works

The future-work discussion points to broader preconditioner integration and improved implementation efficiency. It also frames the current system as a proof of concept and identifies fine-tuning behavior as an open hypothesis.

6. Future Works: Mousse could integrate newer preconditioners, including Kronecker-sum variants and module-specific eigenbasis update frequencies.These directions are presented as potential extensions rather than completed components of the current work.
6.2. Better Engineering Practices: The current Mousse implementation prioritizes algorithmic correctness over hardware efficiency, leaving room for system-level optimization.Suggested alternatives include power iteration with QR decomposition and Newton-Schulz iteration instead of naive eigendecomposition.
6. Future Works: Mousse’s behavior during fine-tuning of AdamW-pretrained models remains a hypothesis tied to its second-order information and Adam-like adaptivity.The passage does not report a completed validation of this proposed transition benefit.

7. Conclusion

Mousse integrates structural curvature information into spectral optimization through whitening and preconditioning, while framing the update within a geometric descent framework. The paper also contrasts this approach with related methods and clarifies the dual-space interpretation of its optimization objective.

Conclusion: Mousse performs Newton-Schulz orthogonalization in a whitened coordinate system to align updates with anisotropic loss-landscape geometry.The method explicitly integrates structural curvature information into spectral optimization.
Conclusion: Mousse formulates optimization through a preconditioner whose curvature approximation determines the resulting geometric spectral-descent algorithm.Different approximations of H yield different algorithms under the geometric spectral-descent framework.
Conclusion: The framework uses a change of variables to express the optimization problem in whitened coordinates before solving for the update.The cited passages describe the transformation Y = √v⊙U and its inverse.
Conclusion: AdaMuon differs from Mousse in accumulating its second-moment scaling from the orthogonalized update rather than the raw gradient or momentum.The paper states that AdaMuon’s empirical success may reflect implicit regularization or stabilized magnitude control rather than curvature-aware optimization.
Conclusion: The natural pairing ⟨G, ΔW⟩ represents a dual-space functional acting on an update, rather than a metric-induced geometric inner product.For matrix parameters, this evaluation becomes a component-wise sum of products computed through the trace operator.

B. Supplementary Empirical Results

Supplementary experiments assess Mousse under the WSD schedule and examine how learning-rate-dependent excessive loss affects optimizer comparisons. The analysis argues that decay-phase loss is a more informative indicator than stable-phase loss when noise floors differ across methods.

WSD schedule: Mousse is evaluated against Muon under the WSD schedule with 10% warmup and 10% warmdown phases.The experiments inherit the main setup and report results in Figure 8.
WSD schedule: Validation-loss comparisons under WSD cover peak learning rates and model sizes from 160M to 800M parameters.Figure 8 reports final validation across the tested learning-rate settings.
Excessive loss: Excessive loss is defined as ΔL = L_stable − L_decay and is used to quantify loss released during learning-rate decay.The paper presents this quantity as a factor that can bias comparisons based only on stable-phase validation loss.
Excessive loss: Higher learning rates sustain a higher noise floor during the stable phase, which is subsequently released upon decay.This positive correlation is reported across SOAP and Mousse at various model scales.
Excessive loss: Decay-phase loss is argued to provide a more accurate indicator of optimizer effectiveness because excessive loss values are not necessarily comparable across methods.The paper notes that stable-phase comparisons can disadvantage larger learning rates by obscuring the depth of the basin found.

B.3. The Stability of Preconditioner

The preconditioner statistics vary substantially across layers, motivating trace normalization for consistent whitening. Although condition numbers stabilize early, occasional module-specific spikes reveal localized instability that remains an open robustness issue.

Preconditioner dynamics: Preconditioner matrices L and R exhibit large differences in spectral magnitude and temporal behavior across layers and architectural modules.Some layers show rapid covariance growth while others remain relatively static or decay.
Normalization: Trace normalization is needed because a fixed damping factor would otherwise affect large- and small-norm layers inconsistently.Without normalization, damping can be negligible in some layers and dominant in others, leading to unstable whitening.
Stability: Preconditioner condition numbers are high during initial statistic accumulation but quickly plateau into a stable range.This early stabilization supports consistent subsequent Newton-Schulz iterations.
Stability: Occasional condition-number spikes in specific modules indicate localized instability or sudden shifts in the gradient landscape.The paper suggests adaptive or per-layer damping as a possible future mitigation.

C.1. Pretraining GPT2

The pretraining study uses GPT-2-style models with several stability-oriented architectural choices and compares optimizers under a controlled training setup. Models span 160M to 800M parameters and are trained with specified normalization, initialization, scheduling, and optimizer configurations.

Architecture: The architecture uses RMSNorm, QK-Norm, RoPE, squared ReLU, and disabled bias terms.These choices are described as targeting stability, efficiency, attention-logit control, or memory efficiency.
Initialization: Weights use variance-scaling initialization based on fan-in and fan-out, while embedding weights use N(0, 1).The embedding initialization is intended to keep activation variances stable across widths.
Model configurations: Models range from 160M to 800M parameters, use a 1024-token context window, and have a vocabulary size of 50,304.The vocabulary is padded from 50,257 for kernel efficiency.
Model configurations: All models use the architecture described in the configuration table.The table is identified as providing precise model architecture configurations.
Optimizer settings: Mousse matches Muon’s β1, weight decay, and NS5 kernel while adopting SOAP’s preconditioner update frequency and whitening-statistics EMA.The listed Mousse settings include T = 10 and β_pc = 0.95.
Training: Training uses 8 Nvidia H200 GPUs, 10,000 steps, a global batch size of 2 Million tokens, and linear warmup followed by cosine annealing.Warmup lasts 1,000 steps before the learning rate decays to zero.