Layer Normalization › PyTorch's LayerNorm

Michael BrenndoerferUpdated June 11, 20252 min read
Language AI HandbookAll chapters
Part 15 of 67Transformer Blocks
Chapter 104Layer Normalization › PyTorch's LayerNorm
Reading section 6 of 16 · PyTorch's LayerNorm · 2 minSection 6 of 16 · 2 min

Book contents

Section 6 of 16 · 2 min read

Layer Normalization › PyTorch's LayerNormOne section · Sections in this chapter6 / 16 ↓
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

PyTorch provides a built-in nn.LayerNorm that handles all these details efficiently. Let's verify our implementation matches PyTorch's behavior.

In[15]:
Code
import torch.nn as nn

# Create PyTorch LayerNorm
pytorch_ln = nn.LayerNorm(hidden_dim, eps=1e-5)

# Initialize with same parameters
with torch.no_grad():
    pytorch_ln.weight.fill_(1.0)  # gamma
    pytorch_ln.bias.fill_(0.0)  # beta

# Compare outputs
pytorch_output = pytorch_ln(x)
our_output, _ = layer_norm_forward(x, gamma, beta)
Out[16]:
Console
Maximum difference between implementations: 1.19e-07
Outputs match: True

The outputs match within floating-point precision, confirming our implementation is correct.

PyTorch exposes the learnable parameters through the weight and bias attributes. weight corresponds to γ\gamma and bias corresponds to β\beta. This naming can be confusing because "weight" typically refers to the parameters of a linear transformation, but here it refers to the per-feature scale. The reason for this naming is historical: when PyTorch first implemented layer normalization, it followed the convention established by batch normalization, where the parameters are called weight and bias to indicate that they scale and shift the normalized values.

When you inspect a trained transformer model and look at the layer normalization parameters, you'll often find that the weight parameters are very close to 1 and the bias parameters are very close to 0, especially in the earlier layers. This suggests that those layers don't need to deviate much from the standard normalized distribution. In later layers, or in attention-specific normalizations, the parameters may drift further from their initialization, indicating that those layers benefit from a different distribution.