Attention and transformers

A transformer¹ is a neural network that turns the embeddings of a text into new vectors, one for each position, with a stack of blocks of the same form. Each block has two parts: attention, which lets each position read the vectors of other positions, and a feed-forward network, which works on each position alone.

For example, the model of Section 7.1 has 32 blocks². Each block turns the 8 vectors³ of the text The ride of 3 km cost $ into 8 new vectors of 4096 numbers. The last vector of the last block gives the logits.

Attention is the only part where the positions exchange numbers. Its cost grows with the square of the length of the text, so long texts are expensive to train on and to read.

8.1Attention

Consider the position of ␣$ in the above text. To predict the price, its vector needs the distance 3, which is three positions earlier. A layer of Section 3.1 changes each vector on its own, so it cannot bring the distance to the position of ␣$.

A straightforward way is to give each position the average of the vectors up to it, as for the text embeddings of Section 6.3. However, the average gives the same weight to each position. For example, 3 gets 1/8 of the average at ␣$, as much as The.

The idea is to take a weighted average, where each position chooses the weights by what it looks for. Attention is the part of a block that gives each position a weighted average of the vectors of the positions that it reads.

Each position calculates three vectors from its own vector: a query, a key and a value. Each of them is the product of a weight matrix and the vector, as in a layer with no bias and no activation function. The query describes what the position looks for, the key what it holds, and the value what it gives to the positions that read it.

Let $q_i$ denote the query of position $i$, and let $k_j$ and $v_j$ denote the key and the value of position $j$. The score of position $j$ for position $i$ is as follows: $$s_{ij} = \frac{q_i \cdot k_j}{\sqrt{d_k}}$$ where $d_k$ is the number of numbers in each query and key. The attention weights $w_{ij}$ are the softmax of the scores of the positions $j$ that $i$ reads (see Section 7.2). The output of position $i$ is $\sum_j w_{ij} v_j$. The attention weights are not parameters: the model calculates them again for each text.

As an example, suppose that $d_k = 4$ and that the query of ␣$ is (1, 1, 1, 1). The following table shows the key of each position, its score and its attention weight:

For example, the score of 3 is $(1 + 1 + 1 + 1)/\sqrt{4} = 2$. The sum in the softmax is $e^2 + e^{1.5} + e^0 + 5 e^{-2} \approx 13.55$, so the attention weight of 3 is $7.39 / 13.55 \approx 0.55$. The output of ␣$ is then 0.55 times the value of 3, plus 0.33 times the value of ␣km, and so on. Thus, more than half of the attention weight goes to the distance.

The three matrices are parameters, and the training changes them together with all other parameters. Nothing in the model says that ␣$ needs the distance. Gradient descent finds a query and a key that match in this way because they make the loss of the next token smaller.

Note that the division by $\sqrt{d_k}$ keeps the scores from growing with $d_k$. Without it, queries and keys with many numbers give large scores, and the softmax gives almost all of the attention weight to one position. Its derivatives are then almost 0, and the training barely changes the attention weights⁴.

8.2Reading earlier positions

Causal attention⁵ is attention in which each position reads only itself and the positions before it, and language models use it. The model sets the scores of the later positions to $-\infty$, so their attention weights are $e^{-\infty} = 0$. For example, the position of ␣ride reads The and itself, but not ␣of.

The following picture shows which positions read which in the above text, where each row is a reading position and its shaded cells are the positions that it reads:

The row of ␣$ has 8 shaded cells, one for each row of the above table. The row of The has only one. In all, 1 + 2 + … + 8 = 36 of the 64 cells are shaded.

The reason for this rule is the training. One pass through a training text gives the loss at every position, where each position predicts the token after it (see Section 7.2). If a position could read the next position, its vector would contain the token that it has to predict.

For example, in the training text The ride of 3 km cost $9., the position of ␣$ predicts 9. With causal attention, it cannot read the position of 9.

Thus, in every block, the vector at a position depends only on the tokens up to that position. When the model writes a new token, the vectors of the earlier positions do not change. Chapter 10 shows how the program reuses them.

An embedding model reads a whole text and writes nothing (see Section 6.3), so it has no next token to hide. In this case, we can let each position read every position, and many embedding models do so⁶. All 64 cells of their picture are shaded. The embedding of the text is often the average of the vectors of the last block⁷.

8.3Positions

Attention by itself does not see the order of the tokens. When each position reads every position, the output of a position depends on its own token and on the keys and values that it reads, but not on their order.

For example, the position of bit reads the same five tokens in the dog bit the man and in the man bit the dog. It gets the same output in both texts. An embedding model would then give the two texts the same embedding, as the average of Section 6.3 does⁸.

A straightforward way to add the order is to give each position its own vector of $d$ numbers and to add this vector to the embedding of the token at that position⁹. The position vectors can be learned by the training, like the embeddings. For example, a model can have learned vectors for 1,024 positions¹⁰. However, such a model has no vector for position 1,025, so it cannot read a longer text.

Rotary position embedding¹¹ is a method that rotates the query and the key of each position by an angle that grows with the position¹². A rotation does not stretch a vector, so the dot product of a query and a key changes only through the angle between them. Thus, for the same two tokens, the score depends on their distance, not on where they are in the text.

As an example, suppose that $d_k = 2$, that the angle grows by 20° per position, and that a query and a key are both (1, 0) before the rotation. The following picture shows the query at position 8 and the key at position 5 on the left, and the query at position 4 and the key at position 1 on the right:

In both cases, the angle between the query and the key is 60°, and their dot product is $\cos 60° = 0.5$. For the distance 1, for example at positions 6 and 5, the angle is 20°. The dot product is then $\cos 20° \approx 0.94$.

Note that a rotation exists for every position, so the method has no last position. However, a model has seen only the distances in its training texts, and it predicts the tokens of longer texts badly¹³. In practice, a model is trained further on long texts before it reads them (see Section 8.5).

8.4The block

The feed-forward network of a block consists of two layers of Section 3.1, which each position goes through on its own. The first layer turns the $d$ numbers of the vector into more numbers and applies an activation function. The second layer turns them back into $d$ numbers. For example, in the model of Section 7.1, the first layer gives 14,336 numbers¹⁴, 3.5 times $d = 4096$.

A block runs attention and then the feed-forward network. Each of the two parts adds its output to the vector that it reads, instead of replacing the vector¹⁵. Each part then only has to learn a change of the vector. The following picture shows a block for the first three positions of the above text, without the additions:

Attention is the only part with lines between the positions. Each line goes from a position to the same position or to a later one. The feed-forward network has one box for each position, because each position goes through it alone.

A head is one of several attentions that a block runs side by side on the same vectors, and each head has its own three matrices and shorter queries, keys and values. For example, the above model has 32 heads with queries of 128 numbers, and 32 · 128 = 4096. Each head divides its scores by $\sqrt{128}$, because $d_k = 128$. A fourth matrix turns the 32 outputs into one vector of 4096 numbers.

The heads learn to read different things. For example, in a model of 2018 whose positions read every position, one head found the object of a verb more than 75% of the time¹⁶.

Thus, a block of the above model has 7 weight matrices. Attention has 4 of them: the matrices of the queries, the keys and the values, each for all heads together, and the matrix that joins the heads¹⁷. The feed-forward network has the other 3. The following table shows their numbers of weights:

The feed-forward network has about four fifths of the weights of a block¹⁸.

8.5The context window

The context window of a model is the largest number of tokens that it can read at once, the prompt and the written tokens together. For example, the model of Section 7.1 has a context window of 8,192 tokens. A later version with blocks of the same sizes has a context window of 131,072 tokens¹⁹. In practice, the largest models of two large providers have context windows of about 1,000,000 tokens²⁰.

An API refuses a request whose prompt alone is longer than the context window. When an answer reaches the end of the window, the API stops it²¹. Chapter 18 shows how well a model uses a long prompt.

Attention at position $i$ calculates $i$ scores, so a text of $n$ tokens needs $1 + 2 + \dots + n = n(n + 1)/2$ scores in each head of each block. For example, the 8 tokens of the above text need the 36 scores of the shaded cells of Section 8.2. A text twice as long needs about 4 times as many scores: 16 tokens need 136.

Thus, the time complexity of attention is $O(n^2)$ for a text of $n$ tokens. The weight matrices take $O(n)$ time, because each position goes through each of them once²².

We can compare the two in the above model. The weight matrices of a block do 218,103,808 multiplications for each position, one for each weight. Attention at position $i$ does $2 \cdot 4096 \cdot i$ multiplications: $4096 \cdot i$ for the scores of the 32 heads, and as many for the weighted averages of the values. The following table shows the share of attention in the multiplications of a block for a text of $n$ tokens:

For example, for 8,192 tokens, attention does 4096 · 8192 · 8193 multiplications and the weight matrices 218,103,808 · 8192, so attention has 13% of the sum. The two shares are equal at about 53,000 tokens. This is the extra cost of a long text that the estimate of 2 operations per parameter in Section 3.3 leaves out.

For this reason, the pretraining (see Section 9.2) can leave long texts to a late stage. For example, a model can be trained on texts of 8,192 tokens first, and then on longer texts until it reads 131,072 tokens²³.

Some models let some or all of their blocks read only the last $\ell$ positions²⁴. Such a block takes $O(n\ell)$ time instead of $O(n^2)$. On the other hand, it cannot read a token that is more than $\ell$ positions back, and the token can reach the position only through several blocks.