Muon Step Certificates

This module connects a backend's orthogonalization contract to the direction, state, and parameter values produced by one executable Muon update. The generic checked-backend theorems are the public proof interface; concrete QR and Newton-Schulz backends instantiate them in their own modules.

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

Evidence that direction is exactly the output of an orthogonalizer on buffer and has column Gram matrix $I$.

• direction_eq : direction = orthogonalizer.apply buffer

  The certified direction is the backend output.

• exact_column_gram : HasExactColumnGram direction

  The certified direction has exact column Gram $I$.

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

Evidence that direction is exactly the output of an orthogonalizer on buffer and has an entrywise column-Gram residual bounded by $\varepsilon$.

• direction_eq : direction = orthogonalizer.apply buffer

  The certified direction is the backend output.

• approx_column_gram : HasApproxColumnGram eps direction

  The certified direction satisfies the requested residual bound.

structure Optim.Muon.ExactCertifiedStep {α : Type} [Context α] {m n : ℕ} (state : State α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (params grads direction : MatrixTensor α m n) : Prop

Certificate for one exact Muon update: the fresh momentum buffer is orthogonalized, the new state stores that buffer, and the parameters move along the certified direction.

• direction_cert : ExactCertifiedDirection state.orthogonalizer (update state params grads).1.buf direction

  Certificate for the direction computed from the fresh momentum buffer.

• state_eq : (update state params grads).1 = { lr := state.lr, momentum := state.momentum, buf := OptimizerUtils.updateMomentumBuf state.buf state.momentum grads, orthogonalizer := state.orthogonalizer }

  Muon changes only the momentum buffer in its optimizer state.

• params_eq : (update state params grads).2 = Spec.Tensor.subSpec params (Spec.Tensor.scaleSpec direction state.lr)

  Parameter equation for the certified update direction.

structure Optim.Muon.ApproxCertifiedStep {α : Type} [Context α] {m n : ℕ} (eps : α) (state : State α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (params grads direction : MatrixTensor α m n) : Prop

The residual-bounded counterpart of ExactCertifiedStep.

• direction_cert : ApproxCertifiedDirection eps state.orthogonalizer (update state params grads).1.buf direction

  Certificate for the direction computed from the fresh momentum buffer.

• state_eq : (update state params grads).1 = { lr := state.lr, momentum := state.momentum, buf := OptimizerUtils.updateMomentumBuf state.buf state.momentum grads, orthogonalizer := state.orthogonalizer }

  Muon changes only the momentum buffer in its optimizer state.

• params_eq : (update state params grads).2 = Spec.Tensor.subSpec params (Spec.Tensor.scaleSpec direction state.lr)

  Parameter equation for the certified update direction.

theorem Optim.Muon.exactCertifiedStep_of_buffer {α : Type} [Context α] {m n : ℕ} (state : State α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (params grads : MatrixTensor α m n) (horth : ExactOrthogonalizesBuffer state.orthogonalizer (update state params grads).1.buf) : ∃ (direction : MatrixTensor α m n), ExactCertifiedStep state params grads direction

A local exact backend fact for the fresh buffer produces a certified Muon step.

theorem Optim.Muon.approxCertifiedStep_of_buffer {α : Type} [Context α] {m n : ℕ} {eps : α} (state : State α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (params grads : MatrixTensor α m n) (horth : ApproxOrthogonalizesBuffer eps state.orthogonalizer (update state params grads).1.buf) : ∃ (direction : MatrixTensor α m n), ApproxCertifiedStep eps state params grads direction

A local residual bound for the fresh buffer produces a certified Muon step.

theorem Optim.Muon.exactCertifiedStep_of_checkedBackend {α : Type} [Context α] {m n : ℕ} (backend : CheckedExactOrthogonalizer α m n) (lr momentum : α) (buf params grads : MatrixTensor α m n) (hsuccess : backend.Success (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := backend.orthogonalizer } params grads).1.buf) : ∃ (direction : MatrixTensor α m n), ExactCertifiedStep { lr := lr, momentum := momentum, buf := buf, orthogonalizer := backend.orthogonalizer } params grads direction

A checked exact backend certifies the concrete direction and equations of one Muon update whenever its success predicate holds on the fresh momentum buffer.

theorem Optim.Muon.approxCertifiedStep_of_checkedBackend {α : Type} [Context α] {m n : ℕ} {eps : α} (backend : CheckedApproxOrthogonalizer α m n eps) (lr momentum : α) (buf params grads : MatrixTensor α m n) (hsuccess : backend.Success (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := backend.orthogonalizer } params grads).1.buf) : ∃ (direction : MatrixTensor α m n), ApproxCertifiedStep eps { lr := lr, momentum := momentum, buf := buf, orthogonalizer := backend.orthogonalizer } params grads direction

A checked approximate backend certifies one Muon update whenever its success predicate establishes the requested Gram-residual bound on the fresh momentum buffer.

theorem Optim.Muon.checkedBackend_updateDirection_hasExactColumnGram {α : Type} [Context α] {m n : ℕ} (backend : CheckedExactOrthogonalizer α m n) (lr momentum : α) (buf params grads : MatrixTensor α m n) (hsuccess : backend.Success (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := backend.orthogonalizer } params grads).1.buf) : HasExactColumnGram (backend.orthogonalizer.apply (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := backend.orthogonalizer } params grads).1.buf)

A checked exact backend gives $Q^\mathsf{T}Q=I$ for the direction used by an update.

theorem Optim.Muon.checkedBackend_updateDirection_hasApproxColumnGram {α : Type} [Context α] {m n : ℕ} {eps : α} (backend : CheckedApproxOrthogonalizer α m n eps) (lr momentum : α) (buf params grads : MatrixTensor α m n) (hsuccess : backend.Success (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := backend.orthogonalizer } params grads).1.buf) : HasApproxColumnGram eps (backend.orthogonalizer.apply (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := backend.orthogonalizer } params grads).1.buf)

A checked approximate backend gives its residual bound for the direction used by an update.

theorem Optim.Muon.ExactCertifiedStep.hasExactColumnGram {α : Type} [Context α] {m n : ℕ} {state : State α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))} {params grads direction : MatrixTensor α m n} (cert : ExactCertifiedStep state params grads direction) : HasExactColumnGram direction

Extract exact column orthogonality from a certified step.

theorem Optim.Muon.ApproxCertifiedStep.hasApproxColumnGram {α : Type} [Context α] {m n : ℕ} {eps : α} {state : State α (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))} {params grads direction : MatrixTensor α m n} (cert : ApproxCertifiedStep eps state params grads direction) : HasApproxColumnGram eps direction

Extract the Gram-residual bound from an approximate certified step.