PaperlayerView exact revision on arXiv

Source-linked AI summary

AdaGrad Meets Muon: Adaptive Stepsizes for Orthogonal Updates

Minxin Zhang, Yuxuan Liu, Hayden Schaeffer

arXiv:2509.02981v2cs.LGmath.OC

TL;DR

Muon’s orthogonalized matrix updates leave effective learning-rate selection unresolved. AdaGO combines norm-based adaptive stepsizes with orthogonalized updates, and experiments on classification and regression show better performance than Muon and Adam.

  • Problem

    The effective learning rate for Muon and other optimizers using orthogonal updates remains unresolved.

    Source: §1 · pp. 1–2

  • Method

    AdaGO combines a norm-based AdaGrad-type stepsize with an orthogonalized update direction while preserving update orthogonality.

    Source: §1.2 · p. 4

  • Results

    Across function regression and CIFAR-10 classification, AdaGO consistently outperforms Adam and Muon.

    Source: §4.2 · p. 8

  • Takeaways & Limitations

    The analysis establishes optimal nonconvex convergence rates in stochastic and deterministic settings under smoothness and unbiased bounded-variance assumptions.

    Source: §3 · p. 5, §5 · p. 8

    The convergence analysis assumes an L-smooth loss and stochastic gradients that are unbiased estimates with uniformly bounded variance.

    Source: §3 · p. 5

Abstract

from arXiv · show

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.

1 Introduction

Muon uses matrix structure through orthogonalized momentum, but effective learning rates for orthogonal updates remain unresolved. The paper introduces AdaGO, combining norm-based adaptive stepsizes with orthogonalized updates.

1 Introduction

  • Muon leverages weight matrices through orthogonalized momentum, but effective learning rates for Muon and other orthogonal-update optimizers remain unresolved.The introduction also connects orthogonalized gradient descent with steepest descent under the spectral norm.Source: §1 · pp. 1–2
  • Large Muon learning rates can yield rapid initial loss decreases followed by plateaus and oscillations, while smaller constant rates converge more slowly but reach lower final loss.The paper motivates gradient-informed adaptive schedules for orthogonal updates.Source: §1 · p. 2

1.1 Related work

  • Prior Muon analyses establish an O(T^-1/4) stationary-point rate with a constant O(T^-3/4) stepsize; other analyses cover practical variants or a spectral-norm constrained view.The cited results are [19] [21] [26] [29], [27], and [7], respectively.Source: §1.1 · p. 3
  • Adaptive Muon variants include element-wise AdaMuon [30] and memory-efficient COSMOS [24], while Shampoo [13], ASGO [1], and PolarGrad [20] have convex-setting convergence analyses.AdaMuon and COSMOS lack theoretical convergence guarantees in the cited discussion.Source: §1.1 · p. 3

1.2 Contributions and organization

  • AdaGO combines norm-based AdaGrad stepsizes with orthogonalized updates, preserves update orthogonality, and uses one additional scalar; the paper reports nonconvex convergence guarantees and better CIFAR-10 and regression results than Muon and Adam.The convergence rates are described as optimal in stochastic and deterministic settings under standard assumptions.Source: §1.2 · pp. 3–4

2 AdaGO: A New Algorithm

AdaGO combines adaptive stepsizes based on past gradients with orthogonalized updates. Its stepsize construction clamps accumulated gradient norms and includes safeguards for stable, vanishing updates.

2 AdaGO: A New Algorithm

  • AdaGO combines stepsizes tuned by past gradients with orthogonalized updates; each parameter step uses orthogonalized momentum and an adaptive stepsize.The algorithm is summarized in Algorithm 2.Source: §2 · p. 4
  • Clamping accumulated squared gradient norms by γ removes AdaGrad-Norm’s uniformly bounded gradient-norm assumption.AdaGrad-Norm scales stochastic gradients by the reciprocal square root of accumulated squared norms.Source: §2 · p. 4
  • Initializing the accumulator with v0 > 0 prevents division by small denominators; γ and v0 affect theoretical error bounds only through logarithmic terms.The paper reports robust empirical performance across a wide range of γ values.Source: §2 · p. 4
  • Scaling orthogonalized momentum by min{∥Gt∥, γ} makes updates decay toward zero at stationary points, while a lower stepsize bound ε ensures convergence at least as fast as Muon with a small constant stepsize.The choice of ε depends on stopping time T.Source: §2 · p. 5
  • The theory permits any matrix norm for gradients, while the implementation uses the Frobenius norm for computational efficiency.Source: §2 · p. 5

3 Convergence Analysis

AdaGO’s convergence analysis assumes an L-smooth loss and unbiased stochastic gradients with bounded variance. It establishes convergence guarantees in stochastic and deterministic settings, including rates described as best possible under the stated assumptions.

3 Convergence Analysis

  • The analysis assumes an L-smooth loss and stochastic gradients that are unbiased estimates of the true gradient with uniformly bounded variance.The smoothness assumption is expressed using the nuclear and spectral norms.Source: §3 · p. 5
  • Under L-smoothness and unbiased bounded-variance stochastic gradients, AdaGO achieves the best-possible O(T^-1/4) convergence rate in the stochastic setting.The theorem uses single-sample batches, ϵ = T^-3/4, 1 − µ = T^-1/2, and η = T^-(3/8+q) for arbitrary q > 0.Source: §3 · p. 5
  • In the deterministic full-batch setting without momentum, AdaGO has an O(1/T) convergence rate, which is best possible for deterministic first-order methods under smoothness.The theorem sets ϵ = T^-1/2 and η = T^-q for arbitrary q > 0.Source: §3 · p. 6
  • Without momentum, AdaGO also converges in the stochastic setting when the batch size increases with iteration, and the result implies adaptation to gradient-noise levels.The supplied theorem passages do not provide a complete batch-size schedule or rate.Source: §3 · p. 6

4 Experiments

The experiments compare Adam, Muon, and AdaGO on function regression and CIFAR-10, with optimizer-specific hyperparameters tuned by grid search. AdaGO outperforms both baselines across the two tasks.

4.1 Experiment Setup

  • The baselines are Adam and Muon; Adam handles scalar and vector parameters, and each optimizer’s learning rate is grid-searched, with AdaGO’s epsilon also tuned.Source: §4.1 · p. 6
  • The study evaluates function regression and CIFAR-10 classification, using a two-layer MLP for regression and a five-layer CNN for classification.Source: §4.1 · p. 7

4.2 Results

  • On function regression, AdaGO reaches lower final training and test losses than Adam and Muon, despite sharp spikes from which it recovers quickly.Source: §4.2 · p. 7, Figure 1 · p. 7
  • On CIFAR-10, AdaGO maintains lower training loss and achieves higher test accuracy than Adam and Muon.Source: §4.2 · p. 8, Figure 2 · p. 8
  • Across function regression and CIFAR-10, AdaGO consistently outperforms Adam and Muon.Source: §4.2 · p. 8

5 Conclusions and Future Work

The paper introduces AdaGO, combining a norm-based AdaGrad-type stepsize with an orthogonalized update direction. It establishes nonconvex convergence rates and reports better performance than Adam and Muon on two tasks.

5 Conclusions and Future Work

  • AdaGO combines a norm-based AdaGrad-type stepsize with an orthogonalized update direction.Source: §5 · p. 8
  • AdaGO preserves orthogonality while adapting stepsizes, and requires only one additional scalar: accumulated squared gradient norms.Source: §5 · p. 8
  • The paper establishes optimal convergence rates for nonconvex functions in stochastic and deterministic settings and reports improved performance over Adam and Muon on CIFAR-10 and function regression.Source: §5 · p. 8
  • Future work includes testing AdaGO on LLM training, relaxing the analysis assumptions, and incorporating practical enhancements.Source: §5 · p. 8

Appendix A. Motivating Example: GD vs. OGD in a Linear Case

The motivating example compares gradient descent and orthogonalized gradient descent in a one-layer linear network. It illustrates why a gradient-norm-adaptive learning-rate schedule is desirable for OGD.

Appendix A. Motivating Example: GD vs. OGD in a Linear Case

  • The example uses a one-layer linear network and assumes full-rank inputs and a unique zero-loss minimizer.Source: Appendix A · p. 11
  • The example concludes that OGD benefits from a learning-rate schedule adaptively tuned to gradient norms.Source: Appendix A · p. 12
  • With constant learning rates, OGD’s large rate initially reduces training loss rapidly but later plateaus and oscillates; a smaller rate converges more slowly to a lower final loss.Source: Appendix A · p. 12, Figure 3 · p. 12

Appendix B. Useful Lemmas

Appendix B collects auxiliary lemmas and proofs used in the analysis, including a result derived from Lipschitz continuity of the gradient.

Appendix B. Useful Lemmas

  • Lemma B.1 states an implication of the Lipschitz continuity of the gradient, and its proof defines a line-segment function between two parameter values.Source: Appendix B · pp. 12–13
  • Lemma B.2 concerns arbitrary nonnegative values, and its proof notes a similar result in [34, Lemma 3.2].Source: Appendix B · p. 13
  • The proof of Lemma B.2 uses the mean value theorem to establish the stated relation for t ≥ 2.Source: Appendix B · p. 13

Appendix C. Proof of Theorem 3.3

The proof bounds the expected change using the orthogonalized direction and gradient-estimation error, then invokes conditional expectation and a recursive argument to complete the analysis.

Appendix C. Proof of Theorem 3.3

  • The argument uses conditional expectation given previous iterates and applies the law of total expectation.It defines the estimation error as M_t − ∇L(Θ_{t−1}).Source: Appendix C · p. 14
  • The proof bounds the update expression using the dual norm of the gradient-estimation error and the gradient norm.It applies successive inequalities to the inner-product expression.Source: Appendix C · p. 14
  • The proof also invokes Cauchy–Schwarz and [34, Lemma 3.2] before completing the large-T argument.Source: Appendix C · pp. 14, 16, 17
  • A recursive formula follows the proof strategy of [8, Theorem], with independently sampled minibatches assumed at each iteration.Source: Appendix C · pp. 14–15

Appendix D. Proof of Theorem 3.4

The supplied proof passage closes after referring to a condition for large T and an expression involving α_t².

Appendix D. Proof of Theorem 3.4

Appendix E. Proof of Theorem 3.5

The proof combines inequalities and considers a stepsize schedule, with an alternative setting that increases batch size over iterations.

Appendix E. Proof of Theorem 3.5

Loading 2509.02981v2…