QR Muon Backend

The real-valued QR orthogonalizer, its positive-pivot condition, and the exact certificates it supplies to Muon updates.

noncomputable def Optim.Muon.qrOrthogonalizer {m n : ℕ} : Orthogonalizer ℝ (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))

QR/Gram-Schmidt orthogonalizer: return the Q factor of the fresh matrix buffer.

def Optim.Muon.HasPositiveQRPivots {m n : ℕ} (buffer : MatrixTensor ℝ m n) : Prop

The success condition for TorchLean's executable QR orthogonalizer.

theorem Optim.Muon.get2_identityTensorSpec_real {n : ℕ} (i j : Fin n) : Spec.get2 (Spec.identityTensorSpec n) i j = if i = j then 1 else 0

theorem Optim.Muon.get2_addSpec_real {m n : ℕ} (A B : MatrixTensor ℝ m n) (i : Fin m) (j : Fin n) : Spec.get2 (Spec.Tensor.addSpec A B) i j = Spec.get2 A i j + Spec.get2 B i j

Entry rule for matrix-shaped tensor addition over $\mathbb{R}$.

theorem Optim.Muon.get2_scaleSpec_real {m n : ℕ} (A : MatrixTensor ℝ m n) (c : ℝ) (i : Fin m) (j : Fin n) : Spec.get2 (Spec.Tensor.scaleSpec A c) i j = Spec.get2 A i j * c

Entry rule for matrix-shaped tensor scaling over $\mathbb{R}$.

theorem Optim.Muon.get2_subSpec_real {m n : ℕ} (A B : MatrixTensor ℝ m n) (i : Fin m) (j : Fin n) : Spec.get2 (Spec.Tensor.subSpec A B) i j = Spec.get2 A i j - Spec.get2 B i j

Entry rule for matrix-shaped tensor subtraction over $\mathbb{R}$.

theorem Optim.Muon.matMul_right_identity_real {m n : ℕ} (A : MatrixTensor ℝ m n) : Spec.matMulSpec A (Spec.identityTensorSpec n) = A

Right multiplication by the identity matrix leaves a real matrix unchanged.

theorem Optim.Muon.add_scaled_three_eq_scale_sum {m n : ℕ} (Q : MatrixTensor ℝ m n) (a b c : ℝ) : ((Spec.Tensor.scaleSpec Q a).addSpec (Spec.Tensor.scaleSpec Q b)).addSpec (Spec.Tensor.scaleSpec Q c) = Spec.Tensor.scaleSpec Q (a + b + c)

Three scaled copies of the same real matrix combine into one scaled copy using the sum of the coefficients.

theorem Optim.Muon.add_scaled_three_eq_self_of_coeff_sum_one {m n : ℕ} (Q : MatrixTensor ℝ m n) (a b c : ℝ) (hsum : a + b + c = 1) : ((Spec.Tensor.scaleSpec Q a).addSpec (Spec.Tensor.scaleSpec Q b)).addSpec (Spec.Tensor.scaleSpec Q c) = Q

If three scaled copies of a matrix are added and the coefficients sum to one, the result is the original matrix.

theorem Optim.Muon.scale_hasExactColumnGram_of_square_eq_one {m n : ℕ} (Q : MatrixTensor ℝ m n) (k : ℝ) (hgram : HasExactColumnGram Q) (hk : k * k = 1) : HasExactColumnGram (Spec.Tensor.scaleSpec Q k)

Scaling an exact-column-orthogonal real matrix by a scalar whose square is one preserves exact column Gram.

theorem Optim.Muon.scale_hasApproxColumnGram_of_exact_column_gram_of_square_error {m n : ℕ} (Q : MatrixTensor ℝ m n) (k eps : ℝ) (hgram : HasExactColumnGram Q) (herr : MathFunctions.abs (k * k - 1) ≤ eps) (heps : 0 ≤ eps) : HasApproxColumnGram eps (Spec.Tensor.scaleSpec Q k)

Scaling an exact-column-orthogonal real matrix gives an approximate Gram certificate whenever $|k^2-1|$ is bounded by the requested tolerance.

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

If $Q^\mathsf{T}Q=I$, then one column-oriented Newton-Schulz step returns $(a+b+c)Q$.

theorem Optim.Muon.newtonSchulzFixedPoint_of_exact_column_gram_of_coeff_sum_one {m n : ℕ} (coeffs : NewtonSchulzCoeffs ℝ) (Q : MatrixTensor ℝ m n) (hgram : HasExactColumnGram Q) (hsum : coeffs.a + coeffs.b + coeffs.c = 1) : NewtonSchulzFixedPoint coeffs Q

For real coefficients whose sum is one, an exact-column-orthogonal matrix is a fixed point of one column-oriented Newton-Schulz step.

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

If $Q^\mathsf{T}Q=I$ and $(a+b+c)^2=1$, then one column-oriented Newton-Schulz step still has exact column Gram.

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

If $Q^\mathsf{T}Q=I$ and $|(a+b+c)^2-1|\leq\varepsilon$, then one column-oriented Newton-Schulz step has entrywise Gram residual bounded by $\varepsilon$.

theorem Optim.Muon.newtonSchulzFixedPointCheckedExact_success_of_coeff_sum_one {m n : ℕ} (coeffs : NewtonSchulzCoeffs ℝ) (steps : ℕ) (buffer : MatrixTensor ℝ m n) (hgram : HasExactColumnGram buffer) (hsum : coeffs.a + coeffs.b + coeffs.c = 1) : (newtonSchulzFixedPointCheckedExactOrthogonalizer coeffs steps).Success buffer

If the Newton-Schulz coefficients sum to one, exact column Gram is enough to satisfy the exact fixed-point backend's success predicate.

theorem Optim.Muon.update_has_exact_certified_step_newtonSchulz_exact_gram_checked {m n : ℕ} (coeffs : NewtonSchulzCoeffs ℝ) (steps : ℕ) (lr momentum : ℝ) (buf params grads : MatrixTensor ℝ m n) (hgram : HasExactColumnGram (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := newtonSchulzOrthogonalizer coeffs steps } params grads).1.buf) (hsum : coeffs.a + coeffs.b + coeffs.c = 1) : ∃ (direction : MatrixTensor ℝ m n), ExactCertifiedStep { lr := lr, momentum := momentum, buf := buf, orthogonalizer := newtonSchulzOrthogonalizer coeffs steps } params grads direction

For real coefficients with $a+b+c=1$, exact column Gram of the fresh momentum buffer is enough to certify a Newton-Schulz Muon update exactly.

theorem Optim.Muon.update_newtonSchulz_exact_gram_direction_has_exact_column_gram_checked {m n : ℕ} (coeffs : NewtonSchulzCoeffs ℝ) (steps : ℕ) (lr momentum : ℝ) (buf params grads : MatrixTensor ℝ m n) (hgram : HasExactColumnGram (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := newtonSchulzOrthogonalizer coeffs steps } params grads).1.buf) (hsum : coeffs.a + coeffs.b + coeffs.c = 1) : HasExactColumnGram ((newtonSchulzOrthogonalizer coeffs steps).apply (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := newtonSchulzOrthogonalizer coeffs steps } params grads).1.buf)

For real coefficients with $a+b+c=1$, exact column Gram of the fresh momentum buffer gives $Q^\mathsf{T}Q=I$ for the actual Newton-Schulz update direction.

theorem Optim.Muon.init_has_exact_certified_step_newtonSchulz_exact_gram_checked {m n : ℕ} (coeffs : NewtonSchulzCoeffs ℝ) (steps : ℕ) (lr momentum : ℝ) (params grads : MatrixTensor ℝ m n) (hgram : HasExactColumnGram (update (init lr momentum (newtonSchulzOrthogonalizer coeffs steps) params) params grads).1.buf) (hsum : coeffs.a + coeffs.b + coeffs.c = 1) : ∃ (direction : MatrixTensor ℝ m n), ExactCertifiedStep (init lr momentum (newtonSchulzOrthogonalizer coeffs steps) params) params grads direction

Initialized version: exact column Gram of the first fresh momentum buffer and $a+b+c=1$ certify the first Newton-Schulz Muon step exactly.

theorem Optim.Muon.qrOrthogonalizer_exact_of_positive_pivots {m n : ℕ} (buffer : MatrixTensor ℝ m n) (hpivots : HasPositiveQRPivots buffer) : ExactOrthogonalizesBuffer qrOrthogonalizer buffer

The QR orthogonalizer satisfies the exact Muon direction contract whenever the executable QR pivots of the input buffer are positive.

noncomputable def Optim.Muon.qrCheckedExactOrthogonalizer {m n : ℕ} : CheckedExactOrthogonalizer ℝ m n

QR packaged as a checked exact Muon backend.

theorem Optim.Muon.update_has_exact_certified_step_qr {m n : ℕ} (lr momentum : ℝ) (buf params grads : MatrixTensor ℝ m n) (hpivots : HasPositiveQRPivots (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := qrOrthogonalizer } params grads).1.buf) : ∃ (direction : MatrixTensor ℝ m n), ExactCertifiedStep { lr := lr, momentum := momentum, buf := buf, orthogonalizer := qrOrthogonalizer } params grads direction

Concrete QR-backed Muon step theorem: if the fresh momentum buffer has positive QR pivots, the executable Muon update has a certified exact step.

theorem Optim.Muon.update_qr_direction_has_exact_column_gram {m n : ℕ} (lr momentum : ℝ) (buf params grads : MatrixTensor ℝ m n) (hpivots : HasPositiveQRPivots (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := qrOrthogonalizer } params grads).1.buf) : HasExactColumnGram (qrOrthogonalizer.apply (update { lr := lr, momentum := momentum, buf := buf, orthogonalizer := qrOrthogonalizer } params grads).1.buf)

Concrete QR-backed direction theorem: if the fresh momentum buffer has positive QR pivots, the actual direction used by the Muon update has column Gram $I$.

theorem Optim.Muon.init_has_exact_certified_step_qr {m n : ℕ} (lr momentum : ℝ) (params grads : MatrixTensor ℝ m n) (hpivots : HasPositiveQRPivots (update (init lr momentum qrOrthogonalizer params) params grads).1.buf) : ∃ (direction : MatrixTensor ℝ m n), ExactCertifiedStep (init lr momentum qrOrthogonalizer params) params grads direction

Initialized QR-backed Muon step theorem: if the first fresh momentum buffer has positive QR pivots, the first initialized Muon update has a certified exact step.

theorem Optim.Muon.init_qr_direction_has_exact_column_gram {m n : ℕ} (lr momentum : ℝ) (params grads : MatrixTensor ℝ m n) (hpivots : HasPositiveQRPivots (update (init lr momentum qrOrthogonalizer params) params grads).1.buf) : HasExactColumnGram (qrOrthogonalizer.apply (update (init lr momentum qrOrthogonalizer params) params grads).1.buf)

Initialized QR-backed direction theorem: if the first fresh momentum buffer has positive QR pivots, the first initialized Muon update direction has column Gram $I$.