Fréchet derivative of BatchNorm

Spec.batchNorm flattens the spatial axes of a channel-first tensor to a [channels, positions] matrix, normalizes every row of that matrix with the row statistics, and applies the per-channel affine parameters. This file proves that, for 0 < ε, the map (x, gamma, beta) ↦ batchNorm is Fréchet differentiable on the flattened vectors and that its derivative is exactly Spec.batchNormJvp.

The calculus is delegated to RowNorm.hasFDerivAt_nrm, the derivative of one normalized row entry. The rest of the file is bookkeeping: reshaping commutes with pointwise operations, the clamps in the spec are inactive because the variance is a mean of squares, and every entry of the JVP is the closed-form row differential.

Reshaping

theorem Proofs.Autograd.BatchNorm.reshapeSpec_map2Spec {s₁ s₂ : Spec.Shape} (f : ℝ → ℝ → ℝ) (a b : TorchLean.Tensor ℝ s₁) (h : s₁.size = s₂.size) : (TorchLean.Tensor.map2Spec f a b).reshapeSpec h = TorchLean.Tensor.map2Spec f (a.reshapeSpec h) (b.reshapeSpec h)

Reshaping commutes with pointwise binary operations.

theorem Proofs.Autograd.BatchNorm.reshapeSpec_mapSpec {s₁ s₂ : Spec.Shape} (f : ℝ → ℝ) (a : TorchLean.Tensor ℝ s₁) (h : s₁.size = s₂.size) : (TorchLean.Tensor.mapSpec f a).reshapeSpec h = TorchLean.Tensor.mapSpec f (a.reshapeSpec h)

Reshaping commutes with pointwise unary operations.

theorem Proofs.Autograd.BatchNorm.reshapeSpec_full {s₁ s₂ : Spec.Shape} (v : ℝ) (h : s₁.size = s₂.size) : (TorchLean.Tensor.full s₁ v).reshapeSpec h = TorchLean.Tensor.full s₂ v

Reshaping a filled tensor fills the target shape.

theorem Proofs.Autograd.BatchNorm.reshapeSpec_reshapeSpec {s₁ s₂ : Spec.Shape} (t : TorchLean.Tensor ℝ s₁) (h : s₁.size = s₂.size) (h' : s₂.size = s₁.size) : (t.reshapeSpec h).reshapeSpec h' = t

Reshaping back and forth is the identity, for any proofs of the size equalities.

theorem Proofs.Autograd.BatchNorm.reshapeSpec_addSpec {s₁ s₂ : Spec.Shape} (a b : TorchLean.Tensor ℝ s₁) (h : s₁.size = s₂.size) : (a.addSpec b).reshapeSpec h = (a.reshapeSpec h).addSpec (b.reshapeSpec h)

Reshaping commutes with tensor addition.

theorem Proofs.Autograd.BatchNorm.reshapeSpec_subSpec {s₁ s₂ : Spec.Shape} (a b : TorchLean.Tensor ℝ s₁) (h : s₁.size = s₂.size) : (a.subSpec b).reshapeSpec h = (a.reshapeSpec h).subSpec (b.reshapeSpec h)

Reshaping commutes with tensor subtraction.

theorem Proofs.Autograd.BatchNorm.reshapeSpec_mulSpec {s₁ s₂ : Spec.Shape} (a b : TorchLean.Tensor ℝ s₁) (h : s₁.size = s₂.size) : (a.mulSpec b).reshapeSpec h = (a.reshapeSpec h).mulSpec (b.reshapeSpec h)

Reshaping commutes with pointwise multiplication.

theorem Proofs.Autograd.BatchNorm.reshapeSpec_divSpec {s₁ s₂ : Spec.Shape} (a b : TorchLean.Tensor ℝ s₁) (h : s₁.size = s₂.size) : (a.divSpec b).reshapeSpec h = (a.reshapeSpec h).divSpec (b.reshapeSpec h)

Reshaping commutes with pointwise division.

theorem Proofs.Autograd.BatchNorm.reshapeSpec_sqrtSpec {s₁ s₂ : Spec.Shape} (a : TorchLean.Tensor ℝ s₁) (h : s₁.size = s₂.size) : a.sqrtSpec.reshapeSpec h = (a.reshapeSpec h).sqrtSpec

Reshaping commutes with the clamped square root.

theorem Proofs.Autograd.BatchNorm.getScalarE_cast {a b : ℕ} (h : a = b) (v : TorchLean.Tensor ℝ [a]) (k : Fin b) : (getScalarE (h ▸ v)).ofLp k = (getScalarE v).ofLp (Fin.cast ⋯ k)

Transporting a vector along a length equality reindexes its Euclidean view.

theorem Proofs.Autograd.BatchNorm.tensorToVec_reshapeSpec {s₁ s₂ : Spec.Shape} (t : TorchLean.Tensor ℝ s₁) (h : s₁.size = s₂.size) (k : Fin s₂.size) : (tensorToVec (t.reshapeSpec h)).ofLp k = (tensorToVec t).ofLp (Fin.cast ⋯ k)

tensorToVec of a reshaped tensor is a reindexing of tensorToVec.

theorem Proofs.Autograd.BatchNorm.get2_sqrtSpec {m n : ℕ} (a : TorchLean.Tensor ℝ [m, n]) (i : Fin m) (j : Fin n) : Spec.get2 a.sqrtSpec i j = √(max (Spec.get2 a i j) 0)

Entries of the clamped square root of a matrix.

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

Row means when the axis is written as rank - 1, as in the spec.

The flattened matrix form

@[reducible, inline]

abbrev Proofs.Autograd.BatchNorm.positions (sSpatial : Spec.Shape) : ℕ

Number of positions per channel.

theorem Proofs.Autograd.BatchNorm.size_flat (channels : ℕ) (sSpatial : Spec.Shape) : (Spec.Shape.dim channels sSpatial).size = (Spec.Shape.dim channels (Spec.Shape.dim (positions sSpatial) Spec.Shape.scalar)).size

Flattening the spatial axes preserves the size.

theorem Proofs.Autograd.BatchNorm.positions_pos {sSpatial : Spec.Shape} (channels : ℕ) [(Spec.Shape.dim channels sSpatial).WellFormed] : 0 < positions sSpatial

The number of positions is positive for a well-formed input shape.

def Proofs.Autograd.BatchNorm.flat {channels : ℕ} {sSpatial : Spec.Shape} (x : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) : TorchLean.Tensor ℝ [channels, positions sSpatial]

The input flattened to a [channels, positions] matrix.

@[reducible, inline]

abbrev Proofs.Autograd.BatchNorm.bAS {channels : ℕ} {sSpatial : Spec.Shape} (v : TorchLean.Tensor ℝ [channels]) : TorchLean.Tensor ℝ [channels, positions sSpatial]

Row broadcast of a channel vector to the flattened matrix shape.

noncomputable def Proofs.Autograd.BatchNorm.batchNormExplicit {channels : ℕ} {sSpatial : Spec.Shape} (x : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (gamma beta : TorchLean.Tensor ℝ [channels]) (ε : ℝ) (hP : 0 < positions sSpatial) : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)

Spec.batchNorm with its axis and reshape evidence spelled out.

theorem Proofs.Autograd.BatchNorm.batchNorm_eq_explicit {channels : ℕ} {sSpatial : Spec.Shape} [(Spec.Shape.dim channels sSpatial).WellFormed] (x : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (gamma beta : TorchLean.Tensor ℝ [channels]) (ε : ℝ) : Spec.batchNorm x gamma beta ε = batchNormExplicit x gamma beta ε ⋯

Spec.batchNorm is batchNormExplicit by unfolding.

theorem Proofs.Autograd.BatchNorm.get2_flat_batchNorm {channels : ℕ} {sSpatial : Spec.Shape} [(Spec.Shape.dim channels sSpatial).WellFormed] (x : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (gamma beta : TorchLean.Tensor ℝ [channels]) (ε : ℝ) (c : Fin channels) (p : Fin (positions sSpatial)) : Spec.get2 ((Spec.batchNorm x gamma beta ε).reshapeSpec ⋯) c p = (Spec.get2 (flat x) c p - Norm.rowMeanE (flat x) c) / √(max (max (Norm.rowVarE (flat x) c) 0 + ε) 0) * gamma.getScalar c + beta.getScalar c

Entries of Spec.batchNorm in the flattened matrix: each row is centered, divided by the clamped standard deviation, scaled and shifted by the channel parameters.

noncomputable def Proofs.Autograd.BatchNorm.batchNormJvpExplicit {channels : ℕ} {sSpatial : Spec.Shape} (x dx : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (gamma dgamma dbeta : TorchLean.Tensor ℝ [channels]) (ε : ℝ) (hP : 0 < positions sSpatial) : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)

Spec.batchNormJvp with its axis and reshape evidence spelled out.

theorem Proofs.Autograd.BatchNorm.batchNormJvp_eq_explicit {channels : ℕ} {sSpatial : Spec.Shape} [(Spec.Shape.dim channels sSpatial).WellFormed] (x dx : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (gamma dgamma beta dbeta : TorchLean.Tensor ℝ [channels]) (ε : ℝ) : Spec.batchNormJvp x dx gamma dgamma beta dbeta ε = batchNormJvpExplicit x dx gamma dgamma dbeta ε ⋯

Spec.batchNormJvp is batchNormJvpExplicit by unfolding.

theorem Proofs.Autograd.BatchNorm.get2_normalizedJvp {C P : ℕ} (hP : 0 < P) (t xHat : TorchLean.Tensor ℝ [C, P]) (invStd gamma dgamma dbeta : TorchLean.Tensor ℝ [C]) (c : Fin C) (p : Fin P) : Spec.get2 (Spec.BatchNorm.normalizedJvp hP t xHat invStd gamma dgamma dbeta) c p = invStd.getScalar c * (Spec.get2 t c p - (∑ k : Fin P, Spec.get2 t c k) / ↑P - Spec.get2 xHat c p * ((∑ k : Fin P, Spec.get2 t c k * Spec.get2 xHat c k) / ↑P)) * gamma.getScalar c + Spec.get2 xHat c p * dgamma.getScalar c + dbeta.getScalar c

Entries of the normalized-data JVP.

Vector form

@[reducible, inline]

abbrev Proofs.Autograd.BatchNorm.Domain (channels : ℕ) (sSpatial : Spec.Shape) : Type

Domain of the vectorized BatchNorm: flattened input and the two channel parameters.

noncomputable def Proofs.Autograd.BatchNorm.bnVec {channels : ℕ} {sSpatial : Spec.Shape} [(Spec.Shape.dim channels sSpatial).WellFormed] (ε : ℝ) (q : Domain channels sSpatial) : Vec (Spec.Shape.dim channels sSpatial).size

Spec.batchNorm on flattened vectors.

def Proofs.Autograd.BatchNorm.rowIdx {channels : ℕ} {sSpatial : Spec.Shape} (k : Fin (Spec.Shape.dim channels sSpatial).size) : Fin channels

Flattened row index of a coordinate of the input.

def Proofs.Autograd.BatchNorm.colIdx {channels : ℕ} {sSpatial : Spec.Shape} (k : Fin (Spec.Shape.dim channels sSpatial).size) : Fin (positions sSpatial)

Flattened column index of a coordinate of the input.

def Proofs.Autograd.BatchNorm.chanIdx {channels : ℕ} {sSpatial : Spec.Shape} (k : Fin (Spec.Shape.dim channels sSpatial).size) : Fin (TapeNodes.Matmul.vecSize channels)

Channel index of a coordinate of the input, as an index into a channel vector.

noncomputable def Proofs.Autograd.BatchNorm.matVec {channels : ℕ} {sSpatial : Spec.Shape} (v : Vec (Spec.Shape.dim channels sSpatial).size) : Vec (TapeNodes.Matmul.matSize channels (positions sSpatial))

The input as a flattened channels × positions matrix vector.

theorem Proofs.Autograd.BatchNorm.cast_eq_idxMN {channels : ℕ} {sSpatial : Spec.Shape} (k : Fin (Spec.Shape.dim channels sSpatial).size) : Fin.cast ⋯ k = TapeNodes.Matmul.idxMN (rowIdx k) (colIdx k)

Every input coordinate is idxMN of its row and column, up to the size cast.

theorem Proofs.Autograd.BatchNorm.get2_flat {channels : ℕ} {sSpatial : Spec.Shape} (t : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (c : Fin channels) (p : Fin (positions sSpatial)) : Spec.get2 (flat t) c p = (matVec (tensorToVec t)).ofLp (TapeNodes.Matmul.idxMN c p)

Entries of a flattened tensor are coordinates of matVec of its vectorization.

theorem Proofs.Autograd.BatchNorm.get2_flat_vecToTensor {channels : ℕ} {sSpatial : Spec.Shape} (v : Vec (Spec.Shape.dim channels sSpatial).size) (c : Fin channels) (p : Fin (positions sSpatial)) : Spec.get2 (flat (vecToTensor v)) c p = (matVec v).ofLp (TapeNodes.Matmul.idxMN c p)

Entries of the flattened input tensor are coordinates of matVec.

theorem Proofs.Autograd.BatchNorm.rowMeanE_flat {channels : ℕ} {sSpatial : Spec.Shape} (t : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (c : Fin channels) : Norm.rowMeanE (flat t) c = RowNorm.rowMean (matVec (tensorToVec t)) c

Row means of a flattened tensor are RowNorm.rowMean of its vectorization.

theorem Proofs.Autograd.BatchNorm.rowVarE_flat {channels : ℕ} {sSpatial : Spec.Shape} (t : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (c : Fin channels) : Norm.rowVarE (flat t) c = RowNorm.rowVar (matVec (tensorToVec t)) c

Row variances of a flattened tensor are RowNorm.rowVar of its vectorization.

theorem Proofs.Autograd.BatchNorm.getScalar_chan {channels : ℕ} (g : TorchLean.Tensor ℝ [channels]) (c : Fin channels) : g.getScalar c = (tensorToVec g).ofLp (Fin.cast ⋯ c)

Entries of a channel vector are coordinates of its vectorization.

noncomputable def Proofs.Autograd.BatchNorm.bnClosed {channels : ℕ} {sSpatial : Spec.Shape} (ε : ℝ) (q : Domain channels sSpatial) : Vec (Spec.Shape.dim channels sSpatial).size

Closed form of the vectorized BatchNorm for positive ε.

theorem Proofs.Autograd.BatchNorm.bnVec_eq_bnClosed {channels : ℕ} {sSpatial : Spec.Shape} [(Spec.Shape.dim channels sSpatial).WellFormed] {ε : ℝ} (hε : 0 < ε) : bnVec ε = bnClosed ε

For positive ε the clamps in Spec.batchNorm are inactive and the vectorized map is its closed form.

theorem Proofs.Autograd.BatchNorm.castCLM_apply {a b : ℕ} (h : a = b) (v : Vec a) : (Graph.castCLM h) v = castVec h v

Graph.castCLM acts by castVec.

noncomputable def Proofs.Autograd.BatchNorm.bnD {channels : ℕ} {sSpatial : Spec.Shape} (X : Vec (TapeNodes.Matmul.matSize channels (positions sSpatial))) (g : Vec (TapeNodes.Matmul.vecSize channels)) (ε : ℝ) : Domain channels sSpatial →L[ℝ] Vec (Spec.Shape.dim channels sSpatial).size

Derivative of the closed form at (X, g, _), as a continuous linear map.

@[simp]

theorem Proofs.Autograd.BatchNorm.bnD_apply {channels : ℕ} {sSpatial : Spec.Shape} (X : Vec (TapeNodes.Matmul.matSize channels (positions sSpatial))) (g : Vec (TapeNodes.Matmul.vecSize channels)) (ε : ℝ) (d : Domain channels sSpatial) (k : Fin (Spec.Shape.dim channels sSpatial).size) : ((bnD X g ε) d).ofLp k = ((RowNorm.nrmJvpCLM X ε) (matVec d.1)).ofLp (TapeNodes.Matmul.idxMN (rowIdx k) (colIdx k)) * g.ofLp (chanIdx k) + RowNorm.nrm X ε (rowIdx k) (colIdx k) * d.2.1.ofLp (chanIdx k) + d.2.2.ofLp (chanIdx k)

Coordinates of the batch-norm differential, split into its three contributions.

Reading the summands left to right: the normalization's own Jacobian applied to the input perturbation and then scaled by γ, the normalized activation times the perturbation of γ, and the perturbation of β. That decomposition is why the three parameter groups can be differentiated independently downstream.

theorem Proofs.Autograd.BatchNorm.hasFDerivAt_bnClosed {channels : ℕ} {sSpatial : Spec.Shape} (hP : 0 < positions sSpatial) {ε : ℝ} (hε : 0 < ε) (q : Domain channels sSpatial) : HasFDerivAt (bnClosed ε) (bnD (matVec q.1) q.2.1 ε) q

The closed form is differentiable for positive ε.

Main theorems

theorem Proofs.Autograd.BatchNorm.hasFDerivAt_batchNorm {channels : ℕ} {sSpatial : Spec.Shape} [(Spec.Shape.dim channels sSpatial).WellFormed] {ε : ℝ} (hε : 0 < ε) (x : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (gamma beta : TorchLean.Tensor ℝ [channels]) : HasFDerivAt (bnVec ε) (bnD (matVec (tensorToVec x)) (tensorToVec gamma) ε) (tensorToVec x, tensorToVec gamma, tensorToVec beta)

Spec.batchNorm is Fréchet differentiable in (x, gamma, beta) for positive ε.

theorem Proofs.Autograd.BatchNorm.bnVec_tensors {channels : ℕ} {sSpatial : Spec.Shape} [(Spec.Shape.dim channels sSpatial).WellFormed] (ε : ℝ) (x : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (gamma beta : TorchLean.Tensor ℝ [channels]) : bnVec ε (tensorToVec x, tensorToVec gamma, tensorToVec beta) = tensorToVec (Spec.batchNorm x gamma beta ε)

The vectorized BatchNorm at packed tensors is the vectorized spec output.

theorem Proofs.Autograd.BatchNorm.get2_flat_batchNormJvp {channels : ℕ} {sSpatial : Spec.Shape} [(Spec.Shape.dim channels sSpatial).WellFormed] {ε : ℝ} (hε : 0 < ε) (x dx : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (gamma dgamma beta dbeta : TorchLean.Tensor ℝ [channels]) (c : Fin channels) (p : Fin (positions sSpatial)) : Spec.get2 ((Spec.batchNormJvp x dx gamma dgamma beta dbeta ε).reshapeSpec ⋯) c p = RowNorm.nrmJvp (matVec (tensorToVec x)) ε (matVec (tensorToVec dx)) c p * gamma.getScalar c + RowNorm.nrm (matVec (tensorToVec x)) ε c p * dgamma.getScalar c + dbeta.getScalar c

Entries of Spec.batchNormJvp in the flattened matrix are the closed-form row differential, for positive ε.

theorem Proofs.Autograd.BatchNorm.bnD_tensors {channels : ℕ} {sSpatial : Spec.Shape} [(Spec.Shape.dim channels sSpatial).WellFormed] {ε : ℝ} (hε : 0 < ε) (x dx : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (gamma dgamma beta dbeta : TorchLean.Tensor ℝ [channels]) : (bnD (matVec (tensorToVec x)) (tensorToVec gamma) ε) (tensorToVec dx, tensorToVec dgamma, tensorToVec dbeta) = tensorToVec (Spec.batchNormJvp x dx gamma dgamma beta dbeta ε)

The derivative of BatchNorm applied to packed tangents is Spec.batchNormJvp.

theorem Proofs.Autograd.BatchNorm.fderiv_batchNorm_eq_batchNormJvp {channels : ℕ} {sSpatial : Spec.Shape} [(Spec.Shape.dim channels sSpatial).WellFormed] {ε : ℝ} (hε : 0 < ε) (x dx : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (gamma dgamma beta dbeta : TorchLean.Tensor ℝ [channels]) : (fderiv ℝ (bnVec ε) (tensorToVec x, tensorToVec gamma, tensorToVec beta)) (tensorToVec dx, tensorToVec dgamma, tensorToVec dbeta) = tensorToVec (Spec.batchNormJvp x dx gamma dgamma beta dbeta ε)

The Fréchet derivative of Spec.batchNorm in (x, gamma, beta) applied to a tangent triple is Spec.batchNormJvp, for positive ε.