Layer Normalization

A common assumption among learners is that a transformer block is simply attention followed by a feed-forward network, with normalization placed wherever it happens to fit. Under that assumption, the order of the operations looks like a detail that can be ignored. The code for a transformer block shows otherwise. The same two components, attention and feed-forward, can be wrapped by normalization in two different ways, and the implementation commits to one of them.

This article corrects that belief in four moves. First you will predict which operation runs first on the input. Then you will see how the block is constructed from a configuration dictionary. Next you will follow the forward pass step by step, including the shortcut connections that add the block input back into its output. Finally you will examine the difference between Pre-LayerNorm and Post-LayerNorm, which is the point where the wrong belief does real damage.

The TransformerBlock class is a PyTorch module that inherits from nn.Module. Its constructor does not hard-code sizes. Instead it reads everything from a configuration dictionary called cfg, which is passed in once when the block is created. The source describes the block as including a multi-head attention mechanism and a feed-forward network, both configured from that dictionary, with the GPT_CONFIG_124M dictionary given as an example of such a configuration.

Inside the constructor, the block creates five components, and each one is stored as an attribute of the block. The attention module, stored as self.att, is a MultiHeadAttention built from the embedding dimension, the context length, the number of heads, the dropout rate, and a flag for query, key, and value bias. The feed-forward module, stored as self.ff, is a FeedForward network built directly from cfg. Two LayerNorm modules, norm1 and norm2, are both sized to the embedding dimension. Finally, a dropout module called drop_shortcut is created from the dropout rate.

The constructor only builds these parts. It does not decide the order in which they run on the data. That job belongs to the forward method, which is why the constructor alone cannot tell you how the block behaves. Two blocks with the same five components could in principle wire them differently, so the forward method is the place to look when you need to know what the block actually does.

The forward method is where the block does its work, and it is where most debugging questions are answered. Reading it in order gives a fixed sequence of actions. The block first saves the input as a shortcut, then normalizes x, runs attention, applies dropout, and adds the saved input back. The same pattern then repeats for the feed-forward half: save the current value as a new shortcut, normalize it, run the feed-forward network, apply dropout, and add the shortcut back.

Read the diagram as two passes over the same data. The first pass takes the block input and produces the attention contribution, which is then added back to that input. The second pass takes that new value and produces the feed-forward contribution, which is again added back to what it started from. The output of the block is therefore the result of the second addition, not the raw output of either component.

The source states that this feature helps gradients flow through the network during training and improves the learning of deep models. It gives a specific reason for the design. A block that stacks attention and feed-forward components with additions back to the input keeps the original signal available alongside the processed signal. The processed path transforms the data, and the shortcut path carries the data through unchanged, so the two are combined at each addition.

Pre-LayerNorm and Post-LayerNorm are not two names for the same arrangement. The source defines Pre-LayerNorm as applying layer normalization before each of the two components, which is what this block does. Older architectures, including the original transformer model, applied layer normalization after the self-attention and feed-forward networks, which is called Post-LayerNorm.

The before-and-after view isolates the single difference that matters for the attention half. The same two operations appear in both arrangements, but their order is reversed. The block in this article is the Pre-LayerNorm form, because norm1 runs before the attention module and norm2 runs before the feed-forward module. A reader who assumes the two designs are interchangeable will misread every line that involves a norm.

Suppose a block produces an output that does not match what a colleague expects. The fastest way to find the cause is to ask a question about each step in the forward pass: is the value at this point what I expect it to be? Because the block has a fixed sequence of steps, the first step where the expected and actual values disagree marks the place where the divergence starts.

This is the same reasoning used when stepping through the execution trace for this lesson. The method works because the forward method has a fixed order. You do not need to guess where the block went wrong. You check each stage in sequence and stop at the first stage that disagrees with what you expect.

A common intuition in machine learning is that a bigger network should be able to learn more. Adding layers gives a network more room to represent complicated relationships. In practice, however, stacking many layers together does not simply make learning easier. Training a deep network can become unreliable, and the network may fail to learn the patterns that would let it make accurate predictions.

Training a neural network is a search. The network has a set of parameters, called weights, and the learning process looks for a particular set of weights that minimizes the loss function. The loss function scores how far the network's outputs are from what they should be. A lower loss means the network is doing a better job on the training data.

The adjustments to the weights during this search are guided by gradients. A gradient carries information about how the loss would change if a weight were nudged slightly. For this article you only need the working idea: gradients are the signals that tell the learning process which way to move the weights. Deep mathematical understanding of gradients is not required to follow the material.

The source names two specific problems that cause this difficulty: vanishing gradients and exploding gradients. Before reading about them, try to reason about the core issue. Imagine a signal that is multiplied by a number, passed through a layer, and then multiplied again by the same kind of number at the next layer. The question is what happens to the size of that signal after many such steps.

The source identifies vanishing or exploding gradients as the problem that can make training deep networks challenging. These two failure modes are named together because they share a cause and a consequence. Both are problems with the size of the gradients that flow through the network during training. Both lead to unstable training dynamics.

The diagram uses an illustrative factor of 0.5 at each layer. It is a simplified picture of the compounding idea, not a measurement of a real network. Each backward step multiplies the gradient by a number, so the gradient reaching earlier layers is the product of many such numbers.

The formula shows the same compounding in both directions. Over ten steps, a factor of 2 grows a signal to about 1024, which is the exploding pattern. A factor of 0.5 shrinks it to roughly one thousandth of its starting size, which is the vanishing pattern. Real networks do not multiply by one fixed number, but the compounding logic is the same, and it explains why depth amplifies whatever tendency a gradient has to shrink or grow.

The source describes the effects of these gradient problems in three connected ways. First, they lead to unstable training dynamics, meaning the training process does not progress in a steady, predictable manner. Second, they make it difficult for the network to effectively adjust its weights. Third, as a result, the learning process struggles to find a set of weights that minimizes the loss function.

The consequence the source emphasizes is a failure to learn the underlying patterns in the data. The network may then be unable to make accurate predictions or decisions. This is a failure of learning itself, not merely a slower version of learning. A network with unstable gradients can end up with weights that do not capture the structure the task requires.

Putting the pieces together gives the motivation for layer normalization. Deep networks can suffer from unstable training, and the gradients that guide weight updates can shrink or grow as they pass through layers. Any remedy has to address that instability so that the network can again adjust its weights and find parameters that minimize the loss. The need is a need for stability in the training signal, and that is the problem layer normalization is introduced to address.