Entrywise formulas for row-normalization tensors

The normalization specs are written with reduceMean, reduceVar, broadcastAfterSum, broadcastTo, and pointwise tensor arithmetic on [m, n] matrices. This file records what each of those operations does to a single entry get2 t i j, so that the LayerNorm and BatchNorm proofs can work with plain real sums. Entries of filled tensors are TorchLean.Tensor.getScalar_full and Spec.get2_full; tensorToVec of a matrix is read at idxMN i j by TapeNodes.Matmul.tensorToVec_idxMN.

Matrix vectorization

theorem Proofs.Autograd.Norm.vec_ext_idxMN {m n : ℕ} {u v : Vec (TapeNodes.Matmul.matSize m n)} (h : ∀ (i : Fin m) (j : Fin n), u.ofLp (TapeNodes.Matmul.idxMN i j) = v.ofLp (TapeNodes.Matmul.idxMN i j)) : u = v

Two flattened matrices agree once they agree at every idxMN i j.

Vector and matrix entries of pointwise operations

@[reducible, inline]

abbrev Proofs.Autograd.Norm.rowBroadcast {m n : ℕ} (v : TorchLean.Tensor ℝ [m]) : TorchLean.Tensor ℝ [m, n]

Row broadcast of a length-m vector to an m × n matrix.

theorem Proofs.Autograd.Norm.get2_broadcastAfterSum_one {m n : ℕ} (v : TorchLean.Tensor ℝ [m]) (i : Fin m) (j : Fin n) : Spec.get2 (TorchLean.Tensor.broadcastAfterSum (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar)) 1 v) i j = v.getScalar i

Entries of a row broadcast: every column of row i reads the vector entry i.

theorem Proofs.Autograd.Norm.get2_broadcastTo_col {m n : ℕ} (h : (Spec.Shape.dim n Spec.Shape.scalar).CanBroadcastTo (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (v : TorchLean.Tensor ℝ [n]) (i : Fin m) (j : Fin n) : Spec.get2 (TorchLean.Tensor.broadcastTo h v) i j = v.getScalar j

Entries of a column broadcast: every row reads the vector entry j.

theorem Proofs.Autograd.Norm.getScalar_reduceSum_one {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (h : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (i : Fin m) : (TorchLean.Tensor.reduceSum 1 x h).getScalar i = ∑ j : Fin n, Spec.get2 x i j

Row sums along the last axis.

theorem Proofs.Autograd.Norm.getScalar_reduceMean_one {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (h : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (i : Fin m) : (TorchLean.Tensor.reduceMean 1 x h).getScalar i = (∑ j : Fin n, Spec.get2 x i j) / ↑n

Row means along the last axis.

theorem Proofs.Autograd.Norm.item_reduceVar_zero {n : ℕ} (row : TorchLean.Tensor ℝ [n]) (h : Spec.Shape.NonemptyAxis 0 (Spec.Shape.dim n Spec.Shape.scalar)) : (TorchLean.Tensor.reduceVar 0 row h).item = (∑ j : Fin n, (row.getScalar j - (∑ k : Fin n, row.getScalar k) / ↑n) * (row.getScalar j - (∑ k : Fin n, row.getScalar k) / ↑n)) / ↑n

Population variance of a vector, as the scalar entry of reduceVar 0.

theorem Proofs.Autograd.Norm.getScalar_reduceVar_one {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (h : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (i : Fin m) : (TorchLean.Tensor.reduceVar 1 x h).getScalar i = (∑ j : Fin n, (Spec.get2 x i j - (∑ k : Fin n, Spec.get2 x i k) / ↑n) * (Spec.get2 x i j - (∑ k : Fin n, Spec.get2 x i k) / ↑n)) / ↑n

Row variances along the last axis (population variance of each row).

Vectors

theorem Proofs.Autograd.Norm.tensorToVec_vec {k : ℕ} (v : TorchLean.Tensor ℝ [k]) (p : Fin (Spec.Shape.dim k Spec.Shape.scalar).size) : (tensorToVec v).ofLp p = v.getScalar (Fin.cast ⋯ p)

tensorToVec of a rank-one tensor reads the corresponding scalar entry.

Pointwise operations on matrices and vectors

theorem Proofs.Autograd.Norm.get2_mulSpec {m n : ℕ} (a b : TorchLean.Tensor ℝ [m, n]) (i : Fin m) (j : Fin n) : Spec.get2 (a.mulSpec b) i j = Spec.get2 a i j * Spec.get2 b i j

Entries of a pointwise matrix product.

theorem Proofs.Autograd.Norm.get2_divSpec {m n : ℕ} (a b : TorchLean.Tensor ℝ [m, n]) (i : Fin m) (j : Fin n) : Spec.get2 (a.divSpec b) i j = Spec.get2 a i j / Spec.get2 b i j

Entries of a pointwise matrix quotient.

theorem Proofs.Autograd.Norm.getScalar_addSpec {k : ℕ} (a b : TorchLean.Tensor ℝ [k]) (i : Fin k) : (a.addSpec b).getScalar i = a.getScalar i + b.getScalar i

Entries of a vector sum.

theorem Proofs.Autograd.Norm.getScalar_divSpec {k : ℕ} (a b : TorchLean.Tensor ℝ [k]) (i : Fin k) : (a.divSpec b).getScalar i = a.getScalar i / b.getScalar i

Entries of a pointwise vector quotient.

theorem Proofs.Autograd.Norm.getScalar_maxSpec {k : ℕ} (a b : TorchLean.Tensor ℝ [k]) (i : Fin k) : (a.maxSpec b).getScalar i = max (a.getScalar i) (b.getScalar i)

Entries of a pointwise vector maximum.

theorem Proofs.Autograd.Norm.getScalar_sqrtSpec {k : ℕ} (a : TorchLean.Tensor ℝ [k]) (i : Fin k) : a.sqrtSpec.getScalar i = √(max (a.getScalar i) 0)

Entries of the clamped square root.

Row statistics of a spec matrix

noncomputable def Proofs.Autograd.Norm.rowMeanE {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (i : Fin m) : ℝ

Mean of row i.

noncomputable def Proofs.Autograd.Norm.rowVarE {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (i : Fin m) : ℝ

Population variance of row i.

theorem Proofs.Autograd.Norm.rowVarE_nonneg {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (i : Fin m) : 0 ≤ rowVarE x i

The row variance is a mean of squares.

theorem Proofs.Autograd.Norm.sum_sub_rowMeanE {m n : ℕ} (hn : 0 < n) (x : TorchLean.Tensor ℝ [m, n]) (i : Fin m) : ∑ j : Fin n, (Spec.get2 x i j - rowMeanE x i) = 0

Centered row entries sum to zero.

Spec.layerNorm entrywise

theorem Proofs.Autograd.Norm.nonemptyAxis_one {m n : ℕ} (hn : 0 < n) : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))

The feature axis of an [m, n] matrix is nonempty when 0 < n.

theorem Proofs.Autograd.Norm.colBroadcast {m n : ℕ} : (Spec.Shape.dim n Spec.Shape.scalar).CanBroadcastTo (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))

Column-broadcast evidence used by Spec.layerNorm for gamma and beta.

noncomputable def Proofs.Autograd.Norm.lnMean {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) : TorchLean.Tensor ℝ [m]

Row means of the LayerNorm input.

noncomputable def Proofs.Autograd.Norm.lnCentered {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) : TorchLean.Tensor ℝ [m, n]

Centered LayerNorm input.

noncomputable def Proofs.Autograd.Norm.lnStd {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (ε : ℝ) : TorchLean.Tensor ℝ [m]

Clamped LayerNorm standard deviation sqrt (max (max var 0 + ε) 0).

noncomputable def Proofs.Autograd.Norm.layerNormMat {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (ε : ℝ) : TorchLean.Tensor ℝ [m, n]

Spec.layerNorm with its axis and broadcast evidence spelled out.

theorem Proofs.Autograd.Norm.layerNorm_eq_layerNormMat {m n : ℕ} (hm : 0 < m) (hn : 0 < n) (ε : ℝ) (x : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) : Spec.layerNorm x gamma beta hm hn ε = layerNormMat x gamma beta ⋯ ε

Spec.layerNorm is layerNormMat by unfolding.

theorem Proofs.Autograd.Norm.getScalar_lnMean {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (i : Fin m) : (lnMean x h1).getScalar i = rowMeanE x i

Entries of the LayerNorm row means.

theorem Proofs.Autograd.Norm.get2_lnCentered {m n : ℕ} (x : TorchLean.Tensor ℝ [m, n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (i : Fin m) (j : Fin n) : Spec.get2 (lnCentered x h1) i j = Spec.get2 x i j - rowMeanE x i

Entries of the centered LayerNorm input.

theorem Proofs.Autograd.Norm.getScalar_lnStd {m n : ℕ} (hn : 0 < n) (x : TorchLean.Tensor ℝ [m, n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (ε : ℝ) (i : Fin m) : (lnStd x h1 ε).getScalar i = √(max (rowVarE x i + ε) 0)

Entries of the LayerNorm standard deviation: the clamp on the variance is inactive.

theorem Proofs.Autograd.Norm.get2_layerNormMat {m n : ℕ} (hn : 0 < n) (x : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) (h1 : Spec.Shape.NonemptyAxis 1 (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))) (ε : ℝ) (i : Fin m) (j : Fin n) : Spec.get2 (layerNormMat x gamma beta h1 ε) i j = (Spec.get2 x i j - rowMeanE x i) / √(max (rowVarE x i + ε) 0) * gamma.getScalar j + beta.getScalar j

Entries of layerNormMat.

theorem Proofs.Autograd.Norm.get2_layerNorm {m n : ℕ} (hm : 0 < m) (hn : 0 < n) (ε : ℝ) (x : TorchLean.Tensor ℝ [m, n]) (gamma beta : TorchLean.Tensor ℝ [n]) (i : Fin m) (j : Fin n) : Spec.get2 (Spec.layerNorm x gamma beta hm hn ε) i j = (Spec.get2 x i j - rowMeanE x i) / √(max (rowVarE x i + ε) 0) * gamma.getScalar j + beta.getScalar j

Entries of Spec.layerNorm: each row is centered, divided by the clamped standard deviation of that row, then scaled and shifted by the per-column parameters.