QR Muon Backend and Real-Valued Certificates

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

The entrywise matrix API (matrix_ext, get2_mat_mul_spec) is stated over ℝ, so this module also holds the real-valued facts about columnGram and about one Newton-Schulz step at an exactly column-orthogonal matrix, which the generic NN.MLTheory.Optimization.Muon.NewtonSchulz module cannot prove.

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

Positive pivots of the real-valued QR specification; no native QR claim is made.

theorem Optim.Muon.add_scaled_three_eq_scale_sum {m n : ℕ} (Q : MatrixTensor ℝ m n) (a b c : ℝ) : ((TorchLean.Tensor.scaleSpec Q a).addSpec (TorchLean.Tensor.scaleSpec Q b)).addSpec (TorchLean.Tensor.scaleSpec Q c) = TorchLean.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.

Column Gram entries

theorem Optim.Muon.get2_columnGram {m n : ℕ} (Q : MatrixTensor ℝ m n) (i j : Fin n) : Spec.get2 (columnGram Q) i j = ∑ r : Fin m, Spec.get2 Q r i * Spec.get2 Q r j

Entries of the column Gram matrix are inner products of columns.

theorem Optim.Muon.get2_columnGram_scaleSpec {m n : ℕ} (Q : MatrixTensor ℝ m n) (k : ℝ) (i j : Fin n) : Spec.get2 (columnGram (TorchLean.Tensor.scaleSpec Q k)) i j = Spec.get2 (columnGram Q) i j * (k * k)

Scaling a matrix by k scales every column Gram entry by k * k.

Scaling exactly column-orthogonal matrices

theorem Optim.Muon.scale_hasExactColumnGram_of_square_eq_one {m n : ℕ} (Q : MatrixTensor ℝ m n) (k : ℝ) (hgram : HasExactColumnGram Q) (hk : k * k = 1) : HasExactColumnGram (TorchLean.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 (TorchLean.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.

Newton-Schulz at an exactly column-orthogonal matrix

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 = TorchLean.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 : ℝ) (momentumBuffer parameters gradients : MatrixTensor ℝ m n) (hgram : HasExactColumnGram (update { lr := lr, momentum := momentum, momentumBuffer := momentumBuffer, orthogonalizer := newtonSchulzOrthogonalizer coeffs steps } parameters gradients).optimizerState.momentumBuffer) (hsum : coeffs.a + coeffs.b + coeffs.c = 1) : ∃ (direction : MatrixTensor ℝ m n), ExactCertifiedStep { lr := lr, momentum := momentum, momentumBuffer := momentumBuffer, orthogonalizer := newtonSchulzOrthogonalizer coeffs steps } parameters gradients 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 : ℝ) (momentumBuffer parameters gradients : MatrixTensor ℝ m n) (hgram : HasExactColumnGram (update { lr := lr, momentum := momentum, momentumBuffer := momentumBuffer, orthogonalizer := newtonSchulzOrthogonalizer coeffs steps } parameters gradients).optimizerState.momentumBuffer) (hsum : coeffs.a + coeffs.b + coeffs.c = 1) : HasExactColumnGram ((newtonSchulzOrthogonalizer coeffs steps).apply (update { lr := lr, momentum := momentum, momentumBuffer := momentumBuffer, orthogonalizer := newtonSchulzOrthogonalizer coeffs steps } parameters gradients).optimizerState.momentumBuffer)

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.

QR certificates

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 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 : ℝ) (momentumBuffer parameters gradients : MatrixTensor ℝ m n) (hpivots : HasPositiveQRPivots (update { lr := lr, momentum := momentum, momentumBuffer := momentumBuffer, orthogonalizer := qrOrthogonalizer } parameters gradients).optimizerState.momentumBuffer) : ∃ (direction : MatrixTensor ℝ m n), ExactCertifiedStep { lr := lr, momentum := momentum, momentumBuffer := momentumBuffer, orthogonalizer := qrOrthogonalizer } parameters gradients 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 : ℝ) (momentumBuffer parameters gradients : MatrixTensor ℝ m n) (hpivots : HasPositiveQRPivots (update { lr := lr, momentum := momentum, momentumBuffer := momentumBuffer, orthogonalizer := qrOrthogonalizer } parameters gradients).optimizerState.momentumBuffer) : HasExactColumnGram (qrOrthogonalizer.apply (update { lr := lr, momentum := momentum, momentumBuffer := momentumBuffer, orthogonalizer := qrOrthogonalizer } parameters gradients).optimizerState.momentumBuffer)

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$.