2. What makes transformers fast

Chapter 1 ended with a model that reverses any string of a length it has seen and none of a length it has not. This chapter makes two changes to it. The first swaps the learned position table for rotary position embeddings (Su et al., 2024), which LLaMA (Touvron et al., 2023) and many open models after it use, and asks whether that fixes the length problem. It mostly does not, and the way it fails is worth understanding. The second change adds a key-value cache to generation and measures what it saves. Everything else, the data, the size, the optimiser and the seeds, is chapter 1's.

The numbers come from 02-tricks.ipynb as executed, 129 seconds in total on the same laptop as chapter 1.

2.1 What a position embedding has to do

Attention is blind to order. Section 1.8 put order back by adding a learned vector for each position to the token at the input, and that vector then had to carry two jobs at once: identify the position and be useful for whatever the layers above wanted to do with it. It also had one row per position, so position 19 in a model trained on sequences of 17 was a row that had never been touched.

There is a different place to put position information. Attention compares a query with a key by a dot product. If the dot product itself could be made to depend on how far apart the two tokens are, the model would get relative position for free, at every layer, without any row for position 19 to be missing. That is what rotary embeddings do.

2.2 Rotary positions by example

Take a query and a key of dimension 2, so each is a point in the plane. Rotate the query by an angle proportional to its position, and the key by an angle proportional to its position. The dot product of two vectors depends only on their lengths and the angle between them, and the angle between the two rotated vectors is the original angle plus the difference of the two position angles. Rotate both by the same extra amount, which is what moving both tokens along the sequence does, and nothing changes.

The notebook checks this directly. A random query at position 5 against a random key at position 2 gives a dot product of 0.6012. The same query at position 40 against the same key at position 37 gives 0.6012. Distance 3 in both cases; the answer does not know where in the sequence it was computed.

Figure 2.1 makes the same check for every pair of positions from 0 to 63 at once. On the left, the dot product of one rotated query with one rotated key forms stripes parallel to the diagonal, which is what a function of $m - n$ alone looks like. On the right, all 2,080 pairs with the key at or before the query fall onto 64 points, one per distance. Pairs at the same distance differ by at most $4.1 \times 10^{-6}$, which is floating-point rounding.

Figure 2.1. One random 32-dimensional query and one random key, each rotated to every position from 0 to 63. Left: the dot product of the query at position $m$ with the key at position $n$. Right: the same values for all pairs with $n \le m$, plotted against the distance $m - n$; darker points are where many pairs overlap.

With more than two dimensions, the vector is split into pairs and each pair is rotated at its own rate, fast for the first pair and slowing geometrically to very slow for the last. Fast pairs resolve small distances; slow pairs can tell position 3 from position 300 without wrapping around.

2.3 The same thing as an equation

For a pair of dimensions $(x_{2i}, x_{2i+1})$ at position $m$, with rotation rate $\theta_i$,

Apply this to the query at position $m$ and the key at position $n$. Because a rotation by $m\theta$ followed by the transpose of a rotation by $n\theta$ is a rotation by $(m-n)\theta$, the dot product $q'_m \cdot k'_n$ is a function of $q$, $k$ and $m-n$ only. The values are not rotated; they carry content, not position.

def apply_rope(x, cos, sin, offset=0):
    # x: (batch, heads, tokens, dim); rotate pairs (x[2i], x[2i+1]) by the angle of their position
    T = x.shape[-2]
    c = cos[offset:offset + T][None, None]
    s = sin[offset:offset + T][None, None]
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack((x1 * c - x2 * s, x1 * s + x2 * c), dim=-1).flatten(-2)

The rotation is applied inside every attention module, to q and k after the projection and before the scores. The position table at the input goes away, which is why the rotary model has 101,644 parameters to the learned model's 103,244: exactly the 25-by-64 table.

2.4 Does it fix the length problem?

Both models were trained on the chapter 1 data (strings of 3 to 8 letters, 1,500 steps of 64) and tested on 500 fresh strings of the training lengths and 500 strings of lengths 9 to 12. Figure 2.2 gives the result by length.

Figure 2.2. Exact-match accuracy by string length. The shaded region is lengths the models never trained on.

On lengths 3 to 8, both models get 500 of 500. On lengths 9 to 12, the learned-position model gets 0 of 500, as in chapter 1. The rotary model gets 14 of 500: 14 of the 122 strings of length 9, and none of the 133, 131 and 114 strings of lengths 10, 11 and 12.

So rotary positions do not fix it. They remove one obstacle, the untrained rows of the table, and expose the next one. The model trained on lengths up to 8 has learned, for each length it saw, where to look; a string of length 9 asks it to look one position further back than it has ever had to, and it manages that about one time in nine, and never for anything longer. Relative position makes the generalisation possible, in that nothing in the architecture now rules it out. It does not make it happen. What makes it happen is training on lengths close to the ones you want, which is what chapter 1's experiments notebook showed, and which is why real models are trained on long contexts rather than extrapolated to them.

That is a smaller result than the usual story about rotary embeddings, and it is the honest one for this task. On natural text, where the useful distances are mostly short and the statistics do not change with position, rotary models do extrapolate somewhat, and there are tricks (scaling the rotation rates after training) that stretch them further. Reversal is a harder case because every output position depends on an exact distance that grows with the length.

2.5 The cost of generating one token

Chapter 1 generated by appending a token and running the whole sequence through the model again. For a prompt of 10 tokens and 10 generated tokens that is 10 forward passes over sequences of 10 to 19 tokens. The waste is that the keys and values at the old positions are recomputed every time and come out identical, because nothing before them has changed and the mask stops anything after them from mattering.

The fix is to keep them. Each attention module stores the keys and values it has computed so far. When one new token arrives, the module computes a query, key and value for that token only, appends the key and value to the store, and attends from the one new query over the whole store. The query is not stored; queries are only ever needed at the moment they are made. Rotary positions make the bookkeeping easy, since the stored keys were rotated by their own positions when they were made and never need touching again; the new token is rotated by its position, passed in as offset.

if cache is not None:
    if cache["k"] is not None:
        k = torch.cat([cache["k"], k], dim=2)
        v = torch.cat([cache["v"], v], dim=2)
    cache["k"], cache["v"] = k, v
scores = q @ k.transpose(-2, -1) / math.sqrt(self.dk)

The mask needs one change: the query block may be shorter than the key block, so the triangle is offset by the difference. That is the diagonal=1 + Tk - Tq in the notebook, and it is the line most likely to be wrong in a first attempt.

2.6 Measured

Correctness first. On 200 random prompts from the reversal task, the trained rotary model produced identical token sequences with and without the cache, 200 of 200. If this check fails, the mask offset is the first thing to look at.

Then time. The reversal model is too small to show anything, so the timing uses an untrained model with 3,290,880 parameters, four blocks, 256-dimensional embedding, generating from a 4-token prompt. Training is irrelevant to the cost of a forward pass. The table and Figure 2.3 give the times.

Figure 2.3. Seconds to generate N tokens, four CPU threads, batch of one.

Without the cache, each new token costs a pass over the whole sequence, so the total grows with the square of the length: doubling from 256 to 512 tokens multiplied the time by 3.9. With the cache, each new token costs one position of work plus one attention over the store, and the total grows close to linearly: the same doubling multiplied the time by 2.2. At 512 tokens the cache is 8.5 times faster, and the gap keeps widening. Every chat model you have used generates this way; the "time to first token" you notice is the prompt being run through once to fill the cache, and the steady stream after it is the per-token path.

What the cache costs is memory. For each layer it holds two tensors of (heads, tokens, dimension per head), and at long contexts in large models that store is bigger than the model's weights. A good deal of the engineering in modern inference, grouped-query attention, paged caches, quantised caches, is about shrinking it.

Where this goes wrong

Rotating the values. The rotation belongs on queries and keys only. Rotate the values too and the content each position hands over changes with where it sits, which is not what you want and will cost accuracy on anything position-sensitive.

Forgetting the offset. With a cache, the new token's query and key must be rotated by its actual position, not by zero. Get this wrong and the cached and uncached paths disagree; the 200-of-200 check in the notebook exists for exactly this.

Mask shape. A square mask of the query length will not fit a key block that is longer. The offset form handles the prompt (many queries, as many keys), the per-token step (one query, many keys), and the uncached case (equal lengths) with one expression. Check all three.

Exercises

1. The rotary model and the learned-position model have different parameter counts. Account for the difference exactly.
2. Rotary positions use a base of 10,000. Retrain with a base of 100 and with 1,000,000 and repeat the length tests. Does the choice matter on this task?
3. The cache stores keys and values but not queries. Explain why queries are never needed again, and work out how much memory the cache holds for the timing model after 512 tokens.
4. In generate_with_cache, the prompt is processed in one call and then tokens one at a time. Time the prompt call alone for prompts of 16, 64 and 256 tokens. What does this say about where the cost of a long prompt goes?
5. Train the rotary model on lengths 3 to 8 and 11 to 12, leaving 9 and 10 out, then test on 9 and 10. Compare with the learned-position result from chapter 1's exercise 5.

Notebook for this chapter: 02-tricks.ipynb, 129 seconds in total, no downloads.

Further reading

• Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo and Yunfeng Liu. 2024. RoFormer: Enhanced transformer with rotary position embedding. Neurocomputing 568, 127063. doi:10.1016/j.neucom.2023.127063.
• Hugo Touvron, Thibaut Lavril, Gautier Izacard, et al. 2023. LLaMA: Open and efficient foundation language models. arXiv:2302.13971.