TorchLean API

NN.Proofs.Analysis.Normalization

Normalization analysis properties #

This file records theorem-level properties of TorchLean's normalization specs. The executable and spec definitions live in NN.Spec.Layers.Normalization; this file belongs under NN.Proofs.Analysis because it proves algebraic facts about those definitions over .

Current focus:

That fact matters because inference-time BatchNorm can be folded into a preceding/following affine layer for verification, bound propagation, and simplification. Training-mode BatchNorm has batch-dependent statistics and is not claimed by this theorem.

BatchNorm inference is affine #

At inference time, BatchNorm uses fixed running mean/variance. The resulting function is affine:

$y=x\odot\operatorname{scale}+\operatorname{bias}$.

theorem Proofs.Normalization.batchNorm_inference_eq_mul_add {channels : } {sSpatial : Spec.Shape} (x : TorchLean.Tensor (Spec.Shape.dim channels sSpatial)) (runningMean runningVar gamma beta : TorchLean.Tensor [channels]) (epsilon : := TorchLean.normalizationEpsilon) :
Spec.batchNormInference x runningMean runningVar gamma beta epsilon = let s := Spec.Shape.dim channels sSpatial; have runningVar := runningVar.maxSpec (TorchLean.Tensor.full (Spec.Shape.dim channels Spec.Shape.scalar) 0); have mean_b := Spec.broadcastChannel sSpatial runningMean; have var_b := Spec.broadcastChannel sSpatial runningVar; have gamma_b := Spec.broadcastChannel sSpatial gamma; have beta_b := Spec.broadcastChannel sSpatial beta; have std := (var_b.addSpec (TorchLean.Tensor.full s epsilon)).sqrtSpec; (x.mulSpec (gamma_b.divSpec std)).addSpec (beta_b.subSpec (mean_b.mulSpec (gamma_b.divSpec std)))

Inference-time BatchNorm is affine in the input x.

This is the public theorem users want for verification and graph simplification:

batchNormInference x runningMean runningVar gamma beta epsilon

is definitionally equal to a pointwise affine map

$x\,\frac{\gamma}{\operatorname{std}} +\left(\beta-\mu\frac{\gamma}{\operatorname{std}}\right)$

after broadcasting channel parameters to the input shape and clamping the running variance exactly as the spec does.