Layer normalization details in GPT-2

Question

I've read that GPT-2 and other transformers use layer normalization before the self-attention and feedforward blocks, but I am still unsure exactly how the normalization works.

Let's say that our context size is 1024 tokens, the embedding size is 768 (so that each token and its subsequent hidden states are represented by vectors of size 768), and we are using 12 multi-attention heads. So in the diagram above, there are 1024 r's and each r has dimensionality 768.

For a given layer in the transformer, how many normalization statistics (sample mean and stdev) are computed? Do we do one normalization per token, for 12x1024 normalizations so that the feature values within each token have mean 0 and std 1? Or do we normalize the values for each feature across tokens, for 12x768 normalizations? Or do we normalize all the feature values for all the tokens together, for 12 normalizations? Do we compute separate normalizations for each context in the minibatch?

I'm also keen to understand intuitively why this normalization is desirable. Assuming that the scheme is to normalize the feature values within each token: let's say one of our tokens is a bland word like "ok" whereas another token is the word "hatred". I would expect that the representation of "hatred" would be spikier, with higher variance among the different feature values. Why is it useful to throw away this information and force the representation for "ok" to be just as spiky? On the other hand, if the normalization scheme is to normalize across feature values, so that if you take feature 1 from all of the tokens in our context, they will have zero mean and stdev 1, doesn't this throw away information when all of the words in our context are very negative, for example in the context "war violence hate fear"?

Separately, with layer normalization it seems like it is optional to re-scale the normalized values through learned bias and gain parameters. Does GPT-2 do this, or does it keep the values normalized to mean 0 and std 1?

Answer 1

The most standard implementation uses PyTorch's LayerNorm which applies Layer Normalization over a mini-batch of inputs. The mean and standard-deviation are calculated separately over the last certain number dimensions which have to be of the shape specified by normalized_shape argument. Most often normalized_shape is the token embedding size.

The paper "On Layer Normalization in the Transformer Architecture" goes into great detail about the topic. The paper proposes "the layer normalization plays a crucial role in controlling the gradient scales." Better behaved gradients help with training.

Answer 2

As per the reference, Layer Normalization is applied 2 times per block (or layer). Once for the hidden states from the output of the attention layer, and once for the hidden states for the output from the feed-forward layer. However, it is

(For hugging-face implementation, you can check out class Block here)

For your example, normalization necessarily doesn't really does not change the order of spikes of various tokens, but changes the magnitude of the spikes. This has shown to increase the model training time, and improve the performance. (paper)

Answer 3

Recently I'm working on NVIDIA's FastTransformer. At least in their implementation of layer norm, normalizations are applied to every input token (i.e. its subsequent hidden states) separately. For your example, it will calculate 12x1024 means and standard deviations.

I think, they choose to implement layer normalization in this way because that, if you normalize all the feature values for all the tokens together, the output tokens for previous input tokens will be affected by succeeding tokens, which needs recalculation in each "feed-forward" iteration. If we apply normalization to every token separately, this won't happen, so we can just store keys & values for previous tokens (NVIDIA calls them k/v cache) instead of running the complete multi-head attention again and again.

Answer 4

Layer normalization (LN) and its implementation is one of the more confusing aspects of transformers.

First, the goal of normalization is to stabilize the gradient descent during training. Specifically, normalization gives the model more control over the distribution of data. This article explains it in the context of batch normalization (BN), but the same principles apply to LN. Note, BN cannot be applied effectively to text because number of tokens vary from one input to the next in a batch, so the need for LN.

In the context of transformers, LN operates on one input text (sequence) at a time. It involves computing one mean and std for each token, i.e., there will be 1024 mean and std values OP's case. Then, typically, each feature/dimension of each token is scaled and shifted using the token's mean, std, and 2 learned parameters called bias and gain. There is one bias and gain value per dimension across all tokens, i.e., input tokens share the learned parameters. See implementations of LN in BERT and GPT-2 as examples. See also Section 5.1 in the original LN paper where it states "They also learn an adaptive bias b and gain g for each neuron after the normalization." Thus, in OP's case, there will be 768 learned bias and gain parameters for each LN layer.

Note, none of the normalization schemes in OP's question align with how normalization is actually done. Normalization statistics are computed per token, but bias and gain are applied across tokens. Additionally, normalization does not necessarily make token embeddings "spiky" or "smooth." Normalization lets the model choose the data distribution.