Muon with Spectral Guidance: Efficient Optimization for Scientific Machine Learning

Binghang Lu; Jiahao Zhang; Guang Lin

cs.LGFeb 18, 2026

Muon with Spectral Guidance: Efficient Optimization for Scientific Machine Learning

Authors: Binghang Lu, Jiahao Zhang, Guang Lin

Organizations: Department of Mathematics, Purdue University, West Lafayette, IN 47907, USA · Department of Mathematics & School of Mechanical Engineering, Purdue University, West Lafayette, IN 47907, USA

Abstract

Physics-informed neural networks and neural operators often suffer from severe optimization difficulties caused by ill-conditioned gradients, multi-scale spectral behavior, and stiffness induced by physical constraints. Recently, the Muon optimizer has shown promise by performing orthogonalized updates in the singular-vector basis of the gradient, thereby improving geometric conditioning. However, its unit-singular-value updates may lead to overly aggressive steps and lack explicit stability guarantees when applied to physics-informed learning. In this work, we propose SpecMuon, a spectral-aware optimizer that integrates Muon's orthogonalized geometry with a mode-wise relaxed scalar auxiliary variable (RSAV) mechanism. By decomposing matrix-valued gradients into singular modes and applying RSAV updates individually along dominant spectral directions, SpecMuon adaptively regulates step sizes according to the global loss energy while preserving Muon's scale-balancing properties. This formulation interprets optimization as a multi-mode gradient flow and enables principled control of stiff spectral components. We establish rigorous theoretical properties of SpecMuon, including a modified energy dissipation law, positivity and boundedness of auxiliary variables, and global convergence with a linear rate under the Polyak-Lojasiewicz condition. Numerical experiments on physics-informed neural networks, DeepONets, and fractional PINN-DeepONets demonstrate that SpecMuon achieves faster convergence and improved stability compared with Adam, AdamW, and the original Muon optimizer on benchmark problems such as the one-dimensional Burgers equation and fractional partial differential equations.

Figures & tables

Figure 1: Training loss (log scale) for the linear regression problem. SpecMuon (red) is compared with Muon, Adam, and AdamW. SpecMuon achieves faster convergence by combining spectral orthogonalization with mode-wise energy regulation.

Optimizer	Hyperparameter	Values Searched	Best Value Used
Adam	Learning Rate ( lr )	{0.01, 5e-3, 1e-3, 5e-4}	5e-3
	Betas ( $\beta_{1},\beta_{2}$ )	–	(0.9, 0.999)
	Weight Decay ( $\lambda$ )	–	0
AdamW	Learning Rate ( lr )	{0.01, 5e-3, 1e-3, 5e-4}	5e-3
	Betas ( $\beta_{1},\beta_{2}$ )	–	(0.9, 0.999)
	Weight Decay ( $\lambda$ )	{0.0, 0.01, 1e-3, 5e-4, 1e-4 }	5e-4

Table 1: Hyperparameter configurations for the PINN experiment on Burgers’ equation. Learning rates were selected via grid search. Other parameters were set to standard defaults unless otherwise specified.

Panel of Figure 2 — Figure 2: PINN training on the one-dimensional Burgers’ equation.

Optimizer	Final Training Loss	MSE	Training Time (s)
Adam	$5.88\pm 12\text{\times}{10}^{-4}$	$1.42\text{\times}{10}^{-3}$	$120\pm 12$
AdamW	$7.46\pm 15\text{\times}{10}^{-4}$	$1.50\text{\times}{10}^{-3}$	$117\pm 11$
Muon	$3.39\pm 5\text{\times}{10}^{-4}$	$1.25\text{\times}{10}^{-3}$	$149\pm 15$
SpecMuon	$1.55\pm 4\text{\times}{10}^{-4}$	$5.06\text{\times}{10}^{-4}$	$180\pm 18$

Table 2: Quantitative comparison for the Burgers’ PINN experiment. Reported metrics include final training loss, mean squared error (MSE), and total training time on an NVIDIA H100 80GB GPU. Best results are shown in bold.

Panel of Figure 3 — Figure 3: Spatiotemporal solution $u(x,t)$ of the Burgers’ equation obtained by PINNs trained with Adam(left), Muon(middle), and SpecMuon(right). SpecMuon produces a more accurate and smoother approximation of the reference solution.

Panel of Figure 4 — Figure 4: Spatial distribution of the $L_{1}$ error for the Burgers’ PINN trained with Adam(left), Muon(middle), and SpecMuon(right). SpecMuon exhibits reduced error magnitude.

Optimizer	Hyperparameter	Values Searched	Best Value Used
Adam	Learning Rate ( lr )	{0.01, 5e-3, 1e-3, 5e-4}	5e-3
	Betas ( $\beta_{1},\beta_{2}$ )	–	(0.9, 0.999)
	Weight Decay ( $\lambda$ )	–	0
AdamW	Learning Rate ( lr )	{0.01, 5e-3, 1e-3, 5e-4}	5e-3
	Betas ( $\beta_{1},\beta_{2}$ )	–	(0.9, 0.999)
	Weight Decay ( $\lambda$ )	{0.0, 0.01, 1e-3, 5e-4, 1e-4}	1e-2

Table 3: Hyperparameter configurations for the DeepONet experiment on Burgers’ equation. Learning rates were selected via grid search. The optimal spectral rank

r_{top}

increases compared with the PINN setting due to richer gradient structure.

Panel of Figure 5 — Figure 5: DeepONet training on the Burgers’ equation.

Optimizer	Final Training Loss	MSE	Training Time (s)
Adam	$1.89\pm 03\text{\times}{10}^{-3}$	$0.0014$	$113\pm 12$
AdamW	$4.71\pm 05\text{\times}{10}^{-3}$	$0.0013$	$114\pm 10$
Muon	$6.38\pm 12\text{\times}{10}^{-4}$	$0.0011$	$204\pm 13$
SpecMuon	$4.82\pm 07\text{\times}{10}^{-4}$	$0.0010$	$259\pm 10$

Table 4: Quantitative comparison for the Burgers’ DeepONet experiment. Metrics include final training loss, MSE, and training time on an NVIDIA H100 80GB GPU. Best results are highlighted in bold.

Optimizer	Hyperparameter	Values Searched	Best Value Used
Adam	Learning Rate ( lr )	{0.01, 5e-3, 1e-3, 5e-4}	5e-3
	Betas ( $\beta_{1},\beta_{2}$ )	–	(0.9, 0.999)
	Weight Decay ( $\lambda$ )	–	0
AdamW	Learning Rate ( lr )	{0.01, 5e-3, 1e-3, 5e-4}	5e-3
	Betas ( $\beta_{1},\beta_{2}$ )	–	(0.9, 0.999)
	Weight Decay ( $\lambda$ )	{0.0, 0.01, 1e-3, 5e-4, 1e-4}	1e-2

Table 5: Hyperparameter configurations for the fPINN–DeepONet experiment on the two-dimensional heat equation. Learning rates were selected via grid search, and spectral rank

r_{top}

was tuned for stability and accuracy.

Panel of Figure 6 — Figure 6: fPINN–DeepONet training on the fractional PDE problem.

Optimizer	Training Loss	MSE	Training Time (s)
Adam	$3.72\pm 04\text{\times}{10}^{-2}$	$0.0268$	$19.41$
AdamW	$3.64\pm 04\text{\times}{10}^{-2}$	$0.0275$	$13.37$
Muon	$3.52\pm 06\text{\times}{10}^{-2}$	$0.0272$	$18.24$
SpecMuon	$2.31\pm 05\text{\times}{10}^{-2}$	$0.0262$	$30.74$

Table 6: Quantitative comparison for the fPINN–DeepONet experiment. Training loss is reported at 500 epochs. Training time was measured on an NVIDIA H100 80GB GPU. Best results are highlighted in bold.

Muon Optimizer

Muon optimizer uses momentum and matrix orthogonalization to shape neural-network updates, especially in large language-model training.

See all 149 papers

Mousse: Rectifying the Geometry of Muon with Curvature-Aware Preconditioning
Mar 10, 2026Yechen Zhang, Shuhao Xing, Junhao Huang +5Deep Learning Optimization Muon Optimizer
Shanghai Jiao Tong University · Shanghai AI Laboratory · Fudan University
Orthogonal Witness Control for Muon Optimization via Sigmoid Spectral Reshaping
Sep 27, 2026Dat Phi Van, Ngo Vu Minh, Tuc Nguyen +3Neural Network Optimization Spectral Optimization
Center for AI Research, VinUniversity, Hanoi, Vietnam · Indiana University, Bloomington, IN, USA · Monash University, Clayton, VIC, Australia
Spectral Flattening Is All Muon Needs: How Orthogonalization Controls Learning Rate and Convergence
May 13, 2026Tien-Phat Nguyen, Truong Nguyen, Minh-Phuc Truong +3Momentum Methods Gradient Descent
Hanoi University of Science and Technology Hanoi, Vietnam · Indiana University 107 S. Indiana Ave, Bloomington, IN 47405, USA · Monash University Clayton, VIC 3800, Australia

Mar 10, 2026J/K move · Enter open · S save

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

Yechen Zhang, Shuhao Xing, Junhao Huang, Kai Lv+4

Shanghai Jiao Tong University · Shanghai AI Laboratory · Fudan University

Deep Learning Optimization Muon Optimizer Stiefel Manifold Optimization Trust-Region Optimization

Open paper AskPDF arXiv

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

\sim

12% reduction in training steps with negligible computational overhead.

Abstract

Muon Optimizer

Explore similar work

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

Orthogonal Witness Control for Muon Optimization via Sigmoid Spectral Reshaping

Spectral Flattening Is All Muon Needs: How Orthogonalization Controls Learning Rate and Convergence