Calculus of row normalization

LayerNorm normalizes every row of a [seqLen, embedDim] matrix; the flattened form of BatchNorm normalizes every row of a [channels, positions] matrix. Both share the map

x ↦ (x - mean(x)) / sqrt (var(x) + ε)

applied to each row, where mean and var are the population statistics of that row. This file proves that map is differentiable on flattened matrices (for 0 < ε) and identifies its derivative with the closed form

dx ↦ invStd * (dx - mean(dx) - xhat * mean(dx * xhat))

used by the normalization JVPs.

noncomputable def Proofs.Autograd.RowNorm.rowMean {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) : ℝ

Mean of row i of a flattened m × n matrix.

noncomputable def Proofs.Autograd.RowNorm.centered {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) (j : Fin n) : ℝ

Centered entry (i, j).

noncomputable def Proofs.Autograd.RowNorm.rowVar {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) : ℝ

Population variance of row i.

noncomputable def Proofs.Autograd.RowNorm.invStd {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (ε : ℝ) (i : Fin m) : ℝ

Inverse standard deviation 1 / sqrt (var + ε) of row i.

noncomputable def Proofs.Autograd.RowNorm.nrm {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (ε : ℝ) (i : Fin m) (j : Fin n) : ℝ

Normalized entry (i, j).

noncomputable def Proofs.Autograd.RowNorm.nrmJvp {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (ε : ℝ) (dX : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) (j : Fin n) : ℝ

Closed-form differential of the normalized entry (i, j) in the direction dX.

theorem Proofs.Autograd.RowNorm.rowVar_nonneg {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) : 0 ≤ rowVar X i

The row variance is a mean of squares.

theorem Proofs.Autograd.RowNorm.rowVar_add_pos {m n : ℕ} {ε : ℝ} (hε : 0 < ε) (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) : 0 < rowVar X i + ε

var + ε is positive when ε is.

theorem Proofs.Autograd.RowNorm.sum_centered {m n : ℕ} (hn : 0 < n) (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) : ∑ j : Fin n, centered X i j = 0

The centered entries of a row sum to zero.

theorem Proofs.Autograd.RowNorm.sum_centered_mul_sub {m n : ℕ} (hn : 0 < n) (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) (d : Fin n → ℝ) (a : ℝ) : ∑ j : Fin n, centered X i j * (d j - a) = ∑ j : Fin n, centered X i j * d j

The centered row is orthogonal to constants, so ∑ c_j (d_j - a) = ∑ c_j d_j.

Derivative building blocks

noncomputable def Proofs.Autograd.RowNorm.projCLM {m n : ℕ} (k : Fin (TapeNodes.Matmul.matSize m n)) : Vec (TapeNodes.Matmul.matSize m n) →L[ℝ] ℝ

Coordinate projection as a continuous linear functional on flattened matrices.

@[simp]

theorem Proofs.Autograd.RowNorm.projCLM_apply {m n : ℕ} (k : Fin (TapeNodes.Matmul.matSize m n)) (v : Vec (TapeNodes.Matmul.matSize m n)) : (projCLM k) v = v.ofLp k

The projection functional reads coordinate k.

noncomputable def Proofs.Autograd.RowNorm.meanD {m n : ℕ} (i : Fin m) : Vec (TapeNodes.Matmul.matSize m n) →L[ℝ] ℝ

Derivative of the row mean.

@[simp]

theorem Proofs.Autograd.RowNorm.meanD_apply {m n : ℕ} (i : Fin m) (dX : Vec (TapeNodes.Matmul.matSize m n)) : (meanD i) dX = rowMean dX i

The row-mean derivative is the row mean of the perturbation, since the mean is already linear.

noncomputable def Proofs.Autograd.RowNorm.centD {m n : ℕ} (i : Fin m) (j : Fin n) : Vec (TapeNodes.Matmul.matSize m n) →L[ℝ] ℝ

Derivative of a centered entry.

@[simp]

theorem Proofs.Autograd.RowNorm.centD_apply {m n : ℕ} (i : Fin m) (j : Fin n) (dX : Vec (TapeNodes.Matmul.matSize m n)) : (centD i j) dX = dX.ofLp (TapeNodes.Matmul.idxMN i j) - rowMean dX i

The centering derivative subtracts the perturbation's own row mean.

Centering is linear too, so its derivative is itself; the nonlinearity in row normalization enters only through the variance and the square root below.

noncomputable def Proofs.Autograd.RowNorm.varD {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) : Vec (TapeNodes.Matmul.matSize m n) →L[ℝ] ℝ

Derivative of the row variance at X.

theorem Proofs.Autograd.RowNorm.varD_apply {m n : ℕ} (hn : 0 < n) (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) (dX : Vec (TapeNodes.Matmul.matSize m n)) : (varD X i) dX = (2 * ∑ k : Fin n, centered X i k * dX.ofLp (TapeNodes.Matmul.idxMN i k)) / ↑n

The variance differential only sees the centered direction.

noncomputable def Proofs.Autograd.RowNorm.sqrtD {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (ε : ℝ) (i : Fin m) : Vec (TapeNodes.Matmul.matSize m n) →L[ℝ] ℝ

Derivative of sqrt (var + ε) at X.

noncomputable def Proofs.Autograd.RowNorm.invD {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (ε : ℝ) (i : Fin m) : Vec (TapeNodes.Matmul.matSize m n) →L[ℝ] ℝ

Derivative of the inverse standard deviation at X.

noncomputable def Proofs.Autograd.RowNorm.nrmD {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (ε : ℝ) (i : Fin m) (j : Fin n) : Vec (TapeNodes.Matmul.matSize m n) →L[ℝ] ℝ

Derivative of the normalized entry (i, j) at X.

theorem Proofs.Autograd.RowNorm.hasFDerivAt_coord {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (k : Fin (TapeNodes.Matmul.matSize m n)) : HasFDerivAt (fun (Y : Vec (TapeNodes.Matmul.matSize m n)) => Y.ofLp k) (projCLM k) X

Coordinates of a flattened matrix are differentiable.

theorem Proofs.Autograd.RowNorm.hasFDerivAt_rowMean {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) : HasFDerivAt (fun (Y : Vec (TapeNodes.Matmul.matSize m n)) => rowMean Y i) (meanD i) X

Derivative of the row mean.

theorem Proofs.Autograd.RowNorm.hasFDerivAt_centered {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) (j : Fin n) : HasFDerivAt (fun (Y : Vec (TapeNodes.Matmul.matSize m n)) => centered Y i j) (centD i j) X

Derivative of a centered entry.

theorem Proofs.Autograd.RowNorm.hasFDerivAt_rowVar {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) : HasFDerivAt (fun (Y : Vec (TapeNodes.Matmul.matSize m n)) => rowVar Y i) (varD X i) X

Derivative of the row variance.

theorem Proofs.Autograd.RowNorm.hasFDerivAt_invStd {m n : ℕ} {ε : ℝ} (hε : 0 < ε) (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) : HasFDerivAt (fun (Y : Vec (TapeNodes.Matmul.matSize m n)) => invStd Y ε i) (invD X ε i) X

Derivative of the inverse standard deviation, for positive ε.

theorem Proofs.Autograd.RowNorm.hasFDerivAt_nrm {m n : ℕ} {ε : ℝ} (hε : 0 < ε) (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) (j : Fin n) : HasFDerivAt (fun (Y : Vec (TapeNodes.Matmul.matSize m n)) => nrm Y ε i j) (nrmD X ε i j) X

The normalized entry is differentiable for positive ε.

theorem Proofs.Autograd.RowNorm.nrmD_apply {m n : ℕ} (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (X : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) (j : Fin n) (dX : Vec (TapeNodes.Matmul.matSize m n)) : (nrmD X ε i j) dX = nrmJvp X ε dX i j

The derivative of the normalized entry is the closed-form JVP.

Whole-matrix normalization

def Proofs.Autograd.RowNorm.rowOf {m n : ℕ} (ip : Fin (TapeNodes.Matmul.matSize m n)) : Fin m

Row index of a flattened matrix coordinate.

def Proofs.Autograd.RowNorm.colOf {m n : ℕ} (ip : Fin (TapeNodes.Matmul.matSize m n)) : Fin n

Column index of a flattened matrix coordinate.

theorem Proofs.Autograd.RowNorm.idxMN_rowOf_colOf {m n : ℕ} (ip : Fin (TapeNodes.Matmul.matSize m n)) : TapeNodes.Matmul.idxMN (rowOf ip) (colOf ip) = ip

Every flattened coordinate is idxMN of its row and column.

theorem Proofs.Autograd.RowNorm.rowOf_idxMN {m n : ℕ} (i : Fin m) (j : Fin n) : rowOf (TapeNodes.Matmul.idxMN i j) = i

Row of the flattened coordinate idxMN i j.

theorem Proofs.Autograd.RowNorm.colOf_idxMN {m n : ℕ} (i : Fin m) (j : Fin n) : colOf (TapeNodes.Matmul.idxMN i j) = j

Column of the flattened coordinate idxMN i j.

theorem Proofs.Autograd.RowNorm.sum_vecSize {n : ℕ} (f : Fin (TapeNodes.Matmul.vecSize n) → ℝ) : ∑ j : Fin (TapeNodes.Matmul.vecSize n), f j = ∑ j : Fin n, f (Fin.cast ⋯ j)

Reindex a sum over Fin (vecSize n) as a sum over Fin n.

theorem Proofs.Autograd.RowNorm.rowMeanCLM_eq {m n : ℕ} (x : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin m) : (TapeNodes.MatrixLinear.rowMeanCLM x).ofLp i = rowMean x i

The tape row-mean map computes rowMean.

theorem Proofs.Autograd.RowNorm.broadcastRowCLM_idxMN {m n : ℕ} (v : Vec m) (i : Fin m) (j : Fin n) : (TapeNodes.MatrixLinear.broadcastRowCLM v).ofLp (TapeNodes.Matmul.idxMN i j) = v.ofLp i

The tape row-broadcast map at idxMN i j reads entry i.

theorem Proofs.Autograd.RowNorm.broadcastColCLM_idxMN {m n : ℕ} (v : Vec n) (i : Fin m) (j : Fin n) : (TapeNodes.MatrixLinear.broadcastColCLM v).ofLp (TapeNodes.Matmul.idxMN i j) = v.ofLp j

The tape column-broadcast map at idxMN i j reads entry j.

noncomputable def Proofs.Autograd.RowNorm.nrmVec {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (ε : ℝ) : Vec (TapeNodes.Matmul.matSize m n)

Normalize every row of a flattened matrix.

@[simp]

theorem Proofs.Autograd.RowNorm.nrmVec_apply {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (ε : ℝ) (ip : Fin (TapeNodes.Matmul.matSize m n)) : (nrmVec X ε).ofLp ip = nrm X ε (rowOf ip) (colOf ip)

Coordinate ip of the normalized matrix is the normalized entry at its row and column.

noncomputable def Proofs.Autograd.RowNorm.nrmJvpCLM {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (ε : ℝ) : Vec (TapeNodes.Matmul.matSize m n) →L[ℝ] Vec (TapeNodes.Matmul.matSize m n)

Closed-form JVP of nrmVec at X, packaged as a continuous linear map.

@[simp]

theorem Proofs.Autograd.RowNorm.nrmJvpCLM_apply {m n : ℕ} (X : Vec (TapeNodes.Matmul.matSize m n)) (ε : ℝ) (dX : Vec (TapeNodes.Matmul.matSize m n)) (ip : Fin (TapeNodes.Matmul.matSize m n)) : ((nrmJvpCLM X ε) dX).ofLp ip = nrmJvp X ε dX (rowOf ip) (colOf ip)

The bundled JVP agrees with the closed-form nrmJvp coordinatewise.

theorem Proofs.Autograd.RowNorm.hasFDerivAt_nrmVec {m n : ℕ} (hn : 0 < n) {ε : ℝ} (hε : 0 < ε) (X : Vec (TapeNodes.Matmul.matSize m n)) : HasFDerivAt (fun (Y : Vec (TapeNodes.Matmul.matSize m n)) => nrmVec Y ε) (nrmJvpCLM X ε) X

Row normalization is differentiable on a nonempty feature axis, with the closed-form JVP as derivative. The coordinatewise chain-rule calculation agrees with that JVP because positive epsilon keeps each row's denominator nonzero.