Chapter 2

Multi-Head Attention

Split the model width into several attention heads, compute them in parallel, and combine their results. The chapter follows every reshape and parameter while separating possible specialization from evidence that a trained head actually specialized.

A single attention head gives each query one distribution over readable positions. Its value vector may contain many coordinates, but every coordinate in that vector is combined using the same attention weights.

Multi-head attention creates several attention computations instead. Each head has its own query, key, and value projections, so it can form a different score matrix, a different normalized weight matrix, and a different value reading. The head results are then concatenated and projected back to the model width.

For head r∈{1,…,h}r\in\{1,\ldots,h\},

Z(r)=softmax⁡last ⁣(Q(r)K(r)⊤dk+M)V(r).Z^{(r)}= \operatorname{softmax}_{last}\!\left( \frac{Q^{(r)}K^{(r)\top}}{\sqrt{d_k}}+M \right)V^{(r)}.

The complete operation is

Z=Concat⁡(Z(1),…,Z(h)),Y=ZWO.Z=\operatorname{Concat}\left(Z^{(1)},\ldots,Z^{(h)}\right), \qquad Y=ZW_O.

We will keep the batch-major convention from Chapter 1. The main new axis is the head axis:

(B,T,dmodel)→(B,h,T,dk) and (B,h,T,dv)→(B,h,T,T)→(B,T,hdv)→(B,T,dmodel).(B,T,d_{model}) \rightarrow(B,h,T,d_k)\text{ and }(B,h,T,d_v) \rightarrow(B,h,T,T) \rightarrow(B,T,hd_v) \rightarrow(B,T,d_{model}).

The chapter begins with the reason for using separate heads rather than one wider value vector. It then defines model width, head count, and per-head widths; calculates two heads independently; and shows how combined projection matrices implement the same operation efficiently. The final lessons cover concatenation, parameter counts, head ablations, empirical evidence about specialization and redundancy, and a complete two-head audit.

After this chapter

  • Track batch, head, sequence, and feature axes through multi-head attention.
  • Calculate projection and output shapes and parameter counts.
  • Compare one-head and multi-head computations without overinterpreting attention patterns.

Lessons

  1. 01
    One Wider Head Still Has One Attention Pattern

    Distinguish value width from the number of attention distributions and state multi-head representational claims without assuming specialization.

    1 exercise
  2. 02
    Model Width, Head Count, and Head Width

    Define d_model, h, d_k, and d_v; derive projection shapes; and examine equal and unequal head-width conventions.

    1 exercise
  3. 03
    Compute Two Heads with Separate Projections

    Compute a complete two-head causal attention example with separate Q, K, and V projections.

    1 exercise
  4. 04
    Combined Projections Introduce a Head Axis

    Convert separate head projections into combined matrices and trace batch, position, head, and feature axes.

    1 exercise
  5. 05
    Concatenate Heads and Project Once

    Join multi-head results, apply W_O, and analyze concatenation axes, head order, and output mixing.

    1 exercise
  6. 06
    Count Parameters and Attention Work

    Count multi-head attention parameters, dense scores, arithmetic, and storage without mixing cost conventions.

    1 exercise
  7. 07
    Specialization, Redundancy, and Head Ablation

    Evaluate evidence for head specialization and redundancy through bounded claims, ablations, pruning, and permutation symmetry.

    1 exercise
  8. 08
    Audit a Two-Head Attention Layer

    Audit separate and combined multi-head attention implementations with numerical, shape, axis, and ablation checks.

    1 exercise

Review and practice

  1. •
    Review

    Review multi-head attention equations, shapes, costs, interventions, evidence, and implementation checks.

  2. •
    Exercises

    Solve multi-head attention shape, calculation, cost, interpretation, permutation, and debugging exercises.

Chapter progress