Computational complexity of attention: Exercises

Definition The symbols we will reuse (read this once)

Everything on this page uses just five letters plus three named matrices. Let us earn each one:

$n$ ::: the sequence length — how many tokens (words/pieces) we feed in at once. Picture a row of $n$ boxes.
$d$ ::: the model dimension — the length of the vector that represents each token. Each box holds a list of $d$ numbers.
$d_{k}$ ::: the per-head dimension — when we split work across heads, each head only looks at $d_{k}$ of those numbers.
$h$ ::: the number of heads — parallel attention "cameras", each seeing a $d_{k}$ -slice.
$O (\cdot)$ ::: "order of"— we count operations but ignore constant multipliers, keeping only the term that grows fastest as $n$ gets huge. $O (n^{2})$ means "grows like $n$ squared".
$Q, K, V$ ::: the three matrices attention builds from the input. Q (query) = "what am I looking for?", K (key) = "what do I offer?", V (value) = "what I actually pass on if chosen". A score is a query-meets-key comparison; the output is a value-weighted sum. Each is made by one linear projection of the input. See 4.1.1-Self-attention-mechanism.

The two costs you will compute over and over: $Time = O (n^{2} d + n d^{2}), Memory = O (n^{2} + n d)$ The first time-term ( $n^{2} d$ ) is the all-pairs attention cost; the second ( $n d^{2}$ ) is the linear projections (Q, K, V, and output). For memory, $n^{2}$ is the score matrix and $n d$ is the stored token-vector activations. See 4.1.1-Self-attention-mechanism for where they come from.

Figure — Computational complexity of attention

Level 1 — Recognition

Goal: can you read off and plug into the formula?

L1.1 — Which term is which?

State, in words, what the two terms $n^{2} d$ and $n d^{2}$ in the time complexity $O (n^{2} d + n d^{2})$ each represent. Name which of $Q, K, V$ contribute to each.

Recall Solution

$n^{2} d$ ::: the attention cost — every one of the $n$ tokens forms a score with every other of the $n$ tokens ( $n^{2}$ pairs), and each score touches a $d$ -length vector. This is the $Q K^{⊤}$ score step and the $V$ -weighting step.
$n d^{2}$ ::: the linear projection cost — for each of the $n$ tokens we multiply its $d$ -vector by a $d \times d$ weight matrix (that costs $d^{2}$ ), and we do this to build Q, K, V and to apply the output projection. The word "quadratic bottleneck" always refers to the first term because of its $n^{2}$ .

L1.2 — Plug in the numbers

For $n = 1000$ and $d = 512$ , compute the two time-terms $n^{2} d$ and $n d^{2}$ (as raw operation counts), and say which dominates. Locate this point on the figure.

Recall Solution

$n^{2} d = 100 0^{2} \times 512 = 1 0^{6} \times 512 = 512, 000, 000 \approx 512 M$ $n d^{2} = 1000 \times 51 2^{2} = 1000 \times 262, 144 = 262, 144, 000 \approx 262 M$ The attention term ( $512$ M) dominates the projection term ( $262$ M). On the figure: $n = 1000$ sits to the right of the yellow crossover line ( $n = d = 512$ ), where the red parabola has already climbed above the blue line — exactly the region attention wins.

L1.3 — Memory of one score matrix

A single attention head stores the full score matrix $S \in R^{n \times n}$ in FP32 (4 bytes each). For $n = 2048$ , how many megabytes is that? (Use $1 MB = 2^{20}$ bytes.)

Recall Solution

Number of floats $= 204 8^{2} = 4, 194, 304$ . Bytes $= 4, 194, 304 \times 4 = 16, 777, 216$ bytes. $\frac{16 , 777 , 216}{2 ^{20}} = \frac{16 , 777 , 216}{1 , 048 , 576} = 16 MB$ So 16 MB per head — this is the $O (n^{2})$ memory term, kept for backpropagation, and it is why long sequences explode memory. See 4.5.1-Memory-optimization-techniques.

Level 2 — Application

Goal: run the formula on realistic model configs.

L2.1 — Crossover point

At what sequence length $n$ does the attention term $n^{2} d$ equal the projection term $n d^{2}$ (treating $d_{k} \approx d$ , single-head form)? Interpret the answer.

Recall Solution

Set them equal. Since $n > 0$ (a sequence has at least one token) we are allowed to divide by $n$ : $n^{2} d = n d^{2} ⟹ \frac{n ^{2} d}{n} = \frac{n d ^{2}}{n} ⟹ n d = d^{2} ⟹ n = d .$ (Dividing by $n$ is legal precisely because $n \neq = 0$ ; likewise $d > 0$ .) Interpretation: the crossover happens exactly at $n = d$ . For $n < d$ (short sequences) projections dominate; for $n > d$ (long sequences) attention dominates. This is why long-context models are the ones that fear the $n^{2}$ term.

L2.2 — GPT-2 Medium FLOPs

GPT-2 Medium uses $d = 1024$ , with $h = 12$ heads. Note that in this model the per-head width is $d_{k} = 64$ , so $h \cdot d_{k} = 768$ — smaller than $d = 1024$ . This is a real design choice (the concatenated head output is $768$ , then a projection maps it back up/into the $1024$ -wide residual stream); it is not the idealised $h d_{k} = d$ we assumed in L2.1. So here we use the parent's exact convention: attention $= h n^{2} d_{k}$ and projections $= 4 n d^{2}$ (with the full $d = 1024$ ). For $n = 512$ , compute (a) attention FLOPs $= h n^{2} d_{k}$ , and (b) projection FLOPs $= 4 n d^{2}$ .

Recall Solution

(a) Attention: $h n^{2} d_{k} = 12 \times 51 2^{2} \times 64 = 12 \times 262, 144 \times 64 = 201, 326, 592 \approx 201 M .$ (b) Projections (use full $d = 1024$ , since projections act on the model dimension, not the per-head width): $4 n d^{2} = 4 \times 512 \times 102 4^{2} = 4 \times 512 \times 1, 048, 576 = 2, 147, 483, 648 \approx 2.1 B .$ Ratio attention/projection $= 201 M /2.1 B \approx 0.094 \approx 0.10 \times$ . Projections still dominate at $n = 512$ — consistent with the parent table.

L2.3 — Doubling the sequence

You measured that attention takes 10 ms at $n = 1024$ . Estimate the attention time at $n = 4096$ , assuming everything else fixed.

Recall Solution

Attention scales as $n^{2}$ . Going from $1024 \to 4096$ multiplies $n$ by $4$ , so time multiplies by $4^{2} = 16$ . $10 ms \times 16 = 160 ms .$ What it looks like: a 4× longer input is a 16× taller point on the red parabola of the figure — this is the whole reason 4.3.2-Sparse-attention-patterns exist.

Level 3 — Analysis

Goal: reason about why a term behaves as it does.

L3.1 — Multi-head asymptotics

Show algebraically that $h$ -head attention (in the idealised case $h d_{k} = d$ ) has the same asymptotic time complexity as single-head attention.

Recall Solution

Single head, full dimension: attention cost $= n^{2} d$ .
$h$ heads, each dimension $d_{k} = d / h$ : each head costs $n^{2} d_{k}$ , and there are $h$ of them: $h \cdot n^{2} d_{k} = h \cdot n^{2} \cdot \frac{d}{h} = n^{2} d .$ The $h$ cancels (legal since $h \geq 1 \neq = 0$ ). Projections are done once over $d$ , giving $n d^{2}$ in both cases. Therefore both are $O (n^{2} d + n d^{2})$ . Splitting into heads reshapes the same work; it does not multiply it.

L3.2 — Memory budget failure

A training run has a GPU with 8 GB free for attention activations. Using the parent's figure of $\approx 16$ MB per head at $n = 2048$ , and a model with 12 heads over 12 layers, does the forward-pass attention memory fit? By how much does it miss/spare?

Recall Solution

Per-head-per-layer $= 16$ MB. Total heads-layers $= 12 \times 12 = 144$ . $16 MB \times 144 = 2304 MB \approx 2.25 GB .$ $2.25$ GB $<$ $8$ GB, so it fits, with $8 - 2.25 = 5.75$ GB to spare — but only for the attention scores; weights, optimizer states, and other activations eat the rest. This is why gradient checkpointing exists.

L3.3 — The other memory term: $O (n d)$ activations

Besides the $n \times n$ score matrix, each layer must also keep the token-vector activations $X \in R^{n \times d}$ (this is the $O (n d)$ memory term). For $n = 2048$ , $d = 1024$ , FP32: (a) how many MB is one such activation tensor? (b) Compare it to the $16$ MB score matrix of L1.3 — which term dominates here, and at what $n$ would they be equal?

Recall Solution

(a) Floats $= n \cdot d = 2048 \times 1024 = 2, 097, 152$ . Bytes $= \times 4 = 8, 388, 608$ . In MB: $8, 388, 608/ 2^{20} = 8$ MB. (b) The score matrix is $16$ MB vs the activation $8$ MB, so the $n^{2}$ score term already dominates at $n = 2048$ . They are equal when $n^{2} = n d$ , i.e. $n = d = 1024$ (using $n > 0$ to divide by $n$ ). For $n < 1024$ the $n d$ activations dominate; for $n > 1024$ the $n^{2}$ scores do. This is the memory version of the L2.1 time crossover — same $n = d$ pivot.

L3.4 — CNN vs attention receptive field

A CNN with kernel $k$ costs $O (nk d^{2})$ and connects tokens only $k$ apart per layer. Self-attention costs $O (n^{2} d)$ but connects any two tokens in one layer. For $n = 8192$ , $d = 1024$ , $k = 3$ , which is cheaper, and what does self-attention buy for the extra cost?

Recall Solution

CNN: $nk d^{2} = 8192 \times 3 \times 102 4^{2} = 8192 \times 3 \times 1, 048, 576 = 25, 769, 803, 776 \approx 25.8 B$ . Attention: $n^{2} d = 819 2^{2} \times 1024 = 67, 108, 864 \times 1024 = 68, 719, 476, 736 \approx 68.7 B$ . The CNN is cheaper ( $\approx 25.8$ B vs $68.7$ B). What attention buys: a depth-1 path between any pair of tokens (global receptive field instantly), whereas the CNN needs many stacked layers to connect distant tokens. You pay ~2.7× more compute for global context.

Level 4 — Synthesis

Goal: combine terms, choose regimes, design.

L4.1 — Total cost with all terms

Write the full per-layer forward FLOP count combining attention and all four projections (use $n^{2} d$ for attention and $4 n d^{2}$ for projections), then evaluate for $n = 2048$ , $d = 1024$ .

Recall Solution

$Total \approx n^{2} d + 4 n d^{2} .$ Attention: $n^{2} d = 204 8^{2} \times 1024 = 4, 194, 304 \times 1024 = 4, 294, 967, 296 \approx 4.29 B$ . Projections: $4 n d^{2} = 4 \times 2048 \times 102 4^{2} = 4 \times 2048 \times 1, 048, 576 = 8, 589, 934, 592 \approx 8.59 B$ . Total $\approx 4.29 B + 8.59 B = 12.88 B$ FLOPs. Projections still lead at $n = 2048$ , but attention is closing in (ratio $\approx 0.5$ ).

L4.2 — When does sparse attention win?

Sparse attention costs $O (n n d)$ instead of $n^{2} d$ . Find the range of $n$ for which sparse attention does less attention-work than dense, given both use the same $d$ . (Ignore constants.)

Recall Solution

Compare $n n d$ vs $n^{2} d$ ; since $n > 0$ and $d > 0$ we may cancel the common $n d$ : $n < n .$ Square both sides (both positive): $n < n^{2}$ , i.e. $1 < n$ . So for every $n > 1$ the sparse form has lower asymptotic order — but the hidden constants in sparse attention are larger, so in practice it only pays off once $n$ is large enough (typically thousands of tokens) that the $n^{2}$ growth overwhelms those constants. See 4.3.2-Sparse-attention-patterns.

L4.3 — Budget-driven design

You are limited to $2 B$ attention FLOPs per layer ( $n^{2} d$ form) and must support $d = 512$ . What is the largest sequence length $n$ you can serve with dense attention? What if you switch to sparse $n n d$ ?

Recall Solution

Dense: $n^{2} d \leq 2 \times 1 0^{9} \Rightarrow n^{2} \leq \frac{2 \times 1 0 ^{9}}{512} = 3, 906, 250 \Rightarrow n \leq 3, 906, 250 \approx 1976$ . So $n_{m a x}^{dense} \approx 1976$ . Sparse: $n^{1.5} d \leq 2 \times 1 0^{9} \Rightarrow n^{1.5} \leq 3, 906, 250 \Rightarrow n \leq 3, 906, 25 0^{2/3} \approx 24, 803$ . Sparse lets you serve roughly 12.5× longer sequences under the same budget. That factor grows with the budget — the whole point of sub-quadratic attention.

Level 5 — Mastery

Goal: derive a novel result / prove a bound.

L5.1 — Prove the crossover is head-independent

Prove that the $n$ at which attention cost equals projection cost does not depend on the number of heads $h$ .

Recall Solution

Attention cost with $h$ heads: $h n^{2} d_{k} = n^{2} d$ (since $h d_{k} = d$ in the idealised case) — the $h$ has already vanished. Projection cost: $n d^{2}$ , applied once over $d$ , independent of $h$ . Set equal and divide by $n$ (legal since $n > 0$ ): $n^{2} d = n d^{2} \Rightarrow n = d$ . Neither side ever contained $h$ , so the crossover $n = d$ is head-independent. $■$ This is the rigorous version of the parent's "multi-head has the same asymptotic complexity."

L5.2 — Amortised memory with checkpointing

Naive training stores all $L$ layers' $n \times n$ score matrices: memory $O (L n^{2})$ . With gradient checkpointing you store only $L$ checkpoints and recompute the rest. Give the checkpointed memory and the extra compute cost, and evaluate memory saving for $L = 144$ layer-heads at $n = 2048$ (16 MB each).

Recall Solution

Checkpointed memory $= O (L n^{2})$ ; extra compute $= O (1)$ extra forward pass, i.e. roughly $2 \times$ the forward FLOPs.
Naive memory $= 144 \times 16 MB = 2304 MB$ .
Checkpointed $= 144 \times 16 MB = 12 \times 16 = 192 MB$ .
Saving factor $= 2304/192 = 12 \times$ (exactly $L$ ). You trade one recomputation for a $L$ -fold memory cut — the standard long-context bargain of 4.5.1-Memory-optimization-techniques.

L5.3 — Positional-encoding scaling sanity check

Attention itself is permutation-blind, so we add positional information whose storage is $O (n d)$ . Show this term never changes the asymptotic time or memory dominated by attention for long sequences, and state the one regime where it could matter.

Recall Solution

Time: attention costs $n^{2} d$ while positional work is $O (n d)$ . Their ratio is $\frac{n ^{2} d}{n d} = n$ (dividing by $n d$ is legal since $n, d > 0$ ). So attention beats positional work by a factor of $n$ , which grows without bound — the $O (n d)$ term is asymptotically absorbed for every $n > 1$ .
Memory: attention scores are $O (n^{2})$ vs positional storage $O (n d)$ . The ratio $n^{2} / (n d) = n / d$ , so once $n > d$ the score term dominates and positional storage is negligible.
The one regime where it matters: very short sequences $n \leq d$ (equivalently $n ≪ d$ ). There $n^{2}$ is small, so the linear-in- $n$ terms (projections and positional storage) lead the cost; positional-encoding design — especially extrapolation to lengths unseen in training — then affects both output quality and the leading constant factor. See 5.2.3-Positional-encoding-scaling. $■$

Recall Quick self-test (close the page, answer these)

The two time-terms of attention are ::: $n^{2} d$ (attention) and $n d^{2}$ (projections). Q, K, V stand for ::: query ("what I look for"), key ("what I offer"), value ("what I pass on"). Attention overtakes projections at $n$ equal to ::: $d$ . The two memory-terms $n^{2} + n d$ cross over at $n$ equal to ::: $d$ (same pivot as time). Doubling $n$ multiplies attention cost by ::: $4$ . Multi-head changes the asymptotic cost by ::: nothing — it stays $O (n^{2} d + n d^{2})$ . Sparse attention's order is ::: $O (n n d)$ .