Inside Transformers, Part 5: Layer normalization and residual connections

The attention blocks in Parts 3 and 4 do not stand alone. Each sub-layer sits inside an Add and Norm block that combines the incoming representation with the sub-layer output and then normalizes the result.

Add and norm

For an input x and a sub-layer Sublayer, the original Transformer uses the post-norm form

y = LayerNorm(x + Sublayer(x))

The addition is the residual connection. It gives the next block both the representation that arrived and the update computed by the attention or feed-forward sub-layer. The residual path can make optimization easier, but it is not a guarantee that gradients never vanish.

Layer normalization

Layer normalization normalizes the configured feature dimensions for each example independently. For a Transformer tensor shaped (batch, sequence, d_model) with normalized_shape=d_model, it computes statistics across the d_model features of each token. It does not combine statistics from different tokens or other examples in the batch.

For one token vector x with d_model features, the operation is

mu    = (1 / d_model) * sum_r x_r
var   = (1 / d_model) * sum_r (x_r - mu)^2
LN(x) = gamma * (x - mu) / sqrt(var + epsilon) + beta

gamma and beta are learned feature-wise scale and shift parameters. epsilon prevents division by zero. The same kind of per-token computation is used during training and evaluation because it does not need batch statistics.

BatchNorm uses a different aggregation pattern. It normally computes statistics for each feature channel over the examples in a mini-batch and, depending on the input shape, over spatial or time positions. The distinction is about which axes supply the statistics. It is not accurate to reduce it to “first dimension versus last dimension” for every tensor layout.

The original paper applies post-norm after each attention and feed-forward sub-layer. Later Transformer implementations often use pre-norm instead:

y = x + Sublayer(LayerNorm(x))

These are different layer arrangements. Neither equation alone guarantees a particular training speed or final performance; the effect depends on the architecture, initialization, optimizer, and task.