Hard End-to-End

Attention Block with RoPE

This block applies rotary position embeddings (RoPE) to projected queries and keys before computing multi-head attention. It then projects the joined head outputs, adds the original input and normalizes each token.

RoPE rotates adjacent feature pairs within each head. Applying position-dependent rotations to both queries and keys makes their dot products sensitive to relative position. The scores still depend on the contents of the queries and keys; position differences alone do not determine them.

Inputs and output

Implement transformer_block_with_rope(x, w_q, w_k, w_v, w_o, num_heads, freqs_cos, freqs_sin). Keep the supplied function name.

x has shape (N, T, d_model).
All four weight matrices have shape (d_model, d_model).
num_heads is positive and divides d_model.
d_head = d_model // num_heads is even.
freqs_cos and freqs_sin each have shape (T, d_head // 2).

All arrays are floating-point, and all dimensions are positive. Return shape (N, T, d_model). The tables already contain the cosine and sine coefficients for each position and pair. Use them directly, broadcasting across batches and heads. Do not calculate new frequencies or apply sine/cosine to the tables.

This is an unmasked attention block with a residual and post-layer normalization. It has no feed-forward layer, projection biases, dropout or learned normalization parameters. Each query can attend to every position in its own sequence.

Project and rotate

Compute Q = x @ w_q, K = x @ w_k and V = x @ w_v. Reshape each to (N, T, num_heads, d_head), then swap the sequence and head axes to obtain (N, num_heads, T, d_head).

For each query or key vector at position t, pair adjacent coordinates (2*k, 2*k+1). Using the original coordinates a and b, compute:

c = freqs_cos[t, k]
s = freqs_sin[t, k]
rotated_even = a * c - b * s
rotated_odd  = a * s + b * c

Interleave the rotated even and odd coordinates back into their original pair order. Rotate Q and K only. This block combines the unrotated V vectors; do not apply the same rotation to V.

For example, coefficients c = 0, s = 1 rotate [2, 3] into [-3, 2]. Coefficients c = 1, s = 0 leave the pair unchanged. With all identity coefficients, the block reduces to the corresponding unrotated attention block.

Attend, add and normalize

Multiply rotated Q by rotated K with only K's last two axes swapped. Divide by sqrt(d_head) to get shape (N, num_heads, T, T).
Apply a stable softmax over the final axis, which contains key positions. Multiply these weights by the unrotated V.
Swap the head and sequence 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 over the last feature axis:

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, with divisor d_model, and add epsilon inside the square root. Normalization follows the residual addition and does not mix batch items or sequence positions.

The visualization shows how position can generate rotation angles. In this exercise that work is already done: the supplied tables are authoritative. Its four pairs are an illustration, not a required head width.

Loading visualization…

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

Give attention information about token order

Self-attention based only on token content has no absolute position input. Position information can be added to token vectors, added to attention scores, or introduced by transforming queries and keys.

Learned absolute embeddings provide a vector for each position. Sinusoidal encodings use sine and cosine pairs at different frequencies. When adding either to token embeddings, the feature widths must match. In a sinusoidal encoding, position determines the angle while the feature-pair index determines the frequency.

Describe the distance between query and key

Relative-position methods represent the relationship between two positions. A learned score table, for example, can use their signed distance as an index. For query position 3 and key position 1, key - query is -2, while query - key is 2. Use the task's sign convention.

Clipping distances to a fixed range makes all more-distant pairs share an entry at the corresponding end of that range.

Rotate vectors or add a score bias

RoPE rotates pairs of query and key coordinates according to position. Rotations preserve vector lengths. Applying the same frequency scheme to queries and keys makes their dot product depend on the relative rotation between their positions. RoPE transforms the vectors rather than adding an embedding to them.

Some implementations pair adjacent features; others pair corresponding features from the two halves of a vector. Use the specified pairing, because it changes the result.

ALiBi adds a linear distance bias to attention scores, with a slope for each head. It affects preferences among keys without creating a position-embedding vector. Causal masking is still needed separately: a finite negative score bias does not give a future key exactly zero attention.

Check position zero, equal positions, and a small offset. With a cache during generation, use the new token's actual sequence position; restarting at zero on every step gives the wrong positional relationship to cached keys.

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…