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\)

Watch one token branch into Q, K, and V

\(\bar x\)

120−1

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?

Retrieval check

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

Practice before moving on

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.

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.

I can project one token into Q/K/V and reshape features into heads.

Ask the teaching agent to give you a fresh \(d_{\text{model}},h\) shape exercise.