Hard End-to-End

Transformer Decoder Block

Implement transformer_decoder as one decoder block: causal self-attention, cross-attention to encoder outputs, then a feed-forward network. Each stage normalizes its input before computing an update and adding a residual.

This is a single-head, unbatched exercise with identity attention projections. No learned query/key/value or attention-output matrices are needed.

Inputs and output

X: (T, D), the decoder sequence.
enc_out: (S, D), the encoder sequence. S and T may differ.
gamma1, beta1, gamma2, beta2, gamma3, beta3: (D,), the three pairs of layer-normalization scale and bias parameters.
W1: (D, F) and b1: (F,), the first feed-forward layer.
W2: (F, D) and b2: (D,), the second feed-forward layer.

T, S, D and F are positive. All inputs share a floating-point dtype and device. Return shape (T, D), preserving dtype and device and leaving all inputs unchanged.

Layer normalization

For each row independently, compute the mean and population variance across its D features. Divide the sum of squared deviations by D, not D-1:

LN(row, gamma, beta) = gamma * (row - mean) / sqrt(variance + 1e-5) + beta

The epsilon is inside the square root. Each stage has its own gamma and beta. Do not normalize across sequence positions.

Three stages

Causal self-attention. Let H1 = LN(X, gamma1, beta1). Use H1 as queries, keys and values. Compute H1 @ H1.T / sqrt(D), with shape (T, T). For query row i, mask key positions j > i with negative infinity before softmax. Position i can attend to itself and earlier positions. Normalize each score row and set X1 = X + attention_weights @ H1.
Cross-attention. Let H2 = LN(X1, gamma2, beta2). Compute H2 @ enc_out.T / sqrt(D), with shape (T, S), and normalize each row over all S encoder positions. Set X2 = X1 + attention_weights @ enc_out. Use enc_out unchanged as both keys and values; do not normalize or causally mask the encoder sequence.
Feed-forward update. Let H3 = LN(X2, gamma3, beta3). Compute FF = ReLU(H3 @ W1 + b1) @ W2 + b2, then return X2 + FF. ReLU replaces negative values with zero.

Use a stable softmax. Each residual adds to the stage's unnormalized input, not to H1, H2 or H3. Apply no final layer normalization, dropout, padding mask or positional encoding.

For a single decoder row, self-attention has only one available key and therefore returns H1 directly. Cross-attention can still choose among multiple encoder positions.

The visualization below illustrates attention visibility. This exercise uses causal visibility for self-attention and all encoder positions for cross-attention.

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…