Why Attention Costs O(n²) in transformers

Module 05 · Transformers · 19 concepts

In depth / 15 lines

Why Attention Costs O(n²), in more detail

Attention scores a query against every key, and in self-attention every token supplies both.

With n tokens that is n times n scores, computed for each head of each layer.

The compute therefore grows with the square of the sequence length, written O(n squared).

Feed-forward layers, by contrast, grow linearly, so attention comes to dominate as context lengthens.

At five hundred tokens the quadratic term is negligible; at a hundred thousand it is everything.

There are two separate costs and they are worth separating, because their fixes differ.

The first is arithmetic: the multiplications needed to build and apply the score matrix.

The second is memory: that matrix is n by n per head and must exist somewhere while softmax runs.

For long sequences the memory cost bites first, and it is what produces out-of-memory errors.

Flash attention attacks exactly that by computing attention in tiles and never materialising the full matrix.

It is mathematically identical, faster in practice, and leaves the quadratic arithmetic untouched.

Sparse and windowed attention change the maths instead, letting each token attend to a subset.

Linear attention approximates the softmax so the operation can be reordered into linear cost.

These approximations trade quality for reach and have never fully displaced exact attention.

Meanwhile the practical answer at inference is the KV cache, which avoids recomputing past keys and values.

Deep dive

What is underneath Why Attention Costs O(n²)

The two terms that matter

\text{time } O(n^2 d) \qquad \text{memory } O(n^2)

$n$ : sequence length.
$d$ : model width.

The time comes from $QK^{\top}$ : an $n \times d$ matrix times a $d \times n$ matrix is $n^2 d$ multiply-adds, and the multiplication by $V$ costs the same again.

The memory comes from the score matrix itself, which is $n \times n$ and has to exist before the softmax can normalise it.

What quadratic actually costs

Doubling the sequence length quadruples both.

n = 1{,}000 \implies 10^6 \text{ scores} \qquad n = 100{,}000 \implies 10^{10} \text{ scores}

At 100,000 tokens the score matrix in fp16 is 20 GB, per head, per layer. That is not a slow model, it is a model that does not run.

This single term is why context windows grew slowly for years while parameter counts grew fast.

Against the alternative

\text{self-attention } O(n^2 d) \qquad \text{recurrent } O(n d^2)

Attention is cheaper than recurrence whenever $n < d$ , which covers most sentences: 512 tokens against a width of 4096.

It is more expensive when the sequence is long, which is exactly the regime everyone now wants.

So the comparison that made attention look cheap in 2017 has inverted, and the field has been chasing the difference ever since.

The three ways out

Make it exact but cheaper in memory, by never storing the full matrix. FlashAttention does this and is not an approximation.
Make it sparse, so each position attends to a subset rather than everything. Sliding windows and dilated patterns take the term to $O(n \sqrt{n})$ or $O(n)$ .
Change the operation, replacing softmax attention with something linear in $n$ . Linear attention and state-space models take this route.
Reduce what is stored rather than what is computed, which is grouped-query attention attacking the cache instead.
Accept the cost and pay it, which is what most production models still do below about 32,000 tokens.

The distinction that gets muddled

Compute and memory are separate problems with separate fixes.

FlashAttention removes the $O(n^2)$ memory but keeps the $O(n^2 d)$ arithmetic, and it is still a large speedup because attention is bandwidth-bound rather than compute-bound.

Sparse attention removes arithmetic and changes what the model can represent. The first is free, the second is a trade.

Test yourself

One question on this card. Every option explains itself once you pick, right or wrong.

Question 1

What happens to the cost of self-attention when the context length doubles?

Transformers

19 concepts in this module, then a graded test with every answer explained.

Study this module →