Exercises

Practice head-axis shapes, two-head calculations, output projection, parameter and score counts, ablation claims, head permutation, and silent implementation errors.

These exercises move from head dimensions and numerical calculations to cost, ablation, and implementation audits. Each prompt states the convention it uses.

Trace Head Dimensions

Q1. Count the attention distributions

A multi-head layer receives a batch with $B=3$ sequences, uses $h=8$ heads, and has sequence length $T=20$ . How many separate attention rows are formed for one query position across the complete batch?

Your answer

Compute it first, then check your number.

Hint

Do not multiply by

T

again; the question fixes one query position.

Solution

There are

Bh=3(8)=24

independently normalized rows for the selected query position.

Q2. Derive a valid output projection

A layer uses $d_{model}=12$ , $h=3$ , and $d_v=5$ . The implementation does not require $hd_v=d_{model}$ . What shape must $W_O$ have?

Choose one

(15, 12)

(12, 15)

(5, 12)

(12, 12)

Select one choice, then check.

Hint

Complete

(B,T,hd_v)(?,?)=(B,T,d_{model})

Solution

W_O

has shape

(hd_v,d_{model})=(15,12)

Calculate Two Heads

Q3. Calculate one head result

For the final query, a head has attention weights $[0.25,0.25,0.50]$ and value vectors $[1,0]$ , $[0,1]$ , and $[1,1]$ . What is the first coordinate of the head result?

Your answer

Compute it first, then check your number.

Hint

Apply the same three weights to the first coordinate of every value.

Solution

The result is

0.25[1,0]+0.25[0,1]+0.50[1,1]=[0.75,0.75]

. Its first coordinate is

0.75

Q4. Place the head axis correctly

Combined query projection produces a tensor with shape $(B,T,hd_k)=(2,7,12)$ . The layer uses $h=3$ equal-width heads. Which shape should enter the batched query-key product?

Choose one

(2, 3, 7, 4)

(2, 7, 3, 4)

(3, 2, 4, 7)

(2, 7, 12)

Select one choice, then check.

Hint

Trace

(B,T,hd_k)\to(B,T,h,d_k)\to(B,h,T,d_k)

Solution

The required shape is

(2,3,7,4)

Q5. Apply the shared output projection

Two heads produce $z^{(1)}=[1,2]$ and $z^{(2)}=[3,4]$ . They are concatenated in that order. Let

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

What is the second coordinate of $[z^{(1)};z^{(2)}]W_O$ ?

Your answer

Compute it first, then check your number.

Hint

The concatenated row is

[1,2,3,4]

Solution

The projected row is

[8,1]

, so the requested coordinate is

1

Count What the Layer Stores and Learns

Q6. Count projection parameters

A bias-free layer uses $d_{model}=16$ , $h=4$ , and $d_k=d_v=4$ . Count all entries in $W_Q$ , $W_K$ , $W_V$ , and $W_O$ .

Your answer

Compute it first, then check your number.

Hint

Under

hd_k=hd_v=d_{model}

, every combined projection is a

d_{model}\times d_{model}

matrix.

Solution

The three input projections and one output projection contain

4d_{model}^2=4(16^2)=1024

entries.

Q7. Count dense score entries

A dense attention layer has $B=2$ , $h=6$ , and $T=50$ . How many score entries does its full score tensor contain before masking?

Your answer

Compute it first, then check your number.

Hint

The score shape is

(B,h,T,T)

Solution

2(6)(50)(50)=30{,}000

entries.

Interpret Interventions Carefully

Q8. Interpret a head-zeroing result

Zeroing head 4 before concatenation reduces validation accuracy by 0.2 percentage points on one task. Which conclusion is supported?

Choose the supported claim

Under this intervention, model, task, and metric, head 4 had a small measured effect

Head 4 is universally unnecessary and may be deleted from every Transformer

The experiment reveals the complete human-readable meaning of head 4

Select one choice, then check.

Hint

Zeroing at inference is not the same experiment as pruning and retraining.

Solution

The bounded claim is supported. The result does not establish a universal role or explain how later components compensate.

Q9. Preserve a layer while permuting heads

All head-result blocks are permuted before concatenation. Which additional change preserves the same layer function?

Choose one

Apply the matching inverse permutation to the input row blocks of W_O

Permute the output columns of W_O in the same order

Permute the query-position rows of the causal mask

No change is needed because head order never affects W_O

Select one choice, then check.

Hint

For row-vector notation,

ZW_O

contracts the coordinates of

Z

with the rows of

W_O

Solution

Apply the matching inverse permutation to the row blocks of

W_O

. Head indices are coordinate labels, not fixed semantic identities.

Integrated Audit

Q10. Locate a silent head-order error

A combined implementation and a verified separate-head implementation agree on each $Z^{(r)}$ . Both final outputs have shape $(B,T,d_{model})$ , but their values differ. The combined implementation concatenates heads in the order $[2,1]$ while using the original $W_O$ . What is the smallest complete repair?

Choose the repair

Restore head order [1, 2], or keep [2, 1] and apply the matching row-block permutation to W_O

Apply softmax over the head axis instead of the key axis

Reshape the final output until its dimensions match

Accept the result because the final shape is correct

Select one choice, then check.

Hint

The error preserves shape and appears only after correct head results are joined.

Solution

Either concatenate in the reference order or permute the input row blocks of

W_O

consistently. Changing softmax or merely reshaping does not repair the semantic coordinate mismatch.

Trace Head Dimensions

Q1. Count the attention distributions

Q2. Derive a valid output projection

Calculate Two Heads

Q3. Calculate one head result

Q4. Place the head axis correctly

Q5. Apply the shared output projection

Count What the Layer Stores and Learns

Q6. Count projection parameters

Q7. Count dense score entries

Interpret Interventions Carefully

Q8. Interpret a head-zeroing result

Q9. Preserve a layer while permuting heads

Integrated Audit

Q10. Locate a silent head-order error

Pause and reflect