Normalize a training batch

Learn it

Words to know

batch

A group of examples that the model processes together in one step.

normalization

Rescaling a group of numbers to a standard size, which keeps training stable.

variance

A measure of spread: the average of the squared distances from the mean.

How it works

1. For one column, find the mean of its values over all rows.
2. Find the variance of that column, dividing by the number of rows N.
3. Subtract the column's mean from each of its values and divide by sqrt(variance + epsilon).
4. In PyTorch this is the layer normalization code with one change: dim=0 in place of dim=-1. Dimension 0 runs down the rows.

Worked example

Input
rows = [[1, 10], [2, 20], [3, 60]]
epsilon = 1e-8

Output
[[-1.22474, -0.92582], [0, -0.46291], [1.22474, 1.38873]]

The second column is far larger than the first, but after normalizing both are on the same scale. Each column now averages 0.

Why it matters

Batch normalization is common in image models. Transformers prefer layer normalization, because it gives the same result for an example whatever else is in the batch.

Your turn

Using PyTorch, normalize each column of a table across its rows, as batch normalization does during training.

rows

a table of numbers with at least one row; each row is one example

epsilon

a positive number added to the variance; tiny in practice

Returns

A table of the same shape, as nested Python lists. Each column has its own mean subtracted and is divided by sqrt(variance + epsilon). For the variance, divide the sum of the squared distances by the number of rows.

This lesson needs PyTorch, which does not run in the browser. The app gives you the exercise as a file with its tests, to run on your own computer.

Input
rows = [[1, 10], [2, 20], [3, 60]]
epsilon = 1e-8

Output
[[-1.22474, -0.92582], [0, -0.46291], [1.22474, 1.38873]]

Input
rows = [[1, 2], [3, 4]]
epsilon = 1e-8

Output
[[-1, -1], [1, 1]]

Input
rows = [[1, 2], [3, 4]]
epsilon = 3

Output
[[-0.5, -0.5], [0.5, 0.5]]

Input
rows = [[5, 5]]
epsilon = 0.00001

Output
[[0, 0]]

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Where to start

   This is the layer normalization exercise turned on its side. The mean and the variance now come from columns.

2. Hint 2: The key step

   Take the mean and the variance along dimension 0, again with keepdim=True and unbiased=False.

3. Hint 3: The outline

   Make x a float64 tensor. Set mean = x.mean(0, keepdim=True) and variance = x.var(0, unbiased=False, keepdim=True). Return ((x - mean) / torch.sqrt(variance + epsilon)).tolist().

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

1. One mean and one variance per column, taken down the rows. Dimension 0 is the one that collapses. What is left has shape [1, columns]: one number per feature.
2. Centre each column on 0 and scale its spread to 1. The [1, columns] means and variances line up under the columns, so each value is normalized with the numbers of its own feature.

The same numbered steps are the comments in the solution.

import torch


def solve(rows, epsilon):
    x = torch.tensor(rows, dtype=torch.float64)
    # 1. One mean and one variance per column, taken down the rows
    mean = x.mean(0, keepdim=True)
    variance = x.var(0, unbiased=False, keepdim=True)
    # 2. Centre each column on 0 and scale its spread to 1
    normalized = (x - mean) / torch.sqrt(variance + epsilon)
    return normalized.tolist()

Watch out

• Mixing up the dimensions gives layer normalization instead. dim=0 goes down the rows, across examples; dim=-1 goes along a row, across features.
• With a batch of one row every result is 0, because each value is its own mean. With a small batch the mean and the variance say little about the feature.
• This is the training behaviour. When the model is used later, PyTorch's batch normalization takes the mean and variance from averages it saved during training, not from the batch.

Check yourself