Hard End-to-End

Multi-Head Attention Block

Multi-head attention computes several attention mixtures in parallel. Each head uses its own slice of the projected queries, keys and values. The heads can assign different weights to the same sequence positions. Their outputs are then joined and mixed by an output projection.

In this exercise, you also add the original input and normalize the result. This ordering is called post-layer normalization: attention comes first, then the residual addition, then layer normalization. Follow that order; normalizing the input before attention defines a different block.

Inputs and output

Implement multi_head_attention(x, w_q, w_k, w_v, w_o, num_heads).

x: floating-point array of shape (N, T, d_model).
w_q, w_k, w_v, w_o: each has shape (d_model, d_model).
num_heads: a positive integer that divides d_model evenly.

N is the batch size and T is the number of positions in each sequence. All dimensions are positive. Return shape (N, T, d_model). There is no causal or padding mask: each query can attend to every position in its own sequence. Keep batches independent. There are no projection biases, dropout, learned normalization parameters or feed-forward layer.

Build the block

Project. Compute Q = x @ w_q, K = x @ w_k and V = x @ w_v. Each result has shape (N, T, d_model). Use the matrices in the supplied orientation; do not transpose them as if they were module weight fields.
Split into heads. Set d_head = d_model // num_heads. Reshape each projection to (N, T, num_heads, d_head), then swap the position and head axes to get (N, num_heads, T, d_head).
Attend within each head. Multiply Q by K with only K's last two axes swapped. Divide by sqrt(d_head), not sqrt(d_model). The score shape is (N, num_heads, T, T). Apply a stable softmax over the last axis (key positions), then multiply by V.
Join heads. Swap the head and position axes back, then reshape to (N, T, d_model). Keep heads in their original order. Reshaping without this axis swap would mix positions and head features.
Project and add the residual. Compute z = joined @ w_o + x.
Normalize each token. Compute the mean and population variance over the last axis of z, keeping the batch and position axes separate:

mean = mean(z, axis=-1, keepdims=True)
variance = mean((z - mean)**2, axis=-1, keepdims=True)
output = (z - mean) / sqrt(variance + 1e-5)

Population variance divides by d_model, not d_model - 1. Add epsilon inside the square root. No learned scale or offset follows normalization.

Shape and value checks

With d_model = 8 and num_heads = 2, each head has width 4. Head 0 uses the first four projected features and head 1 uses the last four. Each head produces a separate T-by-T attention matrix for every batch item.

If w_o is all zeros, the attention branch contributes zero, so the result must be layer normalization of x. If a token's features are all equal, normalization returns zeros because subtracting the mean removes them.

The diagram starts with model width 64. Change the head count to see how the projected feature axis splits; your function must handle other valid model widths too.

Loading visualization…

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Combine attention with a feed-forward network

A transformer block uses attention to combine information across token positions. Its feed-forward network transforms features at each position independently, usually using shared weights across positions. Residual connections add a sublayer's output to its input, and normalization controls the values seen by that sublayer or passed to the next one.

Split attention into heads

Multi-head attention projects inputs into queries, keys, and values for several heads. Each head computes its own attention result. The block concatenates those results and applies an output projection.

For shape (batch, time, heads * head_width), first reshape to (batch, time, heads, head_width), then move the head axis if the attention operation expects (batch, heads, time, head_width). A reshape alone cannot exchange time and head axes while preserving their meanings.

Match masks to the source of the keys

In self-attention, queries, keys, and values come from the same sequence. In encoder-decoder cross-attention, decoder states supply queries and encoder outputs supply keys and values. Output length follows the queries, even when the source sequence has a different length.

A causal mask blocks future keys in decoder self-attention. Cross-attention to an available encoder sequence usually needs a source-padding mask instead. Use the mask appropriate to the information the model is allowed to read.

Keep normalization in the specified place

Pre-normalization computes x + F(Norm(x)). Post-normalization computes Norm(x + F(x)). Moving the normalization changes the function.

The same applies to replacing LayerNorm with RMSNorm or changing the feed-forward activation to SwiGLU. Follow the exercise's equations rather than filling in details from another model with a similar name. Write each intermediate shape, and check residual additions, packed weight order, epsilon, and activations.

Detecting runtime…