Normalize by root mean square

Learn it

Words to know

normalization

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

root mean square

A measure of size: square the numbers, take their mean, then take the square root.

epsilon

A tiny number added before dividing or taking a root, so that the calculation is safe when the other part is 0.

How it works

1. Square every value in the row. Squares are never negative, so positive and negative values both count toward the size.
2. Take the mean of the squares, add epsilon, and take the square root. This is the root mean square, or RMS.
3. Divide each value in the row by that number. Nothing is subtracted, so signs and zeros stay as they were.

Worked example

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

Output
[[0.848528, 1.13137]]

Both results stay positive and keep their ratio of 3 to 4. Layer normalization would turn the same row into -1 and 1.

Why it matters

Several recent open language models, Llama among them, use RMS normalization in place of layer normalization inside each transformer block. It does the same job with less arithmetic.

Your turn

Using PyTorch, divide each row of a table by its root mean square.

rows

a table of numbers; every row has at least one value

epsilon

a positive number added to the mean of squares; tiny in practice

Returns

A table of the same shape, as nested Python lists: each row divided by sqrt(mean(row ** 2) + epsilon).

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 = [[3, 4]]
epsilon = 1e-8

Output
[[0.848528, 1.13137]]

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

Output
[[0, 0]]

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

Output
[[0.848528, 1.13137], [0.848528, 1.13137]]

Input
rows = [[3, 3]]
epsilon = 7

Output
[[0.75, 0.75]]

Hints

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

1. Hint 1: Where to start

   There is no mean to subtract this time. You need one number per row: its root mean square.

2. Hint 2: The key step

   Square the tensor with x * x, then take the mean along the last dimension with keepdim=True.

3. Hint 3: The outline

   Make x a float64 tensor. Set mean_square = (x * x).mean(-1, keepdim=True). Return (x / torch.sqrt(mean_square + 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. The mean of the squared values in each row. Squaring first makes a value of -3 count as much as a value of 3. The result has shape [rows, 1].
2. Divide each row by its root mean square. After this the mean of the squares in each row is 1, or very close to it, whatever size the row had before. Only a row of zeros stays at 0.

The same numbered steps are the comments in the solution.

import torch


def solve(rows, epsilon):
    x = torch.tensor(rows, dtype=torch.float64)
    # 1. The mean of the squared values in each row
    mean_square = (x * x).mean(-1, keepdim=True)
    # 2. Divide each row by its root mean square
    normalized = x / torch.sqrt(mean_square + epsilon)
    return normalized.tolist()

Watch out

• Square first, then average. The square of the mean is a different number: for [3, -3] it is 0, while the root mean square is 3.
• Do not subtract the mean. That would turn this into layer normalization.

Check yourself