Attention's quadratic complexity, and why context length is expensive

The problem, from first principles

[04-seq-modeling/04-attention/01-scaled-dot-product-attention]'s core computation, scores = Q @ K^T, produces a (seq_len, seq_len) matrix: EVERY position's query compared against EVERY position's key. Doubling the sequence length doesn't just double the amount of work this involves, it QUADRUPLES it, since both the number of queries AND the number of keys being compared against doubled simultaneously. This is fundamentally different from [01-transformer-block/04-feedforward-sublayer]'s FFN, which processes each position independently: doubling the sequence length there simply doubles the number of (identical, independent) positions to process, a LINEAR cost.

This asymmetry has real, practical consequences. At the SHORT sequence lengths most models were historically trained at, the FFN's cost usually dominates (it has a large HIDDEN dimension, d_ff, typically 4x d_model, doing real work per position). But as sequence length grows, attention's quadratic term eventually overtakes it, and at the very long context lengths modern LLMs increasingly support (hundreds of thousands of tokens), attention can become THE dominant cost of running the model at all, both in raw compute AND in the memory needed to hold the (seq_len, seq_len) attention matrix itself. This is exactly the pressure [04-flash-attention] and [05-sliding-window-attention], immediately following this question, are each independently designed to relieve.

From theory to code

Implement three cost-estimation functions: attention_compute_cost(seq_len, d_model), the multiply-add count for attention's two big matrix multiplies (Q @ K^T and weights @ V); attention_memory_elements(seq_len, num_heads), the number of scalars in the fully materialized attention weight tensor; and ffn_compute_cost(seq_len, d_model, d_ff), the multiply-add count for the feed-forward sublayer's two linear layers.

Constraints

• attention_compute_cost scales as seq_len^2 * d_model (QUADRATIC in sequence length): 2 (two matrix multiplies of equal size) * seq_len * seq_len * d_model.
• attention_memory_elements scales as seq_len^2 (also quadratic): num_heads * seq_len * seq_len.
• ffn_compute_cost scales as seq_len * d_model * d_ff (LINEAR in sequence length): 2 * seq_len * d_model * d_ff.

Hints