Layer normalization standardizes an input’s features to have approximately zero mean and unit standard deviation. In a Transformer, it sits alongside residual connections in attention and feed-forward sublayers, helping stabilize training. It is similar to batch normalization, but the implementation here calculates its mean and standard deviation over the input’s last dimension. section 9.2.2 section 9.1.2
A small PyTorch implementation makes the steps concrete:
- Compute the mean and standard deviation with `x.mean(-1, keepdim=True)` and `x.std(-1, keepdim=True)`. section 9.2.2
- Subtract the mean and divide by `sqrt(std ** 2 + eps)`; `eps` is a small value included in the denominator. section 9.2.2
- Apply learned parameters `a_2` and `b_2` to scale and shift the standardized values back into the input’s shape. section 9.2.2
Inside a Transformer block
An encoder layer normalizes `x` before passing it to multihead self-attention, adds the original `x` back through a residual connection, and repeats the pattern for its feed-forward network. A GPT-2XL decoder block uses `nn.LayerNorm(config.n_embd)` before causal self-attention and another before its feed-forward network. It also normalizes the stacked blocks’ output before the vocabulary-sized linear head. section 9.2.2 section 11.2.4
class LayerNorm(nn.Module):
def __init__(self, features, eps=1e-6):
super().__init__()
self.a_2 = nn.Parameter(torch.ones(features))
self.b_2 = nn.Parameter(torch.zeros(features))
self.eps = eps
def forward(self, x):
mean = x.mean(-1, keepdim=True)
std = x.std(-1, keepdim=True)
x_zscore = (x - mean) / torch.sqrt(std ** 2 + self.eps)
output = self.a_2*x_zscore+self.b_2
return output
The `LayerNorm` class shows the feature-axis statistics, epsilon-adjusted denominator, and learned scale and shift in one forward pass. section 9.2.2
The same placement carries over to a music Transformer: each of its six decoder blocks pairs normalization and residual connections with causal self-attention and a feed-forward network, followed by a final normalization before predicting music tokens. section 14.3.2
