Hard End-to-End

Transformer Encoder Block

This exercise combines self-attention and a two-layer feed-forward network. Each branch receives a normalized input, and its output is added back to the residual stream. This is a pre-normalization layout.

To keep the calculation focused, attention uses one head with identity query, key and value projections. There is no attention output projection. You will implement the two residual additions and the feed-forward network, not construct framework modules or initialize parameters.

Inputs and output

Implement transformer_encoder(X, gamma1, beta1, gamma2, beta2, W1, b1, W2, b2).

X has shape (seq_len, d_model).
gamma1, beta1, gamma2 and beta2 each have shape (d_model,).
W1 has shape (d_model, d_ff) and b1 has shape (d_ff,).
W2 has shape (d_ff, d_model) and b2 has shape (d_model,).

All arrays are floating-point and all dimensions are positive. Return shape (seq_len, d_model). There is no batch axis, dropout, causal mask or padding mask. Every query can attend to every position, including itself.

Layer normalization

For each row, compute population variance over the d_model features, with no sample-variance correction. Then apply the supplied featurewise scale and offset:

mean = mean(t, axis=-1, keepdims=True)
variance = mean((t - mean)**2, axis=-1, keepdims=True)
layer_norm(t, gamma, beta) = gamma * (t - mean) / sqrt(variance + 1e-5) + beta

The epsilon is inside the square root. The result need not have zero mean or unit variance: gamma and beta change its statistics, and epsilon affects the normalization even before that affine transformation.

Attention branch

Compute norm1 = layer_norm(X, gamma1, beta1). Use norm1 for Q, K and V:

scores = norm1 @ norm1.T / sqrt(d_model)
weights = softmax(scores, axis=-1)
X1 = X + weights @ norm1

Softmax runs across each row, over key positions, and must be numerically stable. Add the attention result to the original X, not to norm1.

Feed-forward branch

Normalize the updated residual with the second set of parameters:

norm2 = layer_norm(X1, gamma2, beta2)
hidden = maximum(norm2 @ W1 + b1, 0)
ffn_out = hidden @ W2 + b2
output = X1 + ffn_out

ReLU is the elementwise maximum with zero. The bias vectors broadcast across positions. Add the feed-forward result to X1 and return it directly; there is no final normalization.

With one token, its attention weight is one, so the attention branch adds norm1 to X. A constant input row normalizes to beta1, which may be nonzero. These are useful checks for both the identity projections and affine normalization.

The comparison below shows allowed connections. This exercise uses the bidirectional pattern, not the causal or cross-attention patterns.

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…