Splitting Heads, Assembling Blocks — The Whole Transformer

Textbook › Part III Deep Learning: what assumption you add

CHAPTER 7

Splitting Heads, Assembling Blocks — The Whole Transformer

★ MEMBER11 min

Chapter 6 looked at a single attention mechanism. Two things remain: why split it into several, and how those parts assemble into one block.

With this chapter Part III's toolkit is complete. Everything built up from Chapter 1 — the machinery for carrying the error down (Chapter 3), where to stop (Chapter 4), assumptions built into structure (Chapter 5), letting the input decide where to look (Chapter 6) — fits into a single diagram.

1. Heads are not split for parameters

First, confirm that nothing grows. Split $d_{\text{model}} = 512$ into $h$ heads of $d_k = 512/h$ and the total of $W_Q, W_K, W_V, W_O$ is:

heads $h$	$d_k$	parameters
1	512	1,048,576
2	256	1,048,576
8	64	1,048,576
64	8	1,048,576

All the same. The width is only being divided, so head count moves nothing. Then what is it for?

The reason is the shape of a softmax. One softmax output sums to 1 (the categorical distribution of Chapter 0-3), so a single head cannot look strongly at $k$ places at once. Aim at two and each gets at most 0.5; at three, 0.333.

targets $k$	$h=1$	$h=2$	$h=4$	$h=8$
1	1.000	1.000	1.000	1.000
2	0.500	1.000	1.000	1.000
3	0.333	0.500	1.000	1.000
4	0.250	0.500	1.000	1.000
8	0.125	0.250	0.500	1.000

What every target is guaranteed is $1/\lceil k/h \rceil$ . Heads are assigned whole, so it is not the average share $h/k$ — at $k=3,\ h=2$ it is 0.500, not 0.667, because one head ends up covering two. Adding heads adds places you can aim at simultaneously.

FIG 1Click positions to choose targets and change the head count. With more targets than heads, one head covers several and the sum-to-1 constraint thins each one. With enough heads every target can carry 1.00. The parameter count is the same at any h

The figure constructs ideal distributions and measures them; nothing is trained. What it shows is the shape of reachable distributions — not a claim that a real model's heads divide the work this neatly. What a trained head attends to cannot be read off its weights, as the end of Chapter 6 showed.

2. The causal mask — checking the future really is hidden

A model generating text left to right must not let position $t$ see $t+1$ onward. The mechanism is simple: set the upper triangle of $QK^{\top}$ to $-\infty$ before the softmax.

Whether it truly cannot see is measurable. Nudge the input at position $j$ and watch the output at position $t$ by central difference (Chapter 0-2, equation (2)):

	largest dependence on the future ( $j>t$ )	on the past ( $j \le t$ )
masked	0.000000	0.6583
unmasked	0.2550	0.3368

Exactly zero. A softmax fed $-\infty$ outputs precisely 0, and no gradient flows past it either. The mask does not politely request that nothing be seen — it removes the path.

3. Assembling the block

A Transformer block is built entirely from parts already seen:

\begin{aligned} x' &= x + \mathrm{MHA}\big(\mathrm{LN}(x)\big) \\ y &= x' + \mathrm{FFN}\big(\mathrm{LN}(x')\big) \end{aligned}

(1)

Equation (1) contains three things.

The residual connection $x +$ . The additive path measured in Exercise 3 of Chapter 0-2, the same shape as Chapter 5's gate. Connect only multiplicatively and the product decays it (Chapter 3); run one additive path and that route passes at a factor of 1. It is the central reason dozens of layers work at all.