Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Layer Normalization › Implementing Layer Normalization from Scratch

Michael BrenndoerferUpdated June 11, 20257 min read

Language AI HandbookAll chapters

Part 15 of 67Transformer Blocks

Chapter 104Layer Normalization › Implementing Layer Normalization from Scratch

Reading section 4 of 16 · Implementing Layer Normalization from Scratch · 7 minSection 4 of 16 · 7 min

Section 4 of 16 · 7 min read

Layer Normalization › Implementing Layer Normalization from ScratchOne section · Sections in this chapter4 / 16 ↓

Sections in this chapter16

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Now that we understand the formula conceptually, let's translate it into code. Building layer normalization from scratch will solidify our understanding and reveal the implementation details that matter in practice.

The Forward PassLink Copied

Our implementation follows the mathematical steps exactly: compute mean, compute variance, normalize, then apply the learnable transformation.

In[6]:

Code

def layer_norm_forward(x, gamma, beta, eps=1e-5):
    """
    Layer normalization forward pass.

    Args:
        x: Input tensor of shape (batch, seq_len, hidden_dim)
        gamma: Scale parameters of shape (hidden_dim,)
        beta: Shift parameters of shape (hidden_dim,)
        eps: Small constant for numerical stability

    Returns:
        Normalized output with same shape as input
    """
    # Step 1: Compute mean across the feature dimension (last axis)
    mu = x.mean(dim=-1, keepdim=True)

    # Step 2: Compute variance across the feature dimension
    var = x.var(dim=-1, keepdim=True, unbiased=False)

    # Step 3: Normalize to zero mean and unit variance
    x_norm = (x - mu) / torch.sqrt(var + eps)

    # Step 4: Apply learnable affine transformation
    out = gamma * x_norm + beta

    return out, (x, x_norm, mu, var, gamma, eps)

The dim=-1 argument tells PyTorch to compute statistics across the last dimension (features), which is exactly what layer normalization requires. The keepdim=True preserves the dimension for broadcasting during subtraction and division. Without keepdim=True, the mean would have shape (batch, seq_len) and could not be subtracted from x with shape (batch, seq_len, hidden_dim) without an explicit unsqueeze operation.

Notice also that we cache the intermediate values (x, x_norm, mu, var, gamma, eps) in the return value. This cache is needed for the backward pass: computing gradients through layer normalization requires the normalized values and statistics from the forward pass. Modern frameworks like PyTorch compute this automatically through autograd, but understanding the forward-backward relationship helps you reason about memory usage in large models.

A Worked ExampleLink Copied

Let's trace through layer normalization with concrete numbers to see exactly what happens at each step.

In[7]:

Code

# Create a simple example: 2 samples, 4 tokens each, 8 features per token

batch_size = 2
seq_len = 4
hidden_dim = 8

# Generate input with non-standard distribution (mean ~2, varied std)
x = torch.randn(batch_size, seq_len, hidden_dim) * 3 + 2

# Initialize gamma=1 and beta=0 (identity affine transform)
gamma = torch.ones(hidden_dim)
beta = torch.zeros(hidden_dim)

# Apply layer normalization
output, cache = layer_norm_forward(x, gamma, beta)

Out[8]:

Console

Input statistics (per token):
  Shape: torch.Size([2, 4, 8])
  First token mean: 1.410, std: 2.718
  Second token mean: 0.350, std: 2.670

Output statistics (per token):
  First token mean: 0.000000, std: 1.069044
  Second token mean: 0.000000, std: 1.069044

The input tokens have varying means (around 2-3) and standard deviations (around 2-4). This reflects the non-standard distribution we created. After layer normalization, each token has mean neededly zero and standard deviation neededly one. The tiny deviations from exactly 0 and 1 are floating-point precision artifacts, not algorithmic issues.

Let's visualize this transformation to see exactly how layer normalization reshapes the feature distribution.

Out[9]:

Visualization

Histogram showing feature values with positive mean around 2-3. — Feature distribution before layer normalization. The 8 features of a single token show varied values with a non-zero mean (dashed line) and non-unit variance. Each feature contributes to the overall distribution.

Histogram showing feature values centered at zero with unit spread. — Feature distribution after layer normalization. The same token's features are now centered at zero with unit variance. The transformation standardizes each token independently.

The histograms make the transformation crystal clear. Before normalization, the feature values are scattered around a positive mean with varied spread. After normalization, they're centered at zero with approximately unit variance. This happens independently for every token in the sequence.

Notice that each token is normalized independently: the first token's statistics don't affect the second token's normalization. This independence is precisely what makes layer normalization suitable for transformers, where tokens must be processed in parallel and sequences have variable lengths.

Numeric Walk-ThroughLink Copied

To build deep intuition, let's work through a minimal example by hand before letting the code do it. Consider a single token with just four features:

\mathbf{x} = [3.0, \; 1.0, \; 5.0, \; 7.0]

Step 1: Compute the mean.

\mu = \frac{3.0 + 1.0 + 5.0 + 7.0}{4} = \frac{16.0}{4} = 4.0

Step 2: Compute the variance.

\sigma^2 = \frac{(3.0 - 4.0)^2 + (1.0 - 4.0)^2 + (5.0 - 4.0)^2 + (7.0 - 4.0)^2}{4}

\sigma^2 = \frac{1.0 + 9.0 + 1.0 + 9.0}{4} = \frac{20.0}{4} = 5.0

So the standard deviation is $\sigma = \sqrt{5.0} \approx 2.236$ .

Step 3: Normalize each feature. Using $\epsilon = 10^{-5} \approx 0$ for simplicity:

\hat{x}_1 = \frac{3.0 - 4.0}{2.236} \approx -0.447

\hat{x}_2 = \frac{1.0 - 4.0}{2.236} \approx -1.342

\hat{x}_3 = \frac{5.0 - 4.0}{2.236} \approx +0.447

\hat{x}_4 = \frac{7.0 - 4.0}{2.236} \approx +1.342

Verify: The mean of $[-0.447, -1.342, 0.447, 1.342]$ is zero. The variance is $\frac{0.200 + 1.800 + 0.200 + 1.800}{4} = 1.0$ . The normalization is correct.

Step 4: Apply learned parameters. If $\gamma = [1, 1, 1, 1]$ and $\beta = [0, 0, 0, 0]$ , the output equals the normalized values. If instead $\gamma = [2, 2, 2, 2]$ and $\beta = [1, 1, 1, 1]$ , the output would be $[-0.894 + 1, -2.684 + 1, 0.894 + 1, 2.684 + 1] = [0.106, -1.684, 1.894, 3.684]$ , a distribution with mean 1 and standard deviation 2. The network can express this transformation freely through the learned parameters.

This walk-through reveals something important: layer normalization is not a mere computational trick; it is a principled recentering and rescaling operation that the network can tune at training time. The initialization puts every token in a standard reference frame; the learned parameters then shift it to wherever the subsequent layer finds most useful.

In[10]:

Code

# Verify the hand-calculated example
x_manual = torch.tensor([[3.0, 1.0, 5.0, 7.0]])  # shape (1, 4)
gamma_manual = torch.ones(4)
beta_manual = torch.zeros(4)

mu_manual = x_manual.mean(dim=-1, keepdim=True)
var_manual = x_manual.var(dim=-1, keepdim=True, unbiased=False)
x_norm_manual = (x_manual - mu_manual) / torch.sqrt(var_manual + 1e-5)

Out[11]:

Console

Manual walk-through verification:
  Input:    [3.0, 1.0, 5.0, 7.0]
  Mean:     4.0000
  Variance: 5.0000
  Std dev:  2.2361
  Normalized: [-0.4472, -1.3416, 0.4472, 1.3416]
  Normalized mean: 0.00000000
  Normalized std:  0.999999

The computed values match the hand calculation exactly (up to floating-point precision from the epsilon term). Each normalized value has the expected sign and magnitude, and the normalized vector has mean zero and variance one.