Skip to content
datarekha
Deep Learning Hard Asked at GoogleAsked at OpenAIAsked at MetaAsked at Anthropic

Why is standard self-attention O(n^2) in sequence length, and how is it addressed?

The short answer

Standard self-attention compares every token with every other token, so its dense score matrix has n squared entries and its compute and memory grow quadratically with sequence length. FlashAttention removes the large intermediate but not the quadratic compute, while sparse, local, low-rank, linear, retrieval, and alternative sequence architectures trade exact all-to-all attention for lower cost.

How to think about it

Standard self-attention is O(n²) in sequence length because each of n tokens compares its query with the key of every one of n tokens, creating pairwise scores. Systems address this by avoiding the large intermediate, restricting which pairs can interact, approximating the attention calculation, or replacing full attention with another sequence architecture.

Why the square appears

Let n be the number of tokens and d_model the width of the model representation for each token. A transformer turns every token representation into three vectors:

  • A query describes what this token is looking for.
  • A key describes what this token offers for matching.
  • A value contains the information passed onward if the match is useful.

For one attention head, the calculation is commonly written as:

Attention(Q, K, V) = softmax(QK^T / sqrt(d_h))V

Here, d_h is the width of one head. Q, K, and V each have shape n × d_h.

The important operation is QK^T. Every row of Q is compared with every row of K. There are n queries and n keys, so the result contains n × n scores. The score at row i, column j answers a specific question: how much should token i attend to token j?

Each score is a dot product of length d_h, so producing the score matrix costs O(n²d_h) operations per head. Across all heads, that is O(n²d_model). The following multiplication by V has the same order of cost. Softmax also touches all scores.

The memory problem comes from the same place. A dense implementation materialises the score matrix, and often the post-softmax attention matrix as well. With h heads, the storage is O(hn²) values. When h is treated as fixed, people shorten that to O(n²).

A causal mask does not change this conclusion. It prevents a token from looking into the future, leaving roughly half the matrix usable: n(n + 1) / 2 entries instead of . Half of a quadratic is still quadratic.

Common misconception: the quadratic part is in sequence length, not that model dimension is irrelevant. The full attention contraction is O(n²d_model), so increasing d_model increases the work too. But it does not make the sequence-length scaling quadratic in d_model. For short sequences and a very wide model, projections such as O(nd_model²) can even dominate. For long sequences, the pairwise attention work and its activations become the problem.

The underlying calculation is the one shown in self-attention.

A concrete 32,000-token example

Imagine a document assistant reading a 32,000-token contract. Use a model with d_model = 4096, 32 heads, and therefore d_h = 128 per head. Consider one request and one transformer layer.

The score matrix contains:

32,000 × 32,000 = 1,024,000,000

entries.

With 32-bit floating-point values, one head’s score matrix needs:

1,024,000,000 × 4 = 4,096,000,000 bytes

That is about 3.81 GiB for one head. Across 32 heads, the dense score tensor needs about 122.1 GiB, before counting the softmax output, masks, Q, K, V, model weights, or workspace. A training run may also retain activations for the backward pass.

The arithmetic is substantial too. The QK^T multiplication alone performs about 4.19 trillion scalar multiply-accumulate terms across all heads. The PV multiplication is another operation of the same order.

A 32-times longer sequence creates 1,024 times as many token pairs. That is why long context becomes painful rather abruptly. A 1,000-token input has 1 million pairs per head. A 32,000-token input has 1.024 billion.

The inference wrinkle: the KV cache

Autoregressive generation has two phases.

During prefill, the model processes the prompt. Every prompt token can attend to the other allowed prompt tokens, so the prompt still incurs the full quadratic attention calculation.

During decode, the model generates one token at a time. It stores previous keys and values in a KV cache, which is memory holding the attention information for earlier tokens. The new token has one query, so it compares that query with t cached keys, where t is the current context length. One decode step is therefore linear in the current context length, not quadratic in that single step.

However, generating many tokens adds those costs together. If the context grows from n to n + m, the decode work includes roughly m × n comparisons plus the comparisons among the generated tokens. Per-token latency tends to worsen as the context gets longer.

The KV cache itself grows linearly with context length. For standard multi-head attention with d_model = 4096, a 32,000-token cache in float16 needs approximately:

2 × 32,000 × 4096 × 2 bytes

for K and V, or about 500 MiB per layer. A 32-layer model would need roughly 16 GiB for that cache alone. Grouped-query attention and multi-query attention reduce this cache by using fewer key and value heads, but they do not remove the quadratic prefill cost of full attention.

How production systems address it

ApproachWhat it changesTypical sequence scalingMain trade-off
FlashAttentionTiles exact attention without storing all scoresQuadratic compute, linear extra memoryDoes not remove pairwise arithmetic
Local or sliding-window attentionEach token sees only nearby tokensO(nw) for window wDistant tokens cannot interact directly
Block-sparse attentionComputes selected blocks and global linksRoughly O(nk) for k attended tokensRequires a useful sparsity pattern
Linear or kernel attentionReorders an approximate attention calculationLinear when feature size is fixedNot generally the same as softmax attention
Low-rank attentionCompresses keys and values into r componentsOften O(nr)Assumes attention has exploitable low rank
Retrieval and chunkingReduces the number of tokens shown to the modelDense attention on a smaller nRetrieval can miss the needed evidence

FlashAttention: exact, but not magically linear

FlashAttention is often the first answer for ordinary long-context transformer serving because it preserves the same mathematical attention, apart from floating-point rounding.

A naive implementation writes the entire n × n score matrix to high-bandwidth GPU memory, reads it back for softmax, then reads it again for the value multiplication. FlashAttention divides Q, K, and V into tiles that fit in fast on-chip memory. It computes a tile, updates a running maximum and normalisation factor for the softmax, and accumulates the output without ever storing the full score matrix.

That reduces memory traffic and removes the quadratic intermediate. The memory needed for the attention workspace becomes linear in n, alongside the n × d_model inputs and output. It can also be faster because moving data to and from GPU memory is often the bottleneck.

It does not reduce the number of query-key pairs. The algorithm still computes dense attention, so its arithmetic remains O(n²). FlashAttention makes exact attention much more practical; it does not make arbitrary 1-million-token full attention cheap.

Local and sparse attention

With a sliding window of 1,024 tokens, a 32,000-token sequence has roughly 32,000 × 1,024, or about 32.8 million, attention pairs instead of 1.024 billion. That is roughly a 31-times reduction.

The price is connectivity. A token at the end of a document cannot directly attend to a definition at the beginning. Multiple layers can pass information across windows, and special global tokens can provide shortcuts, but neither is equivalent to every token seeing every other token in one layer.

This works well when the task is mostly local: nearby code, neighbouring words, or a document whose relevant evidence has already been retrieved. It is risky for questions involving distant references, tables, definitions, and “use clause 2 to interpret clause 47” reasoning.

Linear and low-rank approximations

Linear-attention methods change the algebra. Instead of computing all pairwise softmax scores, they use a feature map φ that approximates the similarity between a query and a key. The weighted sum can then be rearranged conceptually as:

φ(q_i)^T (sum over j of φ(k_j)v_j)

The accumulated key-value state can be updated as the sequence is scanned. For causal attention, this becomes a prefix computation, so cost is linear in n when the feature dimension is fixed.

The catch is important: ordinary softmax attention does not generally factor this way exactly. Approximation errors are most visible when the model needs sharp, selective retrieval rather than broad averaging.

Low-rank methods make a related bet: the full n × n attention pattern can be represented using far fewer landmarks or components. They reduce cost when that assumption fits the data. They can fail when attention contains many independent, sharp relationships.

The senior-level trade-off

There is no universally best replacement. If the required context fits comfortably on the target hardware, exact attention with FlashAttention is usually the safest choice because it preserves model behaviour. If a 32,000-token contract contains only 3,000 relevant tokens, retrieval and chunking may be better than forcing every token to attend to every other token. If the application needs continuous streaming over very long histories, a local, recurrent, or state-space architecture may be a better design, but it brings a different quality profile.

RoPE, or rotary positional embeddings, and ALiBi, or attention with linear biases, do not solve the quadratic problem. They change how position affects the score, but the model still computes the dense pairwise interactions. Quantisation reduces bytes and grouped-query attention reduces KV-cache size; neither changes dense prefill from quadratic in sequence length.

A common production failure appears first as a CUDA out of memory error during prompt prefill, even though one-token decoding worked in a small test. The cause is often a dense attention fallback that materialises scores, a larger-than-expected batch, or an accidentally duplicated mask. Another failure appears as a quality regression rather than an infrastructure alert: a sliding-window model answers a question using nearby text but misses the decisive definition 20 pages earlier.

I would benchmark peak memory, prefill latency, decode tokens per second, and long-range task accuracy on the actual hardware. Big-O tells you the direction. It does not tell you whether a kernel, batch size, cache layout, or retrieval miss will ruin Tuesday afternoon.

What they’ll ask next

Does FlashAttention make attention linear?
No. It makes memory use and memory traffic much better by avoiding the materialised score matrix. Dense pairwise computation remains quadratic.

Does the causal mask remove half the problem?
It removes future-token pairs, leaving about half the entries. That improves the constant factor, but the number of allowed pairs still grows quadratically.

What would you choose for a 100,000-token document?
First, reduce the effective context with retrieval, structure-aware chunking, or summaries. Use FlashAttention for the dense attention that remains. If the task truly requires arbitrary all-to-all interaction across all 100,000 tokens, I would consider sparse or alternative sequence architectures and validate them with cross-document, long-range evaluation rather than relying on short-context scores.

Say this in the interview: Standard attention is quadratic because n queries score against n keys; FlashAttention removes the quadratic memory bottleneck, while sparse, linear, low-rank, retrieval, and alternative architectures reduce or replace the quadratic interaction at the cost of exact global attention.

Learn it properly Self-attention

Keep practising

All Deep Learning questions

Explore further