PaperlayerView exact revision on arXiv

Source-linked AI summary

Muon Optimizes Under Spectral Norm Constraints

Lizhang Chen, Jonathan Li, Qiang Liu

arXiv:2506.15054v2cs.LGmath.OCstat.ML

TL;DR

Muon has promising empirical performance but lacks a clear theoretical foundation. This paper embeds it in Lion-K, identifies its nuclear-norm structure, and shows that decoupled weight decay yields an implicit spectral-norm-constrained problem with convergence to KKT points. The framework also supports generalizations through alternative convex maps.

  • Problem

    Muon has promising empirical performance, but its theoretical underpinnings and connections to broader optimization techniques remain unclear.

  • Method

    The paper analyzes Muon by embedding it in Lion-K and choosing the nuclear norm with the matrix sign function as its subgradient.

  • Results

    With decoupled weight decay, Muon implicitly solves a spectral-norm-constrained optimization problem and converges to its KKT points.

  • Takeaways & Limitations

    The Lion-K perspective provides a theoretically grounded interpretation of Muon and enables generalizations using alternative convex functions with different implicit regularization effects.

  • Takeaways & Limitations

    Further work is needed for broader convex maps, practical conditions such as general learning rates and nonsmooth objectives, and larger models and more diverse tasks.

Abstract

from arXiv · show

The pursuit of faster optimization algorithms remains an active and important research direction in deep learning. Recently, the Muon optimizer [JJB+24] has demonstrated promising empirical performance, but its theoretical foundation remains less understood. In this paper, we bridge this gap and provide a theoretical analysis of Muon by placing it within the Lion-$\mathcal{K}$ family of optimizers [CLLL24]. Specifically, we show that Muon corresponds to Lion-$\mathcal{K}$ when equipped with the nuclear norm, and we leverage the theoretical results of Lion-$\mathcal{K}$ to establish that Muon (with decoupled weight decay) implicitly solves an optimization problem that enforces a constraint on the spectral norm of weight matrices. This perspective not only demystifies the implicit regularization effects of Muon but also leads to natural generalizations through varying the choice of convex map $\mathcal{K}$, allowing for the exploration of a broader class of implicitly regularized and constrained optimization algorithms.

1 Introduction

The paper places Muon within the Lion-K framework to clarify its theoretical basis and implicit regularization. It identifies Muon as a nuclear-norm Lion-K optimizer, derives spectral constraints with decoupled weight decay, establishes convergence, and proposes generalizations through other convex maps.

  • Motivation: Muon introduces orthogonalized gradient-momentum updates, but its theoretical underpinnings and connections to broader optimization techniques remain unclear.The updates use Newton-Schulz iteration and have shown promising empirical results for efficient large-scale training.
  • Contribution: The paper embeds Muon within Lion-K to provide a theoretical explanation and a unified basis for generalizations.Lion-K also unifies techniques including mirror descent, Nesterov momentum, Hamiltonian descent, Frank–Wolfe algorithms, and decoupled weight decay.
  • Contribution: Muon is a special case of Lion-K when K is the nuclear norm and its subgradient is the matrix sign function.This identifies the convex map and subgradient underlying Muon’s update.
  • Contribution: With decoupled weight decay, Muon implicitly solves an optimization problem constraining the singular values of weight matrices.The paper interprets this behavior as spectral regularization.
  • Contribution: Lion-K guarantees yield convergence rates for Muon with deterministic and stochastic gradients, including convergence with decoupled weight decay to KKT points.The deterministic and stochastic cases are stated as Theorems 3 and 4, while the KKT results are Theorems 5 and 6.
  • Contribution: Varying the convex function beyond the nuclear norm generalizes Muon into optimizers with different implicit regularization effects.The paper presents this as a richer family of optimization algorithms and directions for future work.

2 Related work

Prior norm-constraint interpretations cover several optimizers but do not naturally handle momentum or decoupled weight decay. Related work studies decoupled regularization and Muon convergence, while this paper addresses these gaps through the Lion-K framework.

  • Steepest descent under norm constraints: Norm-constraint analyses reinterpret Adam, Shampoo, and Muon as steepest-descent methods under norm constraints.The same perspective has also been applied to a stochastic conditional-gradient approach.
  • Steepest descent under norm constraints: Introducing momentum breaks the strict steepest-descent interpretation, and the perspective does not naturally extend to decoupled weight decay.These are identified as fundamental limitations of the norm-constraint perspective.
  • Decoupled weight decay: Decoupled weight decay separates regularization from adaptive gradient computations and is reported to improve training stability and generalization.The passage describes this approach as standard in modern adaptive optimizers.
  • Concurrent work: Concurrent Muon convergence studies do not consider decoupled weight decay, while this work proves convergence for that variant using Lion-K and decreasing Robbins–Monro step sizes.One concurrent study analyzes decoupled weight decay through a Frank–Wolfe perspective.

3 Main results

The paper identifies Muon as a nuclear-norm instance of Lion-K and shows that, with decoupled weight decay, it implicitly targets spectral-norm-constrained optimization. It establishes convergence to KKT points and explains the dynamics through Lyapunov functions.

  • Generalizations: The Lion-K interpretation generalizes Muon by replacing the nuclear norm with other convex functions, producing optimizers with different implicit regularization effects.This broadens the family of implicitly regularized and constrained optimization algorithms.
  • Muon within Lion-K: Muon is Lion-K with K(X) = ∥X∥tr and subgradient ∇K(X) = msgn(X).The nuclear norm equals the sum of singular values, and msgn is a subgradient of the nuclear norm.
  • Implicit constrained optimization: With decoupled weight decay, Muon implicitly solves a constrained optimization problem rather than minimizing the original loss directly.The corresponding regularized objective uses the convex conjugate K∗ and yields an implicit constraint formulation.
  • Implicit constrained optimization: The induced constraint is on the spectral norm of weight matrices, whose value is the largest singular value and which is dual to the nuclear norm.The spectral norm is defined as ∥X∥op = σ1(X).
  • Convergence analysis: Muon converges to the KKT points of the implicit constrained problem under deterministic and stochastic gradient updates.The paper states an almost-sure convergence result under its conditions, regardless of initialization.
  • Convergence analysis: Outside the constraint ball, VB decreases exponentially and drives iterates rapidly toward the ball; inside it, VK decreases approximately monotonically until the KKT score vanishes.These Lyapunov functions explain the two-stage convergence behavior summarized in Figure 1.

4 Preliminaries

The preliminaries establish the matrix-space notation, convex-analysis tools, norm duality, and assumptions used for deterministic and stochastic optimization analysis.

  • Notation and norms: The paper studies real n × m matrices with the Frobenius inner product and standard matrix norms, including nuclear, Frobenius, and spectral norms.The spectral norm is the dual of the nuclear norm, while the Frobenius norm is self-dual.
  • Convex analysis: For convex functions, ∇K denotes a gradient or selected subgradient, and K* denotes the convex conjugate.Convex conjugacy supports the Fenchel–Young inequality and the characterization of subgradients.
  • Convex analysis: The conjugate of a norm is finite on the unit ball of its dual norm and infinite outside it.This property later converts regularization through K* into a constraint.
  • Assumptions: The analysis assumes that F has a finite attained infimum and is L-smooth.These conditions make the optimization problem well-posed and support the discrete-time analysis.
  • Assumptions: For stochastic optimization, the paper introduces matrix variance and assumes independent identically distributed samples with bounded gradient variance.An additional iteration-wise bounded-variance assumption is also used.
  • Assumptions: The framework requires K to be convex, which also makes it closed and proper.The paper notes that the smoothness and variance assumptions are standard in stochastic optimization analyses.

5 Background on Lion-K

Lion-K extends Lion through a convex-map preconditioner and provides a framework for understanding regularization and constrained optimization induced by decoupled weight decay.

  • Lion-K update: Lion-K combines Polyak momentum, Nesterov momentum, nonlinear preconditioning ∇K, and decoupled weight decay.The two momentum terms are controlled by β2 and β1, respectively.
  • Lion-K update: Convexity makes ∇K a monotone map, legitimizing its use as a nonlinear transformation of momentum.Monotonicity is expressed through a nonnegative inner-product relation between differences of mapped and original variables.
  • Implicit regularization: Decoupled weight decay and ∇K cause Lion-K optimizers to minimize a regularized objective rather than the original loss.A fixed-point argument connects the update condition to stationarity of the regularized objective.
  • Continuous-time analysis: A Lyapunov-function analysis shows that the auxiliary Hamiltonian decreases along Lion-K ODE trajectories until a local minimum is reached.The Hamiltonian combines the regularized objective with a momentum-dependent term.
  • Continuous-time analysis: Minimizing the auxiliary Hamiltonian is equivalent to minimizing the regularized objective through Fenchel–Young equality.Equality occurs when M ∈ ∂K(λX).
  • Constrained optimization: When K* has infinite values, Lion-K effectively solves a constrained optimization problem over dom(K*).Continuous-time dynamics drive λX rapidly into this effective domain and keep it there.

6 Muon meets Lion-K

Muon is identified as a nuclear-norm Lion-K optimizer, and this correspondence explains its spectral-norm constraint, convergence framework, and generalizations.

  • Muon–Lion-K correspondence: Muon corresponds to Lion-K with K(X) equal to the nuclear norm and ∇K(X) equal to the matrix sign function.The matrix sign is a nuclear-norm subgradient and acts as a matrix analogue of the entrywise signum.
  • Muon–Lion-K correspondence: The nuclear and spectral norms are the Schatten 1- and ∞-norms, respectively, corresponding to ℓ1 and ℓ∞ norms on singular values.This norm duality underlies the connection between Muon’s update and a spectral constraint.
  • Spectral constraint: With decoupled weight decay, Muon implicitly solves a constrained problem that bounds the singular values of its weight matrices.Without weight decay, the paper recovers the original unconstrained optimization problem.
  • Spectral constraint: Any bounded-norm update combined with weight decay drives iterates exponentially downward when the corresponding norm constraint is violated.For Muon, the matrix-sign update has spectral norm at most 1, so the proposition applies.
  • Analytical challenge: The nondifferentiability of the trace norm prevents direct use of the continuous-time results for Muon’s convergence guarantees.The paper therefore uses a discrete-time analysis to establish convergence and implicit bias.
  • Generalizations: Replacing the nuclear norm with a general convex spectral function yields updates whose operator norm is bounded by b and constrains ∥X∥op ≤ b/λ.The bound follows when the scalar spectral function has gradient upper-bounded by b.
  • Generalizations: Convexity of K supports convergence of the associated Lion-K optimizers, producing a broad class of provably convergent algorithms parameterized by K.The framework also extends to general convex constraint sets through characteristic functions and linear minimization oracles.

7 Convergence analysis of Muon

The convergence analysis characterizes Muon’s KKT points, derives deterministic and stochastic discrete-time guarantees, and proves convergence to the KKT set under stated assumptions.

  • Discrete-time analysis: The discrete-time analysis extends Lion-K theory to matrix-valued updates and bounds the KKT score function in deterministic and stochastic settings.The analysis uses helper lemmas, a Lyapunov-style argument, and assumptions including bounded variance where required.
  • KKT characterization: A point is a KKT point of the constrained problem exactly when ∥λX*∥op ≤ 1 and the KKT score function vanishes.This equivalence provides the criterion used throughout the convergence proof.
  • Convergence to KKT points: Muon converges almost surely to the set of KKT points under the stated conditions and initialization constraint.LaSalle’s invariance principle links decay of the KKT score to convergence to the KKT set.
  • Assumptions: The analysis assumes ∥λX0∥op ≤ 1, which Proposition 1 then preserves for every iteration.The paper notes that practical implementations can choose the initial point to satisfy this condition.
  • Deterministic gradients: Muon with deterministic gradients admits convergence-rate guarantees under the theorem’s smoothness, momentum, weight-decay, initialization, and bounded-map assumptions.The Muon update is covered as a specialization of the general matrix Lion-K analysis.
  • Stochastic gradients: The stochastic analysis extends the deterministic result under the paper’s stochastic-gradient assumptions.The theorem treats norm-based K and stochastic updates of the matrix Lion-K variables.

8 Experiments

Experiments validate Muon’s spectral constraint enforcement, convergence behavior, and implicit regularization across toy problems, neural architectures, and tasks. They also show that alternative Lion-K convex maps induce different constraints and penalty structures.

  • Toy example: Muon’s singular values quickly enter and remain within constraint regions in a two-dimensional matrix optimization problem.The experiment uses λ = 1.5 and λ = 5.0, with β1 = β2 = 0.95 and η = 0.001.
  • Convergence behavior: Both feasible-region trajectories converge to the feasible optimum, while the Lyapunov function H decreases monotonically despite a nonmonotonic objective spike.The trajectories start from X0 = diag (0.01, 0.75) and X1 = diag (0.7, 0.7); the objective spike occurs near iteration 2000.
  • Constraint enforcement: Approximately 400 training steps are sufficient for initially violating singular values in ResNet-18 on CIFAR-10 to enter the constraint set and remain bounded.Figure 4 studies weight matrices from each ResNet-18 module with λ = 2.0.
  • Constraint enforcement: Across ImageNet and language modeling models, singular values remain below the theoretical upper limit 1/λ for λ = 2.0 and λ = 4.0.The evaluated architectures are ResNet-50, ViT-B/16, Qwen-100M, and LLaMA-300M.
  • Implicit spectral regularization: Muon produces regularized singular-value distributions for LLaMA 0.5B query, key, and value matrices, whereas AdamW shows no spectral regularization.The comparison uses converged weights trained with Muon and AdamW.
  • Generalizations: Varying the Lion-K convex map K produces distinct implicit constraint sets and penalty structures, including soft-thresholding penalties for singular values above threshold e.The experiments use the same matrix objective as Figure 2.

9 Conclusion

The paper identifies Muon as a nuclear-norm Lion-K optimizer and extends the framework to explain its implicit spectral constraint and convergence. It also presents broader generalizations through alternative convex maps, while identifying practical and theoretical directions for future work.

  • Conclusion: Muon is a Lion-K optimizer with K(X) = ∥X∥tr and matrix-sign subgradient ∇K(X) = msgn(X).The nuclear and operator norms are the Schatten 1- and ∞-norms on matrix singular values.
  • Conclusion: The Lion-K perspective generalizes Muon beyond the nuclear norm through alternative convex functions with different implicit regularization effects.The paper presents these generalizations as directions for future work and empirical exploration.
  • Conclusion: With decoupled weight decay, Muon converges to KKT points of a spectral-norm-constrained optimization problem.This gives the optimizer an interpretation as enforcing spectral regularization.
  • Limitations: The paper’s limitations include restricted convex-map coverage, limited practical training conditions, and empirical evaluations that should be scaled to larger and more diverse settings.The authors specifically mention general learning rates, nonsmooth objectives, larger models, and more diverse tasks.

A Continuous-time analysis of Lion-K

The continuous-time analysis extends Lion-K convergence theory to matrix dynamics using a Lyapunov function. It establishes nonnegativity and monotonic decrease of that function, while noting that direct application to Muon is blocked by nuclear-norm nondifferentiability.

  • Continuous-time analysis: The analysis extends Lion-K ODE convergence results to matrix-valued dynamics.The extension is presented as a direct adaptation of prior Lion-K analysis.
  • Lyapunov analysis: Under the stated assumptions, the Lyapunov function H remains nonnegative and has a nonpositive time derivative.Thus H is a Lyapunov function for the analyzed dynamics.
  • Lyapunov analysis: The proof establishes H’s nonnegativity by combining objective nonnegativity with the Fenchel–Young inequality.The argument also uses the convex-conjugate structure of K and K∗.
  • Scope boundary: The continuous-time theorem cannot be applied directly to Muon because the nuclear norm is nondifferentiable.This motivates the separate treatment required for Muon.

B Deferred proofs

The deferred proofs establish KKT characterizations, norm and conjugacy lemmas, and stochastic convergence ingredients used in the Muon analysis. They also support Lyapunov-based conclusions about bounded trajectories and limit sets.

  • KKT characterization: A point is KKT for the constrained problem exactly when ∥λX⋆∥op ≤ 1 and S(X⋆) = 0.This proposition supplies the feasibility and stationarity characterization used in the convergence results.
  • Norm identities: For a norm K, every subgradient satisfies ⟨∇K(X), X⟩ = K(X), while the conjugate evaluated at that subgradient is zero.These identities are established in Lemmas 2 and 4.
  • Muon-specific bounds: For the nuclear norm and matrix-sign map, the proof maintains ∥λX_t∥op ≤ 1 throughout the trajectory.The result is used in the matrix-specific Muon analysis.
  • Stochastic ingredients: The stochastic proof controls random matrix deviations using norm equivalence, Jensen’s inequality, and finite-variance assumptions.One resulting bound is E[∥X − Y∥F] ≤ 2σ + R.
  • Stochastic convergence: LaSalle’s stochastic invariance principle places the ω-limit set inside the zero set of the nonnegative function h under boundedness and step-size conditions.The proof uses supermartingale convergence and establishes convergence of the Lyapunov sequence almost surely.
Loading 2506.15054v2…