BatchNorm operator bounds (IBP + affine)

This file bounds inference-time BatchNorm. Since inference-time BatchNorm is an affine transformation (with frozen statistics), both IBP and affine propagation are exact (componentwise).

At inference time, TorchLean uses y = γ * (x - μ) / sqrt(max(σ², 0) + ε) + β, so the layer reduces to y = scale * x + offset, where scale = γ / sqrt(max(σ², 0) + ε) and offset = β - γ * μ / sqrt(max(σ², 0) + ε).

The max is the same totalization used by Spec.batchNormInference and the IR evaluator. It has no effect on valid nonnegative running variances, while keeping every TorchLean layer aligned on malformed approximate-runtime inputs.

References:

• Ioffe and Szegedy, "Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift", ICML 2015.
• PyTorch analogue: torch.nn.BatchNorm1d/2d/3d in evaluation mode.

structure NN.MLTheory.CROWN.Operators.Batchnorm.BatchNormParams (α : Type) [Context α] : Type

Parameters for BatchNorm layer (frozen at inference).

• dim : ℕ

  Number of channels/features

• running_mean : Spec.Tensor α (Spec.Shape.dim self.dim Spec.Shape.scalar)

  Running mean μ

• running_var : Spec.Tensor α (Spec.Shape.dim self.dim Spec.Shape.scalar)

  Running variance σ²

• gamma : Spec.Tensor α (Spec.Shape.dim self.dim Spec.Shape.scalar)

  Learnable scale γ

• beta : Spec.Tensor α (Spec.Shape.dim self.dim Spec.Shape.scalar)

  Learnable bias β

• eps : α

  Small constant for numerical stability

def NN.MLTheory.CROWN.Operators.Batchnorm.computeScale {α : Type} [Context α] (params : BatchNormParams α) : Spec.Tensor α (Spec.Shape.dim params.dim Spec.Shape.scalar)

Compute the equivalent affine scale: γ / sqrt(max(σ², 0) + ε).

def NN.MLTheory.CROWN.Operators.Batchnorm.computeOffset {α : Type} [Context α] (params : BatchNormParams α) : Spec.Tensor α (Spec.Shape.dim params.dim Spec.Shape.scalar)

Compute the equivalent affine offset: β - γ * μ / sqrt(max(σ², 0) + ε).

def NN.MLTheory.CROWN.Operators.Batchnorm.ibpBatchnorm {α : Type} [Context α] (params : BatchNormParams α) (xB : Box α (Spec.Shape.dim params.dim Spec.Shape.scalar)) : Box α (Spec.Shape.dim params.dim Spec.Shape.scalar)

IBP for BatchNorm. Since BatchNorm is affine, its bounds are exact.

For $y=sx+o$:

• if $s>0$, then $y_{\mathrm{lo}}=s x_{\mathrm{lo}}+o$ and $y_{\mathrm{hi}}=s x_{\mathrm{hi}}+o$;
• if $s<0$, then $y_{\mathrm{lo}}=s x_{\mathrm{hi}}+o$ and $y_{\mathrm{hi}}=s x_{\mathrm{lo}}+o$.

def NN.MLTheory.CROWN.Operators.Batchnorm.affBatchnorm {α : Type} [Context α] {inDim : ℕ} (params : BatchNormParams α) (aff : AffineVec α inDim params.dim) : AffineVec α inDim params.dim

Affine bounds for BatchNorm propagation.

Since BatchNorm is affine, compose the two affine forms:

$$ \begin{aligned} y_{\mathrm{prev}} &= A_{\mathrm{prev}}x_{\mathrm{in}}+c_{\mathrm{prev}},\\ \operatorname{BN}(y) &= sy+o,\\ \operatorname{BN}(y_{\mathrm{prev}}) &= \operatorname{diag}(s)A_{\mathrm{prev}}x_{\mathrm{in}} +(s c_{\mathrm{prev}}+o). \end{aligned} $$

def NN.MLTheory.CROWN.Operators.Batchnorm.derivBatchnorm {α : Type} [Context α] (params : BatchNormParams α) (dB : Box α (Spec.Shape.dim params.dim Spec.Shape.scalar)) : Box α (Spec.Shape.dim params.dim Spec.Shape.scalar)

Derivative bounds for BatchNorm. Since BatchNorm is affine, $\frac{d}{dx}\operatorname{BN}(x)=s$ is constant. Input bounds $[d_{\mathrm{lo}},d_{\mathrm{hi}}]$ therefore become $s[d_{\mathrm{lo}},d_{\mathrm{hi}}]$.

def NN.MLTheory.CROWN.Operators.Batchnorm.secondDerivBatchnorm {α : Type} [Context α] (params : BatchNormParams α) (d2B : Box α (Spec.Shape.dim params.dim Spec.Shape.scalar)) : Box α (Spec.Shape.dim params.dim Spec.Shape.scalar)

Propagate second-derivative bounds through inference-time BatchNorm.

Although the second derivative of the affine map x ↦ scale * x + offset with respect to x is zero, composition with a curve x(t) gives d²/dt² BN(x(t)) = scale * x''(t). Consequently this uses the same signed scaling rule as first-derivative propagation.