BugZoo: normalization state and BatchNorm contracts

BatchNorm is a small formula with a surprisingly large bug surface. Cross-backend testing work found real library bugs around normalization formulas and backend conventions, including epsilon placement in BatchNorm. Model-generation testing also found BatchNormalization failures involving wrong moving statistics and NaN-producing outputs.

References:

• Pham et al., "CRADLE: Cross-Backend Validation to Detect and Localize Bugs in Deep Learning Libraries", ICSE 2019.
• Wang et al., "Deep Learning Library Testing via Effective Model Generation", ISSTA 2020.
• Ioffe and Szegedy, "Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift", ICML 2015.

TorchLean addresses this class in two layers:

• the spec formula is explicit, so epsilon placement is not hidden in backend code;
• inference-time running statistics are explicit inputs, so train/eval state boundaries become part of the checked object rather than ambient mutable framework state.

noncomputable def NN.Examples.BugZoo.NormalizationState.wrongEpsilonOutsideSqrt (x mean variance gamma beta epsilon : ℝ) : ℝ

The buggy BatchNorm pattern reported by cross-backend testing is easy to state: putting epsilon outside the square root changes the formula from

$$ \frac{x-\mu}{\sqrt{\sigma^2+\varepsilon}} $$

to

$$ \frac{x-\mu}{\sqrt{\sigma^2}+\varepsilon}. $$

The separating example below uses zero variance and epsilon four.

noncomputable def NN.Examples.BugZoo.NormalizationState.correctEpsilonInsideSqrt (x mean variance gamma beta epsilon : ℝ) : ℝ

The intended scalar BatchNorm expression. This mirrors the public TorchLean normalization spec: epsilon is added to the variance before the square root.

theorem NN.Examples.BugZoo.NormalizationState.wrongEpsilonOutsideSqrt_ne_correctEpsilonInsideSqrt : wrongEpsilonOutsideSqrt 1 0 0 1 0 4 ≠ correctEpsilonInsideSqrt 1 0 0 1 0 4

At variance zero and epsilon four, the misplaced epsilon divides by four instead of two.

theorem NN.Examples.BugZoo.NormalizationState.normalizeCore_scalar_uses_variance_plus_epsilon (x mean variance gamma beta epsilon : ℝ) : Spec.normalizeCore [] [] [] [] [] epsilon (TorchLean.Tensor.full [] x) (TorchLean.Tensor.full [] mean) (TorchLean.Tensor.full [] variance) (TorchLean.Tensor.full [] gamma) (TorchLean.Tensor.full [] beta) ⋯ ⋯ ⋯ ⋯ = TorchLean.Tensor.full [] ((x - mean) / MathFunctions.sqrt (max (variance + epsilon) 0) * gamma + beta)

Spec-level BatchNorm uses epsilon inside the variance term.

There is one extra implementation detail worth making explicit: sqrtSpec is total, so it computes $\sqrt{\max(\mathrm{variance}+\varepsilon,0)}$. On the usual BatchNorm path, variance is nonnegative and epsilon is positive, so this is the same mathematical formula as $\sqrt{\mathrm{variance}+\varepsilon}$.

structure NN.Examples.BugZoo.NormalizationState.RunningStats (channels : ℕ) : Type

Running statistics are part of the BatchNorm inference contract.

The exact running mean and variance are arguments to the spec, so a reviewer can identify which state the result uses. Their shape does not establish that they are current or came from the intended training run; that provenance remains a separate obligation.

• mean : TorchLean.Tensor ℝ [channels]

  Inference-time running mean, usually learned/updated during training.

• variance : TorchLean.Tensor ℝ [channels]

  Inference-time running variance, clamped by the spec before normalization.

noncomputable def NN.Examples.BugZoo.NormalizationState.batchNormEvalWithStats {channels : ℕ} {sSpatial : Spec.Shape} (x : TorchLean.Tensor ℝ (sSpatial.prependDim channels)) (stats : RunningStats channels) (gamma beta : TorchLean.Tensor ℝ [channels]) (epsilon : ℝ := TorchLean.normalizationEpsilon) : TorchLean.Tensor ℝ (sSpatial.prependDim channels)

Evaluation-time BatchNorm with state packaged as an explicit value.

theorem NN.Examples.BugZoo.NormalizationState.batchNormEvalWithStats_unfolds {channels : ℕ} {sSpatial : Spec.Shape} (x : TorchLean.Tensor ℝ (sSpatial.prependDim channels)) (stats : RunningStats channels) (gamma beta : TorchLean.Tensor ℝ [channels]) (epsilon : ℝ := TorchLean.normalizationEpsilon) : batchNormEvalWithStats x stats gamma beta epsilon = Spec.batchNormInference x stats.mean stats.variance gamma beta epsilon

The packaged-state wrapper is exactly the public inference-time BatchNorm spec.

theorem NN.Examples.BugZoo.NormalizationState.batchNormEvalWithStats_affine {channels : ℕ} {sSpatial : Spec.Shape} (stats : RunningStats channels) (gamma beta : TorchLean.Tensor ℝ [channels]) (epsilon : ℝ := TorchLean.normalizationEpsilon) : ∃ (scale : TorchLean.Tensor ℝ (sSpatial.prependDim channels)) (bias : TorchLean.Tensor ℝ (sSpatial.prependDim channels)), ∀ (x : TorchLean.Tensor ℝ (sSpatial.prependDim channels)), batchNormEvalWithStats x stats gamma beta epsilon = (x.mulSpec scale).addSpec bias

Fixed BatchNorm running statistics determine one scale and bias that work for every input.

The witnesses depend only on the running statistics, affine parameters, and epsilon. This is the uniform affine representation needed when folding inference-time normalization into another layer.

theorem NN.Examples.BugZoo.NormalizationState.batchNormEvalWithStats_is_affine {channels : ℕ} {sSpatial : Spec.Shape} (x : TorchLean.Tensor ℝ (sSpatial.prependDim channels)) (stats : RunningStats channels) (gamma beta : TorchLean.Tensor ℝ [channels]) (epsilon : ℝ := TorchLean.normalizationEpsilon) : ∃ (scale : TorchLean.Tensor ℝ (sSpatial.prependDim channels)) (bias : TorchLean.Tensor ℝ (sSpatial.prependDim channels)), batchNormEvalWithStats x stats gamma beta epsilon = (x.mulSpec scale).addSpec bias

Specialize the shared affine representation to one input.