Skip to main content

Documentation Index

Fetch the complete documentation index at: https://mintlify.com/itsubaki/gpt/llms.txt

Use this file to discover all available pages before exploring further.

The layer package provides two normalization layers — RMSNorm (the default, used in all transformer blocks) and LayerNorm (available but not used by default). Both expose a learnable gamma scaling parameter and apply normalization across the last dimension of the input tensor, keeping all other dimensions intact.

​
RMSNorm

func RMSNorm(embedDim int) *RMSNormT
Constructs a Root Mean Square normalization layer. RMSNorm is the normalization of choice in Block because it omits mean subtraction, making it computationally lighter than standard Layer Normalization while achieving comparable quality in modern LLMs such as LLaMA.
​
embedDim
int
required
The size of the last dimension of the input tensor (the embedding dimension). Determines the shape of the learnable gamma parameter.

​
Parameters

ParameterShapeInitial value
gamma(embedDim,)all ones
Epsilon is hardcoded to 1e-5 and is not exposed as a constructor argument.

​
Formula

rms = sqrt( mean(x²) + ε )
y   = gamma * x / rms

​
Forward computation

x2  := F.Pow(2)(x[0])                         // x²
ms  := F.Reshape(shape...)(F.Mean(last)(x2))   // mean(x²), keepdims
rms := F.Pow(0.5)(F.AddC(l.eps, ms))           // sqrt(mean(x²) + ε)
y   := F.Mul(gamma, F.Div(x[0], rms))
Mean is computed over the last axis and the result is reshaped back to keep the broadcasting dimension, so the division x / rms broadcasts correctly over the full (B, C, E) input.

​
Shape

Input and output shapes are identical. For a typical 3-D tensor:
Shape
Input x(B, C, E)
rms (after keepdims reshape)(B, C, 1)
Output y(B, C, E)

​
Why RMSNorm over LayerNorm?

RMSNorm skips the mean-subtraction step, computing only the root-mean-square of the activations instead of the full mean and variance. This reduces the number of operations per forward pass while empirically matching or exceeding LayerNorm quality in large language models. It is the normalization used in LLaMA, Mistral, and related architectures.

​
LayerNorm

func LayerNorm(embedDim int) *LayerNormT
Constructs a standard Layer Normalization layer. Available in the layer package but not used by Block by default — Block uses RMSNorm for both normalization positions.
​
embedDim
int
required
The size of the last dimension of the input tensor. Determines the shapes of the learnable gamma and beta parameters.

​
Parameters

ParameterShapeInitial value
gamma(embedDim,)all ones
beta(embedDim,)all zeros
Epsilon is hardcoded to 1e-5.

​
Formula

normx = (x - mean(x)) / sqrt(var(x) + ε)
y     = gamma * normx + beta

​
Forward computation

mean     := F.Reshape(shape...)(F.Mean(last)(x[0]))               // keepdims
variance := F.Reshape(shape...)(F.Variance(last)(x[0], mean))     // keepdims

sub   := F.Sub(x[0], mean)
addc  := F.AddC(l.eps, variance)
sqrt  := F.Pow(0.5)(addc)
normx := F.Div(sub, sqrt)

y := F.Add(F.Mul(gamma, normx), beta)
Both mean and variance are reshaped with a keepdims-equivalent operation so they broadcast correctly over the feature dimension.

​
Shape

Shape
Input x(B, C, E)
mean, variance (after keepdims)(B, C, 1)
Output y(B, C, E)
LayerNorm is available in the layer package but is not used by Block by default — Block uses RMSNorm for both normalization positions. New models trained with Block will not produce LayerNorm parameters.

Build docs developers (and LLMs) love

Get started for freeTalk to us