Source-linked AI summary
A Quadratic Lens on Muon's Polar Update: Benefits, Pitfalls, and the Role of Newton-Schulz
Egor Shulgin, Sam Laing, Antonio Orvieto, Peter Richtárik
TL;DR
The paper asks when orthogonalized matrix updates help, a question not resolved by existing analyses that conflate direction effects with step-size choices. It studies PolarGD on matrix quadratics under exact line search and finds that its speed depends on iterate geometry: it can be fast in row-isotropic regimes and slow near rank deficiency.
Problem
Existing analyses do not separate the effect of the orthogonalized direction from step-size scaling, leaving open when that direction improves on gradient descent.
Source: §1 · p. 2Method
The paper analyzes PolarGD on matrix quadratics and compares its update direction with GD and SignGD using exact line search to remove learning-rate tuning.
Source: §1 · p. 2Results
PolarGD’s progress depends on iterate conditioning rather than Hessian conditioning, making it fast for row-isotropic iterates but potentially much slower than GD when the iterate has tiny singular values.
Source: §1 · pp. 2–3Takeaways & Limitations
Near rank deficiency, finite Newton–Schulz softens the polar update at tiny singular values and can converge faster than exact polar orthogonalization.
Source: §5 · p. 12The analysis omits momentum and stochasticity to isolate the effect of the update direction.
Source: §1 · p. 2
Abstract
from arXiv · showhide
Muon updates matrix-valued parameters along approximately orthogonalized directions and often trains large language models faster than AdamW, yet it remains unclear when and why orthogonalization helps. We study PolarGD, the orthogonalized update at the core of Muon, on matrix quadratics, where optimizers can be analyzed in closed form, and compare update directions under line search to remove the effect of learning-rate tuning. Orthogonalization pays off when the gradient is far from low rank and its largest singular directions see high curvature. PolarGD reaches the minimizer in one step whenever the rows of the error are orthogonal with equal norms, for any Hessian, whereas gradient descent needs all Hessian eigenvalues to be equal; and unlike sign descent, PolarGD is unaffected by rotations of the problem. Over many steps, PolarGD acts as gradient descent preconditioned by the current iterate, so its speed is governed by the conditioning of the iterate rather than of the Hessian. This explains fast regimes as well as a slow one, caused by tiny singular values of the iterate, in which PolarGD is much slower than gradient descent. Finally, the few Newton-Schulz iterations used in Muon treat these tiny singular values as gradient descent does, so fewer iterations can converge faster than exact orthogonalization; in a transformer trained with Muon, about a quarter of the singular values of the attention matrices fall in this regime.
1 Introduction
The paper isolates when Muon’s polar update helps by comparing update directions on matrix quadratics under exact line search. It finds that PolarGD’s one-step geometry, rotation equivariance, and iterate-dependent preconditioning explain both fast and slow regimes.
1 Introduction
- Exact line search isolates update-direction effects from learning-rate tuning, while closed-form quadratic decreases expose the tradeoff between gradient gain and curvature along each direction.The paper studies PolarGD, the gradient with its polar factor substituted, alongside GD and SignGD; momentum and stochasticity are omitted to isolate direction effects.Source: §1 · p. 2
- PolarGD reaches the minimizer in one step for any positive-definite Hessian when error rows are orthogonal with equal norms; GD requires an isotropic Hessian for one-step convergence from every point.SignGD has its own condition: equal-magnitude error entries with the Hessian preserving their signs.Source: §1 · p. 2
- PolarGD is orthogonally equivariant, so rotating an ill-conditioned problem preserves its trajectory behavior, whereas SignGD can degrade after rotation.This contrasts with SignGD’s sensitivity to coordinate alignment.Source: §1 · p. 2
- PolarGD acts as gradient descent preconditioned by the current iterate, making its progress depend on iterate conditioning rather than Hessian conditioning.Row-isotropic iterates are fast, while tiny singular values can make PolarGD remain slow even when GD is nearly one-step exact.Source: §1 · pp. 2–3
- A few Newton–Schulz steps soften polar normalization at tiny singular values, where the resulting update can converge faster than exact orthogonalization.The paper therefore treats finite Newton–Schulz as a distinct update behavior, not merely an imperfect polar-factor approximation.Source: §1 · p. 3
2 Notation and Setup
The setup compares GD, PolarGD, and SignGD as directions on matrix objectives, using exact line search to separate direction quality from step-size scaling. Matrix-valued curvature clarifies when polar directions can outperform GD, and experiments show no optimizer dominates across all tested regimes.
2 Notation and Setup
- Exact line search compares directions independently of their scale and learning-rate choice, with quadratic decreases determined by gradient gain and curvature along the direction.For fixed steps, polar and sign directions can be LMO-normalized or dual-adapted; line search treats both forms identically.Source: §2 · p. 3
- Across tested regimes, GD is robust to initialization but degrades under ill-conditioning, SignGD benefits from coordinate-aligned curvature, and PolarGD benefits from row-isotropic initialization and rotation equivariance.No single optimizer dominates across the tested settings.Source: §2 · p. 4
- The polar direction can guarantee more decrease than GD when large singular directions of the gradient encounter high curvature, while scalar curvature bounds never guarantee a larger decrease.With scalar bounds, polar descent guarantees at most GD’s decrease, matching it only when the nonzero gradient singular values are equal.Source: §2 · p. 5
- PolarGD beats GD by orders of magnitude across tested initializations on dense and diagonal Hessians, although its worst-case one-step guarantee is weaker.The one-step comparison is therefore pessimistic about PolarGD’s performance over complete trajectories.Source: §2 · p. 5
3 One-Step Geometry
On matrix quadratics, one-step progress depends jointly on the error geometry, curvature, and coordinate system. PolarGD’s special one-step and rotation properties coexist with trajectory behavior governed by the iterate’s conditioning.
3.1 Isotropic Quadratic Setting
- For an isotropic Hessian, GD reaches the minimizer from every initialization, whereas PolarGD is one-step exact only when the nonzero singular values of the initialization are equal.SignGD is one-step exact when the nonzero entries have equal magnitude.Source: §3.1 · p. 6
3.2 One-Step Convergence Analysis
- For any positive-definite Hessian, PolarGD is one-step exact when the error rows are orthogonal with equal norms; GD requires an isotropic Hessian for one-step exactness from every point.SignGD instead requires equal-magnitude nonzero error entries and sign preservation by the Hessian.Source: §3.2 · p. 6
3.3 Rotation Equivariance and Curvature Adaptation
- GD and PolarGD are orthogonally equivariant, but SignGD is not; consequently, rotating an ill-conditioned problem preserves PolarGD performance while SignGD degrades.The rotation result also holds for Muon’s momentum update with finite Newton–Schulz iterations under deterministic gradients.Source: §3.3 · p. 7
- On Zipf-type diagonal problems, PolarGD behaves similarly to SignGD and, unlike SignGD, retains performance when the Hessian is randomly rotated.Figure 2 keeps the target fixed during the rotation and reports the empirical comparison.Source: §3.3 · p. 7, Figure 2 · p. 7
- PolarGD’s trajectory speed reflects preconditioning by the iterate covariance rather than the Hessian, explaining fast row-isotropic iterates and slow near-rank-deficient ones.Thus one-step geometry and rotation equivariance alone do not explain its full-trajectory behavior.Source: §3.3 · p. 8
4 Beyond One Step: Implicit Preconditioning
PolarGD’s trajectory speed depends on the current iterate’s conditioning, not just the Hessian; tiny singular values can make exact orthogonalization slow. Finite Newton–Schulz softens those modes, and the paper finds this regime in transformer attention matrices.
4.1 PolarGD as an Implicit Preconditioned Method
- PolarGD is a gradient step preconditioned by the current iterate covariance, so its conditioning—and speed—depends on the iterate rather than the Hessian alone.Row-isotropic iterates can make even a badly conditioned Hessian harmless, while tiny singular values can make the preconditioner ill-conditioned.Source: §4 · p. 8, §4.1 · p. 9
4 Beyond One Step: Implicit Preconditioning
- One-step preferences can change along a trajectory, so a direction that wins locally need not yield the better overall method.In a 2×2 diagonal example, the local winner between GD and PolarGD alternates along the GD trajectory.Source: §4 · p. 8
4.1 PolarGD as an Implicit Preconditioned Method
- Near-rank-deficient iterates can make PolarGD much slower than GD, even with exact line search; GD can be nearly one-step exact in this regime.The adverse rate is tight, and exact PolarGD contracts by a fixed factor per step in the example with H = Id.Source: §4.1 · p. 10
- Line search can create tiny singular values: it sends modes near the spectrum’s weighted mean close to zero, while a longer step can reflect small modes to large ones.In the reported setting, the tuned rate keeps p(Xk) near 0.6, whereas line search lets it decay to 0.01 within 150 iterations.Source: §4.1 · pp. 10–11
- Exact PolarGD cannot increase the rank of an iterate, so rank deficiency is treated as an active-subspace regime rather than a failure of the analysis.The rank bound holds for any step size.Source: §4.1 · p. 11
4.2 Near-Rank-Deficiency and Softened Polar Updates
- Finite Newton–Schulz acts like GD below its damping scale and like exact polar above it, so fewer steps can converge faster near rank deficiency.Its damping of tiny singular values decreases geometrically with the number of steps.Source: §4.2 · p. 11
- About a quarter of attention-matrix singular values fall below the five-step damping scale, where Muon acts like GD rather than exact polar.Across all 50 matrices, exact polar would place about 16% of update squared Frobenius norm on these directions, versus about 5% for five Newton–Schulz steps.Source: §4.2 · p. 12
- Attention matrices need 9–18 Newton–Schulz steps to become polar-like on all directions, versus 4–5 for MLP matrices.The paper attributes the difference to the attention matrices’ smaller normalized singular values.Source: §4.2 · p. 12
5 Conclusion
The paper isolates the geometry of the orthogonalized update through PolarGD on matrix quadratics. It identifies both favorable and unfavorable regimes and suggests treating Newton–Schulz steps or damping as optimizer design choices.
5 Conclusion
- PolarGD’s orthogonal equivariance and implicit preconditioning help near row-isotropic iterates, but tiny singular values can make exact polar normalization slow.Finite Newton–Schulz softens those modes and can converge faster than exact polar.Source: §5 · p. 12
- The trajectory results rely on quadratic structure, while equivariance and Newton–Schulz softening hold for any objective, including Muon’s momentum matrix.Transformer measurements show that the regime where exact polar and Newton–Schulz differ occurs in practice.Source: §5 · p. 12
- The number of Newton–Schulz steps or the damping level can be treated as an optimizer design choice, not only an accuracy setting.Source: §5 · p. 13
Where the Main-Text Results Are Proved
The appendix index maps the paper’s main-text results to their proof or detail locations.
- The appendix lists proof locations for the one-step decrease, isotropic case, one-step exactness regimes, and orthogonal equivariance results.Source: Where the Main-Text Results Are Proved · p. 17
- It points to Appendix C for the alternating local-winner example and Appendix E for preconditioning and trajectory-level rates.Source: Where the Main-Text Results Are Proved · p. 17
- Appendices F–H cover softened polar variants, finite Newton–Schulz as damped polar, and the transformer diagnostic.Source: Where the Main-Text Results Are Proved · p. 17
A Experimental Details
The appendix specifies tuning and experimental setups for quadratic and language-modeling-proxy studies, then supplies line-search analyses for isotropic and diagonal quadratics.
A.1 Learning Rate Selection; A.2 Initialization-Hessian Regimes (Figure 1)
- The experiments use isotropic, dense Marchenko–Pastur, and spectrum-matched diagonal Hessians, Gaussian or orthogonal initialization, and a two-stage learning-rate grid; rank(HXk) is tracked for the full-rank assumption.The grid selects rates by final loss among nondivergent runs, then refines around the candidate rate; the regimes use d = 100 and 150 iterations.Source: §A.2 · p. 18
A.3 Zipfian Language-Modeling Proxy (Figure 2)
- The Zipf proxy models next-token prediction with frequencies πk ∝ k^-ζ and a diagonal Hessian, while its rotated-coordinate experiment preserves the spectrum.The target rows are sampled from a Dirichlet distribution around the Zipf frequencies, yielding a full-rank target almost surely.Source: §A.3 · pp. 18–19
- In the rotated Zipf proxy, PolarGD matches its diagonal-case iteration count when Hessian and target rotate together, while SignGD degrades under both tested rotations.SignGD reaches relative loss 10^-8 to 10^-2 after 1,000 iterations, versus 10^-15 in about 200 iterations without rotation.Source: §A.3 · p. 19, Figure 6 · p. 19
A.4 Spectral Recovery Near Rank Deficiency (Figure 4)
- The near-rank-deficiency experiment compares exact PolarGD with Muon-style updates using three, four, or five Newton–Schulz steps under exact line search.It uses an identity Hessian and a near rank-deficient initialization in dimension 10.Source: §A.4 · pp. 19–20
A.5 Line Search versus Tuned Dual-Adapted Steps (Figure 10)
- The line-search comparison tunes dual-adapted PolarGD and GD rates on a dense Hessian, then measures iterations to relative loss 10^-6 across log-spaced Hessian condition numbers.The latter experiment uses three seeds per condition number and caps each run at 30,000 iterations.Source: §A.5 · p. 20
B.1 Isotropic Quadratics
- On isotropic quadratics, line-searched PolarGD subtracts a shared scalar from nonzero singular values and reaches zero loss when those values are equal; SignGD analogously shrinks entries and is exact when nonzero entries have equal magnitudes.GD is one-step exact in this setting.Source: §B.1 · pp. 21–22
B.2 Exact Line-Search Quantities for Diagonal H
- For diagonal curvature, SignGD aligns with the error’s active-row signs, while PolarGD’s curvature depends on row-energy weights of the gradient’s left singular vectors; full row rank makes that curvature iterate-independent.SignGD is one-step exact for a common-magnitude error with a sign pattern preserved by H.Source: §B.2 · p. 23
B.2.1 Local Comparison in the Diagonal Case (Explicit but Initialization-Dependent)
- In the diagonal case, the local PolarGD-versus-GD comparison depends on the gradient’s singular spectrum and left-singular-vector alignment, and the methods coincide on diagonal iterates.Source: §B.2.1 · pp. 23–24
B.3 Proof of Theorem 1
The theorem characterizes when each direction reaches a quadratic minimizer in one exact-line-search step. The accompanying equivariance results show that GD and PolarGD preserve their trajectories under orthogonal coordinate changes, unlike SignGD.
B.3 Proof of Theorem 1
- SignGD is one-step exact when the error’s nonzero entries share a magnitude and sign(HE) matches sign(E), including at zero entries.Strict diagonal dominance with positive diagonal is a sufficient condition for sign preservation for every sign pattern.Source: §B.3 · p. 24
- PolarGD is one-step exact when the error has orthogonal rows of equal norm; GD is exact from every starting point only when H is a positive scalar identity.For a fixed point, GD is exact when the error lies in a single eigenspace of H.Source: §B.3 · p. 24
B.4 Orthogonal Equivariance and Coordinate Dependence
- Under orthogonal changes of basis, GD and PolarGD trajectories and losses correspond exactly, whereas SignGD is not generally rotation-equivariant.This makes diagonalization lossless for GD and PolarGD, but not for SignGD.Source: §B.4 · p. 26
C An Analytic 2 × 2 Example Where the Local Winner Alternates
A two-dimensional quadratic shows that the direction with the better one-step decrease can alternate along a trajectory. The appendix also relates the line-search gap to the evolving iterate shape.
C An Analytic 2 × 2 Example Where the Local Winner Alternates
- Along the exact-line-search GD trajectory from diag(1, √κ), the local PolarGD and GD decrease advantage swaps every iteration, so pointwise wins do not establish a global ranking.The result holds for every κ > 1 in the stated two-cycle and for a broader class of diagonal initializations.Source: C An Analytic 2 × 2 Example Where the… · pp. 26–28
D Line Search versus Tuned Fixed Steps
- In the Figure 10 line-search run, the share of iterate singular values below 10^-2 of the largest grows to about 90%, versus mostly 10%–20% for the tuned step.The passage connects this evolving singular-value profile to the gap between line search and tuned fixed steps.Source: D Line Search versus Tuned Fixed Steps · p. 28
D.1 Proof of Proposition 4
The appendix derives how exact line-search steps scale for GD, PolarGD, and SignGD, then compares fixed rates with line search on a CIFAR-5K MLP.
D.1 Proof of Proposition 4
- Scaling X by a positive constant leaves GD’s exact line-search step unchanged, but scales the PolarGD and SignGD steps by that constant.Source: §D.1 · p. 29
D.2 Line Search on a CIFAR-5K MLP
- On the CIFAR-5K MLP, line search matches or slightly beats the best fixed GD rate, while a fixed PolarGD rate outperforms line search.The comparison uses the three best-performing fixed rates from a grid search.Source: Figure 11 · p. 30, §D.2 · p. 30
E.1 Proofs for the Implicit-Preconditioning Identities
These proofs characterize PolarGD’s dynamics and comparisons with gradient descent, focusing on active curvature and subspaces, iterate singular values, and rank deficiency. They also examine slowdowns and modifications motivated by sensitivity to tiny singular values.
E.1 Proofs for the Implicit-Preconditioning Identities
- Exact line-search PolarGD has progress coefficient p(X) bounded below by λmin(H)/tr(H) and 4κ(X)/(1+κ(X))^2, with one-step exactness iff XX⊤=cI.The squared singular values of the implicit preconditioner A are the eigenvalues of Σ^-1, so its singular values are 1/σ_i(X).Source: §E.1 · pp. 30–32
E.2 A Near-Rank-Deficient Instance Approaching the Pessimistic PolarGD Bound
- 1/[1+(d−1)κ] is the limiting first-step loss-decrease fraction for exact PolarGD as ε→0, while GD is nearly one-step exact on the same family.Tiny nonzero singular modes are normalized to unit-size directions, so the line-search denominator pays for all directions even when most gradient gain comes from one; at the rank-one limit, PolarGD reaches the minimizer in one step.Source: §E.2 · pp. 32–33
E.3 Generic Multiplicative Recursion Under Exact Line Search; E.4 GD with Exact Line Search: Classical Rate
- Exact line search converts any trajectory-wise lower bound on per-step progress into a global geometric rate; GD’s exact-line-search contraction is classical and sharp.The cited GD result is attributed to Nocedal and Wright (2006, Theorem 3.3).Source: §E.3 · p. 34, §E.4 · p. 34
E.5 PolarGD with Exact Line Search: Trajectory Dependence Through Σk
- If κ(X_k)≤K along the trajectory, PolarGD’s progress bound yields the classical exact-line-search GD rate with κ(H) replaced by K.This makes the rate depend on iterate geometry: PolarGD can be fast when the iterate stays close to row-isotropic, but can slow when Σ_k becomes highly anisotropic.Source: §E.5 · p. 34
E.6.1 Active-Subspace Preconditioning
- On the active subspace, κ+(A)≤κ(H)κ+(G)≤κ(H)^2κ+(X); when range(X) is H-invariant, κ+(A)=κ+(X).For full row rank, κ(A)=κ(X), whereas with singular Σ the active preconditioner’s conditioning can also depend on its alignment with H.Source: §E.6.1 · pp. 35–36
E.6.2 Exact Line Search and Progress in the Rank-Deficient Case
- In the rank-deficient case, exact line search uses active curvature tr(HΠ_G), and its progress coefficient is at least λmin(H)/tr(HΠ_G) ≥ 1/[rκ(H)].Thus rank deficiency is not automatically pessimistic: the curvature paid is restricted to the active subspace rather than all of H.Source: §E.6.2 · p. 36
E.6.3 Rank Monotonicity and Rank Drops
- Exact PolarGD cannot increase rank, and in the full-row-rank case a rank drop occurs iff γ∈spec(H^-1/2RH^-1/2).More generally, rank drops iff (R−γH)u=0 for some nonzero u in range(G).Source: §E.6.3 · p. 36
E.6.4 A Rank-r One-Step Exactness Regime
- A nonzero rank-r iterate is one-step exact under PolarGD iff X=√cUV⊤ and range(U) is H-invariant.The theorem also states that any such iterate reaches the minimizer in one exact-line-search step.Source: §E.6.4 · p. 37
- When r = d, the invariance condition is automatic, and Proposition 11 reduces to isotropy on an invariant active subspace.The passage identifies this reduction with Theorem 1(iii).Source: §E.6.4 · p. 37
E.6.5 Active-Subspace Root Dynamics
- For X(γ) = X − γQ, HΣ(γ)H = (R − γH)ΠG(R − γH); with full row rank, its square root is |R − γH|.In the full-row-rank case, ΠG = Id.Source: §E.6.5 · p. 37
E.6.6 Local Comparison with Active Curvature
- The exact local PolarGD–GD comparison extends to rank-deficient gradients using active curvature tr(HΠG).Rank deficiency can lower active curvature relative to tr(H), helping PolarGD, but also narrows the gap between ∥G∥∗ and ∥G∥F; at rank one, PolarGD coincides with GD under exact line search.Source: §E.6.6 · p. 37
F Rank-Truncated and Damped Polar Directions
- Rank-truncated and damped polar directions are introduced as modifications addressing sensitivity to tiny singular values suggested by Q = (GGT)†/2G.Source: F Rank-Truncated and Damped Polar… · p. 38
F.1 Rank-Truncated Polar Direction
Rank truncation and damping limit the influence of tiny singular modes, while finite Newton–Schulz orthogonalization acts as a polynomial filter that behaves like damped polar. These alternatives can improve line-search progress when exact polar normalization is impaired by tiny modes.
F.1 Rank-Truncated Polar Direction
- When X has rapidly decaying singular values, choosing s ≪ d can improve fractional decrease by ignoring tiny singular modes.The rank-s direction maximizes gradient gain over spectral-norm-bounded matrices of rank at most s, attaining the Ky Fan s-norm.Source: §F.1 · p. 38
F.2 Damped Polar Direction
- Damped polar preserves the polar direction’s singular vectors while limiting the weight assigned to tiny singular values.Its preconditioner is uniformly bounded by 1/√ε.Source: §F.2 · p. 38
G Finite Newton–Schulz as Damped Polar
- Finite Newton–Schulz behaves like damped polar: unlike exact polar, it does not apply an unbounded multiplier to tiny nonzero singular values.It is a polynomial map of singular values, not an exact damping or truncation.Source: G Finite Newton–Schulz as Damped Polar · p. 39
G.1 A General Polynomial Newton–Schulz Filter
- Finite Newton–Schulz iterations preserve the singular vectors and apply a scalar polynomial filter to the singular values.The quintic iteration used by Muon has fixed coefficients; step-dependent variants include Polar Express.Source: §G.1 · p. 39
G.2 Comparison with Exact Polar, Damped Polar, and Truncation
- Exact polar, damped polar, rank truncation, and finite Newton–Schulz differ in how their singular-value filters weight small modes.Exact polar uses multiplier 1/σ, unbounded as σ approaches zero; rank truncation hard-filters modes, while finite Newton–Schulz applies a polynomial filter to normalized singular values.Source: §G.2 · p. 39
G.3 Near-Zero Behavior and Damping Equivalence
Finite Newton–Schulz behaves linearly on tiny singular values, matching a damped polar filter rather than exact polar normalization. This bounded amplification can improve convergence near rank deficiency, but it does not restore exact zero singular values.
G.3 Near-Zero Behavior and Damping Equivalence
- For small singular values, q Newton–Schulz steps map σ to approximately ℓqσ, matching damped polar with εq = ℓq^-2.For Muon’s coefficients, ℓq = 3.4445^q; examples include ℓ2 ≈ 11.9, ℓ3 ≈ 40.9, and ℓ5 ≈ 485.Source: §G.3 · p. 39
- The transition scale τq = 1/ℓq separates polar-like behavior above it from an approximately linear, GD-like map below it.For Muon’s coefficients, τq is about 2.1 · 10^-3 at q = 5 and decreases as q increases.Source: §G.3 · p. 40
G.4 Bounded Amplification
- Finite Newton–Schulz has bounded singular-mode amplification, whereas exact polar’s multiplier 1/σ becomes unbounded as σ approaches zero.Near rank deficiency, tiny nonzero singular values are multiplied by at most Lq, which equals ℓq for Muon’s coefficients, so finite Newton–Schulz can converge faster than exact polar.Source: §G.4 · p. 40
- Finite Newton–Schulz leaves exact zero singular values at zero, so it softens near-rank deficiency but does not undo exact low rank.More iterations move the map closer to exact polar, helping away from rank deficiency but potentially hurting near it.Source: §G.4 · p. 40, §G.3 · p. 40
H Transformer Diagnostic
The transformer diagnostic finds the GD-like singular-value regime in attention momentum matrices, but not in MLP, embedding, or LM-head matrices. Its prevalence varies across attention layers, and the momentum spectra remain nearly stationary over training.
Definitions
- The GD-like range is defined by normalized momentum singular values si < τ5 = 1/ℓ5 ≈ 2.1 · 10^-3.The diagnostic normalizes each momentum singular value by the matrix Frobenius norm plus 10^-7.Source: H Transformer Diagnostic · p. 41
Where the GD-like range occurs
- About a quarter of attention singular values are GD-like, while MLP, embedding, and LM-head matrices have none in that range.The share is largest in the first layer, reaching 89% for the query matrix and 60% for the output projection at the final checkpoint.Source: H Transformer Diagnostic · p. 41
- Momentum spectra are nearly stationary over training, while raw gradients have about 40% of attention singular values in the GD-like range at 25–50% of training.By the end of training, raw gradients have about as many as momentum; the authors attribute the momentum spectrum’s stability to momentum averaging.Source: H Transformer Diagnostic · p. 41
Relation to the weights
- GD-like momentum directions are not aligned with the bottom singular directions of attention weights in the reported final-checkpoint comparison.Across 32 matrices, the median mean squared cosine of principal angles was 1.05 (left) and 1.02 (right) times the random-subspace value k/n.Source: H Transformer Diagnostic · p. 41
Extreme directions
- At the final checkpoint, five attention matrices have smin between 5.6 · 10^-9 and 2.7 · 10^-7; exact polar amplifies them by 1.8 · 10^8 to 3.7 · 10^6, versus about 485 for five Newton–Schulz steps.Newton–Schulz becomes polar-like on all directions once ℓq smin ≥ 1; across checkpoints this requires 9–18 steps for attention matrices and 4–5 for MLP matrices.Source: H Transformer Diagnostic · p. 41
Where the GD-like range occurs
- Figure 12 tracks normalized singular-value distributions at four checkpoints and the GD-like-range share over training for attention and all 50 matrices.The shaded band marks the GD-like range.Source: Figure 12 · p. 42
I Exact Dynamics through the Square-Rooted Gradient Covariance
The exact dynamics analysis expresses PolarGD through the square-rooted gradient covariance and yields closed-form trajectory behavior under commuting and two-cluster conditions. Its behavior depends on the spectrum and curvature structure: it can collapse to gradient descent in two dimensions, outperform GD in some 2 × 2 cases, and fail to improve conditioning monotonically.
I Exact Dynamics through the Square-Rooted Gradient Covariance
- For full-row-rank iterates, PolarGD updates the square-rooted gradient covariance through R = (HΣH)^1/2, with exact line-search step γ⋆(Q) = tr(R)/tr(H).The resulting covariance remains a function of Σ, so the commuting structure with H persists along the trajectory.Source: I Exact Dynamics through the… · pp. 42, 44
- When HΣ = ΣH, the singular values evolve by s_i^+ = |s_i − s̄_H|, where s̄_H is the curvature-weighted mean.The corollary assumes H is positive definite and X has full row rank.Source: I Exact Dynamics through the… · pp. 43–44
- With two singular-value clusters, loss contracts by (1 − 2w)^2 per step; balanced curvature mass w = 1/2 reaches the minimizer in two steps, while near-rank deficiency can yield slow contraction.The adverse family has w = 1/(1 + (d − 1)κ), while balanced mass gives the two-step result for any κ(H).Source: I Exact Dynamics through the… · p. 44
- PolarGD need not monotonically improve either the conditioning of Σk or its exact progress coefficient along a trajectory.The cited examples establish that row-covariance conditioning can worsen and the progress coefficient can either increase or decrease between steps.Source: I Exact Dynamics through the… · pp. 44–45
- In the full-rank 2 × 2 case, one exact PolarGD step makes R+ isotropic; thereafter loss contracts at ρ² per step, where ρ = (κ(H) − 1)/(κ(H) + 1), and PolarGD and GD directions coincide.For d > 2, this collapse generally does not hold, and PolarGD need not reduce to GD after one step.Source: I Exact Dynamics through the… · p. 45
I.1 Specialization to the 2 × 2 Case
- In the 2 × 2 line-search comparison, rank-one gradients make PolarGD coincide with GD; for anisotropic H, some full-rank configurations give PolarGD a strictly larger decrease.When H is isotropic and the gradient has equal singular values, GD is at least as good.Source: §I.1 · pp. 45–46
I.2 Optimizer Trajectories in Singular Value Space
- Figure 13 compares optimizer trajectories in singular-value space across initializations, with loss curves in the right column and the von Neumann bound overlaid.PolarGD maintains a 45° direction in singular-value space, while GD optimizes the larger eigen-direction first.Source: Figure 13 · p. 47, §I.2 · p. 47