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

• URL: http://arxiv.org/abs/2603.09697v1
• Date: Tue, 10 Mar 2026 14:03:49 GMT
• Title: Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning
• Abstract summary: Recent advances in spectral optimization, notably Muon, have demonstrated that constraining update steps to the Stiefel manifold can significantly accelerate training and improveimation.<n>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.<n>We propose textbfMousse (textbfMuon textbfOptimization textbfUtilizing textbfShampoo's textbfStruct
• Score: 72.62839712454196
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• From Atoms to Entropy: Optimal Noise Allocation for Diffusion Training in the Convex Regime [52.459513910242485]
  We develop a general statistical framework for studying optimal noise-level allocation in diffusion training.<n>We show that the optimized training schedule admits an atomic minimizer, concentrated on finitely many noise levels.<n>We then evaluate the entropic schedule in larger-scale experiments, where full schedule optimization is currently intractable.
  arXiv  Detail & Related papers  (2026-07-10T20:59:06Z)
• OptMuon: Closed-Loop Orthogonalized Momentum Methods for Stochastic Optimization with Zero-Noise Optimality [23.28384210732827]
  We show that closed-loop scalar momentum.<n>optimization can be combined with Muon-style momentum.<n>optimization while retaining noise adaptivity and zero-noise optimality up to logarithmic factors.<n>These results show that OptMuon-A achieves the noise-rate.<n>(tildemathcal O(T-1/2+1/2T--1/2)) under average smoothness, while OptMuon-I achieves the noise-rate.<n>(tildemathcal O(T-1/2+
  arXiv  Detail & Related papers  (2026-06-07T18:59:24Z)
• Probabilistic Smoothing with Ratio-Monotone Transforms for Global Optimization [12.72057284593005]
  We propose a general smoothing framework that combines flexible symmetric unimodal kernels with monotonic ratio-based transformations.<n> Experiments on high-dimensional benchmarks and blackbox adversarial attacks demonstrate improved robustness and competitive performance.
  arXiv  Detail & Related papers  (2026-05-26T17:25:01Z)
• Extra-Merge: Tracing the Rank-1 Subspace of Model Merging in Language Model Pre-Training [56.323119575322146]
  Extra-Merge consistently outperforms standard merging baselines in experiments across GPT-2 and LLaMA families.<n>It yields consistent zero-shot accuracy gains on Pythia-12B downstream tasks and generalizes effectively to the Muon citepjordan2024muon.
  arXiv  Detail & Related papers  (2026-05-26T02:48:34Z)
• Move on Muon : A Hamiltonian probability gradient flow perspective of Muon optimizer [6.003944904428636]
  We develop a gradient flow on the space of probability measures defined on matrix-valued parameters induced by regularized Muon.<n>We prove an exact Hamiltonian dissipation identity, showing that the Hamiltonian energy decreases monotonically.<n>We extend the formulation to Hilbert-valued feature maps on product matrix spaces, yielding a blockwise Muon probability flow.
  arXiv  Detail & Related papers  (2026-05-22T17:28:53Z)
• Stability and Generalization of Push-Sum Based Decentralized Optimization over Directed Graphs [55.77845440440496]
  Push-based decentralized communication enables optimization over communication networks, where information exchange may be asymmetric.<n>We develop a unified uniform-stability framework for the Gradient Push (SGP) algorithm.<n>A key technical ingredient is an imbalance-aware generalization bound through two quantities.
  arXiv  Detail & Related papers  (2026-02-24T05:32:03Z)
• Muon with Spectral Guidance: Efficient Optimization for Scientific Machine Learning [10.647088281181222]
  SpecMuon is a spectral-aware, multi-mode gradient flow for physics-informed learning.<n>It regulates step sizes according to the global loss energy while preserving Muon's scale-balancing properties.<n>It achieves faster convergence and improved stability compared with Adam AdamW.
  arXiv  Detail & Related papers  (2026-02-18T03:56:20Z)
• Majorization-Minimization Networks for Inverse Problems: An Application to EEG Imaging [4.063392865490957]
  Inverse problems are often ill-posed and require optimization schemes with strong stability and convergence guarantees.<n>We propose a learned Majorization-Minimization (MM) framework for inverse problems within a bilevel optimization setting.<n>We learn a structured curvature majorant that governs each MM step while preserving classical MM descent guarantees.
  arXiv  Detail & Related papers  (2026-01-23T10:33:45Z)
• Error Feedback for Muon and Friends [80.90330715662961]
  We introduce EF21-Muon, the first communication-efficient, non-Euclidean LMO-based with rigorous convergence guarantees.<n>Our theory covers non-Euclidean smooth and the more general $(L0, L1)$-smooth setting, matching best-known Euclidean rates and enabling faster convergence under suitable norm choices.
  arXiv  Detail & Related papers  (2025-10-01T08:20:08Z)
• MPQ-DMv2: Flexible Residual Mixed Precision Quantization for Low-Bit Diffusion Models with Temporal Distillation [74.34220141721231]
  We present MPQ-DMv2, an improved textbfMixed textbfPrecision textbfQuantization framework for extremely low-bit textbfDiffusion textbfModels.
  arXiv  Detail & Related papers  (2025-07-06T08:16:50Z)
• Rao-Blackwell Gradient Estimators for Equivariant Denoising Diffusion [55.95767828747407]
  In domains such as molecular and protein generation, physical systems exhibit inherent symmetries that are critical to model.<n>We present a framework that reduces training variance and provides a provably lower-variance gradient estimator.<n>We also present a practical implementation of this estimator incorporating the loss and sampling procedure through a method we call Orbit Diffusion.
  arXiv  Detail & Related papers  (2025-02-14T03:26:57Z)
• A Riemannian Optimization Perspective of the Gauss-Newton Method for Feedforward Neural Networks [3.48097307252416]
  We analyze the convergence of Gauss-Newton dynamics for training neural networks with smooth activation functions.<n>We show that the Levenberg-Marquardt dynamics with an appropriately chosen damping schedule yields fast convergence rate despite potentially ill-conditioned neural tangent kernel matrices.
  arXiv  Detail & Related papers  (2024-12-18T16:51:47Z)
• Convergence of mean-field Langevin dynamics: Time and space discretization, stochastic gradient, and variance reduction [49.66486092259376]
  The mean-field Langevin dynamics (MFLD) is a nonlinear generalization of the Langevin dynamics that incorporates a distribution-dependent drift. Recent works have shown that MFLD globally minimizes an entropy-regularized convex functional in the space of measures. We provide a framework to prove a uniform-in-time propagation of chaos for MFLD that takes into account the errors due to finite-particle approximation, time-discretization, and gradient approximation.
  arXiv  Detail & Related papers  (2023-06-12T16:28:11Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.