Documentation

LeanMlir.Proofs.Foundation.MuonNewtonSchulz

Newton–Schulz convergence, P1: the iteration is a scalar map in disguise #

The capstone of the Muon-geometry ladder (planning/archive/muon_ns_convergence.md, planning/archive/muon_geometry.md, LeanMlir/Proofs/Foundation/MuonGeometry.lean). L1–L6 proved that the polar factor UVᵀ is the right object — operator-norm steepest descent (L3, von Neumann), the nuclear norm's argmax, Shampoo's single step (L5), the nearest orthogonal matrix to G (L6). What remains is that the implementation actually computes it: Muon's matmul iteration X ↦ aX + b(XXᵀ)X + c(XXᵀ)²X (OptimizerKind.muon, emitMuonUpdate) converges to UVᵀ.

This file is P1 — the spectral-step lemma, the bridge that turns the whole problem scalar. The one idea: a Newton–Schulz step never touches the singular directions, only the singular values. With X = U Σ Vᵀ (U,V orthonormal, Σ = diagonal σ), since XXᵀ = U Σ² Uᵀ,

(XXᵀ)X = U Σ³ Vᵀ,   (XXᵀ)²X = U Σ⁵ Vᵀ   ⟹   nsStep a b c X = U (diagonal (φ ∘ σ)) Vᵀ,

where φ(t) = a t + b t³ + c t⁵ (nsScalar) is applied per singular value, with U,V carried along unchanged (nsStep_spectral). Iterating, nsStep^[k] X = U (diagonal (φ^[k] ∘ σ)) Vᵀ (nsStep_iterate_spectral): matrix convergence to UVᵀ reduces to scalar convergence φ^[k](σᵢ) → 1 per singular value. This is the same U Σ Vᵀ ↦ U f(Σ) Vᵀ motif as L5's conj_diag_pow, now for the polynomial φ. The downstream scalar analysis (P2) and the matrix-continuity assembly (P3) build on these two lemmas. All propext / Classical.choice / Quot.sound-clean.

def Proofs.MuonNewtonSchulz.nsStep {n : } (a b c : ) (X : Matrix (Fin n) (Fin n) ) :
Matrix (Fin n) (Fin n)

The Newton–Schulz step. One iteration of Muon's gradient orthogonalizer: nsStep a b c X = aX + b(XXᵀ)X + c(XXᵀ)²X, the odd matrix polynomial X·p(XᵀX) written via the Gram matrix XXᵀ. The classic inverse-free polar iteration is (a,b,c) = (3/2, −1/2, 0); Muon's tuned quintic is (3.4445, −4.7750, 2.0315). The quintic monomial is associated as (XXᵀ)·((XXᵀ)X) so the spectral collapse threads cleanly.

Equations
Instances For

    The scalar map a Newton–Schulz step induces on each singular value: φ(t) = a t + b t³ + c t⁵. nsStep_spectral shows nsStep a b c acts as nsScalar a b c on the singular values of X, so the matrix iteration's convergence is exactly this scalar iteration's (φ^[k](σᵢ) → 1).

    Equations
    Instances For
      theorem Proofs.MuonNewtonSchulz.nsStep_spectral {n : } (a b c : ) (U V : Matrix (Fin n) (Fin n) ) (σ : Fin n) (hU : U.transpose * U = 1) (hV : V.transpose * V = 1) :
      nsStep a b c (U * Matrix.diagonal σ * V.transpose) = (U * Matrix.diagonal fun (i : Fin n) => nsScalar a b c (σ i)) * V.transpose

      P1 — the spectral-step lemma: a Newton–Schulz step is nsScalar applied per singular value. For X = U (diagonal σ) Vᵀ with U,V orthonormal (UᵀU = VᵀV = 1), nsStep a b c X = U (diagonal (fun i ↦ nsScalar a b c (σ i))) Vᵀ. The singular vectors U,V are carried through untouched; only the singular values move, by the scalar polynomial φ. This is the only matrix-level work in the convergence proof — everything downstream is scalar.

      The proof is pure UᵀU = 1 / VᵀV = 1 collapse algebra (the conj_diag_pow motif of L5): the Gram matrix XXᵀ = U (diagonal σ²) Uᵀ, and each higher monomial (XXᵀ)ᵏX collapses to U (diagonal σ^{2k+1}) Vᵀ because the inner UᵀU contracts to the identity; the three scalar coefficients a, b, c ride through onto the diagonal and sum pointwise to φ.

      theorem Proofs.MuonNewtonSchulz.nsStep_iterate_spectral {n : } (a b c : ) (U V : Matrix (Fin n) (Fin n) ) (σ : Fin n) (hU : U.transpose * U = 1) (hV : V.transpose * V = 1) (k : ) :
      (nsStep a b c)^[k] (U * Matrix.diagonal σ * V.transpose) = (U * Matrix.diagonal fun (i : Fin n) => (nsScalar a b c)^[k] (σ i)) * V.transpose

      P1, iterated: k Newton–Schulz steps act as nsScalar^[k] per singular value. (nsStep a b c)^[k] (U (diagonal σ) Vᵀ) = U (diagonal (fun i ↦ (nsScalar a b c)^[k] (σ i))) Vᵀ. A one-line induction reusing nsStep_spectral at each step: the singular vectors U,V are invariant under the whole orbit, so convergence of the matrix iteration nsStep^[k] X → UVᵀ reduces to the scalar fixed-point convergence (nsScalar a b c)^[k] (σ i) → 1 for each singular value — the entry point for P2 (the cubic monotone argument) and P3 (the matrix-continuity glue).

      theorem Proofs.MuonNewtonSchulz.scalar_iterate_tendsto_one (g : ) (hcont : Continuous g) (hge : ∀ (t : ), 0 tt 1t g t) (hle : ∀ (t : ), 0 tt 1g t 1) (hfix : ∀ (L : ), 0 < LL 1g L = LL = 1) {t₀ : } (h0 : 0 < t₀) (h1 : t₀ 1) :
      Filter.Tendsto (fun (k : ) => g^[k] t₀) Filter.atTop (nhds 1)

      The monotone scalar-convergence engine — convergence is a property of the coefficients, not the degree. Any scalar map g that on [0,1] is a self-map pushing up toward 1 (t ≤ g t and g t ≤ 1) with 1 its only fixed point in (0,1] (hfix) drives every t₀ ∈ (0,1] to 1: g^[k](t₀) → 1. The textbook monotone argument: the orbit is trapped in [t₀,1] (hinv), monotone and bounded above, hence converges to its supremum L ∈ (0,1] (tendsto_atTop_ciSup); continuity makes L a fixed point (isFixedPt_of_tendsto_iterate), and hfix pins L = 1.

      Both the cubic (gCubic_iterate_tendsto_one) and the principled convergent quintic (q5Scalar_iterate_tendsto_one) instantiate this. Muon's tuned quintic (P4) fails g t ≤ 1 — that single broken hypothesis is exactly why it bands instead of converging.

      noncomputable def Proofs.MuonNewtonSchulz.gCubic (t : ) :

      The classic inverse-free cubic g(t) = ½(3t − t³) — the Newton–Schulz scalar map nsScalar (3/2) (−1/2) 0 (gCubic_eq_nsScalar).

      Equations
      Instances For
        theorem Proofs.MuonNewtonSchulz.gCubic_iterate_tendsto_one {t₀ : } (h0 : 0 < t₀) (h1 : t₀ 1) :
        Filter.Tendsto (fun (k : ) => gCubic^[k] t₀) Filter.atTop (nhds 1)

        P2 (cubic) — gCubic^[k](t₀) → 1 on (0,1] (second-order convergent). Instantiates the scalar_iterate_tendsto_one engine: increasing toward 1 on [0,1] (g t − t = t(1−t)(1+t)/2 ≥ 0, 1 − g t = (1−t)²(2+t)/2 ≥ 0) with fixed points g t = t ⟺ t(1−t²)=0 ⟺ t ∈ {0,±1} — only 1 lies in (0,1].

        noncomputable def Proofs.MuonNewtonSchulz.q5Scalar (t : ) :

        The principled convergent quintic q₅(t) = (15t − 10t³ + 3t⁵)/8 — Higham's order-5 Newton–Schulz iteration for the matrix sign / polar function, nsScalar (15/8) (−5/4) (3/8). This is the answer to "does a quintic converge?": yes — if you pick these coefficients. Unlike Muon's tuned quintic (P4), here q₅(1) = 1 is a fixed point and the iteration converges — in fact faster than the cubic (third-order: 1 − q₅(t) = (1−t)³(3t²+9t+8)/8). Convergence is a property of the chosen polynomial, not of its degree.

        Equations
        Instances For
          theorem Proofs.MuonNewtonSchulz.q5Scalar_iterate_tendsto_one {t₀ : } (h0 : 0 < t₀) (h1 : t₀ 1) :

          P2 (principled quintic) — q5Scalar^[k](t₀) → 1 on (0,1] (third-order convergent). The same scalar_iterate_tendsto_one engine as the cubic: q₅ is increasing (q₅′(t) = 15(1−t²)²/8 ≥ 0), q₅ t − t = t(7−3t²)(1−t²)/8 ≥ 0 and 1 − q₅ t = (1−t)³(3t²+9t+8)/8 ≥ 0 on [0,1], and its only fixed point in (0,1] is 1 (q₅ L = L ⟺ L(7−3L²)(1−L)(1+L)=0, and 7−3L² > 0 there).

          theorem Proofs.MuonNewtonSchulz.nsStep_iterate_tendsto_polar {n : } (a b c : ) (g : ) (hbridge : g = nsScalar a b c) (hconv : ∀ (t₀ : ), 0 < t₀t₀ 1Filter.Tendsto (fun (k : ) => g^[k] t₀) Filter.atTop (nhds 1)) (U V : Matrix (Fin n) (Fin n) ) (σ : Fin n) (hU : U.transpose * U = 1) (hV : V.transpose * V = 1) ( : ∀ (i : Fin n), 0 < σ i σ i 1) :

          P3 — the matrix glue: any convergent scalar Newton–Schulz map lifts to nsStep^[k] G → UVᵀ. For a pre-normalized full-rank G = U (diagonal σ) Vᵀ with every singular value σᵢ ∈ (0,1] (the implementation's G / ‖G‖ step), if the scalar map g = nsScalar a b c drives (0,1] → 1 (hconv), then the matmul iterate (nsStep a b c)^[k] G converges to the polar factor U Vᵀ — exactly the object L3–L6 proved optimal (operator-norm steepest descent / nuclear-norm argmax / Shampoo's step / nearest orthogonal matrix). This closes the loop: the thing the hardware computes is the thing the theory says is optimal.

          The §0 spectral reduction cashed out: nsStep_iterate_spectral (P1) makes the matrix iterate U (diagonal (g^[k] ∘ σ)) Vᵀ; each diagonal entry g^[k](σᵢ) → 1 by hconv; pointwise convergence in Fin n → ℝ (tendsto_pi_nhds) plus continuity of d ↦ U (diagonal d) Vᵀ (Continuous.matrix_diagonal/Continuous.matrix_mul) pushes the limit through to U (diagonal 1) Vᵀ = U Vᵀ (Matrix.diagonal_one). The rank hypothesis σᵢ > 0 is what makes the per-value limit 1; σᵢ = 0 would stay 0, giving the partial isometry U (diagonal 1_{σ>0}) Vᵀ.

          theorem Proofs.MuonNewtonSchulz.nsStep_cubic_iterate_tendsto_polar {n : } (U V : Matrix (Fin n) (Fin n) ) (σ : Fin n) (hU : U.transpose * U = 1) (hV : V.transpose * V = 1) ( : ∀ (i : Fin n), 0 < σ i σ i 1) :
          Filter.Tendsto (fun (k : ) => (nsStep (3 / 2) (-1 / 2) 0)^[k] (U * Matrix.diagonal σ * V.transpose)) Filter.atTop (nhds (U * V.transpose))

          The cubic matmul iteration converges to the polar factor UVᵀ — P2 (cubic) through the P3 glue. The classic (3/2, −1/2, 0) Newton–Schulz iteration provably computes Muon's update.

          theorem Proofs.MuonNewtonSchulz.nsStep_q5_iterate_tendsto_polar {n : } (U V : Matrix (Fin n) (Fin n) ) (σ : Fin n) (hU : U.transpose * U = 1) (hV : V.transpose * V = 1) ( : ∀ (i : Fin n), 0 < σ i σ i 1) :
          Filter.Tendsto (fun (k : ) => (nsStep (15 / 8) (-5 / 4) (3 / 8))^[k] (U * Matrix.diagonal σ * V.transpose)) Filter.atTop (nhds (U * V.transpose))

          The principled convergent quintic's matmul iteration also lands on UVᵀ — same polar factor, one degree up, faster (third-order). Convergence is the coefficient choice, not the degree: (15/8, −5/4, 3/8) converges (this theorem), Muon's tuned (3.4445, −4.7750, 2.0315) bands (P4).

          noncomputable def Proofs.MuonNewtonSchulz.qScalar (t : ) :

          Muon's actual tuned Newton–Schulz quintic φ(t) = 3.4445 t − 4.7750 t³ + 2.0315 t⁵ (Jordan 2024 — planning/archive/muon.md). This is not a statement about quintics in general — the principled quintic q5Scalar (15/8, −5/4, 3/8) converges (q5Scalar_iterate_tendsto_one, faster than the cubic even). Convergence is the coefficient choice: Jordan tuned these coefficients for speed to a band near 1 in ~5 steps, deliberately giving up asymptotic convergence. The map straddles 1 (qScalar_one_lt_oneqScalar_half_gt_one) and so oscillates; the theorems below prove the cubic's monotone-convergence hypotheses (scalar_iterate_ tendsto_one) structurally fail here, so one must not state qScalar^[k] → 1.

          Equations
          Instances For

            At the top of the normalized range, the tuned quintic pulls below 1: φ(1) = 0.701 < 1. So 1 is not a fixed point of qScalar (contrast the cubic, whose only relevant fixed point IS 1) — the asymptotic limit the cubic enjoys simply does not exist here.

            In mid-range the tuned quintic overshoots above 1: φ(1/2) ≈ 1.189 > 1. Together with qScalar_one_lt_one this shows qScalar takes values on both sides of 1 — it straddles the target rather than approaching it monotonically.

            theorem Proofs.MuonNewtonSchulz.qScalar_not_le_one :
            ¬∀ (t : ), 0 tt 1qScalar t 1

            The cubic's key bound fails for the tuned quinticqScalar is not ≤ 1 on [0,1] (it overshoots at 1/2). This is exactly the hypothesis g(t) ≤ 1 that powered the monotone-bounded convergence of gCubic_iterate_tendsto_one (P2); its failure is why the clean → 1 proof does not transfer to Muon's quintic, and why claiming qScalar^[k] → 1 would be an overclaim.

            The honest positive statement: a finite-5-step band bound (the form P4 actually supports). Five steps of Muon's tuned quintic from σ = 1/2 land within 0.3 of 1: |qScalar^[5] (1/2) − 1| ≤ 3/10 (the orbit 0.5 → 1.19 → 0.90 → 0.83 → 0.94 → 0.77 oscillates in a band around 1, never reaching it). This matches the implementation's fixed-5-step, "rough is fine — we recompute next optimizer step anyway" design (planning/archive/muon.md): not convergence, a band. The universal interval version ∀ σ ∈ [σ_min, 1], |φ^[5](σ) − 1| ≤ δ is a degree-5⁵ polynomial bound over an interval (genuine interval arithmetic) and is left open by hand.