AdaGrad Meets Muon: Adaptive Stepsizes for Orthogonal Updates

Minxin Zhang, Yuxuan Liu, Hayden Schaeffer

Keywords: AdaGrad, Muon optimizer, orthogonalized momentum, adaptive stepsizes, nonconvex stochastic optimization, convergence rates

TL;DR: We propose AdaGO, an optimizer that combines Muon’s orthogonalized updates with norm-based AdaGrad-type stepsizes, achieving optimal convergence rates and strong empirical performance.

Abstract: The recently proposed optimizer Muon updates weight matrices via orthogonalized momentum and has demonstrated strong empirical success in large language model training. However, it remains unclear how to determine the learning rates for such orthogonalized updates. AdaGrad, by contrast, is a widely used adaptive method that scales stochastic gradients by accumulated past gradients. We propose a new algorithm, AdaGO, which combines a norm-based AdaGrad-type stepsize with an orthogonalized update direction, bringing together the benefits of both approaches. Unlike other adaptive variants of Muon, AdaGO preserves the orthogonality of the update direction, which can be interpreted as a spectral descent direction, while adapting the stepsizes to the optimization landscape by scaling the direction with accumulated past gradient norms. The implementation of AdaGO requires only minimal modification to Muon, with a single additional scalar variable, the accumulated squared gradient norms, to be computed, making it computationally and memory efficient. Optimal theoretical convergence rates are established for nonconvex functions in both stochastic and deterministic settings under standard smoothness and unbiased bounded-variance noise assumptions. Empirical results on CIFAR-10 classification and function regression demonstrate that AdaGO outperforms Muon and Adam.

Submission Number: 99