Lesson 20 of 23 · Transformer block

Heads split features, not the sentence.

Your win: track the token and feature axes through independent heads, concatenation, and the output projection.

12 minutesNeeds: tensor shapesOutcome: reject a common misconception

Mission link: every head applies RoPE and causal attention to the full sequence; heads differ in learned feature projections.

Every head sees all \(T\) positions

For \(d_{\text{model}}=h\,d_h\), the feature width is reorganized into \(h\) heads:

\(B\times T\times d_{\text{model}}\)→\(B\times h\times T\times d_h\)

The \(T\) axis is unchanged. A head is not assigned one word or one sentence segment. Each head forms its own \(T\times T\) attention matrix over the same token sequence.

Every head sees the full token sequence. Heads differ by their learned feature projections, then their outputs are concatenated along the feature axis and mixed by \(W_O\).

Numbers substituted: \(T=3,d_{\text{model}}=8,h=2\)

Quantity	Shape	Meaning
\(Q,K,V\) before split	\(3\times8\)	Three tokens, eight projected features
Each head's \(Q_r,K_r,V_r\)	\(3\times4\)	All three tokens, four head features
Each head's \(P_r\)	\(3\times3\)	Every query position compared with every key position
Each head's \(O_r=P_rV_r\)	\(3\times4\)	One four-feature output per token
\(\operatorname{Concat}(O_1,O_2)\)	\(3\times8\)	Join along the feature axis

Concatenate, then mix heads

\[ O_{\text{attn}}=\operatorname{Concat}(O_1,\ldots,O_h)W_O. \]

For one token with \(d_{\text{model}}=4,h=2,d_h=2\), suppose the head outputs are \(\mathbf o_1=(1,0.2)\) and \(\mathbf o_2=(-0.5,2)\). Concatenation gives

\[ \mathbf c=\begin{bmatrix}1&0.2&-0.5&2\end{bmatrix}. \]

Use this simple output matrix:

\[ W_O=\begin{bmatrix} 1&0&1&0\\ 0&1&0&1\\ 0&0&1&0\\ 0&0&0&1 \end{bmatrix}. \] \[ \mathbf cW_O =\begin{bmatrix}1&0.2&0.5&2.2\end{bmatrix}. \]

The output projection lets coordinates from different heads contribute to the same model feature while restoring the \(d_{\text{model}}\)-wide representation expected by the residual path.

Code checkpoint · join every head

Attention produced one context tensor per head. Transpose restores token-major order, then reshape concatenates head features before the learned output projection.

Show head recombination

# context begins as [B, H, T, Dh]
context = context.transpose(1, 2).contiguous()
# now [B, T, H, Dh]

context = context.view(batch, tokens, hidden_size)
# now [B, T, H * Dh] = [B, T, D]

attention_output = out_proj(context)  # [B, T, D]

Trace it: Why must the output projection return width \(D\) before the residual addition?

Engineering bridge · GQA still returns model width

MiniMind's grouped-query attention stores K/V with H_kv heads, repeats them to match H query heads, and only then transposes into attention-major order. The output projection restores D, regardless of the K/V compression.

Show the layout around repeat_kv

xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
xq, xk, xv = (
    xq.transpose(1, 2),
    repeat_kv(xk, self.n_rep).transpose(1, 2),
    repeat_kv(xv, self.n_rep).transpose(1, 2),
)
# after attention: [B, H, T, Dh] -> [B, T, H*Dh] -> [B, T, D]
output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
output = self.resid_dropout(self.o_proj(output))

Invariant: H = H_kv * n_rep, and H * D_h = D before the residual addition.

Retrieval check

Answer once before reopening the code; use the feedback to correct your explanation.

With \(T=5,d_{\text{model}}=12,h=3\), what does one head receive?

STAGE 4 · PROVE IT

Practice studio · recall → build → verify

Calculate \(d_h\) for \(d_{\text{model}}=24,h=6\).
For \(B=2,T=7,d_{\text{model}}=24,h=6\), give the split Q shape.
Give one head's score-matrix shape for \(T=7\).
Concatenate head outputs \((2,-1)\) and \((0,3)\).
Multiply \((2,-1,0,3)\) by the displayed \(W_O\).
Correct this statement: “Head 1 reads the first half of the sentence and Head 2 reads the second half.”

BUILD CHALLENGE · RUN IT LOCALLY

Engineering challenge · merge GQA heads back to model width

Given Q heads `H`, KV heads `Hkv`, and `n_rep = H // Hkv`, repeat K/V, transpose to token-major order, flatten heads, then apply `o_proj`.

Standard answer · compare after your attempt

assert H % Hkv == 0
key_states = repeat_kv(key_states, n_rep)
context = context.transpose(1, 2).contiguous().view(B, T, H * Dh)
out = o_proj(context)
assert out.shape == (B, T, D)

GQA compresses stored K/V heads, but the residual contract still ends at model width `D`.

Check solutions

\(d_h=4\).
\(2\times6\times7\times4\).
\(7\times7\).
\((2,-1,0,3)\).
\((2,-1,2,2)\).
Both heads read all token positions; their learned projections and feature subspaces differ.

Primary source: Vaswani et al., Section 3.2.2.