Medium End-to-End

Causal Self-Attention Block

Causal self-attention lets a position read itself and earlier positions, while excluding later positions. For query position i, key position j is allowed exactly when j <= i.

For three positions, the allowed query-key pairs are:

           key 0  key 1  key 2
query 0      yes     no     no
query 1      yes    yes     no
query 2      yes    yes    yes

This exercise combines that rule with multiple attention heads, an output projection, a residual addition and layer normalization. Normalization comes after the residual addition, so this is a post-normalization block.

Inputs and output

Implement causal_self_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.

All dimensions are positive. Return shape (N, T, d_model). Keep batches independent. There are no projection biases, dropout, padding masks, learned normalization parameters or feed-forward layer in this exercise.

Build the block

Project Q = x @ w_q, K = x @ w_k and V = x @ w_v.
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).
Compute scores by multiplying Q by K with only its last two axes swapped, then dividing by sqrt(d_head). Scores have shape (N, num_heads, T, T).
Replace scores at j > i with negative infinity. Apply a stable softmax over the last axis, which contains key positions. Multiply the weights by V.
Swap the head and position axes back, then reshape to (N, T, d_model). Keep the original head order when joining features.
Compute z = joined @ w_o + x, then normalize each token over features:

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

Use population variance, dividing by d_model, with epsilon inside the square root. No learned scale or offset follows normalization.

Masking details

Use a boolean (T, T) mask that broadcasts across batches and heads. In PyTorch, create it on the input device. Negative infinity gives excluded positions zero softmax weight when the allowed scores are finite. Each row here keeps at least its diagonal entry, so no row is completely masked.

A finite replacement such as -1e9 is not an exact exclusion: an allowed score can be even smaller, causing a future position to receive weight. Assume the allowed scaled dot products are finite.

Changing later input positions must not change an earlier output position. The residual and normalization preserve this property because they operate separately at each position.

The diagram starts on Causal. Filled cells are allowed connections, not attention weights. Other settings illustrate different patterns; implement the causal pattern shown in the table above.

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…