Newton-Schulz Muon Backends

The column-oriented Newton-Schulz polynomial orthogonalizer, its residual-checked packaging, and the fixed-point condition under which it is an exact backend. Real-valued facts about the polynomial step at an exactly column-orthogonal matrix live in NN.MLTheory.Optimization.Muon.QR, next to the entrywise matrix API they need.

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 HasExactColumnGram. 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.newtonSchulzStep {α : Type} [TorchLean.Storage α] [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) (X : MatrixTensor α m n) : MatrixTensor α m n

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

def Optim.Muon.newtonSchulzIter {α : Type} [TorchLean.Storage α] [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.newtonSchulzOrthogonalizer {α : Type} [TorchLean.Storage α] [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.residualCheckedApproxOrthogonalizer {α : Type} [TorchLean.Storage α] [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} [TorchLean.Storage α] [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} [TorchLean.Storage α] [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} [TorchLean.Storage α] [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} [TorchLean.Storage α] [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.newtonSchulzOrthogonalizer_exactOrthogonalizesBuffer_of_fixedPoint {α : Type} [TorchLean.Storage α] [Context α] {m n : ℕ} (coeffs : NewtonSchulzCoeffs α) (steps : ℕ) (buffer : MatrixTensor α m n) (hgram : HasExactColumnGram buffer) (hfixed : NewtonSchulzFixedPoint coeffs buffer) : ExactOrthogonalizesBuffer (newtonSchulzOrthogonalizer coeffs steps) buffer

An exact-column-orthogonal buffer that is fixed by one Newton-Schulz step is exactly orthogonalized by every finite Newton-Schulz iteration.

def Optim.Muon.newtonSchulzFixedPointCheckedExactOrthogonalizer {α : Type} [TorchLean.Storage α] [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.