Newton-Schulz Muon Backends

Polynomial orthogonalizers, residual checks, and fixed-point conditions used by Muon.

structure Optim.Muon.NewtonSchulzCoeffs (α : Type) : Type

Coefficients for the odd Newton-Schulz polynomial used by Muon-style orthogonalization.

The column-oriented shape is $X\mapsto aX+bX(X^\mathsf{T}X)+cX(X^\mathsf{T}X)^2$, matching the $Q^\mathsf{T}Q$ certificate used below. TorchLean keeps the coefficients explicit so experiments and backend-specific proofs can choose the polynomial they actually use.

• a : α

  Linear coefficient.

• b : α

  Cubic coefficient.

• c : α

  Quintic coefficient.

def Optim.Muon.leftGram {α : Type} [Context α] {m n : ℕ} (X : MatrixTensor α m n) : MatrixTensor α m m

Left Gram matrix $XX^\mathsf{T}$, useful for row-oriented rectangular Newton-Schulz updates.

def Optim.Muon.rightGram {α : Type} [Context α] {m n : ℕ} (X : MatrixTensor α m n) : MatrixTensor α n n

Right/column Gram matrix $X^\mathsf{T}X$, matching TorchLean's column-orthogonality certificate.

def Optim.Muon.newtonSchulzLeftStep {α : Type} [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) (X : MatrixTensor α m n) : MatrixTensor α m n

One row-oriented Newton-Schulz polynomial step using $XX^\mathsf{T}$.

def Optim.Muon.newtonSchulzStep {α : Type} [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) (X : MatrixTensor α m n) : MatrixTensor α m n

One column-oriented Newton-Schulz polynomial step using $X^\mathsf{T}X$.

def Optim.Muon.newtonSchulzLeftIter {α : Type} [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) : ℕ → MatrixTensor α m n → MatrixTensor α m n

Iterate the row-oriented Newton-Schulz polynomial step a fixed number of times.

def Optim.Muon.newtonSchulzIter {α : Type} [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) : ℕ → MatrixTensor α m n → MatrixTensor α m n

Iterate the column-oriented Newton-Schulz polynomial step a fixed number of times.

def Optim.Muon.newtonSchulzLeftOrthogonalizer {α : Type} [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) (steps : ℕ) : Orthogonalizer α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))

Row-oriented Newton-Schulz-shaped orthogonalizer backend.

def Optim.Muon.newtonSchulzOrthogonalizer {α : Type} [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) (steps : ℕ) : Orthogonalizer α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))

Column-oriented Newton-Schulz-shaped orthogonalizer backend.

def Optim.Muon.ResidualApproxSuccess {α : Type} [Context α] {m n : ℕ} (eps : α) (orthogonalizer : Orthogonalizer α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (buffer : MatrixTensor α m n) : Prop

Residual-check success predicate for approximate Muon backends.

This is the lightest sound checker boundary: after a backend returns a direction, prove or check that the direction's Gram residual is bounded by $\varepsilon$.

def Optim.Muon.residualCheckedApproxOrthogonalizer {α : Type} [Context α] {m n : ℕ} (eps : α) (orthogonalizer : Orthogonalizer α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) : CheckedApproxOrthogonalizer α m n eps

Turn any orthogonalizer into a checked approximate backend by using the Gram-residual bound itself as the success predicate.

def Optim.Muon.newtonSchulzResidualCheckedOrthogonalizer {α : Type} [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) (steps : ℕ) (eps : α) : CheckedApproxOrthogonalizer α m n eps

Newton-Schulz packaged as a checked approximate backend.

The backend is the explicit Newton-Schulz tensor program, and the success predicate is the post-check that its returned direction satisfies the requested Gram-residual bound.

def Optim.Muon.NewtonSchulzFixedPoint {α : Type} [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) (buffer : MatrixTensor α m n) : Prop

A buffer is a fixed point of one column-oriented Newton-Schulz step.

This is a deliberately local condition. It does not assert convergence from arbitrary inputs; it records the exact algebraic fact needed when a backend has already reached a stable direction.

theorem Optim.Muon.newtonSchulzIter_fixed_of_step_fixed {α : Type} [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) (steps : ℕ) (buffer : MatrixTensor α m n) (hfixed : NewtonSchulzFixedPoint coeffs buffer) : newtonSchulzIter coeffs steps buffer = buffer

If one Newton-Schulz step fixes a buffer, then any finite number of Newton-Schulz iterations fixes the same buffer.

theorem Optim.Muon.newtonSchulzOrthogonalizer_fixed_apply {α : Type} [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) (steps : ℕ) (buffer : MatrixTensor α m n) (hfixed : NewtonSchulzFixedPoint coeffs buffer) : (newtonSchulzOrthogonalizer coeffs steps).apply buffer = buffer

A Newton-Schulz fixed point is returned unchanged by the corresponding orthogonalizer.

theorem Optim.Muon.newtonSchulzFixedPoint_exactOrthogonalizesBuffer {α : Type} [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) (steps : ℕ) (buffer : MatrixTensor α m n) (hgram : HasExactColumnGram buffer) (hfixed : NewtonSchulzFixedPoint coeffs buffer) : ExactOrthogonalizesBuffer (newtonSchulzOrthogonalizer coeffs steps) buffer

If a buffer already has exact column Gram and is fixed by one Newton-Schulz step, then the finite-iteration Newton-Schulz backend exactly orthogonalizes that buffer.

def Optim.Muon.newtonSchulzFixedPointCheckedExactOrthogonalizer {α : Type} [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) (steps : ℕ) : CheckedExactOrthogonalizer α m n

Newton-Schulz packaged as a checked exact backend for already-stable directions.

The success predicate says that the input buffer already has exact column Gram and is a fixed point of one Newton-Schulz step. Under that explicit condition, every finite Newton-Schulz iteration returns the same exact-column-orthogonal direction.