Batch Normalization

Generic channel-first BatchNorm semantics together with the JVP and VJP used by TorchLean's concrete 2D graph operator. Training-time statistics and inference-time running statistics remain separate mathematical operations.

def Spec.batchNorm {α : Type} [Context α] {channels : ℕ} {sSpatial : Shape} (x : Tensor α (Shape.dim channels sSpatial)) (gamma beta : Tensor α (Shape.dim channels Shape.scalar)) (epsilon : α := Numbers.epsilon) [(Shape.dim channels sSpatial).WellFormed] : Tensor α (Shape.dim channels sSpatial)

Stateless BatchNorm for channel-first tensors of shape .dim 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.

def Spec.instanceNorm {α : Type} [Context α] {channels : ℕ} {sSpatial : Shape} (x : Tensor α (Shape.dim channels sSpatial)) (gamma beta : Tensor α (Shape.dim channels Shape.scalar)) (epsilon : α := Numbers.epsilon) [(Shape.dim channels sSpatial).WellFormed] : Tensor α (Shape.dim channels sSpatial)

Alias: per-sample normalization over spatial axes ("InstanceNorm-style").

The spec-level batchNorm* operators model the N=1 case (no explicit batch axis and no running statistics update). Many ML codebases refer to that behavior as instance normalization.

These aliases make that intent explicit without changing the existing API surface.

def Spec.batchNorm2d {α : Type} [Context α] {channels height width : ℕ} (x : Tensor α (Shape.dim channels (Shape.dim height (Shape.dim width Shape.scalar)))) (gamma beta : Tensor α (Shape.dim channels Shape.scalar)) (h_c : channels > 0 := by norm_num) (h_h : height > 0 := by norm_num) (h_w : width > 0 := by norm_num) (epsilon : α := Numbers.epsilon) : Tensor α (Shape.dim channels (Shape.dim height (Shape.dim width Shape.scalar)))

batchNorm specialized to a single channel-first image (C,H,W).

def Spec.batchNorm2dJvp {α : Type} [Context α] {channels height width : ℕ} (x tangent : Tensor α (Shape.dim channels (Shape.dim height (Shape.dim width Shape.scalar)))) (gamma dgamma _beta dbeta : Tensor α (Shape.dim channels Shape.scalar)) (_h_c : channels > 0 := by norm_num) (_h_h : height > 0 := by norm_num) (_h_w : width > 0 := by norm_num) (epsilon : α := Numbers.epsilon) : Tensor α (Shape.dim channels (Shape.dim height (Shape.dim width Shape.scalar)))

Forward-mode JVP for batchNorm2d.

TorchLean's stateless BatchNorm2d computes one set of statistics per channel over the spatial grid. The input tangent therefore uses the same closed-form normalization differential as LayerNorm, but with the mean taken over (height,width) for each channel:

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

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

def Spec.batchNorm2dBackward {α : Type} [Context α] {channels height width : ℕ} (x : Tensor α (Shape.dim channels (Shape.dim height (Shape.dim width Shape.scalar)))) (gamma : Tensor α (Shape.dim channels Shape.scalar)) (grad_output : Tensor α (Shape.dim channels (Shape.dim height (Shape.dim width Shape.scalar)))) (_h_c : channels > 0 := by norm_num) (_h_h : height > 0 := by norm_num) (_h_w : width > 0 := by norm_num) (epsilon : α := Numbers.epsilon) : Tensor α (Shape.dim channels (Shape.dim height (Shape.dim width Shape.scalar))) × Tensor α (Shape.dim channels Shape.scalar) × Tensor α (Shape.dim channels Shape.scalar)

Backward/VJP for batchNorm2d.

Returns (dx, dGamma, dBeta). This matches the shape of gradients you expect from a PyTorch-style BatchNorm2d, but note that our forward is the per-image variant (no explicit batch dimension and no running statistics).

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} [Context α] {channels : ℕ} {sSpatial : Shape} (x : Tensor α (Shape.dim channels sSpatial)) (runningMean runningVar gamma beta : Tensor α (Shape.dim channels Shape.scalar)) (epsilon : α := Numbers.epsilon) : Tensor α (Shape.dim channels sSpatial)

Inference-time BatchNorm for channel-first tensors of shape .dim 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).

At inference time, (μ, σ², γ, β) are constants, so this is an affine map in x. See NN.Proofs.Analysis.Normalization.batchNorm_inference_eq_mul_add.