Skip to content

← Normalization step 2 of 6

Medium Primitives

Implement Layer Normalization

Layer normalization rescales each example using statistics from that example's features. For a matrix with examples in rows, each row gets its own mean and variance. Changing another row does not affect those statistics.

For the row [1, 2, 3], the mean is 2. The squared deviations are [1, 0, 1], so the variance is 2 / 3. Subtract the mean and divide by sqrt(2 / 3 + eps). With eps = 1e-5, the result before scaling and shifting is approximately [-1.224736, 0, 1.224736].

The final scale gamma and shift beta each contain one value per feature. For feature j, return gamma[j] * normalized[j] + beta[j].

Your task

Implement layer_norm(x, gamma, beta, eps=1e-5):

  • x is a floating-point tensor of shape (N, D), with positive dimensions.
  • gamma and beta have shape (D,) and broadcast across rows.
  • eps is a positive float. Add it to the variance inside the square root.
  • Return a tensor of shape (N, D).

Compute row means and population variance: the mean squared deviation, dividing by D, not D-1. Keep the reduced dimension so the statistics have shape (N, 1) and broadcast over each row.

Use basic tensor reductions and arithmetic rather than a built-in LayerNorm operation. A constant row, including a row with D=1, normalizes to zero before the affine step, so its output is beta. Positive epsilon means the pre-affine variance is not generally exactly 1; gamma and beta can also change the output mean and variance.

The diagram compares normalization groups. For this 2D exercise, each column is one feature and there are no extra spatial dimensions. With LayerNorm selected, entries in one row share their statistics.

Loading visualization…