Layer normalization details in GPT-2
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 &quot;On Layer Normalization in the Transformer Architecture&quot; goes into great detail about the topic. The paper proposes &quot;the layer normalization plays a crucial role in controlling the gradient scales.&quot; Better behaved gradients help with training.