Batch Normalization Backward Correctness #

The training-time BatchNorm differential couples all positions in a channel through the channel mean and variance. This file proves the reverse-mode rule used by TorchLean for an arbitrary spatial shape. The proof is stated as JVP/VJP adjointness: pairing an input and parameter tangent with the forward differential gives the same scalar as pairing the upstream gradient with Spec.batchNormBackward.

The theorem is independent of a memory layout. The spatial shape is flattened only inside the BatchNorm definition, so the result covers vectors, images, volumes, and higher-rank tensors with one statement.

The companion module NN.Proofs.Autograd.Tape.Ops.Norm.BatchNormFDeriv proves that Spec.batchNormJvp is the Fréchet derivative of Spec.batchNorm in (x, gamma, beta) whenever 0 < ε (hasFDerivAt_batchNorm, fderiv_batchNorm_eq_batchNormJvp). Together with the adjointness below this identifies Spec.batchNormBackward as the adjoint of that derivative.

source

theorem Proofs.Autograd.BatchNorm.normalizedJvp_normalizedBackward_adjoint {channels positions : ℕ} (hPositions : 0 < positions) (tangent gradOutput xHat : TorchLean.Tensor ℝ [channels, positions]) (invStd gamma dgamma dbeta : TorchLean.Tensor ℝ [channels]) :

Spec.dot (Spec.BatchNorm.normalizedJvp hPositions tangent xHat invStd gamma dgamma dbeta) gradOutput = have backward := Spec.BatchNorm.normalizedBackward hPositions gradOutput xHat invStd gamma; Spec.dot tangent backward.1 + Spec.dot dgamma backward.2.1 + Spec.dot dbeta backward.2.2

The normalized BatchNorm reverse rule is adjoint to its forward differential.

source

theorem Proofs.Autograd.BatchNorm.batchNormJvp_batchNormBackward_adjoint {channels : ℕ} {sSpatial : Spec.Shape} (x tangent gradOutput : TorchLean.Tensor ℝ (Spec.Shape.dim channels sSpatial)) (gamma dgamma beta dbeta : TorchLean.Tensor ℝ [channels]) (epsilon : ℝ := TorchLean.normalizationEpsilon) [(Spec.Shape.dim channels sSpatial).WellFormed] :

Spec.dot (Spec.batchNormJvp x tangent gamma dgamma beta dbeta epsilon) gradOutput = have backward := Spec.batchNormBackward x gamma gradOutput epsilon; Spec.dot tangent backward.inputGradient + Spec.dot dgamma backward.scaleGradient + Spec.dot dbeta backward.biasGradient

The rank-general BatchNorm reverse rule is adjoint to its forward differential.

This is the correctness theorem for the public channel-first operator. The spatial shape is arbitrary; flattening is an implementation detail used only to compute each channel's statistics. The conclusion accounts for the input tangent and both affine parameters.