"What happens to the cost of your model when you double the context length?" is one of the most common questions in a systems interview, and the two most common answers are both wrong. "It quadruples, because attention is quadratic" ignores that most of a Transformer's compute is not attention. "It doubles, because the parameters are the same" ignores that attention's share grows with length until, at long enough contexts, it is almost everything. This section counts the costs exactly, for real model sizes, so you can give the right answer with numbers.

Then we turn from what attention costs to what it buys. We look at the attention patterns trained models actually show, build an induction head by hand (the two-head circuit behind much of in-context learning), and review what is proven about the power and the limits of attention. The section ends with a map of the ways people have made attention cheaper, each of which is a later chapter.

Counting FLOPs, one matrix product at a time

The convention throughout the book is that multiplying an $m\times k$ matrix by a $k\times n$ matrix costs $2mkn$ floating-point operations (FLOPs), one multiply and one add per term of each of the $mn$ dot products. Take one multi-head attention layer of width $d$, with $h$ heads of width $d_k = d/h$, and one sequence of $T$ tokens. The forward pass does six products.

The total is

In plain words, the projections cost the same per token at any length, while the score and value products cost per token in proportion to the length.

The first term is linear in $T$, like every other weight multiplication in the model. Only the second, the attention core, is quadratic. With a causal mask, only the $T(T+1)/2$ scores on and below the diagonal are needed, and kernels such as FlashAttention skip the masked blocks, so the core costs about $2T^2d$. (A naive implementation computes all $T^2$ scores and throws half away.)

To see where the compute goes, divide by $T$ to get per-token costs and compare with the rest of a Transformer block. Every attention layer is followed by a feed-forward network applied to each token separately, which §5.3 builds. In its classic form it is two matrices, $d\times 4d$ and $4d\times d$, costing $2 \times 2 \times 4d^2 = 16d^2$ FLOPs per token. So per token and per layer,

The ratio of the core to everything else is $4Td/24d^2 = T/(6d)$, or $T/(12d)$ with a causal mask. With a causal mask, the core is a minor cost while $T \ll 12d$ and the dominant one when $T \gg 12d$. This is the criterion Kaplan et al. (2020) use to drop the context term from their compute estimate, which is fine when $d$ is larger than $T/12$. (Their per-token count for this term, $2n_{\text{layer}}n_{\text{ctx}}d_{\text{attn}}$, is our causal count.) For training, multiply everything by 3, since the backward pass costs about twice the forward pass (§4.3, and §5.6 for the whole model).

Now real numbers. The cell below takes three widths with the classic block, GPT-2 small ($d = 768$), a width of 4,096, and GPT-3 175B ($d = 12{,}288$), and counts causal attention at several context lengths.

At GPT-2's context of 1,024 tokens the attention core is about a tenth of the compute. At width 4,096 it is about 14% at 8,192 tokens and 73% at 128K, with the crossover at $12d = 49{,}152$ tokens. GPT-3 at its training context of 2,048 spent about 1% there. So the honest answer to "what happens when you double the context?" depends on where you start. Well below the crossover, cost per token barely changes and total cost roughly doubles. Well above it, cost per token doubles too, and total cost roughly quadruples. Real models change the weight term a little. Llama 3 8B ($d = 4096$) shares its keys and values across heads and uses a wider, gated feed-forward network, which moves its crossover to about 53,000 tokens and puts about 71% of its forward FLOPs in the core at 128K. §5.6 counts whole models, unembedding and backward pass included.

def attention_flops(T, d, h, causal=False):
    raise NotImplementedError

Memory, the cost that bites first

FLOPs tell you how long a job takes on a perfect machine. Memory tells you whether it runs at all, and for attention, memory is usually the first wall you hit. Three quantities matter.

The score matrices. A naive implementation materializes $\mathbf{S}$ and $\mathbf{A}$, which are $hT^2$ numbers per layer and per sequence, and keeps $\mathbf{A}$ for the backward pass (§4.3). Take a model with 32 layers of 32 heads, the shape of Llama 3 8B. At $T = 8{,}192$ in bfloat16 (2 bytes), one layer's attention weights take $32 \times 8192^2 \times 2$ bytes $= 4$ GiB, and all 32 layers take 128 GiB, for a single sequence, more than any single GPU holds. At 128K tokens one layer would need a full TiB. Korthikanti et al. (2022) count the activation memory of a whole layer as about $sbh(34 + 5as/h)$ bytes (sequence length $s$, batch $b$, width $h$, $a$ heads, in their notation), and the $5as/h$ term is exactly the score-sized tensors (the softmax output, 2 bytes per entry, plus the dropout mask and the dropped-out weights, 1 and 2 bytes). At long contexts it dwarfs everything else.

Everything else is linear. The other activations of a layer (inputs, $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$, outputs) are $O(Td)$, so they grow linearly with length.

The cache at inference. Generating text one token at a time keeps every past key and value, so that each new token does not recompute them. That is $2Ld$ numbers per token for $L$ layers of ordinary multi-head attention, a cache that grows linearly with the context and, for long conversations, can outgrow the weights. §6.5 builds it and counts it for real models, and §12.5 shows how sharing keys and values across heads shrinks it (to 128 KiB per token for Llama 3 8B in bfloat16).

The first of these is a problem of the algorithm, not of attention. FlashAttention computes exactly the same output without ever storing $\mathbf{S}$ or $\mathbf{A}$, keeping only $O(T)$ statistics per head, by working in tiles and using the online softmax (§12.3). Its FLOPs are the same as the naive algorithm's (slightly more, because of recomputation in the backward pass), yet it is much faster. The reason is that the naive algorithm is limited by memory traffic, not arithmetic. Writing and rereading $hT^2$ numbers from GPU memory takes longer than the matrix products that produce them (§12.1, and the roofline model of §11.2).

The quadratic column is why nobody trains long-context models with naive attention, and the linear one is what FlashAttention keeps instead.

What trained heads look like

So much for cost. What does the model get for it? Looking at the attention weights of trained models reveals a small zoo of recurring patterns. These are empirical findings, seen across many models, and the widget below draws each one on a short sequence that repeats itself.

• Previous-token heads put nearly all weight on position $t-1$, a bright line just below the diagonal. You built one in Exercise 4.2b.
• Positional and local heads attend to a fixed offset or a small window, like a learned convolution.
• Sink heads put most of their weight on the first token, the "attend to nothing" behavior of §4.4. Xiao et al. (2023) found this in most layers of the language models they studied.
• Broad heads spread their weight evenly over the past, computing something like a running average.
• Duplicate-token heads attend to earlier occurrences of the current token.
• Induction heads attend to the token after an earlier occurrence of the current token, which shows up as a stripe parallel to the diagonal once text starts repeating.

Induction heads, built by hand

The last pattern deserves a closer look, because it implements an algorithm. Suppose the text so far contains "... Harry Potter ... Harry" and the model must guess the next token. A good guess is "Potter", because when a sequence repeats, the token that followed last time is likely to follow again. Formally, the rule is if the current token is A, find an earlier position where A appeared, and predict the token that came right after it, written $[A][B] \dots [A] \to [B]$. Olsson et al. (2022) called heads that implement this induction heads, and gave several lines of evidence that they are a major mechanism of in-context learning in Transformers. They form during a sudden change early in training, visible as a bump in the loss curve, and the model's ability to use earlier context improves sharply at the same moment.

No single attention head can do this. To attend to "the token after A", the head would need the key at that position to say "the token before me was A", and a token's own vector does not contain its predecessor. The trick is two heads in two layers. Elhage et al. (2021) describe it precisely.

1. In layer 1, a previous-token head copies each token's identity into the next position. After it runs, position $j$ carries both "I am $w_j$" and "the token before me was $w_{j-1}$".
2. In layer 2, the induction head uses the current token as its query ("I am A, who came after an A?") and reads the "previous token" information as its keys. The score $\mathbf{q}_t^\top\mathbf{k}_j$ is high exactly when $w_{j-1} = w_t$. Its value copies the token identity $w_j$, which is the prediction.

The second head's keys are built from the output of the first head. Elhage et al. call this K-composition, and it is the simplest example of heads cooperating to compute something neither can compute alone. We can build the whole circuit with explicit weights. The residual stream is the concatenation of three blocks of coordinates, a one-hot token id, a one-hot position, and a slot where layer 1 writes the previous token's id. Each head reads from some blocks and writes to others, which is how real models use the residual stream too (§5.1).

import numpy as np

def induction_weights(V, T, c=20.0):
    d = 2 * V + T
    return np.zeros((d, V)), np.zeros((d, V)), np.zeros((d, V))   # placeholder: predicts nothing

Here is the complete circuit running, with accuracy measured position by position on 200 random repeated sequences.

A few notes on what this toy does and does not show. Real induction heads use rotary or learned positions rather than one-hot ones, live in models that also have MLPs, and work in superposition with many other features, so they are messier than ours. But the mechanism is the same, and it explains something important. Nothing in the weights stores "Harry Potter". The circuit copies whatever followed the previous occurrence, so it generalizes to sequences it has never seen, which is the essence of in-context learning. Finding such circuits in trained models, by patching activations and tracing which heads feed which, is the program of mechanistic interpretability.

What attention can and cannot compute

Here is what has been proven, stated carefully, because these results are often quoted loosely.

Attention can approximate anything, given room. Yun et al. (2020) proved that Transformers with positional encodings are universal approximators of continuous sequence-to-sequence functions on a compact domain, and that without positional encodings they can approximate any continuous permutation-equivariant function, but nothing else (the symmetry of §4.6). Like every universal approximation theorem, this says nothing about how large the model must be or whether training finds it.

With unbounded precision and unbounded steps, even Turing completeness. Pérez et al. (2021) showed that a Transformer with hard attention, arbitrary-precision numbers, and as many decoding steps as needed can simulate a Turing machine. The assumptions are far from practice, but the result frames the next one.

In one forward pass, with realistic precision, attention is shallow. Merrill and Sabharwal (2023) proved that Transformers whose numbers have precision logarithmic in the input length can be simulated by uniform $\mathsf{TC}^0$ circuits, constant-depth circuits with threshold gates. If, as most complexity theorists believe, $\mathsf{TC}^0 \ne \mathsf{NC}^1$, then such a Transformer cannot solve $\mathsf{NC}^1$-complete problems for arbitrary lengths in a single forward pass. Evaluating Boolean formulas, and tracking the state of certain finite automata over a long input (for example, composing a long sequence of permutations of five objects), are examples. The intuition is that a fixed number of layers gives a fixed number of sequential steps, however wide each step is.

Chain of thought adds serial steps. Merrill and Sabharwal (2024) showed that letting a decoder generate intermediate tokens before answering increases its power, by an amount that depends on how many it may generate. A linear number of steps lets it recognize every regular language, and a polynomial number lets it solve exactly the problems solvable in polynomial time. This is one theoretical reason why thinking out loud helps models reason (§18.1).

Single-token sensitivity is a real limitation. Hahn (2020) showed that, if the scores stay bounded, changing one input token changes a soft-attention layer's output by only $O(1/T)$, so a fixed Transformer cannot compute functions such as parity (is the number of 1s odd?) with confidence that stays bounded away from chance as inputs grow. Chiang and Cholak (2022) showed that this is a limit on confidence, not on accuracy. They built a Transformer that recognizes parity exactly, whose confidence does shrink toward chance as length grows, and then used layer normalization to restore it, bringing the cross-entropy arbitrarily close to zero. They also proposed multiplying the scores by $\log T$, which helped length generalization in their experiments.

That last fix points to a practical issue, softmax dilution. Suppose one key is relevant and $T-1$ are not, and the relevant one has a score higher by $\Delta$. Its weight is

In plain words, as the context grows, a head needs ever larger score gaps to keep its focus, about one extra unit of logit for each factor of $e$ in context length.

A head trained at 4K tokens that runs at 128K sees its weights spread over 32 times as many competitors, which is one reason long-context extension needs care (§12.7). Several long-context methods scale scores with $\log T$ for exactly this reason, such as the "log-n scaling" used in Qwen (Bai et al. 2023) and the "scalable softmax" of Nakanishi (2025).

import numpy as np

def gap_for_weight(T, p):
    raise NotImplementedError

def weight_for_gap(T, gap):
    raise NotImplementedError

Two more limits round out the picture. Dong et al. (2021) proved that a stack of pure attention layers, with no residual connections or MLPs, loses rank doubly exponentially fast with depth, so all tokens collapse toward the same vector. Residual connections and MLPs are what prevent it (§5.1). And for exact attention, quadratic time is probably unavoidable. Keles et al. (2023) proved that computing self-attention, even approximately, needs time quadratic in $T$ unless the Strong Exponential Time Hypothesis fails, and Alman and Song (2023) located the boundary precisely. With head width $O(\log T)$ and every entry of $\mathbf{Q}$, $\mathbf{K}$, $\mathbf{V}$ bounded by $B$, attention can be approximated to within $1/\mathrm{poly}(T)$ in almost linear time when $B = o(\sqrt{\log T})$, and cannot be in truly subquadratic time once $B$ is of order $\sqrt{\log T}$, under the same hypothesis.

The roads to cheaper attention

Every limitation above has prompted a line of work. Chapter 12 covers most of it, including splitting one sequence across GPUs, and later chapters the alternatives to softmax attention. Here is the map.

The last row deserves a sentence of its own. If the softmax of $\mathbf{q}^\top\mathbf{k}$ is replaced by a product of feature maps $\phi(\mathbf{q})^\top\phi(\mathbf{k})$, the sum $\sum_j\phi(\mathbf{k}_j)\mathbf{v}_j^\top$ can be accumulated left to right, and causal attention becomes an RNN with a fixed-size state (Katharopoulos et al. 2020). That idea, refined into selective state-space models such as Mamba, is the main alternative to attention for very long sequences, at the price of the exact, content-addressed lookup this chapter has been building.

In PyTorch

You can measure the quadratic cost directly. The timings below come from PyTorch's CPU kernels, so the absolute numbers depend on your machine, but the trend does not.

Read more

• Kaplan et al. (2020). Scaling Laws for Neural Language Models. Table 1 counts parameters and FLOPs per operation, including the attention context term.
• Korthikanti et al. (2022). Reducing Activation Recomputation in Large Transformer Models. Section 4 derives the activation memory of a layer, including the attention score term.
• Olsson et al. (2022). In-context Learning and Induction Heads. Six lines of evidence that induction heads drive in-context learning, and the phase change in which they form.
• Elhage et al. (2021). A Mathematical Framework for Transformer Circuits. QK and OV circuits, head composition, and the first account of induction heads.
• Merrill & Sabharwal (2023). The Parallelism Tradeoff and (2024), The Expressive Power of Transformers with Chain of Thought. What one forward pass can compute, and what chain of thought adds.
• Hahn (2020). Theoretical Limitations of Self-Attention in Neural Sequence Models and Chiang & Cholak (2022), Overcoming a Theoretical Limitation of Self-Attention. Parity, sensitivity, and log-length scaling.
• Keles et al. (2023). On The Computational Complexity of Self-Attention and Alman & Song (2023), Fast Attention Requires Bounded Entries. Why exact attention is probably quadratic, and when approximation helps.
• Dao et al. (2022). FlashAttention. The IO-aware algorithm that makes the quadratic memory disappear, the subject of Chapter 12.

Watch next

▶

▶

▶