AdaGrad Meets Muon: Adaptive Stepsizes for Orthogonal Updates
This work addresses a specific optimization bottleneck for training large language models, but it is incremental as it builds on existing Muon and AdaGrad methods.
The paper tackled the problem of determining learning rates for the Muon optimizer's orthogonalized updates by proposing AdaGO, which combines AdaGrad-type adaptive stepsizes with orthogonalized momentum, resulting in improved performance over Muon and Adam on CIFAR-10 classification and function regression tasks.
The recently proposed Muon optimizer 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.