PaperlayerView exact revision on arXiv

Source-linked AI summary

Preconditioning Benefits of Spectral Orthogonalization in Muon

Jianhao Ma, Yu Huang, Yuejie Chi, Yuxin Chen

arXiv:2601.13474v1cs.LGcs.AImath.OCstat.ML

TL;DR

The paper addresses the limited end-to-end evidence explaining Muon’s advantages over classical optimizers. It analyzes simplified Muon on matrix factorization and linear-transformer in-context learning, proving condition-number-independent linear convergence and identifying spectral decoupling as a preconditioning mechanism.

  • Problem

    Few existing results provide end-to-end, rigorous analyses demonstrating Muon’s advantages over classical optimizers in concrete applications.

  • Method

    The paper studies simplified Muon on matrix factorization and in-context learning of linear transformers, analyzing its spectral dynamics and comparing them with GD and SignGD.

  • Results

    Muon achieves linear convergence with iteration complexity independent of the relevant condition number in both problems, while GD and SignGD have condition-number-dependent complexity.

  • Takeaways & Limitations

    Spectral gradient normalization acts as an adaptive preconditioner by decoupling Muon’s dynamics into independent scalar sequences with comparable convergence behavior.

  • Takeaways & Limitations

    The theoretical analysis is limited to two simple matrix optimization problems, leaving a general theory for practically relevant conditions as future work.

Abstract

from arXiv · show

The Muon optimizer, a matrix-structured algorithm that leverages spectral orthogonalization of gradients, is a milestone in the pretraining of large language models. However, the underlying mechanisms of Muon -- particularly the role of gradient orthogonalization -- remain poorly understood, with very few works providing end-to-end analyses that rigorously explain its advantages in concrete applications. We take a step by studying the effectiveness of a simplified variant of Muon through two case studies: matrix factorization, and in-context learning of linear transformers. For both problems, we prove that simplified Muon converges linearly with iteration complexities independent of the relevant condition number, provably outperforming gradient descent and Adam. Our analysis reveals that the Muon dynamics decouple into a collection of independent scalar sequences in the spectral domain, each exhibiting similar convergence behavior. Our theory formalizes the preconditioning effect induced by spectral orthogonalization, offering insight into Muon's effectiveness in these matrix optimization problems and potentially beyond.

1 Introduction

Muon is a matrix-structured optimizer whose spectral orthogonalization has shown empirical promise, but its theoretical advantages remain incompletely understood. This paper studies simplified Muon on two matrix problems and proves condition-number-independent linear convergence relative to GD and SignGD.

  • 1.1 The Muon algorithm and prior theory: Muon uses matrix sign updates based on spectral orthogonalization, distinguishing it from per-coordinate preconditioning in Adam and AdamW.The matrix sign can be computed from a compact SVD and is efficiently approximated with Newton-Schulz iterations.
  • 1.1 The Muon algorithm and prior theory: Few existing results provide end-to-end, rigorous demonstrations of Muon’s advantages over classical optimizers in concrete applications.This motivates analyzing simplified Muon in fundamental matrix optimization settings.
  • 1.2 This paper: preconditioning with Muon: The paper analyzes simplified Muon for matrix factorization and in-context learning of linear transformers.These two case studies are chosen as concrete matrix optimization problems rather than as a general treatment of all settings.
  • 1.2 This paper: preconditioning with Muon: Muon converges linearly for matrix factorization with iteration complexity independent of the target matrix’s condition number, unlike GD and SignGD.The result covers both exactly parameterized and over-parameterized settings, while GD and SignGD require complexities scaling at least linearly with the condition number.
  • 1.2 This paper: preconditioning with Muon: Muon also converges linearly for linear-transformer in-context learning with iteration complexity independent of target covariance conditioning, unlike GD and SignGD lower bounds.The lower bounds for GD and SignGD scale polynomially with the relevant condition number.
  • 1.2 This paper: preconditioning with Muon: The analyses decouple Muon’s dynamics into independent spectral-domain scalar sequences with comparable convergence behavior.The paper interprets gradient-spectrum normalization as a preconditioning mechanism and supports the theory with numerical experiments.

2 Main results: two case studies

The paper analyzes simplified Muon on matrix factorization and in-context learning for linear transformers, proving linear convergence independent of relevant condition numbers. Its spectral orthogonalization decouples dynamics into independently evolving scalar sequences, while GD and SignGD/Adam retain condition-number-dependent slowdowns.

  • Scope and assumptions: The theory for the regime r ≤ k < d requires more restrictive conditions than the regime k ≥ d, and the paper leaves relaxing them to future work.The authors state that these restrictions may not be fundamental.
  • Matrix factorization: Simplified Muon converges rapidly in both exactly parameterized and over-parameterized matrix factorization settings, with iteration complexity independent of the target condition number.The result applies when k = r and when k > r, including k > d.
  • Preconditioning interpretation: The analysis attributes Muon’s advantage to spectral preconditioning that normalizes gradient spectra and produces better-conditioned optimization trajectories than GD and Adam.The paper frames this as the preconditioning effect of gradient orthogonalization.
  • Matrix factorization: GD and SignGD cannot generally achieve condition-number-free convergence for matrix factorization, with lower bounds showing at least linear dependence on κ for SignGD.The paper also notes that GD cannot attain condition-number-free guarantees without sacrificing linear convergence.
  • Matrix factorization: Muon’s matrix-factorization dynamics decouple into independent spectral scalar sequences, each converging at a comparable rate despite different eigenvalue magnitudes.This spectral behavior explains how gradient orthogonalization mitigates imbalance between large and small eigenvalues.
  • Numerical evidence: Numerically, Muon reaches machine-level precision within hundreds to thousands of iterations across large condition numbers and severe rank over-specification, while GD and SignGD slow substantially.The experiments use an exponentially decaying learning-rate schedule and initialize Muon with α = 0.1.
  • In-context learning with linear transformers: Simplified Muon converges linearly for in-context learning with one-layer linear transformers, using iteration complexity independent of the target covariance condition number.The theorem uses a geometrically decaying learning-rate schedule and establishes ε-accuracy under the stated initialization and schedule conditions.
  • In-context learning with linear transformers: For the linear-transformer problem, GD requires iteration complexity scaling with √κ, whereas SignGD has a condition-number-dependent lower bound and therefore lacks condition-number-free convergence.The comparison is established through classical GD lower bounds and a dedicated SignGD lower-bound theorem.

3 Analysis for matrix factorization (proof of Theorem 1)

The analysis reduces simplified Muon’s matrix-factorization dynamics to scalar recursions in the spectral domain, first for a scalar problem and then for aligned and general matrix settings. This structure yields linear convergence under stated initialization and learning-rate conditions.

  • Scalar dynamics: The proof begins with scalar Muon dynamics, where a randomly sampled learning-rate prefactor prevents iterates from reaching the undesirable stationary point 0 with probability 1.The prefactor Cη is uniformly sampled from [1, 2].
  • Scalar dynamics: Scalar Muon converges linearly at rate ρ when 1/2 ≤ρ < 1, and this analysis exposes the spectral dynamics used in the general case.The scalar result is the foundation for the subsequent matrix analysis.
  • Aligned initialization: With perfectly aligned initialization, every Muon iterate preserves the target singular subspace and decomposes into diagonal spectral components.The iterate has the form Ut = V⋆ΣtR⊤ for a diagonal Σt.
  • Aligned initialization: The aligned matrix dynamics decouple into r independent scalar sequences, each resembling the scalar Muon recursion and inheriting its convergence behavior.The resulting convergence bound requires |σi,0| ≤η0 for every component.
  • General initialization: For general matrix factorization, the proof applies the aligned analysis after establishing initialization-dependent subspace alignment and controlling the resulting spectral components.The treatment distinguishes the regimes k ≥d and r ≤k < d, with the latter requiring a separate alignment argument.
  • General initialization: In the undercomplete regime r ≤k < d, the first Muon iteration is analyzed through a leading rank-r gradient component and a residual component to establish alignment with the target eigenspace.The decomposition of G0 separates its leading r singular components from the remaining k −r components.

4 Analysis for linear transformers (proof of Theorem 3)

The analysis diagonalizes the target covariance matrix and shows that simplified Muon iterates evolve through independent scalar recursions in the spectral domain. Each sequence converges at a comparable rate, yielding an ε-accuracy iteration bound expressed through the contraction parameters and minimum singular value.

  • The target matrix is decomposed as S = V⋆Λ⋆V⋆⊤, with Λ⋆ containing eigenvalues and V⋆ orthonormal eigenvectors.
  • The simplified Muon iterate admits the decomposition Q_t = V⋆Θ_tV⋆⊤ for a diagonal Θ_t.
  • The diagonal dynamics follow Θ_t+1 = Θ_t − η_t diag-sign(Λ⋆Θ_t − I), with diag-sign applied entrywise to diagonal matrices.
  • The Muon dynamics decompose into d independent scalar sequences, each associated with one eigenvalue and evolving without cross-sequence interaction.
  • T ≥ (1/(1−ρ)) log(Cη/(σmin(S)ε)) iterations suffice to ensure ε-accuracy under the stated recurrence conditions.

5 Discussion

The discussion interprets spectral orthogonalization as adaptive preconditioning that produces independent spectral dynamics and faster convergence in the two studied matrix problems. It also emphasizes that the theory remains limited in scope and calls for broader analyses and applications.

  • Muon’s spectral orthogonalization acts as adaptive preconditioning, transforming the dynamics into independent spectral scalar sequences with comparable convergence rates.
  • The theoretical and empirical studies report faster convergence than GD and Adam on the two analyzed matrix optimization problems.
  • The theoretical analysis is limited to two simple problems, leaving generalization to other matrix-structured tasks as an open direction.
  • A broader theory should address practically relevant conditions such as gradient Lipschitz continuity and explain Muon’s empirical advantage under less idealized assumptions.

A Connection between Muon and ScaledGD for matrix factorization

The appendix relates simplified Muon to ScaledGD through their preconditioners and update rules. Under an idealized local condition, the two updates coincide up to learning-rate scaling, while experiments show structurally similar block-diagonal preconditioners.

  • ScaledGD achieves condition-number-independent convergence, motivating a comparison with simplified Muon’s update rule.
  • Muon constructs its preconditioner from the gradient, whereas ScaledGD constructs it from the iterate.
  • When the relevant gradient product is proportional to the identity, the Muon update coincides with the ScaledGD update up to learning-rate scaling.
  • Figure 3 compares the two preconditioners during matrix-factorization training and shows a consistent block-diagonal pattern.
  • The local equivalence argument is heuristic, and a fully rigorous analysis of the Muon–ScaledGD connection is left for future work.

B.1 Proof of Lemma 4

The proof controls scalar update quantities one index at a time, then combines the resulting bounds across all indices using a union bound. It uses assumptions on singular values and obtains high-probability intermediate guarantees.

  • A lower bound on σ_r(Q) is combined with earlier inequalities to obtain a probability guarantee exceeding 0.995.
  • The argument assumes positive minimum singular values for the relevant gradients and uses triangle inequalities and prior bounds to control iterate size.
  • The proof first controls g_i,t+1 for a single index before extending the result to all indices with a union bound.
  • Conditional on past randomness, the updated scalar eσ_i,t+1 is uniformly distributed over an interval determined by the learning-rate range.
  • When eσ_i,t = 0, the proof obtains g_i,t+1 = |eσ_i,t+1|^3 and uses this relation to bound the update quantity.

B.4 Proof of Lemma 7

The proof derives tail bounds by converting conditional moment-generating-function bounds into Chernoff inequalities, treating the cases k=r and k>r separately. It also establishes a high-probability singular-value bound for a Gaussian-derived matrix.

  • MGF and tail bounds: The proof first relates MGFs to tail probabilities and applies this connection to the process X_t through a conditional tail bound.The argument invokes Lemma 6 to control σ_min(∇f(eU_t)).
  • Case k=r: When k=r, the first term vanishes, enabling an MGF bound that is iterated recursively before applying Markov’s inequality and choosing θ=1/2.This yields P(S_T≥u)≤δ and completes Part (i) of Lemma 7.
  • Case k>r: When k>r, the proof writes Π_T=e^{S_T} and controls S_T using conditional MGFs, a Chernoff argument, and θ=1/6.The resulting bound establishes Part (ii) of Lemma 7.
  • Singular-value bound: For the Gaussian-derived orthogonal matrix, the first r columns share the distribution G(G^⊤G)^−1/2, reducing the analysis to a lower bound on σ_r(V⋆⊤G(G^⊤G)^−1/2).The bound holds with probability at least 0.995.
  • Scalar-update extension: The proof of the scalar-update extension uses induction over two cases, showing that the recurrence remains valid under the stated learning-rate bounds.The two cases compare Δ_t with η_t and conclude the claim at time t+1.

C.1 Proof of Lemma 12

The proof shows that SignGD becomes a coordinate-selective process in a rotated eigenbasis: each step updates exactly one coordinate according to the relative magnitudes of the two coordinates. A carefully chosen initialization then creates a learning-rate barrier that prevents rapid reduction of the slow coordinate.

  • Rotated-basis dynamics: In the rotated eigenbasis, SignGD has only four possible sign patterns and therefore updates exactly one coordinate at each step.The update direction is determined by the signs of the gradient components.
  • Coordinate-selection condition: The updated coordinate is selected by comparing |κẽz_1,t| with |ẽz_2,t|.The two strict inequalities produce separate update rules, while equality permits arbitrary choice.
  • Learning-rate barrier: A non-increasing learning-rate sequence admits an initialization for which x_2,t<κε can occur only after η_t<4ε.The construction keeps the second coordinate fixed while recursively choosing the first-coordinate initialization.
  • Learning-rate barrier: The initialization remains in the prescribed interval by induction, so the update rule continues acting only on the first coordinate while x_2,t stays at κε.This establishes the barrier through the time range where η_t≥4ε.
  • Transfer to SignGD: The constructed sequence is embedded back into the rotated SignGD iterates, transferring the coordinate-wise barrier to the original quadratic problem.The proof then concludes using the per-step change bound and the corresponding objective dynamics.

C.2 Proof of Lemma 13

The proof reduces SignGD on a structured matrix-factorization instance to coordinate-selective dynamics in a two-dimensional eigenspace. Invariance of the structured set and a local objective lower bound transfer the scalar learning-rate barrier to a matrix-factorization lower bound.

  • Invariance of S: The structured matrix set S is invariant under gradient and entrywise-sign updates, so trajectories initialized in S remain in S.Products of matrices in S remain in S, which preserves the gradient and SignGD update structure.
  • Eigenvalue representation: Within S, the matrix is parameterized by two induced eigenvalues, and the gradient becomes diagonal in the rotated basis.The eigenvalues are λ_1=a+b and λ_2=a−b, with gradient components given by g_1(λ) and g_2(λ).
  • Matching dynamics: In a local region around the target eigenvalues, the gradient-sign dynamics can be expressed through perturbations δ_1,t and δ_2,t and matched to the coordinate-selective dynamics of Lemma 12.The proof uses the same initialization construction with a transformed learning-rate sequence and target accuracy.
  • Objective lower bound: The objective admits an exact eigen-form on S, and a local lower bound implies that achieving objective value at most ε requires the perturbations to satisfy a corresponding accuracy condition.The local-region assumption controls the higher-order terms used in the lower bound.

D Derivation of the training objective in Section 2.2

The section derives a tractable in-context-learning objective for a simplified linear transformer on fixed-design linear tasks. After specifying the prompt distribution and model parameterization, it shows that training reduces exactly to a quadratic optimization problem.

  • In-context-learning setup: The in-context-learning setup samples tasks, support examples, and a query, with the training goal of predicting w^⊤x_q from the prompt embedding.The task family is specialized to noiseless linear regression with fixed support inputs.
  • Task and prompt construction: The fixed-design task distribution uses inputs with empirical covariance S, isotropic task parameters, noiseless labels, and a query sampled uniformly from the support set.The prompt embeds the support inputs, labels, and a zero query-label slot.
  • Simplified linear transformer: The simplified transformer removes softmax, merges query and value weights into W_KQ, and fixes the normalization factor to γ=N.The resulting model is parameterized by Q, whose output is read from the bottom-right entry.
  • Training objective: Direct calculation gives a closed-form prediction expression whose expected squared risk defines the in-context-learning training objective.Randomness in the objective comes from the task parameter and query across prompts.
  • Optimization equivalence: Minimizing the in-context-learning objective is exactly equivalent to solving the stated quadratic optimization problem.This equivalence supplies the optimization formulation used for the subsequent analysis.

E Lower bounds for SignGD in ICL (Proof of Theorem 4)

The proof constructs a two-dimensional covariance instance whose invariant matrix slice reduces Muon’s induced SignGD dynamics to a known ill-conditioned recursion. This yields an iteration lower bound proportional to the condition parameter.

  • Instance construction: κ(S)^3 = κ for the constructed covariance matrix, establishing the target condition-number scaling.The construction uses a two-dimensional instance and sets σ1 = κ^1/3 and σ2 = 1.
  • Invariant slice: The matrix slice S is invariant under the update because commutation, gradient structure, and the entrywise sign map preserve S.Thus, iterates initialized in S remain in the two-dimensional parameterized family Q(a,b).
  • Quadratic reduction: In spectral coordinates, the objective becomes a quadratic function with minimizer z = 0 and gradient ∇g(z) = 2Hz.The proof introduces error coordinates from the diagonalized representation to express the dynamics in these coordinates.
  • Dynamics: The induced Muon update is z_t+1 = z_t − η_t sign(H z_t), exactly matching the SignGD recursion analyzed in Lemma 12.This identification transfers the SignGD lower-bound construction to the matrix iterates.
  • Lower bound: (κ−1)/4 iterations are necessary to achieve ∥Q_t−Q⋆∥F ≤ ε for a suitable initialization.The proof translates the lower bound from error coordinates back to the original matrix objective and iterates.
  • Technical lemmas: The appendix also collects perturbation, Gaussian singular-value, and orthogonal-completion lemmas used elsewhere in the analysis.These include bounds for matrix signs, rank-r approximations, Gaussian matrices, and augmenting orthonormal bases.

F.1 Proof of Lemma 22

The proof of Lemma 22 compares two r-dimensional subspaces through principal angles, selects controlled orthogonal complements, and completes the bases to square orthogonal matrices. It then bounds the distance between the completed matrices.

  • Subspace distance: The subspaces spanned by O1 and O2 are characterized by their principal angles, including the maximum angle θmax.The argument also considers the orthogonal complements of the two subspaces.
  • Subspace distance: Classical perturbation theory provides a bound on the distance between orthonormal bases Q1 and Q2 in terms of the subspaces they span.This supplies the starting comparison for constructing compatible complements.
  • Complement selection: Orthonormal bases R1 and R2 for the orthogonal complements can be chosen with controlled distance.The construction uses the relationship between the principal angles of the original subspaces and their complements.
  • Orthogonal completion: Ai = [Oi, Ri] forms a square orthogonal matrix for each i.The completed matrices combine the original orthonormal columns with the selected orthogonal-complement bases.
  • Final bound: The bounds for the original bases and complements combine to control the distance between A1 and A2.The proof concludes by applying norm inequalities to the completed orthogonal matrices.
Loading 2601.13474v1…