BugZoo: LayerNorm on a one-feature axis

LayerNorm has a sharp degenerate case: if the normalized axis has length one, then the mean is the single input value and the variance is zero. The normalized value is therefore exactly zero, so the affine output is the bias and the forward result is independent of both the input and the scale.

PyTorch analogy:

torch.nn.functional.layer_norm(x, normalized_shape=(1,), weight=w, bias=b)

For every finite scalar x, the mathematical contract is:

$$ \begin{aligned} \operatorname{mean}[x] &= x,\\ \operatorname{var}[x] &= 0,\\ \left(\frac{x-\operatorname{mean}}{\sqrt{\operatorname{var}+\varepsilon}}\right) \operatorname{weight}+\operatorname{bias} &= \operatorname{bias}. \end{aligned} $$

So reverse mode must report zero gradient for weight and zero input gradient. This file keeps the contract small: the real-valued theorems record the algebra, and the concrete definitions below are the public TorchLean spec terms used by the Python reproducer notes.

The algebraic statements use real arithmetic with totalized division and square root. They do not assert finite native results for every epsilon or verify an external backward implementation.

Run python3 scripts/verification/normalization_contract_probe.py --device cpu to sweep input magnitudes and compare PyTorch forward and backward residuals in float32 and float64. --device cuda probes the GPU implementation when available; observed numbers depend on the PyTorch build and hardware.

@[reducible, inline]

abbrev NN.Examples.BugZoo.LayerNormDegenerateAxis.OneMat (α : Type) [TorchLean.Storage α] : Type

A one-by-one matrix: a single batch element with a single feature.

@[reducible, inline]

abbrev NN.Examples.BugZoo.LayerNormDegenerateAxis.OneVec (α : Type) [TorchLean.Storage α] : Type

A length-one vector, the shape LayerNorm's γ and β take here.

def NN.Examples.BugZoo.LayerNormDegenerateAxis.oneMat {α : Type} [TorchLean.Storage α] (x : α) : OneMat α

Build a one-by-one matrix from a scalar.

def NN.Examples.BugZoo.LayerNormDegenerateAxis.oneVec {α : Type} [TorchLean.Storage α] (x : α) : OneVec α

Build a length-one vector from a scalar.

theorem NN.Examples.BugZoo.LayerNormDegenerateAxis.one_feature_layernorm_scalar_contract (x gamma beta epsilon : ℝ) : (x - x) / MathFunctions.sqrt (max (0 + epsilon) 0) * gamma + beta = beta

The scalar algebra behind one-feature LayerNorm: normalization contributes zero, so the affine result is the bias.

theorem NN.Examples.BugZoo.LayerNormDegenerateAxis.one_feature_layernorm_scale_grad_contract (x dy epsilon : ℝ) : dy * ((x - x) / MathFunctions.sqrt (max (0 + epsilon) 0)) = 0

The scale/weight gradient is zero because the normalized one-feature value is zero.

theorem NN.Examples.BugZoo.LayerNormDegenerateAxis.one_feature_layernorm_input_grad_contract (dy gamma invStd : ℝ) : invStd * (dy * gamma - dy * gamma - 0) = 0

The input gradient is zero because the one-feature LayerNorm forward is constant in the input.

def NN.Examples.BugZoo.LayerNormDegenerateAxis.reproLayerNormForward : Float

TorchLean spec value for the public PyTorch repro: forward output.

def NN.Examples.BugZoo.LayerNormDegenerateAxis.reproLayerNormDWeight : Float

TorchLean spec value for the public PyTorch repro: gradient with respect to weight.

def NN.Examples.BugZoo.LayerNormDegenerateAxis.reproLayerNormDX : Float

TorchLean spec value for the public PyTorch repro: gradient with respect to input.