ML coding · Deep learning layers
Layer normalisation
Given X (N × D), and gamma, beta of length D, return
Y with Yn = γ ⊙ (Xn − μn) / √(σ²n + eps) + β, where μn
and σ²n are the mean and the biased variance (divide by D) of row n.
Examples
Input X = [[1, 2, 3]], gamma = [1, 1, 1], beta = [0, 0, 0], eps = 0
Output [[-1.2247, 0, 1.2247]]
Mean 2, variance 2/3: the row becomes (−1.2247, 0, 1.2247).
Input X = [[5, 5, 5, 5]], gamma = [1, 1, 1, 1], beta = [0.5, 0.5, 0.5, 0.5], eps = 1.000e-5
Output [[0.5, 0.5, 0.5, 0.5]]
A constant row: ε keeps the division finite and the output is β.
Input X = [[1, 2], [10, 30]], gamma = [2, 1], beta = [0, 1], eps = 0
Output [[-2, 2], [-2, 2]]
Each row is normalised on its own, whatever its scale.
Constraints
- 1 ≤ N ≤ 512, 1 ≤ D ≤ 1024
- eps ≥ 0; when eps = 0 no row is constant
- Checked to a relative tolerance of 10⁻⁹
Hints
Hint 1
Take the mean and variance along axis 1, the features, with keepdims=True.
Hint 2
np.var divides by n by default, which is the convention here. Do not pass ddof=1.
Hint 3
Add eps to the variance before the square root, then multiply by gamma and add beta.
Where candidates go wrong
- Normalising over axis 0, the batch. That is batch norm.
- Using the unbiased variance (ddof=1). PyTorch's LayerNorm divides by D.
- Adding eps to the standard deviation instead of to the variance.
The idea
Compute each row's mean and biased variance with keepdims=True, so that they broadcast back against X. Subtract, divide by √(σ² + ε), then apply γ and β, which broadcast along rows. Three vectorised lines, no loop.
Time O(ND) · Space O(ND)
Reference solution
import numpy as np
def layer_norm(X, gamma, beta, eps):
X = np.asarray(X, dtype=float)
mu = X.mean(axis=1, keepdims=True)
var = X.var(axis=1, keepdims=True) # biased: divide by D
return (X - mu) / np.sqrt(var + eps) * np.asarray(gamma, dtype=float) + np.asarray(beta, dtype=float)Follow-ups an interviewer asks
- RMSNorm, used in Llama: drop the mean subtraction and β, divide by the root mean square.
- Write the backward pass. The gradient with respect to X has three terms; derive them.
- Why does a transformer apply layer norm before attention (pre-norm) rather than after?
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
It is short, and it is in every transformer block, so it is a natural opener in an ML engineering screen. The interviewer is checking the axis, the biased variance and where ε goes.
The theory behind it is in the lesson Attention, normalisation and convolution from scratch.
Keep going
Layer normalisation
Normalise each row over its features, then scale by gamma and shift by beta.
Open the problem → Medium15 minBatch normalisation in training mode
Normalise each feature over the batch and update the running mean and variance the way PyTorch does: biased variance to normalise, unbiased to track..
Open the problem → Medium20 minScaled dot-product attention with a causal mask
Implement softmax(QKᵀ/√d)V with an optional causal mask in NumPy.
Open the problem → Medium20 min2-D convolution with stride and padding
Implement a 2-D convolution (cross-correlation) with zero padding and stride, first with loops, then with sliding_window_view and einsum..
Open the problem →