Compute Two Heads with Separate Projections

Calculate two causal heads from the same three-token input using separate projection matrices. The trace makes each score, softmax row, value reading, and difference between the two heads numerically visible.

Separate-head notation makes the calculation easiest to inspect. Each head has its own projection matrices and runs the single-head operation from Chapter 1.

Use one three-position sequence with model width 4:

X=\begin{bmatrix} 1&0&1&0\\ 0&1&0&1\\ 1&1&-1&1 \end{bmatrix}.

There are two heads, with $d_k=d_v=2$ . Head 1 selects the first two input coordinates, and head 2 selects the final two:

P_1=\begin{bmatrix}1&0\\0&1\\0&0\\0&0\end{bmatrix}, \qquad P_2=\begin{bmatrix}0&0\\0&0\\1&0\\0&1\end{bmatrix}.

For this trace, set

W_Q^{(1)}=W_K^{(1)}=W_V^{(1)}=P_1,

W_Q^{(2)}=W_K^{(2)}=W_V^{(2)}=P_2.

Head 1 Reads the First Coordinate Pair

The projected rows are

Q^{(1)}=K^{(1)}=V^{(1)}= \begin{bmatrix}1&0\\0&1\\1&1\end{bmatrix}.

For the final query $[1,1]$ , the scaled scores are

\frac{[1,1,2]}{\sqrt2} \approx[0.707,0.707,1.414].

Stable softmax gives approximately

a_3^{(1)}=[0.248,0.248,0.503].

The reading is

z_3^{(1)} \approx0.248[1,0]+0.248[0,1]+0.503[1,1] \approx[0.752,0.752].

Head 2 Reads the Final Coordinate Pair

The second projection gives

Q^{(2)}=K^{(2)}=V^{(2)}= \begin{bmatrix}1&0\\0&1\\-1&1\end{bmatrix}.

For the final query $[-1,1]$ , the scaled scores are

\frac{[-1,1,2]}{\sqrt2} \approx[-0.707,0.707,1.414].

Its weights are approximately

a_3^{(2)}=[0.074,0.306,0.620],

and its reading is

z_3^{(2)} \approx0.074[1,0]+0.306[0,1]+0.620[-1,1] \approx[-0.546,0.926].

The two heads received the same input rows and used the same causal visibility rule. Their projections produced different queries, keys, values, weights, and readings.

Concatenate the Results

Joining the head features in head order gives

z_3=\operatorname{Concat}(z_3^{(1)},z_3^{(2)}) \approx[0.752,0.752,-0.546,0.926].

For now, let $W_O=I$ , so this vector is also the final layer output. Lesson 5 will use a non-identity $W_O$ and show why head order is meaningful only together with the corresponding blocks of that projection.

Trace two heads separately

The program projects, scores, normalizes, and reads each head independently. Change one projection and compare both weight rows.

Q1. Calculate one head coordinate

Head 1 has weights approximately $[0.248,0.248,0.503]$ and values $[1,0]$ , $[0,1]$ , and $[1,1]$ . What is the first coordinate of its reading to three decimal places?

Your answer

Compute it first, then check your number.

Hint

Only the first and third values have first coordinate 1.

Solution

Using the displayed rounded weights, the coordinate is

0.248(1)+0.248(0)+0.503(1)=0.751

. The unrounded calculation is approximately 0.752.

Compare Heads Before Averaging Them Away

When inspecting multi-head attention, retain the head axis long enough to compare its scores, weights, and values. Averaging attention maps across heads can hide the very differences the architecture permits.