Building a PyTorch module for layer normalization
Layer normalization is a technique for this purpose.
Why Normalize Activations
Neural network training works best when the values flowing through the network stay in a predictable range. Layer normalization is a technique for this purpose. Its main idea is to adjust the activations, which are the outputs of a neural network layer, so that they have a mean of 0 and a variance of 1. The source describes this adjustment as speeding up the convergence to effective weights and as supporting consistent, reliable training.
In GPT-2 and modern transformer architectures, layer normalization is typically applied before and after the multi-head attention module, and before the final output layer. This article follows the source's approach: first we build a small layer that produces outputs, then we measure those outputs, and then we normalize them. Each of these three stages is something you can run and inspect yourself, and the rest of the article traces the values through them one stage at a time.
Recreating the Layer Outputs
We start by recreating a layer that takes two input examples, each with five features, and produces six outputs per example. The layer is a Linear layer followed by a ReLU activation. Every value we normalize later comes from this output, so it is worth running the code first and looking at the printed tensor before doing anything else.
torch.randn(2, 5)creates two rows of five features each; the random values are illustrativenn.Linear(5, 6)maps five features to six outputs per examplenn.ReLU()is applied after the Linear layer, so it acts on the Linear outputprint(out)shows the six layer outputs for each of the two inputs
The printed tensor has two rows. The first row lists the six layer outputs for the first input, and the second row lists the six layer outputs for the second input. Notice the zeros scattered through both rows. They come from the activation function, which is the next topic.
How ReLU Shapes the Inputs
The layer we coded consists of a Linear layer followed by a nonlinear activation function called ReLU, short for rectified linear unit. ReLU is a standard activation function in neural networks. It thresholds negative inputs to 0, so a layer that uses it outputs only positive values or zero. This is why the printed layer output contains no negative values.
The diagram shows the order in which the values are produced. ReLU is the step that removes negatives, so by the time layer normalization runs, its input already contains only zeros and positive values. The source notes that a more sophisticated activation function is used later in GPT. For this article, the important point is only that the layer outputs we normalize contain zeros and positive values and no negatives, because ReLU has already been applied to them.
Measuring Mean and Variance per Row
Before we normalize, we need to measure the layer outputs. The mean of a tensor is the average of its elements along a specified dimension. The variance is the average of the squared differences from the mean. For our two-row output, the dimension we want is the feature dimension, which is dim=-1 (the last dimension). Taking the statistics along dim=-1 means each row is summarized by its own numbers.
out.mean(dim=-1, keepdim=True)averages across the six features in each row- keepdim=True keeps the reduced dimension as size 1, so the result has shape (2, 1) instead of (2,)
- ((out - mean) ** 2).mean(...) computes the average squared difference from the mean, which matches the source's definition of variance
- Printing the shape confirms there is one mean value per row
The first row of the mean tensor holds the mean for the first input row, and the second row holds the mean for the second input row. Each input row therefore gets its own mean and its own variance. This per-row behavior is the key to what layer normalization does.
Each statistic sits in the row it summarizes. Because keepdim=True keeps the reduced axis as size 1, the means form a column with shape (2, 1). When that column is combined with the (2, 6) output, each row's single mean lines up with that row's six features.
What do you think happens?
If keepdim=True is removed, what shape does the mean have for a (2, 6) tensor reduced along dim=-1?
Reveal answer
(2,)
Without keepdim=True, the reduced dimension is removed, leaving one value per row with shape (2,). That shape does not line up with the (2, 6) output for subtraction in the same way, which is why the source keeps the dimension with keepdim=True.
Applying the Two Normalization Steps
Now we apply layer normalization to the layer outputs. The operation has two steps: subtract the mean, then divide by the square root of the variance, which is also known as the standard deviation. Written as a single expression, with x standing for a layer output, mu for its row mean, and sigma squared for its row variance, the operation is shown below.
- out - mean subtracts each row's single mean from all six of that row's values
torch.sqrt(var)takes the square root of each row's variance, giving the standard deviation- Dividing applies the same row-wise statistic across every feature in that row through broadcasting
out # ReLU outputs: zeros and positive values only
No negative values, mean and spread set by the layer
norm = (out - mean) / torch.sqrt(var)Values rescaled and shifted; mean near 0, variance 1
The subtraction and division each use one number per row, but they apply to all six features in that row. This is the broadcasting that keepdim=True makes possible: the (2, 1) mean and standard deviation stretch across the (2, 6) outputs. After this step, the normalized outputs contain negative values, which the ReLU outputs did not.
Checking the Result
The check is the step that turns a formula into evidence. Measure the normalized tensor the same way you measured the original, row by row, and confirm that the mean is extremely close to 0 and the variance is 1.
- The same dim=-1, keepdim=True pattern from the earlier measurement is reused
- The normalized means should print as values very close to 0
- The normalized variances should print as values very close to 1
Change the layer so it has eight outputs instead of six, rerun the code, and confirm that the mean and variance of each normalized row are still close to 0 and 1. Then explain which dimension you measured along and why the statistic shape stays (2, 1).
Hints
- Only the second argument of nn.Linear changes; the input features stay at five
- The statistics use dim=-1, which refers to the feature dimension, so the number of outputs per row does not change the shape of the statistics
- Print
norm_mean.shape to confirm the shape is (2, 1)
Mistakes Beginners Make
Common mistake Taking the mean over all elements instead of over each row
What it looks likeout.mean() returns one number for the whole tensor.
Why it failsThe source's point is that each input row gets its own mean and variance. A single global mean cannot normalize each row independently.
Do this insteadUse out.mean(dim=-1, keepdim=True) so that each row gets its own statistic.
Common mistake Dropping keepdim=True and expecting the statistic to broadcast the same way
What it looks likemean = out.mean(dim=-1) gives shape (2,) instead of (2, 1).Why it failsWithout keepdim=True the reduced dimension disappears, so the statistic no longer lines up with the rows of the output in the way the normalization step expects.
Do this insteadKeep keepdim=True, and print the shape of the statistic to confirm it is (batch, 1).
Common mistake Expecting the normalized output to be non-negative because the layer output was
What it looks likeAfter ReLU every value is 0 or positive, but after normalization the printed tensor contains negatives.
Why it failsReLU only guarantees non-negative values before normalization. Subtracting the mean introduces negative values.
Do this insteadExpect negatives after normalization, and check the mean and variance rather than the sign of the values.
Common mistake Expecting the mean to print as exactly 0
What it looks likeA result of -5.9605e-08 looks like a bug.
Why it failsThe source attributes this to finite numerical precision, so a very small nonzero value is expected.
Do this insteadTreat values near 0 as correct, and use sci_mode=False if the scientific notation is hard to read.
Common mistake Checking the schematic numbers instead of the printed tensor when verifying variance
Key Takeaways
- Layer normalization adjusts a layer's activations so each row has a mean of 0 and a variance of 1, which the source says speeds convergence and supports reliable training.
- The layer outputs being normalized come from a Linear layer followed by ReLU, so they contain zeros and positive values but no negatives before normalization.
- Mean and variance are computed along dim=-1 with keepdim=True, giving each row its own statistics with shape (batch, 1) that broadcast across the row's features.
- Normalization subtracts the row mean and divides by the square root of the row variance; check the result by measuring the normalized tensor, and expect a mean very close to 0 rather than exactly 0.