Batch Normalization

Generic channel-first BatchNorm semantics together with its JVP and VJP. All spatial axes are flattened only for the reduction, so the definitions apply uniformly at every tensor rank. Training-time statistics and inference-time running statistics remain separate operations.

def Spec.broadcastChannel {α : Type} [TorchLean.Storage α] [Context α] {channels : ℕ} (sSpatial : Shape) (x : TorchLean.Tensor α [channels]) : TorchLean.Tensor α (Shape.concat [channels] sSpatial)

Repeat a channel vector over every position of a spatial shape.

def Spec.BatchNorm.normalizedJvp {α : Type} [TorchLean.Storage α] [Context α] {channels positions : ℕ} (hPositions : 0 < positions) (tangent xHat : TorchLean.Tensor α [channels, positions]) (invStd gamma dgamma dbeta : TorchLean.Tensor α [channels]) : TorchLean.Tensor α [channels, positions]

Differential of normalized, affine channel data after statistics have been computed.

def Spec.BatchNorm.normalizedBackward {α : Type} [TorchLean.Storage α] [Context α] {channels positions : ℕ} (hPositions : 0 < positions) (gradOutput xHat : TorchLean.Tensor α [channels, positions]) (invStd gamma : TorchLean.Tensor α [channels]) : TorchLean.Tensor α [channels, positions] × TorchLean.Tensor α [channels] × TorchLean.Tensor α [channels]

Reverse rule adjoint to normalizedJvp, including affine-parameter gradients.

def Spec.batchNorm {α : Type} [TorchLean.Storage α] [Context α] {channels : ℕ} {sSpatial : Shape} (x : TorchLean.Tensor α (Shape.concat [channels] sSpatial)) (gamma beta : TorchLean.Tensor α [channels]) (epsilon : α := TorchLean.normalizationEpsilon) [(Shape.concat [channels] sSpatial).WellFormed] : TorchLean.Tensor α (Shape.concat [channels] sSpatial)

Stateless BatchNorm for channel-first tensors with shape [channels] ++ sSpatial.

This computes per-channel mean/variance over the sSpatial axes and applies:

y = ((x - mean) / sqrt(var + eps)) * gamma + beta.

PyTorch analogy: torch.nn.BatchNorm{1,2,3}d in training mode on an input with batch size N=1. TorchLean does not model the running-statistics update here. A singleton spatial domain is accepted and has zero variance; PyTorch training BatchNorm requires more than one value per channel. The usual normalization interpretation assumes positive epsilon, which this raw spec does not validate.

def Spec.batchNormJvp {α : Type} [TorchLean.Storage α] [Context α] {channels : ℕ} {sSpatial : Shape} (x tangent : TorchLean.Tensor α (Shape.concat [channels] sSpatial)) (gamma dgamma _beta dbeta : TorchLean.Tensor α [channels]) (epsilon : α := TorchLean.normalizationEpsilon) [(Shape.concat [channels] sSpatial).WellFormed] : TorchLean.Tensor α (Shape.concat [channels] sSpatial)

Forward-mode JVP for batchNorm.

Stateless BatchNorm computes one set of statistics per channel over every spatial position. The input tangent therefore uses the same closed-form normalization differential as LayerNorm, with the mean taken over the flattened spatial shape for each channel:

dxhat = inv_std * (dx - mean(dx) - xhat * mean(dx*xhat)).

Affine tangents contribute xhat * dgamma + dbeta channel-wise.

def Spec.batchNormBackward {α : Type} [TorchLean.Storage α] [Context α] {channels : ℕ} {sSpatial : Shape} (x : TorchLean.Tensor α (Shape.concat [channels] sSpatial)) (gamma : TorchLean.Tensor α [channels]) (gradOutput : TorchLean.Tensor α (Shape.concat [channels] sSpatial)) (epsilon : α := TorchLean.normalizationEpsilon) [(Shape.concat [channels] sSpatial).WellFormed] : NormalizationGradients α (Shape.concat [channels] sSpatial) [channels]

Backward/VJP for batchNorm.

Statistics and affine-parameter gradients are reduced over every spatial axis, independently for each channel.

BatchNorm (inference-time, running statistics)

PyTorch distinction:

• training: normalize using batch statistics (and update running mean/variance);
• inference: normalize using the stored running mean/variance.

TorchLean keeps things pure and explicit: inference-time BatchNorm takes the running statistics as arguments.

def Spec.batchNormInference {α : Type} [TorchLean.Storage α] [Context α] {channels : ℕ} {sSpatial : Shape} (x : TorchLean.Tensor α (Shape.concat [channels] sSpatial)) (runningMean runningVar gamma beta : TorchLean.Tensor α [channels]) (epsilon : α := TorchLean.normalizationEpsilon) : TorchLean.Tensor α (Shape.concat [channels] sSpatial)

Inference-time BatchNorm for channel-first tensors with shape [channels] ++ sSpatial, using fixed running statistics.

Formula (per channel c):

y = ((x - μ) / sqrt(σ² + eps)) * γ + β

This matches the standard evaluation-time behavior of torch.nn.BatchNorm{1,2,3}d (no batch-statistics computation, no running-statistics update).

Over real arithmetic, (μ, σ², γ, β) are constants, so this is an affine map in x. Floating-point evaluation still rounds the individual operations. See Proofs.Normalization.batchNorm_inference_eq_mul_add.