Muon Certificate Examples

These examples show how to obtain and consume the exact and approximate certificates attached to Muon steps. Runtime code configures Muon through TorchLean.optim.runtimeMuon; proofs about the algorithm use the canonical Optim.Muon namespace.

theorem NN.Examples.Optimization.MuonCertificates.qr_update_step_direction_has_exact_gram {m n : ℕ} (lr momentum : ℝ) (buf params grads : Optim.Muon.MatrixTensor ℝ m n) (hpivots : Optim.Muon.HasPositiveQRPivots (Optim.Muon.update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := Optim.Muon.qrOrthogonalizer } params grads).1.buf) : ∃ (direction : Optim.Muon.MatrixTensor ℝ m n), Optim.Muon.HasExactColumnGram direction ∧ (Optim.Muon.update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := Optim.Muon.qrOrthogonalizer } params grads).2 = Spec.Tensor.subSpec params (Spec.Tensor.scaleSpec direction lr)

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} [Context α] {m n : ℕ} {eps : α} (coeffs : Optim.Muon.NewtonSchulzCoeffs α) (steps : ℕ) (lr momentum : α) (buf params grads : Optim.Muon.MatrixTensor α m n) (hresidual : Optim.Muon.ResidualApproxSuccess eps (Optim.Muon.newtonSchulzOrthogonalizer coeffs steps) (Optim.Muon.update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := Optim.Muon.newtonSchulzOrthogonalizer coeffs steps } params grads).1.buf) : ∃ (direction : Optim.Muon.MatrixTensor α m n), Optim.Muon.HasApproxColumnGram eps direction ∧ (Optim.Muon.update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := Optim.Muon.newtonSchulzOrthogonalizer coeffs steps } params grads).2 = Spec.Tensor.subSpec params (Spec.Tensor.scaleSpec direction lr)

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.

theorem NN.Examples.Optimization.MuonCertificates.newtonSchulz_fixed_update_step_direction_has_exact_gram {α : Type} [Context α] {m n : ℕ} (coeffs : Optim.Muon.NewtonSchulzCoeffs α) (steps : ℕ) (lr momentum : α) (buf params grads : Optim.Muon.MatrixTensor α m n) (hsuccess : (Optim.Muon.newtonSchulzFixedPointCheckedExactOrthogonalizer coeffs steps).Success (Optim.Muon.update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := Optim.Muon.newtonSchulzOrthogonalizer coeffs steps } params grads).1.buf) : ∃ (direction : Optim.Muon.MatrixTensor α m n), Optim.Muon.HasExactColumnGram direction ∧ (Optim.Muon.update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := Optim.Muon.newtonSchulzOrthogonalizer coeffs steps } params grads).2 = Spec.Tensor.subSpec params (Spec.Tensor.scaleSpec direction lr)

If the fresh momentum buffer is already exact-column-orthogonal and fixed by one Newton-Schulz step, the fixed-point checked backend upgrades Newton-Schulz from an approximate residual-checked path to an exact certified step.

theorem NN.Examples.Optimization.MuonCertificates.newtonSchulz_step_scales_exact_gram_matrix {m n : ℕ} (coeffs : Optim.Muon.NewtonSchulzCoeffs ℝ) (Q : Optim.Muon.MatrixTensor ℝ m n) (hgram : Optim.Muon.HasExactColumnGram Q) : Optim.Muon.newtonSchulzStep coeffs Q = Spec.Tensor.scaleSpec Q (coeffs.a + coeffs.b + coeffs.c)

At an exact-column-orthogonal matrix, one real Newton-Schulz step acts by the scalar $a+b+c$. This is the algebra behind the fixed-point shortcut below.

theorem NN.Examples.Optimization.MuonCertificates.newtonSchulz_step_preserves_exact_gram_when_sum_squares_one {m n : ℕ} (coeffs : Optim.Muon.NewtonSchulzCoeffs ℝ) (Q : Optim.Muon.MatrixTensor ℝ m n) (hgram : Optim.Muon.HasExactColumnGram Q) (hsquare : (coeffs.a + coeffs.b + coeffs.c) * (coeffs.a + coeffs.b + coeffs.c) = 1) : Optim.Muon.HasExactColumnGram (Optim.Muon.newtonSchulzStep coeffs Q)

The same algebra shows exact column Gram is preserved when the coefficient sum squares to one, including the sign-flip case.

theorem NN.Examples.Optimization.MuonCertificates.newtonSchulz_step_has_approx_gram_from_coeff_square_error {m n : ℕ} (coeffs : Optim.Muon.NewtonSchulzCoeffs ℝ) (Q : Optim.Muon.MatrixTensor ℝ m n) (eps : ℝ) (hgram : Optim.Muon.HasExactColumnGram Q) (herr : MathFunctions.abs ((coeffs.a + coeffs.b + coeffs.c) * (coeffs.a + coeffs.b + coeffs.c) - 1) ≤ eps) (heps : 0 ≤ eps) : Optim.Muon.HasApproxColumnGram eps (Optim.Muon.newtonSchulzStep coeffs Q)

If the coefficient sum is only approximately square-one, the same one-step algebra gives the entrywise residual bound used by the approximate Muon certificate.

theorem NN.Examples.Optimization.MuonCertificates.newtonSchulz_coeff_sum_update_step_direction_has_exact_gram {m n : ℕ} (coeffs : Optim.Muon.NewtonSchulzCoeffs ℝ) (steps : ℕ) (lr momentum : ℝ) (buf params grads : Optim.Muon.MatrixTensor ℝ m n) (hgram : Optim.Muon.HasExactColumnGram (Optim.Muon.update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := Optim.Muon.newtonSchulzOrthogonalizer coeffs steps } params grads).1.buf) (hsum : coeffs.a + coeffs.b + coeffs.c = 1) : ∃ (direction : Optim.Muon.MatrixTensor ℝ m n), Optim.Muon.HasExactColumnGram direction ∧ (Optim.Muon.update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := Optim.Muon.newtonSchulzOrthogonalizer coeffs steps } params grads).2 = Spec.Tensor.subSpec params (Spec.Tensor.scaleSpec direction lr)

Over ℝ, the common coefficient condition $a+b+c=1$ means an already exact-column-orthogonal fresh buffer is automatically a fixed point, so the Newton-Schulz Muon update is exactly certified.

theorem NN.Examples.Optimization.MuonCertificates.qr_initialized_step_direction_has_exact_gram {m n : ℕ} (lr momentum : ℝ) (params grads : Optim.Muon.MatrixTensor ℝ m n) (hpivots : Optim.Muon.HasPositiveQRPivots (Optim.Muon.update (Optim.Muon.init lr momentum Optim.Muon.qrOrthogonalizer params) params grads).1.buf) : ∃ (direction : Optim.Muon.MatrixTensor ℝ m n), Optim.Muon.HasExactColumnGram direction ∧ (Optim.Muon.update (Optim.Muon.init lr momentum Optim.Muon.qrOrthogonalizer params) params grads).2 = Spec.Tensor.subSpec params (Spec.Tensor.scaleSpec direction lr)

The initialized QR theorem has the same proof shape as the stateful update theorem, but starts from the state produced by Optim.Muon.init.

theorem NN.Examples.Optimization.MuonCertificates.qr_initialized_step_state_eq {m n : ℕ} (lr momentum : ℝ) (params grads : Optim.Muon.MatrixTensor ℝ m n) (hpivots : Optim.Muon.HasPositiveQRPivots (Optim.Muon.update (Optim.Muon.init lr momentum Optim.Muon.qrOrthogonalizer params) params grads).1.buf) : (Optim.Muon.update (Optim.Muon.init lr momentum Optim.Muon.qrOrthogonalizer params) params grads).1 = have __src := Optim.Muon.init lr momentum Optim.Muon.qrOrthogonalizer params; { lr := __src.lr, momentum := __src.momentum, buf := Optim.OptimizerUtils.updateMomentumBuf (Optim.Muon.init lr momentum Optim.Muon.qrOrthogonalizer params).buf (Optim.Muon.init lr momentum Optim.Muon.qrOrthogonalizer params).momentum grads, orthogonalizer := __src.orthogonalizer }

The same QR certificate also exposes the whole next-state equation, not only the direction certificate.

theorem NN.Examples.Optimization.MuonCertificates.newtonSchulz_initialized_step_direction_has_approx_gram {α : Type} [Context α] {m n : ℕ} {eps : α} (coeffs : Optim.Muon.NewtonSchulzCoeffs α) (steps : ℕ) (lr momentum : α) (params grads : Optim.Muon.MatrixTensor α m n) (hresidual : Optim.Muon.ResidualApproxSuccess eps (Optim.Muon.newtonSchulzOrthogonalizer coeffs steps) (Optim.Muon.update (Optim.Muon.init lr momentum (Optim.Muon.newtonSchulzOrthogonalizer coeffs steps) params) params grads).1.buf) : ∃ (direction : Optim.Muon.MatrixTensor α m n), Optim.Muon.HasApproxColumnGram eps direction ∧ (Optim.Muon.update (Optim.Muon.init lr momentum (Optim.Muon.newtonSchulzOrthogonalizer coeffs steps) params) params grads).2 = Spec.Tensor.subSpec params (Spec.Tensor.scaleSpec direction lr)

The initialized residual-checked Newton-Schulz theorem is the training-path version: initialize the Muon state, run one update, then consume the certified step.

theorem NN.Examples.Optimization.MuonCertificates.newtonSchulz_initialized_step_state_eq {α : Type} [Context α] {m n : ℕ} {eps : α} (coeffs : Optim.Muon.NewtonSchulzCoeffs α) (steps : ℕ) (lr momentum : α) (params grads : Optim.Muon.MatrixTensor α m n) (hresidual : Optim.Muon.ResidualApproxSuccess eps (Optim.Muon.newtonSchulzOrthogonalizer coeffs steps) (Optim.Muon.update (Optim.Muon.init lr momentum (Optim.Muon.newtonSchulzOrthogonalizer coeffs steps) params) params grads).1.buf) : (Optim.Muon.update (Optim.Muon.init lr momentum (Optim.Muon.newtonSchulzOrthogonalizer coeffs steps) params) params grads).1 = have __src := Optim.Muon.init lr momentum (Optim.Muon.newtonSchulzOrthogonalizer coeffs steps) params; { lr := __src.lr, momentum := __src.momentum, buf := Optim.OptimizerUtils.updateMomentumBuf (Optim.Muon.init lr momentum (Optim.Muon.newtonSchulzOrthogonalizer coeffs steps) params).buf (Optim.Muon.init lr momentum (Optim.Muon.newtonSchulzOrthogonalizer coeffs steps) params).momentum grads, orthogonalizer := __src.orthogonalizer }

The initialized residual-checked Newton-Schulz certificate also exposes the whole next-state equation.