Muon with Spectral Guidance: Efficient Optimization for Scientific Machine Learning

TL;DR

This paper introduces SpecMuon, a spectral-aware optimizer that enhances physics-informed neural network training by adaptively controlling spectral components, leading to faster convergence and improved stability.

Contribution

The paper proposes SpecMuon, integrating spectral decomposition with a relaxed scalar auxiliary variable mechanism, providing theoretical guarantees and superior performance over existing optimizers.

Findings

SpecMuon converges faster than Adam, AdamW, and Muon.

It offers improved stability in physics-informed neural network training.

Theoretical analysis confirms global convergence and energy dissipation properties.

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.