Spectral reshaping improves neural network optimization with Muon
Orthogonal Witness Control for Muon Optimization via Sigmoid Spectral Reshaping
Machine Learning
Summary
Training large neural networks involves adjusting many parameters efficiently. Previous methods like Muon try to speed up this process by using special math that treats neural network updates as matrices but lose some important info about their strength. The authors propose Soren, a new method that smoothly adjusts these updates’ strengths without losing key info, improving learning stability and speed. They also create a way to approximate this approach efficiently without heavy computation and show it works well in different training tasks.
What this means in practice
- •For machine learning engineers: Improve training speed and stability of large neural networks using the Soren optimizer for better parameter updates.
- •For ai model fine-tuning teams: Enhance supervised fine-tuning and preference optimization workflows by incorporating Soren to handle gradient updates more effectively.
Authors
Dat Phi Van, Ngo Vu Minh, Tuc Nguyen, Thin Nguyen, Ngoc-Thanh Dinh, Trung Le
Abstract
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 \emph{Soren} (\textbf{S}pectral \textbf{O}rthogonal \textbf{Re}shapi\textbf{n}g), 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.