Improving matrix optimization for faster AI training with curvature insights
Beyond the Matrix Sign: Quadratic Spectral Descent
Machine LearningArtificial Intelligence
Summary
Training large AI models involves adjusting complex matrix calculations to find better solutions efficiently. The authors explore a method called Muon, which simplifies updates by focusing on key directions in the data but treats changes in a straightforward way. They propose a new technique called Quadratic Spectral Descent (QSD) that considers curvature—how the problem bends locally—leading to smarter updates. By approximating curvature cleverly and optimizing step-by-step, QSD improves training results and speeds up the process compared to prior methods.
matrix optimizationspectral normsingular value decompositionquadratic approximationlocal curvatureFrank–Wolfe algorithmKronecker factorizationgradient descentAI model training
Authors
Qiaozhe Zhang, Jun Sun, Yingzhuang Liu
Abstract
Muon can be interpreted as optimizing a linear local objective over a spectral-norm ball. This gives a matrix-sign update that preserves the singular directions of the gradient and assigns the same magnitude to all active singular modes. We ask whether these two properties remain optimal when local curvature is taken into account. To answer this question, we keep Muon's spectral-norm constraint unchanged and replace the linear local model with a quadratic one. We call the resulting method \emph{Quadratic Spectral Descent} (QSD). We show that curvature can change both the singular values and the singular directions of the optimal update. To make QSD practical, we approximate curvature with Kronecker-factored statistics and solve the constrained quadratic with a small number of Frank--Wolfe steps, each of which has a closed-form matrix-sign subproblem. We further provide an optimality certificate, a comparison with Muon under the same quadratic surrogate, and an $O(1/K)$ convergence rate for the inner solver. Experiments on GPT pre-training show that QSD consistently improves validation loss over Muon and recent Muon variants, and reduces wall-clock training time by up to $8.49\%$ at matched validation loss.