Neural Network Layers: theory

Embedding lookup and its backward pass, LayerNorm and RMSNorm, inverted dropout, and the weight initializations that keep deep networks stable.

Besides attention, a transformer is built from a handful of simple layers:

an embedding table that turns token ids into vectors;
a normalization layer before every block, which keeps the numbers at a steady scale;
dropout in some models, which randomly switches off parts of the network during training;
weights that must start at the right size (initialization).

Each one is a few lines of NumPy, but each has a detail that silently breaks training when it is wrong. This page explains each layer with a small example, and points out that detail.

Embeddings: a lookup table that learns

A language model works on token ids: integers from $0$ to $V - 1$ , where $V$ is the size of the vocabulary. But a network computes with vectors, not with ids. An embedding gives every id its own vector of $D$ numbers. All these vectors are stored in a matrix $E$ with $V$ rows and $D$ columns: the vector of token $i$ is simply row $i$ .

The numbers in $E$ are parameters, learned like any other weight. Bengio et al. (2003) introduced this idea for language models: during training, words with similar roles end up with similar rows.

Example. A vocabulary of $V = 3$ tokens, vectors of $D = 2$ numbers: $E$ = [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]. The text with ids [2, 0, 2] becomes rows 2, 0 and 2: [[0.5, 0.6], [0.1, 0.2], [0.5, 0.6]].

The forward pass is just indexing, E[ids], for an array of ids of any shape. It gives the same result as multiplying a one-hot matrix by $E$ , without ever building that matrix (see the one-hot exercise in the NumPy section).

The backward pass is where the bug hides. Each output row came from one row of $E$ , so its gradient goes back to that row; rows that were not used get 0. And a row used twice gets both gradients added (gradients add up, see section 1). A batch of text repeats tokens all the time: "the", spaces, punctuation.

Example. Same ids = [2, 0, 2], with upstream gradients [1, 1], [2, 2] and [3, 3] for the three output rows. The gradient of $E$ must be:

row 0: [2, 2] (used once, by the second output);

row 1: [0, 0] (not used);

row 2: [1, 1] + [3, 3] = [4, 4] (used by the first and third outputs).

In NumPy, dE[ids] += grad returns [3, 3] for row 2: it keeps only the last of the repeated updates. np.add.at(dE, ids, grad) returns the correct [4, 4].

LayerNorm

Deep networks train badly when the size of the numbers drifts from layer to layer: some layers end up working with values in the thousands, others with values close to 0. LayerNorm (layer normalization; Ba, Kiros and Hinton, 2016) fixes the scale of every token's vector on its own, in three steps:

compute the mean $\mu$ of the vector's $D$ numbers, and their variance $\sigma^2$ : the average squared distance from the mean, a measure of how spread out they are;
subtract the mean and divide by the standard deviation $\sqrt{\sigma^2}$ : now the vector has mean 0 and variance 1;
multiply by a learned gain $\gamma$ and add a learned bias $\beta$ , so the network can still choose its own scale and shift.

As formulas, first the mean and the variance:

\mu = \frac{1}{D}\sum_j x_j

\sigma^2 = \frac{1}{D}\sum_j (x_j - \mu)^2

Then the normalized vector and the output:

\hat x = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}, \qquad y = \gamma \odot \hat x + \beta

where:

$x$ is one token's vector, with entries $x_j$ ;
$\mu$ (mu) is its mean and $\sigma^2$ (sigma squared) its variance;
$\hat x$ ("x hat") is the normalized vector: mean 0 and variance 1;
$\gamma$ (gamma) and $\beta$ (beta) are learned vectors of $D$ numbers;
$\odot$ means elementwise multiplication;
$\epsilon$ (epsilon) is a small constant that avoids dividing by zero.

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

Mean: $\mu = (1 + 2 + 3 + 4)/4 = 2.5$ . Variance: the distances from the mean are $(-1.5, -0.5, 0.5, 1.5)$ , their squares $(2.25, 0.25, 0.25, 2.25)$ , average $\sigma^2 = 1.25$ .

Normalized: $(x - 2.5) / \sqrt{1.25} \approx (-1.34, -0.45, 0.45, 1.34)$ .

With $\gamma = (2, 2, 2, 2)$ and $\beta = (1, 1, 1, 1)$ : $y \approx (-1.68, 0.11, 1.89, 3.68)$ .

The vector $x = (10, 20, 30, 40)$ , ten times larger, gives exactly the same $\hat x$ : the scale of the input is gone, which is the point.

Two details match PyTorch's LayerNorm: the variance divides by $D$ (not $D - 1$ ), and $\epsilon$ (default $10^{-5}$ ) goes inside the square root.

Unlike batch normalization, nothing depends on the other examples in the batch. So LayerNorm behaves identically at training and inference time, and works with a batch of one.

The backward pass is worth deriving once, because every entry of $\hat x$ depends on every entry of $x$ , through $\mu$ and $\sigma$ . With $g = \gamma \odot \partial \mathcal{L}/\partial y$ (the upstream gradient times the gain) and $s = \sqrt{\sigma^2 + \epsilon}$ :

\frac{\partial \mathcal{L}}{\partial x} = \frac{1}{s}\left(g - \overline{g} - \hat x \, \overline{g \odot \hat x}\right)

where a bar means the mean over the $D$ features. In words: the two subtracted terms remove the parts of the gradient that would change the mean and the variance of $x$ , because normalization cancels those changes anyway. For example, adding the same number to every entry of $x$ does not change $\hat x$ , so the gradient must not push in that direction: subtracting $\overline{g}$ makes its entries add up to 0.

RMSNorm

RMSNorm (root mean square layer normalization; Zhang and Sennrich, 2019) keeps only the rescaling: it does not subtract the mean, and has no bias.

\mathrm{RMS}(x) = \sqrt{\frac{1}{D}\sum_j x_j^2 + \epsilon}

y = \gamma \odot \frac{x}{\mathrm{RMS}(x)}

where $\mathrm{RMS}(x)$ , the RMS (root mean square) of the vector, is the square root of the mean of its squared entries: a measure of its typical size, like the standard deviation but measured from 0 instead of from the mean.

Example. The same $x = (1, 2, 3, 4)$ : the squares are $(1, 4, 9, 16)$ , their mean is $7.5$ , and $\mathrm{RMS}(x) = \sqrt{7.5} \approx 2.74$ . With $\gamma = 1$ the output is $x / 2.74 \approx (0.37, 0.73, 1.10, 1.46)$ . Compare with LayerNorm above: the mean is not removed, so all the entries stay positive.

It is cheaper (one reduction over the vector instead of two), and Zhang and Sennrich report quality comparable to LayerNorm. LLaMA (Large Language Model Meta AI) uses it, applied to the input of every sub-layer (pre-normalization).

Dropout

Dropout (Hinton et al., 2012; Srivastava et al., 2014) sets each activation to zero with probability $p$ during training, so the network cannot rely on any single unit and has to spread what it learns. To keep the average value of each activation unchanged, the survivors are multiplied by $1/(1-p)$ ("inverted" dropout). At inference time dropout does nothing at all.

y = \frac{m \odot x}{1 - p}, \qquad m_i \sim \mathrm{Bernoulli}(1 - p)

where:

$p$ is the probability of dropping an entry;
$m$ is the mask: each $m_i$ is 1 with probability $1 - p$ and 0 otherwise (that is what " $\sim \mathrm{Bernoulli}(1 - p)$ " means);
dividing by $1 - p$ makes up for the dropped entries.

Example. $p = 0.5$ , $x = (1, 1, 1, 1)$ , and the random mask comes out as $(1, 0, 1, 1)$ . Then $y = (1, 0, 1, 1) / 0.5 = (2, 0, 2, 2)$ .

Each entry survives half of the time, with value 2, and is 0 the other half: on average it is still 1. Without the division, the average would drop to 0.5 during training and jump back to 1 at inference, and the next layer would see different numbers in the two cases.

The backward pass reuses the same mask: $\partial \mathcal{L}/\partial x = m \odot \partial \mathcal{L}/\partial y \,/\, (1 - p)$ . When a model sees each training token only about once, there is little to overfit: nanoGPT defaults to dropout = 0.0 for pre-training and suggests 0.1 or more for fine-tuning.

Initialization

Before training, every weight is set to a random number. The size of those random numbers matters a lot. Take a linear layer $y = Wx$ with $n_{\text{in}}$ inputs. Suppose the weights and the inputs are independent, have mean zero, and have variances $\mathrm{Var}(w)$ and $\mathrm{Var}(x)$ . Each output adds up $n_{\text{in}}$ products, so

\mathrm{Var}(y) = n_{\text{in}} \, \mathrm{Var}(w) \, \mathrm{Var}(x)

where $\mathrm{Var}$ is the variance, the average squared distance from the mean. In words: one layer multiplies the variance of the signal by $n_{\text{in}} \mathrm{Var}(w)$ .

Stack 50 layers and that factor is raised to the 50th power. If it is 1.5, the activations grow by $1.5^{50} \approx 6 \times 10^8$ ; if it is 0.5, they shrink to almost nothing.

Example. A random input passed through 10 layers of width 512, each a matrix product followed by ReLU (rectified linear unit), measured with NumPy. The standard deviation of the output:

weights with standard deviation 1: about $8 \times 10^{11}$ , the numbers explode;

weights with standard deviation 0.01: about $8 \times 10^{-9}$ , the signal vanishes;

weights with standard deviation $\sqrt{2/512} = 0.0625$ (He initialization, below): about 0.9, it stays steady.

So the weights must start with a variance tied to the layer's size, chosen to make the factor about 1:

Xavier/Glorot (Glorot and Bengio, 2010): $\mathrm{Var}(w) = 2 / (n_{\text{in}} + n_{\text{out}})$ , a compromise between keeping the forward activations and the backward gradients stable. It is derived for activations that are roughly linear around 0, such as tanh.
He/Kaiming (He et al., 2015): ReLU zeroes half its inputs, which halves the variance, so it uses $\mathrm{Var}(w) = 2 / n_{\text{in}}$ : the factor 2 makes up for the lost half.

Models in the style of GPT-2 (Generative Pre-trained Transformer 2) use a simpler rule: every weight is drawn from $\mathcal{N}(0, 0.02^2)$ , a normal distribution with mean 0 and standard deviation $0.02$ . The projections that write into the residual stream are scaled down further, by $1/\sqrt{2L}$ for $L$ layers, because $2L$ of them are added together. This is what nanoGPT's model.py does.

References

Y. Bengio, R. Ducharme, P. Vincent, C. Jauvin, A Neural Probabilistic Language Model, JMLR 3, 1137–1155, 2003. Learned word embeddings.
J. L. Ba, J. R. Kiros, G. E. Hinton, Layer Normalization, arXiv, 2016.
B. Zhang, R. Sennrich, Root Mean Square Layer Normalization, NeurIPS, 2019.
H. Touvron et al., LLaMA: Open and Efficient Foundation Language Models, arXiv, 2023. Section 2.2: pre-normalization with RMSNorm.
G. E. Hinton, N. Srivastava, A. Krizhevsky, I. Sutskever, R. R. Salakhutdinov, Improving neural networks by preventing co-adaptation of feature detectors, arXiv, 2012.
N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, R. Salakhutdinov, Dropout: A Simple Way to Prevent Neural Networks from Overfitting, JMLR 15(56), 1929–1958, 2014.
X. Glorot, Y. Bengio, Understanding the difficulty of training deep feedforward neural networks, AISTATS, 2010.
K. He, X. Zhang, S. Ren, J. Sun, Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification, ICCV, 1026–1034, 2015.
PyTorch documentation: torch.nn.LayerNorm and torch.nn.RMSNorm.