Documentation

LeanMlir.Proofs.Foundation.MuonGeometry

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‖_*.

optimizerstep normd⋆here
SGDEuclidean ‖·‖₂g/‖g‖, value ‖g‖₂steepest_l2_*
sign / Adam-ish‖·‖∞sign(g), value ‖g‖₁steepest_linf_*
Muonoperator normpolar factor UVᵀ, value nuclear Σσᵢmuon_polar_steepest
Shampoo (1-step)Kronecker-factored(GGᵀ)^{-1/4}G(GᵀG)^{-1/4} = UVᵀ = Muonshampoo_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_orthogonalUVᵀ 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.

Euclidean steepest ascent is bounded by ‖g‖. Over the unit ‖·‖₂ ball, no direction beats the gradient: ⟨g,d⟩ ≤ ‖g‖. Pure Cauchy–Schwarz — the geometry behind plain SGD.

…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‖₂.

theorem Proofs.MuonGeometry.steepest_linf_bound {n : } (g d : Fin n) (hd : ∀ (i : Fin n), |d i| 1) :
i : Fin n, g i * d i i : Fin n, |g i|

‖·‖∞-steepest ascent is bounded by Σ|gᵢ|. Over the box |dᵢ| ≤ 1, the pairing Σ gᵢ dᵢ ≤ Σ|gᵢ| = ‖g‖₁. The geometry behind sign and Adam-style coordinate updates.

theorem Proofs.MuonGeometry.steepest_linf_attained {n : } (g : Fin n) :
(∀ (i : Fin n), |if 0 g i then 1 else -1| 1) (∑ i : Fin n, g i * if 0 g i then 1 else -1) = i : Fin n, |g i|

…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‖₁.

def Proofs.MuonGeometry.fInner {n : } (A B : Matrix (Fin n) (Fin n) ) :

Frobenius inner product ⟨A,B⟩_F = tr(Aᵀ B) = Σᵢⱼ AᵢⱼBᵢⱼ — the inner product the update ⟨∇L, D⟩ is taken in.

Equations
Instances For
    theorem Proofs.MuonGeometry.muon_polar_achieves_nuclear {n : } (U V : Matrix (Fin n) (Fin n) ) (s : Fin n) (hU : U.transpose * U = 1) (hV : V.transpose * V = 1) :
    fInner (U * Matrix.diagonal s * V.transpose) (U * V.transpose) = i : Fin n, s i

    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.)

    theorem Proofs.MuonGeometry.muon_polar_is_max {n : } (U V D : Matrix (Fin n) (Fin n) ) (s : Fin n) (hU : U.transpose * U = 1) (hV : V.transpose * V = 1) (hs : ∀ (i : Fin n), 0 s i) (hD : ∀ (x : Fin n), D.mulVec x ⬝ᵥ D.mulVec x x ⬝ᵥ x) :
    fInner (U * Matrix.diagonal s * V.transpose) D i : Fin n, s i

    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 Σ.

    theorem Proofs.MuonGeometry.muon_polar_steepest {n : } (U V : Matrix (Fin n) (Fin n) ) (s : Fin n) (hU : U.transpose * U = 1) (hV : V.transpose * V = 1) (hs : ∀ (i : Fin n), 0 s i) :
    (∀ (x : Fin n), (U * V.transpose).mulVec x ⬝ᵥ (U * V.transpose).mulVec x x ⬝ᵥ x) fInner (U * Matrix.diagonal s * V.transpose) (U * V.transpose) = i : Fin n, s i ∀ (D : Matrix (Fin n) (Fin n) ), (∀ (x : Fin n), D.mulVec x ⬝ᵥ D.mulVec x x ⬝ᵥ x)fInner (U * Matrix.diagonal s * V.transpose) D i : Fin n, s i

    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 Σσᵢ:

    • feasibleUVᵀ 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 D has ⟨G,D⟩_F ≤ Σσᵢ (muon_polar_is_max). Compare the SGD/sign rungs (steepest_l2_*, steepest_linf_*): same bound+attained shape, one norm up. This is why Muon's den = UVᵀ Newton–Schulz update is steepest descent.
    theorem Proofs.MuonGeometry.svd_of_isUnit {n : } (G : Matrix (Fin n) (Fin n) ) (hG : IsUnit G) :
    ∃ (U : Matrix (Fin n) (Fin n) ) (V : Matrix (Fin n) (Fin n) ) (s : Fin n), U.transpose * U = 1 V.transpose * V = 1 (∀ (i : Fin n), 0 s i) G = U * Matrix.diagonal s * V.transpose

    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).

    theorem Proofs.MuonGeometry.muon_polar_achieves_nuclear_of_isUnit {n : } (G : Matrix (Fin n) (Fin n) ) (hG : IsUnit G) :
    ∃ (U : Matrix (Fin n) (Fin n) ) (V : Matrix (Fin n) (Fin n) ) (s : Fin n), U.transpose * U = 1 V.transpose * V = 1 (∀ (i : Fin n), 0 s i) G = U * Matrix.diagonal s * V.transpose fInner G (U * V.transpose) = i : Fin n, s i

    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.)

    theorem Proofs.MuonGeometry.conj_diag_pow {n : } (W : Matrix (Fin n) (Fin n) ) (d : Fin n) (hWtW : W.transpose * W = 1) (k : ) :
    (W * Matrix.diagonal d * W.transpose) ^ k = (W * Matrix.diagonal fun (i : Fin n) => d i ^ k) * W.transpose

    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.

    theorem Proofs.MuonGeometry.shampoo_eq_muon {n : } (U V : Matrix (Fin n) (Fin n) ) (s : Fin n) (hU : U.transpose * U = 1) (hV : V.transpose * V = 1) (hs : ∀ (i : Fin n), 0 < s i) :
    ((V * Matrix.diagonal fun (i : Fin n) => ((s i))⁻¹) * V.transpose) ^ 4 * ((U * Matrix.diagonal s * V.transpose).transpose * (U * Matrix.diagonal s * V.transpose)) = 1 ((U * Matrix.diagonal fun (i : Fin n) => ((s i))⁻¹) * U.transpose) ^ 4 * (U * Matrix.diagonal s * V.transpose * (U * Matrix.diagonal s * V.transpose).transpose) = 1 (U * Matrix.diagonal fun (i : Fin n) => ((s i))⁻¹) * U.transpose * (U * Matrix.diagonal s * V.transpose) * ((V * Matrix.diagonal fun (i : Fin n) => ((s i))⁻¹) * V.transpose) = U * V.transpose

    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) = 1R really is (GᵀG)^{-1/4} (GᵀG = V (diagonal s²) Vᵀ, so R⁴ inverts it);
    • L⁴ · (GGᵀ) = 1L really is (GGᵀ)^{-1/4};
    • L · G · R = U Vᵀthe jewel. The collapse is the scalar identity s^{-1/2} · s · s^{-1/2} = 1 applied to each singular value (conj_diag_pow turns the matrix fourth-roots into these pointwise powers). Cf. muon_polar_achieves_nuclear (the same UVᵀ, now reached from Shampoo's side instead of the nuclear-norm side).
    theorem Proofs.MuonGeometry.shampoo_eq_muon_of_isUnit {n : } (G : Matrix (Fin n) (Fin n) ) (hG : IsUnit G) :
    ∃ (U : Matrix (Fin n) (Fin n) ) (V : Matrix (Fin n) (Fin n) ) (s : Fin n), U.transpose * U = 1 V.transpose * V = 1 (∀ (i : Fin n), 0 < s i) G = U * Matrix.diagonal s * V.transpose ((V * Matrix.diagonal fun (i : Fin n) => ((s i))⁻¹) * V.transpose) ^ 4 * (G.transpose * G) = 1 ((U * Matrix.diagonal fun (i : Fin n) => ((s i))⁻¹) * U.transpose) ^ 4 * (G * G.transpose) = 1 (U * Matrix.diagonal fun (i : Fin n) => ((s i))⁻¹) * U.transpose * G * ((V * Matrix.diagonal fun (i : Fin n) => ((s i))⁻¹) * V.transpose) = U * V.transpose

    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ᵀ.

    theorem Proofs.MuonGeometry.muon_polar_orthogonal {n : } (U V : Matrix (Fin n) (Fin n) ) (hU : U.transpose * U = 1) (hV : V.transpose * V = 1) :

    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.

    theorem Proofs.MuonGeometry.muon_polar_nearest_orthogonal {n : } (U V Q : Matrix (Fin n) (Fin n) ) (s : Fin n) (hU : U.transpose * U = 1) (hV : V.transpose * V = 1) (hs : ∀ (i : Fin n), 0 s i) (hQ : Q.transpose * Q = 1) :

    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.