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.

@[reducible, inline]

abbrev NN.Examples.BugZoo.LayerNormDegenerateAxis.OneMat (α : Type) : Type

@[reducible, inline]

abbrev NN.Examples.BugZoo.LayerNormDegenerateAxis.OneVec (α : Type) : Type

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

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

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.