Audit a Two-Head Attention Layer

Implement the same two-head layer in separate and combined forms, then compare every intermediate. The audit checks masks, row sums, split widths, axis order, concatenation, W_O, head zeroing, and failures that retain valid shapes.

Multi-head attention adds transformations around a familiar single-head core. The audit must therefore check two kinds of correctness:

each head performs valid attention;
split, transpose, concatenate, and output projection preserve the intended axes and head order.

A final tensor with shape $(B,T,d_{model})$ does not prove either one.

Compare Two Equivalent Implementations

The clearest equivalence test calculates the same layer in two forms:

separate form: project and run each head with its own matrices;
combined form: concatenate projection matrices, project once, split the final feature axis into heads, and run those heads.

Before $W_O$ , the corresponding $Z^{(r)}$ tensors must agree. After the same concatenation order and $W_O$ , the final outputs must agree.

Audit separate and combined two-head attention

The program reconstructs the chapter trace in both forms and checks head outputs, row sums, causal zeroes, joined width, and the output projection.

The expected final rows are approximately

z_3=[0.752,0.752,-0.546,0.926],

y_3=[0.479,1.215,-0.170,1.302].

Failures to Introduce Deliberately

Once the verified implementation passes, a useful audit changes one operation at a time:

Broken change	Expected evidence
omit $(B,T,h,d_k)\to(B,h,T,d_k)$ transpose	score axes contract or broadcast incorrectly
split $hd_k$ with the wrong widths	head projections disagree with the separate form
concatenate along position instead of features	$T$ changes or token records interleave
give both heads the same projected slice	their Q/K/V and weights become identical in this trace
swap head results without swapping $W_O$ row blocks	joined shape remains valid but $Y$ changes

The separate-head implementation acts as an oracle only because it has already been checked against the hand calculation. Two implementations can agree and still share the same conceptual error.

Head Zeroing Needs a Defined Location

To test head 2's current contribution, set $Z^{(2)}=0$ before concatenation and keep $W_O$ fixed. Zeroing an attention map before multiplying values is a different intervention. Removing projection parameters and retraining is also different. Record the intervention point beside every ablation result.

Q1. Diagnose a valid-shape head-order bug

Two head results are concatenated in reverse order, but the row blocks of the trained $W_O$ remain unchanged. The final output still has the expected shape. What is the most direct test?

Choose one

Compare ordered pre-W_O head blocks and the final output with a separate-head reference

Accept the result because its final shape is correct

Remove causal masking and see whether the shape changes

Select one choice, then check.

Hint

Ask what

W_O

expects in its first

d_v

input coordinates.

Solution

Compare the ordered concatenated blocks and

Y

with the verified separate-head form. Reversing heads requires reversing the corresponding row blocks of

W_O

to preserve the function.

Audit in Two Nested Loops

First audit projections, masks, row sums, and values inside every head. Then audit the transformations across heads: split, transpose, concatenate, head order, and $W_O$ . This nested order localizes a failure before it is hidden by the final projection.