Normalization layers (spec layer)

This file collects a few normalization operators used throughout TorchLean's spec/model code.

The common pattern is:

• compute per-axis statistics (mean / variance or RMS),
• normalize with an epsilon for numerical stability,
• optionally apply an affine transform (gamma, beta) like PyTorch does.

The familiar normalization and differential interpretations require a positive epsilon and suitable real-number laws; the raw scalar-polymorphic definitions do not validate that parameter. For floating-point contexts, the forward, JVP, and VJP are separate rounded programs. Their closed-form differential formulas do not assert a derivative of IEEE rounding or bitwise equality with a native backend. In particular, LayerNorm computes reduceVar of already centered data, which centers again; simplifying that second centering changes floating-point execution.

References (papers + PyTorch behavior)

• LayerNorm: Ba et al., "Layer Normalization" (2016): https://arxiv.org/abs/1607.06450

• BatchNorm: Ioffe, Szegedy, "Batch Normalization" (2015): https://arxiv.org/abs/1502.03167

• GroupNorm: Wu, He, "Group Normalization" (2018): https://arxiv.org/abs/1803.08494

• RMSNorm: Zhang, Sennrich, "Root Mean Square Layer Normalization" (2019): https://arxiv.org/abs/1910.07467

• WeightNorm: Salimans, Kingma, "Weight Normalization" (2016): https://arxiv.org/abs/1602.07868

• PyTorch LayerNorm: https://docs.pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html

• PyTorch BatchNorm modules: https://docs.pytorch.org/docs/stable/nn.html#normalization-layers

structure Spec.NormalizationGradients (α : Type) [TorchLean.Storage α] (inputShape parameterShape : Shape) : Type

Named reverse-mode result shared by affine normalization operators.

• inputGradient : TorchLean.Tensor α inputShape

  Gradient with respect to the normalized input.

• scaleGradient : TorchLean.Tensor α parameterShape

  Gradient with respect to the learned multiplicative scale.

• biasGradient : TorchLean.Tensor α parameterShape

  Gradient with respect to the learned additive bias.

def Spec.instReprNormalizationGradients.repr {α✝ : Type} {inst✝ : TorchLean.Storage α✝} {inputShape✝ parameterShape✝ : Shape} [Repr α✝] : NormalizationGradients α✝ inputShape✝ parameterShape✝ → ℕ → Std.Format

@[instance_reducible]

instance Spec.instReprNormalizationGradients {α✝ : Type} {inst✝ : TorchLean.Storage α✝} {inputShape✝ parameterShape✝ : Shape} [Repr α✝] : Repr (NormalizationGradients α✝ inputShape✝ parameterShape✝)

def Spec.normalizeCore {α : Type} [TorchLean.Storage α] [Context α] (s sMean sVar sGamma sBeta : Shape) (epsilon : α) (x : TorchLean.Tensor α s) (mean : TorchLean.Tensor α sMean) (variance : TorchLean.Tensor α sVar) (gamma : TorchLean.Tensor α sGamma) (beta : TorchLean.Tensor α sBeta) (cbMean : sMean.CanBroadcastTo s) (cbVar : sVar.CanBroadcastTo s) (cbGamma : sGamma.CanBroadcastTo s) (cbBeta : sBeta.CanBroadcastTo s) : TorchLean.Tensor α s

Core normalization routine with explicit broadcast proofs.

This is the shared “math step” behind normalization layers:

y = ((x - mean) / sqrt(variance + ε)) * gamma + beta.

def Spec.layerNorm {α : Type} [TorchLean.Storage α] [Context α] {seqLen embedDim : ℕ} (x : TorchLean.Tensor α [seqLen, embedDim]) (gamma beta : TorchLean.Tensor α [embedDim]) (h_seq_pos : seqLen > 0 := by norm_num) (h_embed_pos : embedDim > 0 := by norm_num) (epsilon : α := TorchLean.normalizationEpsilon) : TorchLean.Tensor α [seqLen, embedDim]

LayerNorm over the last dimension of a (seqLen, embedDim) tensor.

Uses epsilon (default TorchLean.normalizationEpsilon) for numerical stability in the denominator. The default can round to zero in tiny formats and has no fallback. For those formats, pass a representable positive, finite epsilon explicitly; a constant row otherwise produces a zero denominator. This raw scalar-polymorphic operation does not validate the argument.

def Spec.layerNormBackward {α : Type} [TorchLean.Storage α] [Context α] {seqLen embedDim : ℕ} (sequenceLengthPositive : seqLen > 0) (embeddingWidthPositive : embedDim > 0) (input : TorchLean.Tensor α [seqLen, embedDim]) (scale : TorchLean.Tensor α [embedDim]) (outputGradient : TorchLean.Tensor α [seqLen, embedDim]) (epsilon : α := TorchLean.normalizationEpsilon) : NormalizationGradients α [seqLen, embedDim] [embedDim]

Backward/VJP for layerNorm, with named input, scale, and bias gradients.

def Spec.layerNormJvp {α : Type} [TorchLean.Storage α] [Context α] {seqLen embedDim : ℕ} (h_seq_pos : seqLen > 0) (h_embed_pos : embedDim > 0) (x tangent : TorchLean.Tensor α [seqLen, embedDim]) (gamma dgamma _beta dbeta : TorchLean.Tensor α [embedDim]) (epsilon : α := TorchLean.normalizationEpsilon) : TorchLean.Tensor α [seqLen, embedDim]

Forward-mode JVP for layerNorm.

For each sequence position, LayerNorm is the map y = gamma ⊙ xhat + beta with xhat = (x - mean(x)) / sqrt(var(x)+eps). The input tangent is normalized by the standard closed form

dxhat = invStd ⊙ (dx - mean(dx) - xhat ⊙ mean(dx ⊙ xhat)),

and affine-parameter tangents contribute xhat ⊙ dgamma + dbeta. This is the forward-mode counterpart of the closed-form VJP above and follows the same clamped-variance convention as the forward pass.

Group normalization

def Spec.groupNorm {α : Type} [TorchLean.Storage α] [Context α] {batch channels groups : ℕ} {spatial : Shape} (x : TorchLean.Tensor α (Shape.concat [batch, channels] spatial)) (gamma beta : TorchLean.Tensor α [channels]) (hGroups : groups > 0 := by norm_num) (hGroupsLe : channels ≥ groups) (hDiv : channels % groups = 0) (epsilon : α := TorchLean.normalizationEpsilon) [(Shape.concat [batch, channels] spatial).WellFormed] : TorchLean.Tensor α (Shape.concat [batch, channels] spatial)

Normalize each sample over groups of channels and every spatial position.

The spatial domain is an arbitrary Shape. Channels are split into groups contiguous groups; each group is flattened together with the spatial axes, normalized, and then transformed by the per-channel gamma and beta parameters.

def Spec.normalizeAlongDim {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} (x gamma beta : TorchLean.Tensor α s) (dim : ℕ) (h_valid : Shape.HasNonemptyAxis dim s) (_h_wf : s.WellFormed) (epsilon : α := TorchLean.normalizationEpsilon) : TorchLean.Tensor α s

Normalize along a chosen axis dim of a tensor x, using per-element affine parameters gamma and beta of the same shape as x.

This is a "generic building block" that is handy in specs; it is closer to the raw math than to a single PyTorch module. Most named normalizations (LayerNorm, GroupNorm, BatchNorm) are special cases of this pattern with a specific choice of axis set and parameter shape.

def Spec.rmsNorm {α : Type} [TorchLean.Storage α] [Context α] {seqLen embedDim : ℕ} (x : TorchLean.Tensor α [seqLen, embedDim]) (gamma : TorchLean.Tensor α [embedDim]) (h_seq_pos : seqLen > 0 := by norm_num) (h_embed_pos : embedDim > 0 := by norm_num) (epsilon : α := TorchLean.normalizationEpsilon) : TorchLean.Tensor α [seqLen, embedDim]

RMSNorm over the last dimension of a (seqLen, embedDim) tensor.

Compared to LayerNorm, RMSNorm skips subtracting the mean and normalizes by:

rms(x) = sqrt(mean(x^2) + eps).

This shows up in many Transformer-style models as a cheaper alternative to LayerNorm.

def Spec.weightNorm {α : Type} [TorchLean.Storage α] [Context α] {inDim outDim : ℕ} (weight : TorchLean.Tensor α [outDim, inDim]) (gamma : TorchLean.Tensor α [outDim]) (h_out_pos : outDim > 0 := by norm_num) (h_in_pos : inDim > 0 := by norm_num) (epsilon : α := TorchLean.normalizationEpsilon) : TorchLean.Tensor α [outDim, inDim]

WeightNorm for a dense weight matrix (outDim, inDim).

This implements the "normalize weight vectors then scale" idea:

• normalize each output row by its L2 norm,
• then rescale by gamma (one scalar per output row).

PyTorch analogy: weight normalization is typically applied as a parametrization of a module's weights rather than as a standalone tensor operator.