The Quadratic Cost of Attention

Where the square comes from

For a sequence of $n$ tokens, each attention head computes the scores $QK^\top$: every query against every key, an $n \times n$ matrix (self-attention). The softmax runs over each row and the result multiplies $V$. The projections and the feed-forward layer handle each token separately, so they grow only linearly with $n$.

Dao and colleagues open the FlashAttention paper with this: Transformers are slow and memory-hungry on long sequences because the time and memory of self-attention are quadratic in sequence length Established.

Tiny example. At 8,192 tokens, one head's score matrix has $8{,}192^2 \approx 67$ million entries, 134 MB in 16 bits, for one head in one layer. At 32,768 tokens it is 16 times that. Storing all of them for the backward pass is how attention used to run out of memory.

Training versus decoding

During training and prefill, all $n$ positions are processed together and the full matrix is computed. During decoding with a KV cache, the new token's single query meets $n$ cached keys: one row of the matrix, linear in $n$ per token. Summed over a long generation the quadratic total comes back, and the cache itself grows with $n$.

Three ways out

1. Don't store the matrix. Compute exact attention in tiles that never write the $n \times n$ scores to slow memory: FlashAttention.
2. Store less per token. Shrink the cache that long contexts fill: grouped-query attention, quantization.
3. Change attention. Attend to a local window, a sparse pattern or a compressed summary. The FlashAttention authors observe that approximate attention methods reduce arithmetic but often don't achieve a wall-clock speed-up Established.