Normalization Layers

These layers package normalization programs with their learned parameters and, for BatchNorm, their running statistics. LayerNorm and RMSNorm operate on the final axis. BatchNorm, InstanceNorm, and GroupNorm take inputs of shape [batch, channels, ...spatial], where spatial can describe a sequence, image, volume, or any other collection of spatial axes.

Each constructor accepts a positive eps, added under the square root. We store it as a rational and convert it with Context.ofRat when the forward program chooses its scalar type. This keeps the layer configuration usable with both real-valued specifications and floating-point execution. Floating-point execution requires the converted epsilon to remain positive and finite; the static configuration check only enforces positivity of the rational input. In tiny formats the default can round to zero. Neither eager execution nor typed-graph lowering substitutes Context.defaultEpsilon; a constant LayerNorm input can therefore produce NaNs. Pass an eps whose converted value is positive and finite before lowering the model.

def Runtime.Autograd.Model.Layers.Internal.validateEpsilon (kind : String) (eps : ℚ) : Except String Unit

Require a positive epsilon before constructing runtime state.

A constant input has zero variance. Adding eps under the square root gives that input a positive denominator in the real-valued normalization formula.

def Runtime.Autograd.Model.Layers.Internal.affineNormalization (kind : String) (shape parameterShape : Spec.Shape) (affine bias : Bool) (validateConfig : Except String Unit) (normalize : {α : Type} → [inst : TorchLean.Storage α] → [inst_1 : Context α] → Program α [parameterShape, parameterShape, shape] shape) : Layer shape shape

Build a layer around a normalization program that expects scale, bias, and input references.

With affine := true, the scale starts at one and the optional bias starts at zero. Setting bias := false leaves only the scale in model state; setting affine := false leaves neither parameter. The forward program supplies constant ones and zeros for the omitted arguments, so the same normalization program can run eagerly or be recorded in a typed graph.

def Runtime.Autograd.Model.Layers.layerNorm (leading : Spec.Shape) (width : ℕ) {hWidth : width > 0} (eps : ℚ := 1e-5) (affine bias : Bool := true) : Layer (leading.appendDim width) (leading.appendDim width)

Layer normalization over the final axis.

For an input of shape [batch, tokens, width], each (batch, token) position has its own mean and variance over the width entries. We subtract that mean, divide by sqrt(variance + eps), then apply the learned scale and bias. The variance divides by width.

Scale and bias each have shape [width] and are shared across the leading axes. The scale starts at one and the bias at zero. Setting bias := false keeps only the scale; setting affine := false removes both parameters. Training and evaluation use the same input statistics.

def Runtime.Autograd.Model.Layers.rmsNorm (leading : Spec.Shape) (width : ℕ) {hWidth : width > 0} (eps : ℚ := 1e-5) (affine : Bool := true) : Layer (leading.appendDim width) (leading.appendDim width)

Root-mean-square normalization over the final axis.

Each row is divided by sqrt(mean(x * x) + eps) and multiplied by a learned scale of shape [width]. RMSNorm squares the entries directly, without first subtracting the row's mean. Its only learned parameter is the scale.

The scale starts at one. Setting affine := false removes it from model state and uses a constant scale of one. The default eps is 1e-5 for every scalar type; pass a different value when the model you are reproducing uses another epsilon.

def Runtime.Autograd.Model.Layers.batchNorm (batch channels : ℕ) (spatial : Spec.Shape) (hWellFormed : (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)).wellFormed) (momentum : Float := 0.1) (eps : ℚ := 1e-5) : Layer (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)) (Spec.Shape.dim batch (Spec.Shape.dim channels spatial))

Batch normalization over a batch, channel axis, and arbitrary spatial shape.

For each channel, training computes the mean and variance over the batch and every spatial position. The forward pass divides the variance sum by batch * spatial.size, then normalizes with sqrt(variance + eps) and applies the learned scale and bias.

Running statistics are updated separately with next = (1 - momentum) * running + momentum * batch. The running variance uses the unbiased estimate when there is more than one sample; for a single sample it keeps the finite biased value. Evaluation reads these stored statistics and leaves them unchanged.

Model state contains scale, bias, running mean, running variance, and momentum in that order. Only scale and bias receive gradients. Scale and running variance start at one; bias and running mean start at zero.

def Runtime.Autograd.Model.Layers.instanceNorm (batch channels : ℕ) (spatial : Spec.Shape) (hWellFormed : (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)).wellFormed) (eps : ℚ := 1e-5) (affine bias : Bool := true) : Layer (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)) (Spec.Shape.dim batch (Spec.Shape.dim channels spatial))

Instance normalization over the spatial axes of each sample and channel.

For an image tensor [batch, channels, height, width], each (batch, channel) pair has its own mean and variance over height * width entries. Both training and evaluation compute these statistics from the current input; this layer has no running-statistics buffers.

The optional scale and bias have shape [channels] and are shared across samples and spatial positions. Both are enabled by default. Setting bias := false keeps only the scale, while affine := false removes both. The variance divides by spatial.size, and eps is added before taking its square root.

def Runtime.Autograd.Model.Layers.groupNorm (batch channels groups : ℕ) (spatial : Spec.Shape) (hWellFormed : (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)).wellFormed) (hGroups : groups > 0) (hGroupsLe : channels ≥ groups) (hDiv : channels % groups = 0) (eps : ℚ := 1e-5) (affine bias : Bool := true) : Layer (Spec.Shape.dim batch (Spec.Shape.dim channels spatial)) (Spec.Shape.dim batch (Spec.Shape.dim channels spatial))

Group normalization over channel groups and their spatial positions, separately for each sample.

Channels are split into groups equal, contiguous groups. For example, six channels with groups := 2 form two groups of three channels. Each group's mean and variance include all three channels and all their spatial positions. The variance divides by the number of entries in that group, and eps is added before taking its square root.

Scale and bias still have one entry per channel, shared across samples and spatial positions. Setting bias := false keeps only the scale; setting affine := false removes both parameters. Training and evaluation use the same input statistics.