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
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 ( β1,β2 )
–
(0.9, 0.999)
Weight Decay ( λ )
–
0
AdamW
Learning Rate ( lr )
{0.01, 5e-3, 1e-3, 5e-4}
5e-3
Betas ( β1,β2 )
–
(0.9, 0.999)
Weight Decay ( λ )
{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.
Figure 2: PINN training on the one-dimensional Burgers’ equation.
Optimizer
Final Training Loss
MSE
Training Time (s)
Adam
5.88±12\times10−4
1.42\times10−3
120±12
AdamW
7.46±15\times10−4
1.50\times10−3
117±11
Muon
3.39±5\times10−4
1.25\times10−3
149±15
SpecMuon
1.55±4\times10−4
5.06\times10−4
180±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.
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.
Figure 4: Spatial distribution of the L1 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 ( β1,β2 )
–
(0.9, 0.999)
Weight Decay ( λ )
–
0
AdamW
Learning Rate ( lr )
{0.01, 5e-3, 1e-3, 5e-4}
5e-3
Betas ( β1,β2 )
–
(0.9, 0.999)
Weight Decay ( λ )
{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 rtop increases compared with the PINN setting due to richer gradient structure.
Figure 5: DeepONet training on the Burgers’ equation.
Optimizer
Final Training Loss
MSE
Training Time (s)
Adam
1.89±03\times10−3
0.0014
113±12
AdamW
4.71±05\times10−3
0.0013
114±10
Muon
6.38±12\times10−4
0.0011
204±13
SpecMuon
4.82±07\times10−4
0.0010
259±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 ( β1,β2 )
–
(0.9, 0.999)
Weight Decay ( λ )
–
0
AdamW
Learning Rate ( lr )
{0.01, 5e-3, 1e-3, 5e-4}
5e-3
Betas ( β1,β2 )
–
(0.9, 0.999)
Weight Decay ( λ )
{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 rtop was tuned for stability and accuracy.
Figure 6: fPINN–DeepONet training on the fractional PDE problem.
Optimizer
Training Loss
MSE
Training Time (s)
Adam
3.72±04\times10−2
0.0268
19.41
AdamW
3.64±04\times10−2
0.0275
13.37
Muon
3.52±06\times10−2
0.0272
18.24
SpecMuon
2.31±05\times10−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.
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.
Yechen Zhang, Shuhao Xing, Junhao Huang +5
Shanghai Jiao Tong University · Shanghai AI Laboratory · Fudan University
Matrix-valued optimizers such as Muon exploit the spectral structure of neural network updates through Newton--Schulz orthogonalization, but their near-flattening of the singular spectrum discards relative magnitude information across gradient modes. We introduce Soren (Spectral Orthogonal Reshaping), a matrix-valued optimizer that preserves the singular subspaces of the gradient while applying a bounded, monotone sigmoid transformation to its singular values. This smoothly compresses dominant modes without fully flattening the spectrum. We interpret Soren as a positive-definite preconditioned gradient method and establish convergence guarantees under relative smoothness and metric Polyak--Łojasiewicz geometry. To avoid explicit singular value decomposition, we further develop a finite-depth Soft Newton--Schulz (SNS) polynomial realization of the sigmoid spectral map and characterize how its spectral approximation affects the induced convergence geometry. Experiments across LLM pre-training, supervised fine-tuning, and direct preference optimization demonstrate the effectiveness and robustness of Soren against established optimizers.
Dat Phi Van, Ngo Vu Minh, Tuc Nguyen +3
Center for AI Research, VinUniversity, Hanoi, Vietnam · Indiana University, Bloomington, IN, USA · Monash University, Clayton, VIC, Australia
Muon orthogonalizes the momentum buffer before each update, replacing its singular values with ones via Newton-Schulz iterations. This simple change lets Muon tolerate far larger learning rates and converge faster than other optimizers, but why? We show that the mechanism is spectral flattening, and develop two results around it. First, we prove that Muon's maximal stable step size scales with the average singular value of the gradient rather than the largest, which bottlenecks standard gradient descent. Second, we recast Muon as a preconditioned gradient method and show, under a Kronecker-factored curvature model, that it improves the effective convergence factor, with the improvement controlled by the spectrum of the gradient covariance. Extensive experiments validate both results: Muon remains stable at learning rates that cause SGD to diverge within the first few iterations, and reaches accuracy milestones several epochs earlier even at identical step sizes. Taken together, our results offer a principled, geometric explanation for Muon's empirical success.
Hanoi University of Science and Technology Hanoi, Vietnam · Indiana University 107 S. Indiana Ave, Bloomington, IN 47405, USA · Monash University Clayton, VIC 3800, Australia