Notes
🔍

Pre-LN vs Post-LN

A transformer block repeatedly updates a running representation called the residual stream. Pre-LN and Post-LN are two ways to place Layer Normalization around each attention or feed-forward sublayer. The placement looks like a small wiring choice, but it changes how easily gradients travel through a deep network.

The two layouts
For a sublayer F, such as multi-head attention, the common forms are:

  • Post-LN: y = LayerNorm(x + F(x)). The sublayer updates x, the residual connection adds the original input back, and normalization happens afterward.
  • Pre-LN: y = x + F(LayerNorm(x)). The input is normalized before entering the sublayer, while the residual path itself remains an unmodified addition. A final LayerNorm is usually applied at the end of the full stack.

Why training behaves differently
In a Pre-LN transformer, gradients have a clean identity route through each residual addition: they can pass backward through the x term without repeatedly passing through LayerNorm. This makes very deep stacks substantially easier to optimize and reduces dependence on delicate initialization and long learning-rate warmup. Modern decoder-style transformers therefore commonly use Pre-LN. Post-LN, used in the original Transformer, forces the backward signal through normalization after every residual update. It can train well, but deep versions are more sensitive: an aggressive learning rate can make the loss spike or diverge in the first epochs. Post-LN has also been associated with stronger final performance in some carefully stabilized settings, so it has not disappeared.

Practical consequences
A PyTorch-style Pre-LN block typically resembles x = x + attention(norm1(x)), then x = x + mlp(norm2(x)). If a deep model stalls, becomes unstable early, or needs unusually cautious warmup, switching a Post-LN design to Pre-LN gives optimization a more reliable path. The trade-off is that Pre-LN can allow residual activations to accumulate across depth, which is why the final normalization and sensible initialization remain important. Both designs normalize each token’s feature vector; neither changes the model’s basic attention or feed-forward computation.

Pre-LN and Post-LN describe whether layer normalization is applied before or after a transformer sublayer and its residual addition. Pre-LN normalizes the residual stream before attention or feed-forward computation; Post-LN normalizes after adding the sublayer output. Pre-LN provides more direct gradient paths through residual connections, making deep transformers more stable to optimize, while Post-LN can require stricter initialization and learning-rate control.

Imagine an assembly line where each worker either tidies an item before working on it or afterward. Pre-LN and Post-LN describe those two choices in a transformer.

Layer normalization is a kind of “scale-setting”: it keeps the information flowing through the network from becoming wildly uneven. In Pre-LN, the tidying happens before a block does its main work. In Post-LN, it happens afterward. Both aim to keep learning steady, but Pre-LN is often easier to train in very deep transformers. The choice can affect how reliably a large AI model learns from examples.