Visual walkthrough — Computational complexity of attention

Before we start, three plain-English words we will use constantly. No symbol appears until it is drawn.

So $X$ is just a table: one row per word, one column per feature. Every figure below starts from this table.

---

Step 1 — What a matrix multiply costs

WHAT. Everything in attention is matrix multiplication. So before anything else we must agree on the price of ONE multiply. Take a grid with $a$ rows and $b$ columns, times a grid with $b$ rows and $c$ columns. The rule of matrix multiply: to fill one output cell, you slide along a full row of the left grid and a full column of the right grid, multiplying pairs and adding them up.

WHY this tool. We count multiplications because on a computer/GPU the multiply-add is the unit of work — a "FLOP" (floating-point operation). Counting FLOPs is how we compare architectures fairly, independent of which chip runs them.

PICTURE. In the figure, one output cell (coral) needs one row (lavender) dotted with one column (mint). That row/column each have length $b$, so one cell costs $\approx b$ multiply–adds. There are $a\times c$ cells to fill.

Multiply those two: total cost $=a\cdot b\cdot c$. Memorise this; every later step is just plugging numbers into $a\cdot b\cdot c$.

---

Step 2 — Turning tokens into Q, K, V

WHAT. From the input table $X$ (size $n\times d$) we make three new tables by multiplying $X$ with three learned weight grids $W_Q,W_K,W_V$. In multi-head attention each of these is a single $d\times d$ block that internally holds all $h$ heads side by side — column-block $g$ is head $g$'s slice of width $d_k=d/h$. So one matrix multiply $X{W}_Q$ produces every head's queries at once; we never do $h$ separate small multiplies. Here $d_k$ is the head dimension — the width each token gets inside one attention head.

$Q=X{W}_Q,K=X{W}_K,V=X{W}_V$

• $X$ ::: input, $n$ rows (tokens) by $d$ columns (features)
• $W_Q,W_K,W_V$ ::: fixed learned grids, $d$ tall by $d$ wide — but read as $h$ side-by-side $d\times d_k$ blocks, one per head
• $Q,K,V$ ::: full outputs $n\times d$; slicing off head $g$'s columns gives that head's $n\times d_k$ query/key/value table

WHY. A query is "what am I looking for?", a key is "what do I offer?", a value is "what I'll hand over if chosen". We need all three before any comparison happens. This step is self-attention's setup, and the block view is exactly how multi-head attention packs heads into one multiply.

PICTURE. Three arrows fan out of $X$, each producing an $n\times d$ table sliced into $h$ head-blocks of width $d_k$.

Cost, using Step 1's rule with $a=n,b=d,c=d$ (the full block does all heads): $\text{one projection}=n\cdot d\cdot d,\text{three of them}=3n{d}^2.$

Notice: this grows once with $n$. Double the words, double the work. That is linear — gentle. Hold that thought; it is NOT where the pain comes from. (If you only look at a single head, its slice costs $nd{d}_k$; summing over all $h$ heads gives $h\cdot nd{d}_k=n{d}^2$ again — same number, since $h{d}_k=d$.)

---

Step 3 — Every token meets every token: $S=Q{K}^T$

WHAT. Now the heart, done inside one head. We multiply that head's $Q$ (size $n\times d_k$) by $K^T$ — the "T" means transpose, flipping $K$ on its side so its shape becomes $d_k\times n$.

$S=Q{K}^T,S\in {\mathbb{R}}^{n\times n}$

• $Q$ ::: $n$ query rows, each of length $d_k$
• $K^T$ ::: keys flipped, so now $d_k$ tall by $n$ wide
• $S$ ::: the score matrix, $n\times n$ — cell $S_{ij}$ = how much token $i$ cares about token $j$

WHY transpose? We want the dot product of query $i$ with key $j$ (a dot product measures agreement — big when two vectors point the same way). To line up "row $i$ of $Q$" against "row $j$ of $K$" as a matrix multiply, $K$ must be turned sideways so its rows become columns. That is exactly what $K^T$ does.

PICTURE. This is the moment the output stops being tall-thin and becomes a full $n\times n$ square — one cell for every ordered pair of tokens.

Cost, Step 1 rule with $a=n,b=d_k,c=n$: $n\cdot d_k\cdot n=n^2{d}_k\text{per head}.$

There it is — the quadratic term $n^2{d}_k$. Why unavoidable? The output has $n^2$ cells because we insisted every token gets a score against every other token. The square shape is the cost.

---

Step 4 — Softmax: turning scores into weights

WHAT. Each row of $S$ becomes a set of proper attention weights. For row $i$: $A_{ij}=\frac{\exp (S_{ij}/\sqrt{d_k})}{\sum _{k=1}^n\exp (S_{ik}/\sqrt{d_k})}$

• $\exp (\cdot )$ ::: makes every score positive (bigger score → much bigger value)
• $\sqrt{d_k}$ ::: shrinks the scores first so the exponentials don't blow up (a stability scale)
• denominator sum ::: adds up the whole row, so the weights end up summing to $1$
• $A_{ij}$ ::: final weight — the share of attention token $i$ gives token $j$

WHY. Raw scores can be any size, positive or negative. We need each token's attention to be a clean set of shares (all $\ge 0$, summing to $1$). Softmax is the tool that answers "turn these numbers into a probability split".

PICTURE. A single square row lights up: big-score cells become bright, tiny-score cells fade, and the whole row's brightness sums to one bar of length 1.

Cost. For each of the $n^2$ cells we do a constant amount of work: one division by $\sqrt{d_k}$, one $\exp$, and one accumulation into its row-sum, then one final divide. That is $O(1)$ per cell — crucially, the head dimension $d_k$ does not appear here (we already collapsed it away in Step 3). So softmax costs $\Theta (n^2\cdot 1)=\mathcal{O}(n^2)$: the same square, touched a constant number of times, with no factor of $d_k$ and no new asymptotic term.

---

Step 5 — Collecting the answer: $\text{Output}=AV$

WHAT. Multiply the weight square $A$ (size $n\times n$) by the value table $V$ (size $n\times d_k$).

$\text{Output}=AV\in {\mathbb{R}}^{n\times d_k}$

• $A$ ::: attention weights, $n\times n$
• $V$ ::: value rows, $n\times d_k$
• output row $i$ ::: a blend of all value rows, mixed by row $i$'s weights

WHY. Now that token $i$ knows how much to trust each other token (that's $A$), it builds its new representation as a weighted average of everyone's values. This is the payoff of the whole mechanism.

PICTURE. Each output row is a weighted stack of value rows; heavy weights pull their value rows in strongly.

Cost, Step 1 rule with $a=n,b=n,c=d_k$: $n\cdot n\cdot d_k=n^2{d}_k.$

The second quadratic term — same size as Step 3. So both "compare" and "collect" cost $n^2{d}_k$.

---

Step 6 — The output projection $W_O$ (the term people forget)

WHAT. After the heads are glued back into a full $d$-wide table $Z$ (size $n\times d$), one last multiply with $W_O$ (size $d\times d$) mixes the heads' information: $\text{Out}=Z{W}_O\in {\mathbb{R}}^{n\times d}.$

• $Z$ ::: concatenated heads, $n\times d$
• $W_O$ ::: learned mixing grid, $d\times d$
• $\text{Out}$ ::: final attention output, $n\times d$

WHY. Each head worked in its own little $d_k$-slice. $W_O$ lets the heads talk to each other and returns to full model width. See multi-head attention for what the heads are.

PICTURE. Head-slices concatenated into one wide bar, then squeezed through the square $W_O$.

Cost, Step 1 rule with $a=n,b=d,c=d$: $n\cdot d\cdot d=n{d}^2.$

Together with the Q, K, V projections of Step 2 (each $n{d}^2$), the linear part is about $4n{d}^2$, i.e. $\mathcal{O}(n{d}^2)$.

---

Step 7 — Adding it all up (and the multi-head twist)

WHAT. Sum every step's price (attention terms summed over all $h$ heads).

Step	What	Cost
2	Q, K, V projections (all heads, one multiply each)	$3n{d}^2$
3	scores $Q{K}^T$ (per head $\times h$)	$h{n}^2{d}_k=n^{2}d$
4	softmax (per entry $O(1)$, no $d_k$)	$\Theta (n^2)$
5	collect $AV$ (per head $\times h$)	$h{n}^2{d}_k=n^{2}d$
6	output projection	$n{d}^2$

WHY the two survivors. Softmax's $\Theta (n^2)$ is lower-order than $n^{2}d$ (no $d_k$ factor), so it drops out. Keep only the biggest terms: $\boxed{\mathcal{O}(n^{2}d+n{d}^2)}$

Now the multi-head twist. With $h$ heads, each head uses $d_k=d/h$. The attention work across all heads is $h\cdot n^2{d}_k=h\cdot n^2\cdot \frac{d}{h}=n^{2}d.$ The $h$ cancels! And the projections are done once over the full $d$ (recall Step 2: a single $d\times d$ multiply already outputs every head), not once per head. So: $\boxed{\text{Total}=\mathcal{O}(n^{2}d+n{d}^2)}$

PICTURE. A bar chart: for small $n$ the $n{d}^2$ bar towers; as $n$ grows the $n^{2}d$ bar overtakes it — the crossover that forces long-context models toward sparse attention and memory-optimization tricks.

---

The one-picture summary

Read it left to right: the tall-thin $X$ becomes tall-thin $Q,K,V$ (cheap, linear), those crash together into a fat $n\times n$ square (the quadratic villain), softmax paints the square, the square collects values back into a tall-thin output, and $W_O$ mixes it. The square is the whole story.

Recall Feynman retelling — say it in plain words

Imagine a room of $n$ people (tokens). First everyone writes three cards: a "want" card, an "offer" card, and a "give" card. Writing cards is quick — it scales with how many people there are, one round of writing. Then comes the expensive part: every person compares their want-card against everyone's offer-card. That's $n$ people each doing $n$ comparisons — an $n\times n$ grid of handshakes. Double the crowd and the handshakes quadruple. After comparing, each person mixes everyone's give-cards in proportion to how much they liked them — another full $n\times n$ pass. Finally a group huddle ($W_O$) blends what the different sub-groups (heads) found. The huddle and the card-writing are linear (cheap); the two $n\times n$ handshake grids are the quadratic cost that makes long documents so painful.

Recall Quick self-check

Where does the $n^2$ term come from? ::: From $S=Q{K}^T$ producing an $n\times n$ score for every token pair, and again from $AV$ walking that same square. Why don't $h$ heads multiply the cost? ::: Each head uses $d/h$ width, so $h\cdot n^2(d/h)=n^{2}d$ — the $h$ cancels; projections run once over full $d$ in a single $d\times d$ multiply. Which term dominates for very long sequences? ::: $n^{2}d$ (the attention square), because $n^2$ outgrows the linear $n{d}^2$ once $n>d$. Why doesn't softmax add a $d_k$ factor? ::: It does $O(1)$ work per already-computed score cell, so it is $\Theta (n^2)$ — lower-order than $n^2{d}_k$.

Related scaling concern for long inputs: positional encoding scaling.