Muon geometry: the optimizer as steepest descent under a norm #
The geometric motivation for Muon (planning/archive/muon_geometry.md), in the unifying frame every
optimizer is steepest descent under a choice of norm: the update direction is the dual-norm
maximizer d⋆ = argmax_{‖d‖≤1} ⟨g,d⟩, with optimal value the dual norm ‖g‖_*.
| optimizer | step norm | d⋆ | here |
|---|---|---|---|
| SGD | Euclidean ‖·‖₂ | g/‖g‖, value ‖g‖₂ | steepest_l2_* |
| sign / Adam-ish | ‖·‖∞ | sign(g), value ‖g‖₁ | steepest_linf_* |
| Muon | operator norm | polar factor UVᵀ, value nuclear Σσᵢ | muon_polar_steepest |
| Shampoo (1-step) | Kronecker-factored | (GGᵀ)^{-1/4}G(GᵀG)^{-1/4} = UVᵀ = Muon | shampoo_eq_muon |
The SGD and sign rungs are the framework, proven outright. The Muon rung is now proved both ways:
the polar factor UVᵀ of G = UΣVᵀ attains the nuclear norm Σσᵢ (muon_polar_achieves_nuclear,
trace algebra) and is the maximum — von Neumann's trace inequality, ⟨G,D⟩_F ≤ Σσᵢ for every
contraction D (muon_polar_is_max, per-singular-vector Cauchy–Schwarz). muon_polar_steepest
packages them: UVᵀ is feasible (an isometry), attains, and is unbeatable — Muon's update IS the
operator-norm steepest-ascent direction, the same bound+attained shape as the SGD/sign rungs.
L4 (this layer): the SVD is now constructed, not hypothesized — for an invertible (full-rank) G,
svd_of_isUnit builds U, V orthogonal and s ≥ 0 with G = U (diagonal s) Vᵀ out of Mathlib's
spectral theorem of GᵀG: V = eigenvector basis, sᵢ = √λᵢ the singular values, U = G V Σ⁻¹.
No matrix square root is needed — only the spectral decomposition, scalar √, and diagonal inverses
(invertibility makes every λᵢ > 0, so Σ⁻¹ exists). Composing with the achievability half gives
muon_polar_achieves_nuclear_of_isUnit: for any invertible G, the polar factor UVᵀ (Muon's
update direction) pairs with G to the nuclear norm Σσᵢ — the SVD hypothesis fully discharged.
L5 (this layer): the Shampoo = Muon jewel. Single-step Shampoo preconditions the gradient by the
inverse fourth-roots of its Gram matrices, G ↦ (GGᵀ)^{-1/4} G (GᵀG)^{-1/4}, and shampoo_eq_muon
proves this equals Muon's polar factor UVᵀ — two optimizers, one geometry. Reusing the L4 SVD
pieces V, Σ: the fourth-roots are spectral ((GᵀG)^{-1/4} = V (diagonal s^{-1/2}) Vᵀ), the helper
conj_diag_pow turns the matrix fourth-power into pointwise scalar powers, and the whole thing
collapses by s^{-1/2}·s·s^{-1/2} = 1. shampoo_eq_muon_of_isUnit makes it unconditional for any
invertible G.
L6 (manifold view): the polar factor lands on O(n), and is the nearest orthogonal matrix to
G. muon_polar_orthogonal — UVᵀ is orthogonal (a point of the Stiefel manifold);
muon_polar_nearest_orthogonal — ‖G − UVᵀ‖_F ≤ ‖G − Q‖_F for every orthogonal Q, the projection
of the gradient onto O(n). The latter reuses the von Neumann bound: minimizing Frobenius distance
to O(n) is maximizing ⟨G,·⟩_F over it, so "steepest" and "nearest orthogonal" are the same fact.
The only remaining layer is the singular G case (the orthonormal completion of U, which would
drop the invertibility hypothesis from the _of_isUnit capstones). All
propext / Classical.choice / Quot.sound-clean.
…and the normalized gradient attains it. g/‖g‖ is a unit vector with ⟨g, g/‖g‖⟩ = ‖g‖,
so the SGD direction g/‖g‖ is the steepest-ascent maximizer and the dual norm is ‖g‖₂.
…and sign(g) attains it. The box-corner dᵢ = ±1 = sign(gᵢ) is feasible and gives
Σ gᵢ·sign(gᵢ) = Σ|gᵢ|, so the sign update is the ‖·‖∞-steepest direction, dual norm ‖g‖₁.
Muon's update is the steepest ascent in operator-norm geometry — the achievability half.
Given an SVD G = U Σ Vᵀ (U,V orthogonal, Σ = diagonal s, s ≥ 0), the polar factor
U Vᵀ — exactly Muon's update direction — pairs with G to give the nuclear norm Σσᵢ:
⟨G, UVᵀ⟩_F = Σ sᵢ. (UVᵀ is on the operator-norm sphere, and by von Neumann's trace
inequality Σσᵢ is the max of ⟨G,·⟩ over that ball — the upper half, next layer.)
No feasible direction beats the polar factor — von Neumann's trace inequality. Over the
operator-norm unit ball, the gradient pairing is bounded by the nuclear norm:
⟨G, D⟩_F ≤ Σσᵢ for every contraction D. Here ‖D‖op ≤ 1 is spelled elementarily as the
Euclidean contraction (D x)·(D x) ≤ x·x (*ᵥ = mulVec, ⬝ᵥ = dotProduct), which avoids a
matrix operator-norm instance while saying exactly that.
Proof: with G = U Σ Vᵀ, cyclic trace gives ⟨G,D⟩_F = Σᵢ sᵢ Mᵢᵢ for M = Uᵀ D V, and each
diagonal entry Mᵢᵢ = uᵢ · (D vᵢ) is bounded by 1 — Cauchy–Schwarz (‖uᵢ‖ = 1) then the
contraction (‖D vᵢ‖ ≤ ‖vᵢ‖ = 1), i.e. Mᵢᵢ² ≤ (uᵢ·uᵢ)((Dvᵢ)·(Dvᵢ)) ≤ 1. Since sᵢ ≥ 0,
Σ sᵢ Mᵢᵢ ≤ Σ sᵢ. This is L1's per-singular-vector Cauchy–Schwarz, summed against Σ.
Muon's update UVᵀ IS the steepest-ascent direction under the operator norm — the L3 claim,
both halves now proved. For an SVD G = U Σ Vᵀ, the polar factor UVᵀ is the argmax of
⟨G,·⟩_F over the operator-norm unit ball, with optimal value the dual (nuclear) norm Σσᵢ:
- feasible —
UVᵀis an isometry (‖UVᵀ x‖ = ‖x‖), hence a contraction, so it lies in the ball; - attains —
⟨G, UVᵀ⟩_F = Σσᵢ(muon_polar_achieves_nuclear); - unbeatable — every contraction
Dhas⟨G,D⟩_F ≤ Σσᵢ(muon_polar_is_max). Compare the SGD/sign rungs (steepest_l2_*,steepest_linf_*): samebound+attainedshape, one norm up. This is why Muon'sden = UVᵀNewton–Schulz update is steepest descent.
The SVD of an invertible matrix, constructed. For invertible G, there are orthogonal
U, V (UᵀU = VᵀV = 1) and nonnegative singular values s with G = U (diagonal s) Vᵀ.
The build is spectral, not a black box: A := GᵀG is symmetric positive definite (positive
definite ⇐ G invertible), so the spectral theorem gives an orthogonal eigenbasis V and
eigenvalues λ with A = V (diagonal λ) Vᵀ, all λᵢ > 0. Set the singular values
sᵢ := √λᵢ and U := G V Σ⁻¹ (Σ⁻¹ = diagonal (1/sᵢ), which exists because λᵢ > 0). Then
UᵀU = Σ⁻¹ (Vᵀ A V) Σ⁻¹ = Σ⁻¹ (diagonal λ) Σ⁻¹ = 1 and U Σ Vᵀ = G V Vᵀ = G. No matrix
square root is needed — only the spectral decomposition, the scalar √, and diagonal
inverses. This discharges the SVD hypothesis of muon_polar_achieves_nuclear for full-rank G
(the singular case needs the orthonormal completion of U, the remaining layer).
Muon's update is the steepest ascent in operator-norm geometry — unconditionally, for any
invertible G. Combining the constructed SVD (svd_of_isUnit) with the achievability half
(muon_polar_achieves_nuclear): for invertible G there exist orthogonal U, V and singular
values s ≥ 0 with G = U (diagonal s) Vᵀ, whose polar factor U Vᵀ — exactly Muon's
update direction — pairs with G to the nuclear norm ⟨G, UVᵀ⟩_F = Σ sᵢ. The SVD is no
longer a hypothesis: it is built from the spectral theorem. (Von Neumann's trace inequality —
that Σσᵢ is the max over the operator-norm ball, not merely achieved — is the next layer.)
Powers of a diagonal conjugation become pointwise powers — the spectral-calculus workhorse.
For orthonormal W (WᵀW = 1), (W (diagonal d) Wᵀ)^k = W (diagonal dᵏ) Wᵀ: conjugating a
diagonal commutes with raising to a power, sending a matrix power to a scalar power of each
diagonal entry. This is what lets the Shampoo preconditioners' inverse fourth-roots
(GᵀG)^{-1/4}, (GGᵀ)^{-1/4} collapse to diagonal arithmetic in shampoo_eq_muon.
The Shampoo = Muon jewel. Single-step Shampoo preconditions the gradient G by the inverse
fourth-roots of its two Gram matrices: G ↦ (GGᵀ)^{-1/4} G (GᵀG)^{-1/4}. This equals Muon's
update — the polar factor UVᵀ of G = UΣVᵀ. Two famous optimizers, one geometry.
Given the SVD G = U (diagonal s) Vᵀ (U,V orthonormal, s > 0), the inverse fourth-roots are
spectral: with sᵢ^{-1/2} = (√sᵢ)⁻¹, take R := V (diagonal s^{-1/2}) Vᵀ and
L := U (diagonal s^{-1/2}) Uᵀ. The three conjuncts are:
R⁴ · (GᵀG) = 1—Rreally is(GᵀG)^{-1/4}(GᵀG = V (diagonal s²) Vᵀ, soR⁴inverts it);L⁴ · (GGᵀ) = 1—Lreally is(GGᵀ)^{-1/4};L · G · R = U Vᵀ— the jewel. The collapse is the scalar identitys^{-1/2} · s · s^{-1/2} = 1applied to each singular value (conj_diag_powturns the matrix fourth-roots into these pointwise powers). Cf.muon_polar_achieves_nuclear(the sameUVᵀ, now reached from Shampoo's side instead of the nuclear-norm side).
The Shampoo = Muon jewel, unconditional for any invertible G. Composing the constructed
SVD (svd_of_isUnit) with shampoo_eq_muon: for invertible G there are orthonormal U, V and
singular values s > 0 (strict, since diagonal s = Uᵀ G V is a unit) with G = U (diagonal s) Vᵀ, such that the two factors R = V (diagonal s^{-1/2}) Vᵀ, L = U (diagonal s^{-1/2}) Uᵀ are
the inverse fourth-roots of GᵀG, GGᵀ (R⁴·GᵀG = L⁴·GGᵀ = 1) and Shampoo's preconditioned
gradient is Muon's polar factor: L · G · R = U Vᵀ.
Muon's update lands on the orthogonal group. The polar factor UVᵀ is orthogonal —
(UVᵀ)ᵀ(UVᵀ) = (UVᵀ)(UVᵀ)ᵀ = 1 — i.e. a point of O(n) (the Stiefel manifold of orthonormal
frames). This is the geometric content of "Muon orthogonalizes the gradient": the update is not
a vector in flat weight space but a point on the manifold of orthogonal maps, and the
implementation's Newton–Schulz iteration is the retraction that computes this projection.
Muon's update is the nearest orthogonal matrix to G — the projection of the raw gradient
onto O(n) in Frobenius distance: ‖G − UVᵀ‖_F ≤ ‖G − Q‖_F for every orthogonal Q (stated in
squared fInner form to avoid √). This is why the polar factor is "the orthogonalized
gradient", and it is the ladder's punchline reusing its own prize: expanding
‖G − Q‖_F² = ‖G‖_F² − 2⟨G,Q⟩_F + n (orthogonal Q has ‖Q‖_F² = tr(QᵀQ) = n), minimizing the
distance is maximizing ⟨G,Q⟩_F over O(n) ⊆ {contractions} — exactly the von Neumann bound
muon_polar_is_max, attained at UVᵀ. The same inequality that makes UVᵀ the steepest
direction makes it the nearest orthogonal matrix.