MATH7005 Chap.9 Multi-Head and Causal Attention
Multi-Head and Causal Attention
Parallel subspaces with controlled time access
Multi-head attention divides the model channel width across several learned query, key and value subspaces. Each head attends across the token sequence, then the outputs are concatenated and projected back to model width.
Causal masking prevents any autoregressive position from using later tokens.
What to verify
Label batch, head, time and channel axes before reshaping. Test a tiny lower-triangular mask, confirm every permitted attention row sums to one and distinguish a causal mask from a padding mask.
Reshaping must preserve semantic axes
A combined projection begins with batch, time and three times model width.
Splitting Q, K and V and exposing heads changes how the channel dimension is interpreted without changing the number of elements. A transpose then places head and time where the attention multiplication expects them.
Record each shape instead of trusting a successful reshape.
Causality is an invariance test
Two sequences with the same prefix and different futures should produce identical outputs throughout the shared prefix when the decoder runs in evaluation mode. Pair this negative control with a changed-prefix case that should alter later outputs.
The combination proves both no leakage and test sensitivity.
Mask families solve different problems
A causal mask blocks future positions by order. A padding mask blocks placeholders with no content. Cross-attention may need source padding without a decoder triangle. Test a row containing both kinds of blocked cell and apply the combined restriction before softmax so permitted weights still sum to one.
What this chapter covers
- 01
Combined QKV projection
- 02
Head width
- 03
Tensor reshape
- 04
Per-head attention
- 05
Causal mask
- 06
Cross-attention lengths
Worked example · free
Audit multi-head tensor shapes
- 2Find the per-head width.
- 3Write the Q shape.
- 3Write the score shape.
Key terms
- Attention Head
- One learned query, key and value subspace operating across the sequence.
- Head Width
- The model channel width allocated to one attention head.
- Causal Mask
- A triangular restriction that hides future positions before attention normalisation.
- Padding Mask
- A restriction hiding placeholder tokens that contain no sequence content.
- Output Projection
- The learned transformation applied after concatenating all head outputs.
Multi-Head and Causal Attention FAQ
Does each head receive different tokens?
No. Each head usually sees the same sequence positions but projects their channel representations into different learned subspaces. The time axis remains intact while the model-width axis is divided.
How can I test a causal mask?
Use a very short sequence and inspect every row. The first position may see only itself, the next sees the first two positions, and the final row sees the complete available prefix.
How is cross-attention shaped differently?
Queries may have one sequence length while keys and values share another. The score matrix uses query length by key length, and the returned output keeps the query length.
How do I prove that future tokens are hidden?
Run two evaluation-mode sequences with the same prefix and different futures. Outputs throughout the shared prefix should match. Then change a permitted prefix token as a positive control and confirm that later prefix outputs can respond.
Why test masks at several sequence lengths?
A mask sliced incorrectly from a precomputed maximum may work at one length and fail elsewhere. Exercise short, boundary and product lengths while checking shapes, blocked cells and row sums after normalisation.
Why must model width divide by the head count?
Equal heads partition the model channels into whole per-head widths. If divisibility fails, the intended reshape cannot preserve every element consistently, so the architecture or head count must change before attention is computed.
When is a padding mask different from a causal mask?
A padding mask hides placeholder positions that contain no real token, while a causal mask hides future positions even when they contain valid tokens. A decoder batch may need both restrictions combined before softmax.
Assessment move
Pick a model width divisible by several heads and write the combined QKV shape, separated tensors, exposed-head layout, score matrix and concatenated output. Track the element count through every reshape and transpose. Draw a four-token causal mask and state the permitted positions in each row, then combine it with a padding example.
Build two sequences sharing a prefix and changing only the future; verify identical prefix outputs in evaluation mode. Add a changed-prefix positive control to prove the test is sensitive. For cross-attention, use different query and key lengths and write the rectangular score shape. Keep these assertions in the project code so refactoring cannot silently swap time and head axes.
Working through Multi-Head and Causal Attention in MATH7005? Sia is AskSia’s AI Computer Science tutor — ask any MATH7005 Multi-Head and Causal Attention question and get a clear, step-by-step explanation grounded in how MATH7005 is taught and assessed. Read this chapter free, then take your hardest questions to Sia.