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:

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

The theorems below connect those contracts to the executable Muon step: if the backend satisfies one of these contracts, the direction used in the parameter update is certified at the same level.

@[reducible, inline]

abbrev Optim.Muon.MatrixTensor (α : Type) (m n : ℕ) : Type

A matrix-shaped TorchLean tensor.

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

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

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

Exact column orthogonality for a matrix-shaped update direction.

def Optim.Muon.ExactOrthogonalizesBuffer {α : Type} [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} [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} [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} [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.

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

An exact matrix Muon orthogonalizer maps every momentum buffer to a direction whose columns have Gram matrix $I$.

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

An approximate matrix Muon orthogonalizer maps every momentum buffer to a direction whose Gram residual is entrywise bounded by $\varepsilon$.

Certified backend records

structure Optim.Muon.ExactCertifiedOrthogonalizer (α : Type) [Context α] (m n : ℕ) : Type

Unconditionally certified exact Muon backend.

Use this when the orthogonalizer is known to return an exact $Q^\mathsf{T}Q=I$ direction for every buffer of a fixed matrix shape.

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

  Executable orthogonalizer used by Muon.

• certified : ExactMatrixOrthogonalizer self.orthogonalizer

  The backend returns an exactly orthogonalized direction for every buffer.

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

Unconditionally certified approximate Muon backend.

Use this when the orthogonalizer is known to return a direction whose Gram residual is bounded by $\varepsilon$ for every buffer of a fixed matrix shape.

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

  Executable orthogonalizer used by Muon.

• certified : ApproxMatrixOrthogonalizer eps self.orthogonalizer

  The backend returns an approximately orthogonalized direction for every buffer.

structure Optim.Muon.CheckedExactOrthogonalizer (α : Type) [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 below is the first instance: its success predicate is positive executable R pivots.

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