PaperlayerView exact revision on arXiv

Source-linked AI summary

Nonsmooth Optimization via Orthogonalized Momentum

Lexiao Lai, Tianyi Lin, Jiayu Zhang

arXiv:2609.13677v1math.OCcs.LG

TL;DR

Matrix-aware orthogonalized momentum is empirically successful, but its reliability for locally Lipschitz objectives is unclear. This paper analyzes that setting with conservative fields and proposes coupled momentum schedules and MAGD. Fixed momentum can fail, while the proposed approaches yield asymptotic convergence and convex finite-time guarantees.

  • Problem

    Whether orthogonalized momentum remains reliable for locally Lipschitz objectives is unclear, since existing convergence theory described here is tied to smoothness or generalized smoothness.

    Source: §1 · p. 4

  • Method

    The paper analyzes orthogonalized momentum with conservative fields compatible with backpropagation and proposes coupling momentum to the step size and combining the Muon direction with a progress-weighted gradient branch in MAGD.

    Source: §1 · pp. 2–3

  • Results

    For every fixed β in [0, 1), Muon can fail on a convex Lipschitz objective, whereas coupled momentum restores asymptotic convergence and MAGD achieves O(min{m, n}ϵ^-2) convex complexity with optimal dimension dependence.

    Source: §1 · p. 3, §5 · pp. 26–27

  • Takeaways & Limitations

    The results distinguish fixed from coupled momentum schedules, and experiments on synthetic problems, image classification, and LLM pretraining show MAGD safeguards orthogonalized updates while retaining empirical performance.

    Source: §5 · pp. 26–27

    The asymptotic analysis requires vanishing polar approximation error, while practical implementations use a fixed number of iterations.

    Source: §5 · p. 27

Abstract

from arXiv · show

Modern real application problems involve matrix-valued parameters, yet conventional optimizers treat them as vectors, thereby motivating matrix-aware methods that exploit input-output geometry, such as Muon which orthogonalizes the momentum matrices before parameter updates. Its empirical success raises a conceptual question: can orthogonalized momentum remain effective beyond smooth optimization? This paper studies this question for locally Lipschitz functions using a generalized derivative framework compatible with backpropagation. Our first contribution is to identify a key limitation: for every fixed momentum factor $\beta\in[0,1)$, Muon can fail to approach the global optimal solution of a convex Lipschitz objective from almost every initialization, when step sizes adapt to the full gradient history. The failure can occur even along bounded iterates. Our example is inspired by the one of Parshakova et al. which only covers $\beta\in[0,\frac{1}{2})$. Then, we show that the obstruction lies in fixed momentum rather than orthogonalization. Indeed, when the momentum factor is adaptive and approaches 1 together with a vanishing step size, Muon recovers asymptotic convergence for nonconvex nonsmooth optimization under the boundedness and regularity conditions. Moreover, we propose MAGD, which combines orthogonalized momentum with gradient, weighted based on their relative progress. MAGD retains asymptotic convergence in nonconvex settings and achieves an $O(\min\{m,n\}\epsilon^{-2})$ rate in convex settings. A lower bound shows the optimal dimension dependence. Experiments on synthetic problems, image classification, and LLM pretraining show MAGD is a simple and practical alternative to Muon. Together, our results characterize when orthogonalized momentum fails without smoothness and how it can be made reliable and we hope that the analysis may be useful more broadly.

1 Introduction

The introduction frames orthogonalized momentum as a matrix-aware alternative to vectorized first-order methods, then asks whether it remains reliable for locally Lipschitz objectives. It distinguishes failure under fixed momentum from convergence under coupled adaptive schedules and introduces MAGD for finite-time convex guarantees.

  • Matrix-aware methods use operator-norm geometry to respect matrices’ input-output structure, unlike vectorized Euclidean or coordinate-wise methods.Neural-network layers are a prominent application, while the underlying problem is broader matrix-norm optimization.Source: §1 · p. 1
  • Muon forms a momentum matrix, orthogonalizes it, and uses the result for the parameter update.Inexact orthogonalization can preserve singular subspaces while pushing nonzero singular values toward 1.Source: §1 · p. 2
  • For locally Lipschitz objectives, abrupt changes in first-order information and magnitude removal by orthogonalization make convergence unclear.Existing convergence analyses cited here rely on smoothness or generalized smoothness.Source: §1 · pp. 2, 4
  • Coupling momentum to a vanishing step size restores asymptotic convergence for inexact Muon on convex objectives under controlled approximation.The coupled updates evolve on the same time scale, yielding a dissipative Lyapunov structure.Source: §1 · p. 3
  • For every fixed β in [0, 1), Muon can fail to approach a convex Lipschitz objective’s global optimum from almost every initialization, even with step sizes adapted to full gradient history.The failure can occur along bounded iterates.Source: §1 · p. 3
  • MAGD combines orthogonalized momentum with a progress-weighted gradient branch, with an O(min{m, n}ϵ^-2) convex rate and matching dimension dependence.Experiments cover synthetic problems, image classification, and LLM pretraining; the lower bound uses hidden directions in operator-norm geometry.Source: §1 · pp. 3–4

2 Preliminaries

The preliminaries establish the nonsmooth calculus, matrix geometry, inexact Muon update, and oracle model used in the analysis. They specify conditions that accommodate backpropagation and formalize convex optimization complexity in operator-norm geometry.

2.1 Nonsmooth analysis and conservative fields

  • Conservative fields capture generalized first-order information compatible with backpropagation, and their calculus supplies a chain rule for locally Lipschitz objectives.The analysis assumes a convex-valued conservative field satisfying the weak Sard condition; the cited construction covers standard neural-network architectures.Source: §2.1 · pp. 5–6

2.2 Muon and operator-norm geometry

  • Muon is studied in operator-norm geometry, where the nuclear norm is its dual under the Frobenius inner product.The inexact recursion updates momentum from a field value, normalizes it, approximately orthogonalizes it, and steps the parameter; the exact update minimizes a linear model over an operator-norm ball.Source: §2.2 · p. 6
  • Inexact orthogonalization approximates the polar update while preserving nonzero singular subspaces and modifying singular values.Newton-Schulz and PolarExpress are described as finite polynomial-iteration subroutines.Source: §2.2 · p. 6

2.3 Oracle complexity of convex nonsmooth optimization

  • The black-box model counts local first-order oracle queries needed to return an ϵ-optimal point for a convex Lipschitz objective.A distributional complexity lower bound also lower-bounds worst-case complexity by Yao’s minimax principle.Source: §2.3 · p. 7

3 Main Results

This section analyzes convergence guarantees and limitations for Muon and MAGD, including how adaptive momentum and vanishing step sizes affect asymptotic convergence. It also establishes convex oracle-complexity bounds in operator-norm geometry and shows that their dimension dependence is optimal.

3.1 Convergence behavior of Muon

  • For every fixed β ∈ [0, 1), Muon can remain suboptimal from almost every initialization on a convex Lipschitz objective, even with adaptive step sizes based on full gradient history.The failure example has bounded iterates whenever the step sizes are bounded.Source: §3.1 · pp. 8, 10
  • With βt = 1 − τηt, vanishing ηt and polar error, and bounded iterates, inexact Muon has stationary cluster points and convergent objective values.The theorem also assumes a nonsummable step-size sequence and bounded σt; under Assumption 2.3, the objective converges to its global infimum.Source: §3.1 · pp. 10, 12

3.2 Muon advised by gradient descent

  • MAGD maintains Muon and gradient-descent branches, mixes their iterates using linearized-loss decreases, and favors the branch with larger predicted decrease.The method is designed to stay close to Muon when the Muon branch performs better.Source: §3.2 · p. 13
  • In two regimes, MAGD has stationary cluster points and convergent objective values; regime (ii) additionally yields Mt → 0.Regime (i) requires the gradient-descent branch to asymptotically dominate, while regime (ii) allows neither branch to dominate if Muon momentum approaches one.Source: §3.2 · p. 13
  • MAGD’s objective values converge, and every cluster point W satisfies 0 ∈ Df(W).The proof also establishes that every cluster point of the momentum matrices is zero.Source: §3.2 · pp. 15–16
  • Theorem 3.5 guarantees that MAGD attains an ϵ-optimal function value in O(min{m, n}(GR/ϵ)^2) subgradient evaluations under Assumption 2.3.Source: §3.3 · p. 18

3.3 Oracle complexity in operator-norm geometry

  • The O(min{m, n}) dimension factor is unavoidable for first-order methods in operator-norm geometry when max{m, n} is sufficiently large.The lower-bound argument reduces optimization of a max-of-affine function encoding a hidden sign vector to string guessing [Braun et al., 2017, Oracle III.2 and Proposition III.3].Source: §3.3 · pp. 18–19
  • O(min{m,n}(GR/epsilon)^2) upper and matching Ω(min{m,n}(GR/epsilon)^2) lower bounds establish the oracle complexity in operator-norm geometry.Together, the bounds give the optimal min{m,n} dimension dependence.Source: §3.3 · p. 21
  • The lower bound reduces optimization to sign guessing: any local-oracle algorithm induces a procedure with error at most 1/8, forcing Ω(min{m,n}δ^-2) expected queries.The hard family has M proportional to min{m,n}δ^-2, and the sign-guessing lower bound yields the query bound.Source: §3.3 · p. 21
  • Rescaling transfers the unit-radius lower bound to general G and R, with an ϵ-optimal output corresponding to an (ϵ/(GR))-optimal unit-scale output.The reduction maps matrices in the unit operator-norm ball to radius-R matrices and rescales oracle replies.Source: §3.3 · p. 22

4 Experiments

Experiments compare Muon and MAGD with other methods on synthetic hard instances, quadratic optimization, CNN classification, and LLM pretraining, using grid-searched hyperparameters and matched Muon schedules where specified.

4.1 Nonsmooth hard instance

  • For δ = 0, fixed-momentum Muon fails to reach an objective gap below 10^-6 on all runs, while adaptive Muon, GD, and MAGD perform better.Source: §4.1 · p. 23, Table 1 · p. 23

4.2 Quadratic optimization

  • MAGD beats the stronger of GD and Muon on five of seven quadratic spectra, while retaining some stronger-baseline gains on the two single-outlier spectra.Source: §4.2 · pp. 24–25, Figure 2 · p. 24

4.3 Convolutional neural network

  • On the nonsmooth CNN task, MAGD is competitive with Muon, improves with more accurate approximation, and both outperform SGDM.Source: §4.3 · p. 25, Table 3 · p. 25

4.4 LLM pretraining

5 Conclusion

The conclusion distinguishes fixed from coupled momentum schedules for nonsmooth optimization and presents MAGD as a method with asymptotic and finite-time guarantees. It also identifies open gaps between the theory and practical orthogonalization.

5 Conclusion

  • For every fixed β ∈ [0, 1), Muon can fail to approach a convex Lipschitz objective’s global optimum from almost every initialization, whereas β_t = 1 − τη_t restores asymptotic convergence under boundedness and vanishing polar error.Source: §5 · p. 26
  • MAGD combines inexact Muon and gradient-descent branches, retains stationary cluster points and convergent objective values for nonconvex problems, and reaches an ε-optimal convex iterate in O(min{m, n}ε^-2) iterations.Source: §5 · pp. 26–27
  • The convex iteration bound has optimal dependence on min{m, n} and ε, as shown by the paper’s lower bound.Source: §5 · p. 27
  • The asymptotic analysis requires vanishing polar approximation error, unlike practical implementations that use a fixed number of orthogonalization iterations.Source: §5 · p. 27
Loading 2609.13677v1…