Layer Norm, Backwards

Every block of every transformer you have used runs this twice. It is also the gradient people get wrong, because the mean and the variance both depend on every element of the row, so changing one input changes the normalised value of all of its neighbours.

Write layer_norm_backward(x, gamma, grad_out, eps=1e-5) returning (grad_x, grad_gamma, grad_beta).

The forward pass normalises over the last axis:

mu = x.mean(-1, keepdims=True)
var = x.var(-1, keepdims=True)
xhat = (x - mu) / sqrt(var + eps)
y = gamma * xhat + beta

x is (..., D), gamma and beta are (D,), grad_out has x's shape. grad_gamma and grad_beta are (D,): summed over every axis except the last.

Your gradient will be checked against finite differences to 1e-6. There is nowhere to hide.