What a Muon Step Certificate Says

A Muon update first forms a momentum buffer, then chooses a matrix direction Q, and finally updates parameters by P' = P - learningRate * Q. An exact step certificate proves that the chosen direction has orthonormal columns (QᵀQ = I) and is the direction actually used in that update. It also records the new momentum state.

Start with Concrete below. Its one-column direction is Q = (3/5, 4/5)ᵀ, so the entire Gram condition is (3/5)² + (4/5)² = 1. Lean proves this equation, rejects (1, 1)ᵀ, constructs a real Optim.Muon.ExactCertifiedStep, and derives the updated parameters (97/50, 73/25)ᵀ.

This example uses exact real arithmetic and an already-normalized buffer. The identity backend is sufficient for this one buffer; it does not orthogonalize arbitrary inputs. The later QR and Newton–Schulz examples show which additional backend hypotheses a general application must supply. None of these step certificates proves convergence, lower loss, speed, or agreement with CUDA.

Build with lake build NN.Examples.Optimization. This is a proof tutorial, with no CLI or training run. Runtime configuration uses TorchLean.optim.muon.optimizer; theorem statements use Optim.Muon.

A certificate with actual numbers

noncomputable def NN.Examples.Optimization.MuonCertificates.Concrete.column (x y : ℝ) : TorchLean.Tensor ℝ [2, 1]

A two-row, one-column matrix containing x above y.

@[simp]

theorem NN.Examples.Optimization.MuonCertificates.Concrete.get2_column (x y : ℝ) (i : Fin 2) (j : Fin 1) : Spec.get2 (column x y) i j = if i = 0 then x else y

Reading the column exposes the two numbers used in the scalar calculations below.

noncomputable def NN.Examples.Optimization.MuonCertificates.Concrete.direction : TorchLean.Tensor ℝ [2, 1]

The 3–4–5 triangle supplies a unit column, Q = (3/5, 4/5)ᵀ.

theorem NN.Examples.Optimization.MuonCertificates.Concrete.direction_has_exact_gram : Optim.Muon.HasExactColumnGram direction

The single entry of QᵀQ is (3/5)² + (4/5)² = 1.

theorem NN.Examples.Optimization.MuonCertificates.Concrete.unnormalized_direction_rejected : ¬Optim.Muon.HasExactColumnGram (column 1 1)

The column (1, 1)ᵀ fails the certificate: its squared length is 2, not 1.

Now use Q as the gradient, start from parameters (2, 3)ᵀ, and set momentum to zero. The fresh buffer is therefore 0 * oldBuffer + Q = Q. Passing it through the identity backend leaves the already-proved unit column unchanged. A non-unit gradient would require a different backend or would fail the exact Gram obligation, as the preceding counterexample shows.

noncomputable def NN.Examples.Optimization.MuonCertificates.Concrete.parameters : TorchLean.Tensor ℝ [2, 1]

Parameters before the step.

noncomputable def NN.Examples.Optimization.MuonCertificates.Concrete.state : Optim.Muon.State ℝ [2, 1]

Learning rate 1/10, zero momentum, and a backend valid for this already-unit direction.

theorem NN.Examples.Optimization.MuonCertificates.Concrete.fresh_buffer_eq_direction : (Optim.Muon.update state parameters direction).optimizerState.momentumBuffer = direction

Zero momentum makes the fresh buffer exactly the supplied gradient.

theorem NN.Examples.Optimization.MuonCertificates.Concrete.step_certified : Optim.Muon.ExactCertifiedStep state parameters direction direction

A complete exact step certificate for these concrete inputs, with no backend hypothesis left open. The first field proves the actual direction and its unit Gram matrix; the remaining fields tie that direction to the next optimizer state and parameter update.

theorem NN.Examples.Optimization.MuonCertificates.Concrete.updated_parameters_eq : (Optim.Muon.update state parameters direction).parameters = column (97 / 50) (73 / 25)

Consuming the certificate gives (2, 3)ᵀ - (1/10) * (3/5, 4/5)ᵀ = (97/50, 73/25)ᵀ.

Supplying a general orthogonalization backend

For multiple columns, QᵀQ = I says that each column has length one and distinct columns have inner product zero. The QR example requires positive pivots. The approximate Newton–Schulz example requires a proved entrywise residual bound on QᵀQ - I; choosing an iteration count alone does not establish that bound.

These are conditional API examples. Unlike Concrete.step_certified, they leave the backend's success obligation to the caller and then extract the useful fields of the resulting certificate.

theorem NN.Examples.Optimization.MuonCertificates.qr_update_step_direction_has_exact_gram {m n : ℕ} (learningRate momentum : ℝ) (momentumBuffer parameters gradients : TorchLean.Tensor ℝ [m, n]) (hpivots : Optim.Muon.HasPositiveQRPivots (Optim.Muon.update { learningRate := learningRate, momentum := momentum, momentumBuffer := momentumBuffer, orthogonalizer := Optim.Muon.qrOrthogonalizer } parameters gradients).optimizerState.momentumBuffer) : ∃ (direction : TorchLean.Tensor ℝ [m, n]), Optim.Muon.HasExactColumnGram direction ∧ (Optim.Muon.update { learningRate := learningRate, momentum := momentum, momentumBuffer := momentumBuffer, orthogonalizer := Optim.Muon.qrOrthogonalizer } parameters gradients).parameters = parameters.subSpec (direction.scaleSpec learningRate)

Using the QR checked backend, a positive-pivot proof gives a certified step; from that step we can recover both the exact Gram certificate for the direction and the parameter-update equation.

theorem NN.Examples.Optimization.MuonCertificates.newtonSchulz_update_step_direction_has_approx_gram {α : Type} [TorchLean.Storage α] [Context α] {m n : ℕ} {eps : α} (coeffs : Optim.Muon.NewtonSchulzCoeffs α) (steps : ℕ) (learningRate momentum : α) (momentumBuffer parameters gradients : TorchLean.Tensor α [m, n]) (hresidual : Optim.Muon.ApproxOrthogonalizesBuffer eps (Optim.Muon.newtonSchulzOrthogonalizer coeffs steps) (Optim.Muon.update { learningRate := learningRate, momentum := momentum, momentumBuffer := momentumBuffer, orthogonalizer := Optim.Muon.newtonSchulzOrthogonalizer coeffs steps } parameters gradients).optimizerState.momentumBuffer) : ∃ (direction : TorchLean.Tensor α [m, n]), Optim.Muon.HasApproxColumnGram eps direction ∧ (Optim.Muon.update { learningRate := learningRate, momentum := momentum, momentumBuffer := momentumBuffer, orthogonalizer := Optim.Muon.newtonSchulzOrthogonalizer coeffs steps } parameters gradients).parameters = parameters.subSpec (direction.scaleSpec learningRate)

Using the residual-checked Newton-Schulz backend, a residual proof gives a certified approximate step; from that step we can recover the residual-bound certificate and the parameter equation.