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.

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 BatchNorm2d: https://docs.pytorch.org/docs/stable/generated/torch.nn.BatchNorm2d.html

def Spec.normalizeCore {α : Type} [Context α] (s s_mean s_var s_gamma s_beta : Shape) (epsilon : α) (x : Tensor α s) (mean : Tensor α s_mean) (variance : Tensor α s_var) (gamma : Tensor α s_gamma) (beta : Tensor α s_beta) (cb_mean : s_mean.CanBroadcastTo s) (cb_var : s_var.CanBroadcastTo s) (cb_gamma : s_gamma.CanBroadcastTo s) (cb_beta : s_beta.CanBroadcastTo s) : 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} [Context α] {seqLen embedDim : ℕ} (x : Tensor α (Shape.dim seqLen (Shape.dim embedDim Shape.scalar))) (gamma beta : Tensor α (Shape.dim embedDim Shape.scalar)) (h_seq_pos : seqLen > 0 := by norm_num) (h_embed_pos : embedDim > 0 := by norm_num) (epsilon : α := Numbers.epsilon) : Tensor α (Shape.dim seqLen (Shape.dim embedDim Shape.scalar))

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

Uses epsilon (default Numbers.epsilon) for numerical stability in the denominator.

def Spec.layerNormBackward {α : Type} [Context α] {seqLen embedDim : ℕ} (h_seq_pos : seqLen > 0) (h_embed_pos : embedDim > 0) (x : Tensor α (Shape.dim seqLen (Shape.dim embedDim Shape.scalar))) (gamma _beta : Tensor α (Shape.dim embedDim Shape.scalar)) (grad_output : Tensor α (Shape.dim seqLen (Shape.dim embedDim Shape.scalar))) (epsilon : α := Numbers.epsilon) : Tensor α (Shape.dim seqLen (Shape.dim embedDim Shape.scalar)) × Tensor α (Shape.dim embedDim Shape.scalar) × Tensor α (Shape.dim embedDim Shape.scalar)

Backward/VJP for layerNorm (returns (dx, dGamma, dBeta)).

def Spec.layerNormJvp {α : Type} [Context α] {seqLen embedDim : ℕ} (h_seq_pos : seqLen > 0) (h_embed_pos : embedDim > 0) (x tangent : Tensor α (Shape.dim seqLen (Shape.dim embedDim Shape.scalar))) (gamma dgamma _beta dbeta : Tensor α (Shape.dim embedDim Shape.scalar)) (epsilon : α := Numbers.epsilon) : Tensor α (Shape.dim seqLen (Shape.dim embedDim Shape.scalar))

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 = inv_std ⊙ (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.

def Spec.groupNorm {α : Type} [Context α] {batchSize height width channels groups : ℕ} (x : Tensor α (Shape.dim batchSize (Shape.dim height (Shape.dim width (Shape.dim channels Shape.scalar))))) (gamma beta : Tensor α (Shape.dim channels Shape.scalar)) (h_b : batchSize > 0 := by norm_num) (h_h : height > 0 := by norm_num) (h_w : width > 0 := by norm_num) (h_c : channels > 0 := by norm_num) (h_g : groups > 0 := by norm_num) (h_ge : channels ≥ groups) (h_div : channels % groups = 0) (epsilon : α := Numbers.epsilon) : Tensor α (Shape.dim batchSize (Shape.dim height (Shape.dim width (Shape.dim channels Shape.scalar))))

GroupNorm for channel-last tensors (batch, height, width, channels).

PyTorch analogy: torch.nn.GroupNorm(num_groups=groups, num_channels=channels) applied per sample. The mean/variance are computed over both the spatial dimensions and the channels within each group, then an affine transform is applied per channel via gamma and beta.

This operator is useful in settings where BatchNorm's dependence on batch statistics is awkward (small batches, verification, or when you want purely per-sample behavior).

def Spec.normalizeAlongDim {α : Type} [Context α] {s : Shape} (x gamma beta : Tensor α s) (dim : ℕ) (h_valid : Shape.valid_axis_inst dim s) (h_wf : s.WellFormed) (epsilon : α := Numbers.epsilon) : 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} [Context α] {seqLen embedDim : ℕ} (x : Tensor α (Shape.dim seqLen (Shape.dim embedDim Shape.scalar))) (gamma : Tensor α (Shape.dim embedDim Shape.scalar)) (h_seq_pos : seqLen > 0 := by norm_num) (h_embed_pos : embedDim > 0 := by norm_num) (epsilon : α := Numbers.epsilon) : Tensor α (Shape.dim seqLen (Shape.dim embedDim Shape.scalar))

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} [Context α] {inDim outDim : ℕ} (weight : Tensor α (Shape.dim outDim (Shape.dim inDim Shape.scalar))) (gamma : Tensor α (Shape.dim outDim Shape.scalar)) (h_out_pos : outDim > 0 := by norm_num) (h_in_pos : inDim > 0 := by norm_num) (epsilon : α := Numbers.epsilon) : Tensor α (Shape.dim outDim (Shape.dim inDim Shape.scalar))

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.