LayerNorm Forward & Backward

The problem

LayerNorm (layer normalization) keeps the numbers inside a network at a steady scale. It takes each token's vector of $D$ numbers on its own, shifts it to mean 0, rescales it to variance 1, then applies a learned gain $\gamma$ and bias $\beta$ (both vectors of $D$ numbers). The theory page explains why this helps.

You will write the forward pass and the backward pass.

Forward

First the mean and the variance of the vector:

Then normalize, scale and shift:

where:

• $x$ is one token's vector and $x_j$ its entry $j$;
• $\mu$ (mu) is its mean and $\sigma^2$ (sigma squared) its variance, the average squared distance from the mean;
• $\hat x$ ("x hat") is the normalized vector, with mean 0 and variance 1;
• $\epsilon$ (epsilon) is eps, which avoids dividing by zero;
• $\odot$ is elementwise multiplication.

Example. $x = (1, 2, 3, 4)$, $\gamma = (2, 2, 2, 2)$, $\beta = (1, 1, 1, 1)$.

1. $\mu = 2.5$. Distances from the mean $(-1.5, -0.5, 0.5, 1.5)$, squared $(2.25, 0.25, 0.25, 2.25)$, so $\sigma^2 = 1.25$.
2. $\hat x = (x - 2.5) / \sqrt{1.25} \approx (-1.34, -0.45, 0.45, 1.34)$.
3. $y = 2 \hat x + 1 \approx (-1.68, 0.11, 1.89, 3.68)$.

Backward

The backward pass receives grad_y, the gradient of the loss with respect to $y$, and returns three gradients.

For $\beta$ and $\gamma$: $\beta_j$ is added to entry $j$ of every token's output, and $\gamma_j$ multiplies $\hat x_j$ of every token. So:

• the gradient of $\beta$ is grad_y summed over all tokens;
• the gradient of $\gamma$ is grad_y * x_hat summed over all tokens.

Example. Two tokens with grad_y = [[1, 0, 0, 0], [1, 0, 0, 0]]: the gradient of $\beta$ is [2, 0, 0, 0]. If both tokens had the $\hat x$ above, the gradient of $\gamma$ would be $1 \cdot (-1.34) + 1 \cdot (-1.34) = -2.68$ in its first entry, and 0 elsewhere.

For $x$: every entry of $\hat x$ depends on every entry of $x$, through $\mu$ and $\sigma$. The theory page derives the result. With $g = \gamma \odot$ grad_y and $s = \sqrt{\sigma^2 + \epsilon}$:

where:

• a bar means the mean over the $D$ entries of the token;
• $\overline{g}$ removes the part of the gradient that would only shift all entries together (normalization undoes any shift);
• $\hat x \, \overline{g \odot \hat x}$ removes the part that would only stretch the vector (normalization undoes any stretch).

What to implement

• layernorm_forward(x, gamma, beta, eps=1e-5): x has shape $(\dots, D)$. Return (y, cache), where cache is anything you need for the backward pass (for example $\hat x$, $s$ and $\gamma$).
• layernorm_backward(grad_y, cache): return (dx, dgamma, dbeta), with the shapes of x, gamma and beta.

Requirements

• Match PyTorch's LayerNorm: the variance divides by $D$ (np.var does this by default), and eps is added inside the square root.
• x can have any number of leading axes: $(N, D)$, $(B, T, D)$, ... Normalize over the last axis only. dgamma and dbeta sum over all the other axes.
• No Python loops.

Hints

• Use keepdims=True for the means and the variance over the last axis, so they line up with x.
• To sum over every axis except the last, pass a tuple: axis=tuple(range(x.ndim - 1)).