Learn: Multi-Head Attention & the Transformer Block

What you're learning

How to go from one scaled dot-product attention (lesson 01-scaled-dot-product-attention) to the full transformer layer that LLMs actually stack. By the end you should be able to: explain why splitting attention into $h$ heads buys you more than one big head; write the full MHA formula and count its parameters; assemble the block (MHA → residual → LayerNorm, then FFN → residual → LN) and say why pre-LN beat post-LN for deep stacks; and derive the per-layer FLOPs/params that feed the $C\approx 6ND$ training-compute rule.

This lesson builds directly on 01-scaled-dot-product-attention — if the $\text{softmax}(Q{K}^{\top }/\sqrt{d})V$ mechanism and the $\Theta (N^{2}d)$ cost of one head aren’t automatic, reread it first.

---

1. Learning map

graph TD
    A["Scaled dot-product attention<br/>(lesson 01): one head"] --> B["Why more than one head?<br/>representation subspaces"]
    B --> C["Multi-head attention<br/>split d into h heads, concat, W^O"]
    C --> D["Sublayer wrapper<br/>residual + LayerNorm"]
    D --> E["FFN / MLP sublayer<br/>2 linears + GELU, d_ff = 4d"]
    D --> F["Pre-LN vs Post-LN<br/>(Xiong et al. 2020)"]
    C --> G["Transformer block<br/>attention sublayer + MLP sublayer"]
    E --> G
    F --> G
    G --> H["Params ~12d² / layer<br/>FLOPs ~24bsd² + 4bs²d"]
    H --> I["C ≈ 6ND training rule<br/>(cross-link roofline)"]

    style C fill:#44a,color:#fff
    style G fill:#4a4,color:#fff
    style H fill:#a83,color:#fff
    style I fill:#a44,color:#fff

---

2. Intuition — why many heads beat one big head

A single attention head produces one set of attention weights per query: one softmax distribution over the keys. That means one relationship, computed in one projected subspace. But language (and vision, and code) has many relationships holding simultaneously between the same two tokens: a verb agrees with its subject (syntax), a pronoun binds to its antecedent (coreference), a token attends to its immediate left neighbor (local/positional), a closing bracket looks back to its opener (structure). Forcing all of that through one softmax makes the head average these patterns together — and an average of a “look at the subject” distribution and a “look at the previous token” distribution is a blurry distribution that does neither well.¹

Multi-head attention’s move is embarrassingly simple: run $h$ independent scaled dot-product attentions in parallel, each in its own low-dimensional projected subspace, then concatenate their outputs and mix them with a learned projection. Each head gets its own $W^Q,W^K,W^V$, so each is free to specialize: one head can learn the projection that makes “attend to the syntactic head” easy, another the projection for “attend to the previous token.” Empirically, trained heads do fall into interpretable families — positional/previous-token heads, syntactic heads, rare-token and coreference heads — though the specialization is soft and heads are partly redundant.¹

Key insight — heads are cheap because they're narrow

The heads don’t cost extra: we take the same $d_{model}$ and split it into $h$ slices of width $d_h=d_{model}/h$. So $h$ heads of dim $d_h$ cost (almost) exactly the same parameters and FLOPs as one head of dim $d_{model}$ — you get parallel, specialized views of the sequence essentially for free. Multi-head is a reshape, not an enlargement.

Common misconception — "more heads = more capacity"

Adding heads while holding $d_{model}$ fixed does not add parameters or FLOPs; it repartitions the same budget into more, narrower subspaces. Push $h$ too high and each $d_h$ becomes so small the per-head softmax loses expressiveness (too few dimensions to separate keys). There’s a sweet spot; typical LLMs use $d_h=64$ or $128$. (This is also exactly the tension GQA/MQA later exploit — sharing K/V across heads to shrink the KV-cache. Out of scope here.)

---

3. Math — multi-head attention and the block

3.1 Multi-head attention

Let the input be $X\in {\mathbb{R}}^{s\times d}$ ($s$ = sequence length, $d=d_{model}$). Choose $h$ heads, each of dimension $d_h=d/h$. For head $i\in \{1,\dots ,h\}$, project into the head’s subspace with learned matrices $W_i^Q,W_i^K,W_i^V\in {\mathbb{R}}^{d\times d_h}$:

$$Q_i=X{W}_i^Q,K_i=X{W}_i^K,V_i=X{W}_i^V\in {\mathbb{R}}^{s\times d_h}.$$

Each head runs the same scaled dot-product attention from lesson 01, but at the reduced dimension $d_h$ (note the scaling uses $d_h$, not $d$):

$${\text{head}}_i=\text{Attention}(Q_i,K_i,V_i)=\text{softmax}\left(\frac{Q_i{K}_i^{\top }}{\sqrt{d_h}}\right)V_i\in {\mathbb{R}}^{s\times d_h}.$$

Concatenate the $h$ head outputs back to width $d$ and apply the output projection $W^O\in {\mathbb{R}}^{d\times d}$:

$$\boxed{\text{MHA}(X)=\text{Concat}({\text{head}}_1,\dots ,{\text{head}}_h)W^O\in {\mathbb{R}}^{s\times d}.}$$

In practice the $h$ per-head projections are packed into single matrices $W^Q,W^K,W^V\in {\mathbb{R}}^{d\times d}$ (one big GEMM each), the result reshaped to $(s,h,d_h)$, attention done batched over the head axis, then reshaped back — but the math above is what that computes.

What $W^O$ is for

Concatenation just stacks the heads side by side; nothing has mixed information across heads yet. $W^O$ is the learned mixer that lets the layer combine “what head 3 found” with “what head 7 found” into each output coordinate. Dropping $W^O$ (or making it block-diagonal) would leave the heads siloed. It also maps the concatenated $d$-vector back into the residual stream’s basis.

3.2 The FFN / MLP sublayer

After attention mixes information across positions, the position-wise feed-forward network mixes it across features, independently at each position. Two linear layers with a nonlinearity, expanding to an inner width $d_{ff}$ (canonically $d_{ff}=4d$):¹

$$\text{FFN}(x)=\varphi \left(x{W}_1+b_1\right)W_2+b_2,W_1\in {\mathbb{R}}^{d\times d_{ff}},W_2\in {\mathbb{R}}^{d_{ff}\times d}.$$

The original transformer used ReLU for $\varphi$; modern LLMs use GELU² (a smooth, probabilistically-motivated gate, $\text{GELU}(x)=x\Phi (x)$ with $\Phi$ the standard-normal CDF) or gated variants (SwiGLU). The FFN holds most of a transformer’s parameters — $8d^2$ of the $\approx 12d^2$ per layer (see §5).

3.3 Wrapping each sublayer: residual + LayerNorm

Each of the two sublayers (MHA, FFN) is wrapped in a residual connection and layer normalization. Where the LN goes defines the two conventions:

Post-LN (original Transformer¹) — normalize after adding the residual:

$$x^{\prime }=\text{LayerNorm}(x+\text{Sublayer}(x)).$$

Pre-LN (now standard in deep LLMs³) — normalize the input to the sublayer, add the raw residual:

$$x^{\prime }=x+\text{Sublayer}(\text{LayerNorm}(x)).$$

The full block is these two wrapped sublayers in sequence: attention first, then FFN.

Why pre-LN won for deep stacks

Xiong et al. (2020) showed that in post-LN, the expected gradient magnitude at initialization grows with depth, so training deep transformers requires a learning-rate warmup stage to avoid blowing up early — and it’s still delicate. In pre-LN, the LayerNorm sits inside the residual branch, so the residual path is a clean identity: gradients flow straight through the sum without being renormalized at every layer. This bounds the gradient scale, makes deep stacks train stably, and largely removes the need for warmup and heavy LR tuning.³ For the 100+-layer stacks in frontier models, that stability is decisive — which is why pre-LN (usually with RMSNorm) is the default. The tradeoff: pre-LN can leave the residual stream’s variance growing with depth, so a final LN before the output head is standard.

Don't confuse the two orderings

Pre-LN vs post-LN is not cosmetic — a deep post-LN model trained without warmup often diverges, while the same architecture in pre-LN trains fine. When reproducing an old paper, check which it used; swapping the ordering silently changes the optimization landscape.

---

4. Illustration — the transformer block

The diagram shows one pre-LN block: the two sublayers, each with its LayerNorm on the input branch and a residual add. Data flows top to bottom; the straight-down arrows are the residual highways.

graph TD
    X["Input x  (s × d)"] --> LN1["LayerNorm"]
    LN1 --> MHA["Multi-Head Attention<br/>h heads, d_h = d/h<br/>Q,K,V,O projections"]
    X -->|residual| ADD1(("+"))
    MHA --> ADD1
    ADD1 --> LN2["LayerNorm"]
    LN2 --> FFN["FFN / MLP<br/>d → d_ff=4d → d<br/>GELU"]
    ADD1 -->|residual| ADD2(("+"))
    FFN --> ADD2
    ADD2 --> OUT["Output  (s × d)"]

    style MHA fill:#44a,color:#fff
    style FFN fill:#4a4,color:#fff
    style ADD1 fill:#a83,color:#fff
    style ADD2 fill:#a83,color:#fff

Figure 1 — One pre-LN transformer block. The attention sublayer mixes information across positions; the FFN sublayer mixes across features per position. Each sublayer normalizes its input and adds a raw residual, so the vertical path (the residual stream) is an unobstructed identity from input to output. Post-LN would instead place each LayerNorm after its residual add, on the trunk.

---

5. FLOPs & params per layer

Let $b$ = batch, $s$ = sequence length, $d=d_{model}$, and $d_{ff}=4d$. Use the convention that a GEMM multiplying $(m\times k)\cdot (k\times n)$ costs $2mkn$ FLOPs (the 2 = one multiply + one add). Interactive: Attention FLOPs & KV-cache calculator.

5.1 Attention sublayer

Q, K, V, O projections. Each is a $(bs\times d)\cdot (d\times d)$ GEMM = $2bs{d}^2$ FLOPs. Four of them (three input projections + $W^O$):

$$\text{proj FLOPs}=4\times 2bs{d}^2=8bs{d}^2\approx 4bs{d}^2\text{ MACs.}$$

(The "$\approx 4bs{d}^2$" in the brief counts multiply-adds; $\times 2$ for FLOPs gives $8bs{d}^2$.) Note this is independent of $h$ — splitting into heads just reshapes these same matrices.

Attention scores + values ($Q{K}^{\top }$ and $PV$). Summed over all heads, the head axis collapses ($h\cdot d_h=d$), so:

$$Q{K}^{\top }:2b{s}^{2}d,PV:2b{s}^{2}d\Rightarrow \text{attn-matrix FLOPs}=4b{s}^{2}d.$$

This is the $\Theta (s^2)$ term — it dominates once $s\gtrsim d$, and it’s the term FlashAttention (see 06-memory-wall-and-flashattention) makes memory-cheap without changing its FLOP count.

5.2 FFN sublayer

Two GEMMs, $(bs\times d)\cdot (d\times 4d)$ and $(bs\times 4d)\cdot (4d\times d)$:

$$\text{FFN FLOPs}=2bsd(4d)+2bs(4d)d=16bs{d}^2.$$

5.3 Totals and the $6ND$ rule

$${\text{FLOPs}}_{\text{layer}}\approx \underset{\text{QKVO}}{\underbrace{8bs{d}^2}}+\underset{\text{scores}}{\underbrace{4b{s}^{2}d}}+\underset{\text{FFN}}{\underbrace{16bs{d}^2}}=24bs{d}^2+4b{s}^{2}d.$$

Parameters per layer (weights, ignoring biases/LN):

$$\underset{W^Q,W^K,W^V,W^O}{\underbrace{4d^2}}+\underset{\text{FFN }{W}_1,W_2}{\underbrace{2\cdot 4d^2}}=12d^2.$$

Now connect to scaling. For the whole model, non-embedding params $N\approx 12d^2\cdot n_{layers}$. Per token, the forward pass does $\approx 2N$ FLOPs (the "$24bs{d}^2$" term divided by $bs$ tokens is $24d^2\approx 2\cdot 12d^2$ per token per layer — matching $2N$ across the stack, when the $s^2$ term is subdominant). Backprop costs $\approx 2\times$ the forward, so training costs $\approx 3\times$ forward:

$$\boxed{C\approx 6ND}\text{(total training FLOPs; }D=\text{tokens seen).}$$

This is the workhorse estimate behind compute budgeting and scaling laws — derived straight from the per-layer count above. It ignores the $4b{s}^{2}d$ attention term, which is a good approximation only while $s\ll d/(\text{something})$; at long context the $s^2$ term is not negligible and $6ND$ undercounts. See the roofline/perf-modeling treatment in 03-performance-modeling-roofline for when that matters.

The two regimes in one glance

Per layer, the $d^2$ terms ($24bs{d}^2$) scale linearly in $s$; the attention-matrix term ($4b{s}^{2}d$) scales quadratically. Their crossover is around $s\approx 6d$: below it the projections+FFN dominate (model is “compute-in-the-weights”); above it attention’s $s^2$ takes over. That crossover is exactly why long-context work obsesses over the attention term.

---

6. Retrieval practice

Quiz 1: Splitting $d_{model}$ into $h$ heads of width $d_h=d/h$ — does this increase the parameter count and FLOPs of the projections versus a single head of width $d$? Why or why not?

Answer

No (to first order). The per-head projections $W_i^{Q/K/V}\in {\mathbb{R}}^{d\times d_h}$ stacked over $h$ heads form matrices of total size $d\times (h{d}_h)=d\times d$ — identical to a single full-width projection. So Q/K/V/O together are still $4d^2$ params and $8bs{d}^2$ FLOPs regardless of $h$. Multi-head is a reshape of the same budget into $h$ narrower, independent subspaces, not an enlargement. What changes is expressiveness/inductive bias (parallel specialized attention patterns), not cost. (The $Q{K}^{\top }$/$PV$ FLOPs are also $h$-independent since $h\cdot d_h=d$.)

Quiz 2: In pre-LN, where exactly does the LayerNorm sit relative to the residual add, and why does that make deep transformers trainable without warmup?

Answer

Pre-LN normalizes the input to the sublayer, inside the residual branch: $x^{\prime }=x+\text{Sublayer}(\text{LayerNorm}(x))$. The residual trunk $x\to x^{\prime }$ is therefore a clean identity — gradients flow through the additive skip without being renormalized at each layer, so their magnitude stays bounded with depth. Post-LN instead normalizes the trunk after the add ($\text{LayerNorm}(x+\text{Sublayer}(x))$), which Xiong et al. showed makes expected gradients grow with depth at init, requiring a learning-rate warmup to train stably.³ Pre-LN removes that fragility, which is why deep LLM stacks use it (typically with a final LN before the output head to control residual-stream variance growth).

Quiz 3: A transformer layer has roughly $12d^2$ parameters. Break that down by component, and identify where most of the parameters live.

Answer

Attention: $W^Q,W^K,W^V,W^O$, each $d\times d$ → $4d^2$. FFN: $W_1$ ($d\times 4d$) and $W_2$ ($4d\times d$) → $4d^2+4d^2=8d^2$. Total $=12d^2$ (ignoring biases and LN params, which are $O(d)$ and negligible). Most parameters (~two-thirds) live in the FFN, not attention — a fact that motivates a lot of MoE and FFN-sparsity work, since the FFN is the fattest single target.

Quiz 4: Why does the $C\approx 6ND$ rule undercount training FLOPs at very long context, and which term is being dropped?

Answer

$6ND$ comes from the $2N$ forward FLOPs/token (× 3 for fwd+bwd), where $N\approx 12d^2{n}_{layers}$ captures only the weight GEMMs (projections + FFN), which scale linearly in $s$. It drops the attention-matrix term $4b{s}^{2}d$ per layer, which scales quadratically in $s$ and has no weights (so it isn’t in $N$). While $s\ll d$ that term is negligible; once $s$ grows past $\sim 6d$ it dominates the per-layer cost, so $6ND$ increasingly under-predicts actual FLOPs. That’s precisely the regime where efficient-attention and FlashAttention economics matter (see 06-memory-wall-and-flashattention).

---

7. Practice problems

Practice problem 1: Count params and FLOPs of a concrete block

Problem: A pre-LN transformer block with $d=4096$, $h=32$ (so $d_h=128$), $d_{ff}=4d=16384$. Batch $b=1$, sequence $s=2048$, BF16. (a) Parameters in the block. (b) Forward FLOPs, split into the weight-GEMM part and the attention-matrix part. (c) Which dominates here, and at what $s$ would they be equal?

Worked solution:
(a) Params. $12d^2=12\cdot 4096^2=12\cdot 1.678\times 10^7\approx 2.01\times 10^8$ ≈ 201 M params/layer (attention $4d^2\approx 67$M, FFN $8d^2\approx 134$M).
(b) FLOPs. Weight GEMMs: $24bs{d}^2=24\cdot 1\cdot 2048\cdot 4096^2=24\cdot 2048\cdot 1.678\times 10^7\approx 8.24\times 10^{11}$ ≈ 0.82 TFLOP. Attention matrix: $4b{s}^{2}d=4\cdot 1\cdot 2048^2\cdot 4096=4\cdot 4.19\times 10^6\cdot 4096\approx 6.87\times 10^{10}$ ≈ 0.069 TFLOP.
(c) Dominance. Weight term is ~12× the attention term here → projections+FFN dominate at $s=2048$. They’re equal when $24bs{d}^2=4b{s}^{2}d\Rightarrow s=6d=6\cdot 4096=24,576$. Above ~24k tokens, attention’s $s^2$ term overtakes. (Matches the ”$s\approx 6d$ crossover” tip.)

Practice problem 2: Reason about the $6ND$ estimate

Problem: A model has $N=7\times 10^9$ non-embedding params and is trained on $D=2\times 10^{12}$ tokens at short context. (a) Estimate total training FLOPs via $6ND$. (b) If instead trained mostly at $s=32768>6d$ with $d\approx 4096$, is $6ND$ an over- or under-estimate, and roughly why?

Worked solution:
(a) $C\approx 6ND=6\cdot 7\times 10^9\cdot 2\times 10^{12}=8.4\times 10^{22}$ FLOPs (84 ZettaFLOP). (For scale, ~1000 H100-days at ~1 PFLOP/s sustained.)
(b) Under-estimate. $6ND$ omits the attention-matrix term $4b{s}^{2}d$, which is quadratic in $s$ and grows relative to the weight GEMMs once $s>6d\approx 24,576$. At $s=32768$ (past the crossover) attention contributes a non-trivial extra fraction of per-layer FLOPs that $6ND$ ignores, so true compute exceeds the estimate. The correction grows with $s$; for careful long-context budgeting, add the ${\sum }_{layers}4b{s}^{2}d\times 3$ (fwd+bwd) term explicitly. See 03-performance-modeling-roofline.

---

8. Primary source & going deeper

• Primary source (read this): Vaswani et al., “Attention Is All You Need” — arXiv:1706.03762.¹ Read §3.2 (scaled dot-product + multi-head attention) and §3.3 (position-wise FFN); the block structure and $d_{ff}=4d$, $h=8$, $d_h=64$ defaults are here. Note it uses post-LN.
• Pre-LN analysis: Xiong et al., “On Layer Normalization in the Transformer Architecture” — arXiv:2002.04745³ — the gradient-at-init argument for why pre-LN trains without warmup.
• GELU: Hendrycks & Gimpel, “Gaussian Error Linear Units (GELUs)” — arXiv:1606.08415.²

Ask the teacher

This is a compressed treatment. If any step is fuzzy — the $W^O$ mixing argument, the pre-LN gradient bound, the $s\approx 6d$ crossover derivation, or how MHA batches over the head axis in real code — ask for a deeper drill or a worked code walkthrough. Follow-up questions are the point.

---

9. What’s next

1. 03-positional-information (next) — attention as defined here is permutation-equivariant: it has no notion of token order. The next lesson adds positional information (sinusoidal, learned, and RoPE/ALiBi), the missing ingredient that makes the block sequence-aware.
2. 01-scaled-dot-product-attention (back) — revisit the single-head mechanism now that you’ve seen how $h$ of them compose; the $\sqrt{d_h}$ scaling and softmax saturation arguments there apply per head.

Further threads: GQA/MQA & the KV-cache (sharing K/V across heads for cheaper inference), MoE FFNs (sparsifying the $8d^2$ that dominates params), and normalization variants (RMSNorm, and where the final LN goes).

---

Reference topic: attention-mechanisms | Glossary: attention-glossary | Filed: 2026-09-10

---

References

Footnotes

1. [established] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Łukasz Kaiser, Illia Polosukhin, “Attention Is All You Need” — arXiv:1706.03762 (NeurIPS 2017). arxiv.org/abs/1706.03762. Source of the multi-head attention formulation (§3.2), the position-wise FFN with $d_{ff}=4d$ (§3.3), the $h=8$, $d_h=64$ defaults, post-LN residual wrapping, and the observation that heads attend in different representation subspaces. ↩ ↩² ↩³ ↩⁴ ↩⁵

2. [established] Dan Hendrycks, Kevin Gimpel, “Gaussian Error Linear Units (GELUs)” — arXiv:1606.08415 (2016). arxiv.org/abs/1606.08415. Defines $\text{GELU}(x)=x\Phi (x)$, the smooth activation used in the FFN of most modern transformers. ↩ ↩²

3. [established] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, Tie-Yan Liu, “On Layer Normalization in the Transformer Architecture” — arXiv:2002.04745 (ICML 2020). arxiv.org/abs/2002.04745. Shows post-LN has depth-growing expected gradients at init (needs LR warmup), while pre-LN bounds them and trains deep stacks stably without warmup. ↩ ↩² ↩³ ↩⁴