Muon Orthogonalization Contracts

Muon is more than momentum with a different name: after the momentum buffer is updated, the parameter direction is supposed to be an orthogonalized matrix direction. The runtime optimizer therefore takes an explicit orthogonalizer backend, while this file records the proof contract that such a backend should satisfy.

There are two useful levels, each stated for one momentum buffer:

• ExactOrthogonalizesBuffer says the backend returns a matrix $Q$ with $Q^\mathsf{T}Q=I$.
• ApproxOrthogonalizesBuffer eps says the Gram residual $Q^\mathsf{T}Q-I$ is entrywise bounded by $\varepsilon$.

CheckedExactOrthogonalizer and CheckedApproxOrthogonalizer package a backend with the precondition under which it meets one of them.

NN.MLTheory.Optimization.Muon.Certificates connects those contracts to the executable Muon step: if the backend satisfies one of them on the fresh momentum buffer, the direction used in the parameter update is certified at the same level. Concrete backends live in NN.MLTheory.Optimization.Muon.QR (exact, under positive QR pivots) and NN.MLTheory.Optimization.Muon.NewtonSchulz (residual-checked).

@[reducible, inline]

abbrev Optim.Muon.MatrixTensor (α : Type) [TorchLean.Storage α] (m n : ℕ) : Type

A matrix-shaped TorchLean tensor.

def Optim.Muon.columnGram {α : Type} [TorchLean.Storage α] [Context α] {m n : ℕ} (Q : MatrixTensor α m n) : MatrixTensor α n n

The column Gram matrix $Q^\mathsf{T}Q$.

def Optim.Muon.HasExactColumnGram {α : Type} [TorchLean.Storage α] [Context α] {m n : ℕ} (Q : MatrixTensor α m n) : Prop

Exact column orthogonality for a matrix-shaped update direction.

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

A backend exactly orthogonalizes one specified momentum buffer.

def Optim.Muon.columnGramResidual {α : Type} [TorchLean.Storage α] [Context α] {m n : ℕ} (Q : MatrixTensor α m n) : MatrixTensor α n n

Residual matrix $Q^\mathsf{T}Q-I$, used for approximate orthogonalization certificates.

def Optim.Muon.HasApproxColumnGram {α : Type} [TorchLean.Storage α] [Context α] {m n : ℕ} (eps : α) (Q : MatrixTensor α m n) : Prop

Entrywise approximate column orthogonality.

For an exact backend use HasExactColumnGram. For Newton-Schulz or CUDA implementations, this is the certificate shape we want the backend to establish or export: every entry of $Q^\mathsf{T}Q-I$ is bounded by $\varepsilon$.

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

A backend approximately orthogonalizes one specified momentum buffer.

This is also the lightest sound checker boundary we can put on an approximate backend: after it returns a direction, prove or check that the direction's Gram residual is bounded by $\varepsilon$, and nothing more. residualCheckedApproxOrthogonalizer in NN.MLTheory.Optimization.Muon.NewtonSchulz uses it exactly that way.

Certified backend records

structure Optim.Muon.CheckedExactOrthogonalizer (α : Type) [TorchLean.Storage α] [Context α] (m n : ℕ) : Type

Checked exact Muon backend with a per-buffer success predicate.

This is the practical interface for algorithms whose correctness has preconditions on the matrix being orthogonalized. The QR backend in NN.MLTheory.Optimization.Muon.QR is the first instance: its success predicate is positive pivots of the real-valued QR specification.

• orthogonalizer : Orthogonalizer α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))

  Executable orthogonalizer used by Muon.

• Success : MatrixTensor α m n → Prop

  Backend-specific success/certification condition for an input buffer.

• certified (buffer : MatrixTensor α m n) : self.Success buffer → ExactOrthogonalizesBuffer self.orthogonalizer buffer

  Whenever Success buffer holds, the backend exactly orthogonalizes that buffer.

structure Optim.Muon.CheckedApproxOrthogonalizer (α : Type) [TorchLean.Storage α] [Context α] (m n : ℕ) (eps : α) : Type

Checked approximate Muon backend with a per-buffer success predicate.

This is the intended proof shape for Newton-Schulz/CUDA-style backends: the kernel may be fast and approximate, but the exported proof or checker must establish Success buffer, which then gives the entrywise $Q^\mathsf{T}Q-I$ bound.

• orthogonalizer : Orthogonalizer α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))

  Executable orthogonalizer used by Muon.

• Success : MatrixTensor α m n → Prop

  Backend-specific success/certification condition for an input buffer.

• certified (buffer : MatrixTensor α m n) : self.Success buffer → ApproxOrthogonalizesBuffer eps self.orthogonalizer buffer

  Whenever Success buffer holds, the backend approximately orthogonalizes that buffer.