Normalization analysis properties

This file proves algebraic properties of normalization and cosine loss over ℝ. The executable definitions live in NN.Spec.Layers.Normalization and NN.Spec.Layers.Loss.

Current focus:

• BatchNorm inference is affine in the input x when the running statistics are fixed.
• Scaled cosine loss and its selected gradient preserve the epsilon-clipped real formulas.

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. The folding identity is over real arithmetic; it does not justify reordering finite-precision division, multiplication, and addition without a separate rounding analysis.

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.

batchNormInference x runningMean runningVar gamma beta epsilon

equals the 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.

Scaled cosine normalization

theorem Proofs.Normalization.cosineSimilarity_eq {s : Spec.Shape} (x t : TorchLean.Tensor ℝ s) (epsilon : ℝ) (he : 0 < epsilon) : Spec.cosineSimilaritySpec x t epsilon = 1 - (x.mulSpec t).sumSpec / (max (√x.squareSpec.sumSpec) epsilon * max (√t.squareSpec.sumSpec) epsilon)

Scaling before normalization preserves the epsilon-clipped real cosine loss. This identity does not assert bitwise equivalence of the two floating-point schedules.

theorem Proofs.Normalization.cosineSimilarityDeriv_eq {s : Spec.Shape} (x t : TorchLean.Tensor ℝ s) (epsilon : ℝ) (he : 0 < epsilon) : Spec.cosineSimilarityDerivSpec x t epsilon = have xn := √x.squareSpec.sumSpec; have tn := √t.squareSpec.sumSpec; have xd := max xn epsilon; have td := max tn epsilon; TorchLean.Tensor.map2Spec (fun (p q : ℝ) => (if xn > epsilon then (x.mulSpec t).sumSpec / (xd ^ 3 * td) * p else 0) - q / (xd * td)) x t

The scaled pullback preserves the clipped real formula, including its zero radial term at and below the predicted-norm floor. No rounding-error bound is asserted here.