Lesson 17 of 23 · Transformer block

One token becomes three views.

Your win: multiply one normalized token by explicit \(W_Q,W_K,W_V\) matrices, then reshape features into attention heads.

15 minutesNeeds: matrix multiplicationOutcome: track Q/K/V shapes

Mission link: RoPE is applied after the Q/K projections and after head structure is known, but before query-key scores are formed.

Three learned linear projections

\[ Q=\bar XW_Q,\qquad K=\bar XW_K,\qquad V=\bar XW_V. \]

In ordinary multi-head attention, each projection has total output width \(h\,d_h=d_{\text{model}}\). Before reshaping:

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

Numbers substituted: \(d_{\text{model}}=4,h=2,d_h=2\)

Take one normalized token \(\bar{\mathbf x}=\begin{bmatrix}1&2&0&-1\end{bmatrix}\) and choose simple projection matrices:

\[ W_Q=I,\quad W_K=\begin{bmatrix} 1&0&0&0\\ 0&1&0&0\\ 0&0&0&1\\ 0&0&1&0 \end{bmatrix},\quad W_V=\begin{bmatrix} 2&0&0&0\\ 0&1&0&0\\ 0&0&1/2&0\\ 0&0&0&-1 \end{bmatrix}. \]

\[ \mathbf q=\begin{bmatrix}1&2&0&-1\end{bmatrix},\quad \mathbf k=\begin{bmatrix}1&2&-1&0\end{bmatrix},\quad \mathbf v=\begin{bmatrix}2&2&0&1\end{bmatrix}. \]

Reshape by feature chunks:

Head	\(\mathbf q\)	\(\mathbf k\)	\(\mathbf v\)
1	\((1,2)\)	\((1,2)\)	\((2,2)\)
2	\((0,-1)\)	\((-1,0)\)	\((0,1)\)

Where RoPE enters

\[Q'=\operatorname{RoPE}(Q),\qquad K'=\operatorname{RoPE}(K),\qquad V'=V.\]

RoPE rotates adjacent pairs inside each head. With \(d_h=2\), each head above contains exactly one rotation plane. Every head still receives all \(T\) token positions.

Code checkpoint · project, split heads, rotate Q and K

Three linear maps create Q, K, and V. Reshaping exposes the head axis; RoPE then changes Q and K while V stays untouched.

Show projection and head splitting

def split_heads(x):
    batch, tokens, _ = x.shape
    return x.view(batch, tokens, num_heads, head_dim).transpose(1, 2)

q = split_heads(q_proj(x))  # [B, H, T, Dh]
k = split_heads(k_proj(x))  # [B, H, T, Dh]
v = split_heads(v_proj(x))  # [B, H, T, Dh]

q = apply_rope(q, rope_cos, rope_sin)
k = apply_rope(k, rope_cos, rope_sin)

Trace it: If x is [2, 5, 12] and there are 3 heads, what is each projected tensor's shape after split_heads?

Engineering bridge · grouped-query attention

The teaching block gives Q, K, and V the same head count. MiniMind keeps more query heads than key/value heads, then repeats each K/V head so it can serve a small group of queries. This reduces the cached memory footprint without changing the residual width.

Q: \(B\times T\times H\times D_h\)→K/V: \(B\times T\times H_{\mathrm{kv}}\times D_h\)→ repeat\(B\times T\times H\times D_h\)

Show MiniMind's head-count contract

self.n_local_heads = config.num_attention_heads
self.n_local_kv_heads = config.num_key_value_heads
self.n_rep = self.n_local_heads // self.n_local_kv_heads

q = self.q_proj(x).view(bsz, seq_len, self.n_local_heads, self.head_dim)
k = self.k_proj(x).view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
v = self.v_proj(x).view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
q = self.q_norm(q)
k = self.k_norm(k)
k = repeat_kv(k, self.n_rep).transpose(1, 2)
v = repeat_kv(v, self.n_rep).transpose(1, 2)

Invariant: \(H\) must be divisible by \(H_{\mathrm{kv}}\). After repetition, Q, K, and V agree on the head axis before score formation.

Retrieval check

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

What does \(d_{\text{model}}=8,h=2\) imply for \(d_h\)?

STAGE 4 · PROVE IT

Practice studio · recall → build → verify

For \(d_{\text{model}}=12,h=3\), calculate \(d_h\).
Reshape \((1,2,3,4,5,6)\) into \(h=3\) feature chunks.
If \(X\) has shape \(2\times5\times12\), state the Q shape before and after splitting into three heads.
Using the matrices above, recompute \(\mathbf k\) from \(\bar{\mathbf x}\).
Which tensors receive RoPE, and which does not?
Explain why the projection weights mix features but do not mix token positions.

BUILD CHALLENGE · RUN IT LOCALLY

Engineering challenge · split heads without losing tokens

Implement the reshape from `[B, T, D]` to `[B, H, T, Dh]`, then apply RoPE to Q and K only. Add assertions that `D == H * Dh` and V remains untouched.

Standard answer · compare after your attempt

B, T, D = x.shape
assert D % H == 0
Dh = D // H
q = q.view(B, T, H, Dh).transpose(1, 2)
k = k.view(B, T, H, Dh).transpose(1, 2)
v = v.view(B, T, H, Dh).transpose(1, 2)
q, k = apply_rope(q, k, cos, sin)
assert q.shape == k.shape == v.shape == (B, H, T, Dh)

The reshape changes feature organization; it never assigns a head a subset of token positions.

Check solutions

\(d_h=4\).
\((1,2),(3,4),(5,6)\).
Before: \(2\times5\times12\). After: \(2\times3\times5\times4\).
\((1,2,-1,0)\).
Q and K receive RoPE; V does not.
The same matrix multiplies each token row independently along the feature axis.

Primary source: Vaswani et al., Sections 3.2.1–3.2.2; RoPE placement follows Su et al., Section 3.2.