Hard Framework

Causal Attention with Framework Operations

Compute one head of causal attention using framework operations. Every query may use its own position and earlier positions, but not later ones. This version uses unbatched matrices and negative infinity for masked scores.

Build the scores and mask

Calculate scores = (Q @ K.T) / sqrt(d_k). Row i contains the scores for query i, and column j refers to key j.

Create a Boolean mask for the strict upper triangle: True where j > i. Here, True means blocked. Keep the diagonal unmasked so a query can use its own key. Replace blocked scores with negative infinity before applying softmax across the last axis. Then multiply the weights by V.

For Q = K = [[1], [2]] and V = [[10], [20]], the masked scores are:

[[1, -inf],
 [2,    4]]

The first row has weights [1, 0] and returns 10. The second row has weights approximately [0.119203, 0.880797] and returns 18.807971.

Use framework operations

In PyTorch, use torch.triu(..., diagonal=1), masked_fill and torch.softmax(..., dim=-1). Build the Boolean mask on the scores' device, for example from torch.ones_like(scores, dtype=torch.bool).

In JAX, use jnp.triu(..., k=1), jnp.where and jax.nn.softmax(..., axis=-1). Softmax gives blocked positions zero weight when each row has finite allowed scores. The diagonal ensures an allowed position here; do not extend this rule blindly to rows blocked by padding.

This implementation constructs a dense square score matrix. Masking changes which values contribute, but does not by itself avoid the quadratic score storage or make this a sparse attention algorithm.

Inputs

Implement causal_attention(Q, K, V) following these steps.

Q and K: floating-point tensors of shape (seq_len, d_k).
V: a floating-point tensor of shape (seq_len, d_v).

Dimensions are positive, and the inputs produce finite unmasked scores. Use compatible dtypes and, in PyTorch, the same device. The value width may differ from the query/key width. No batching, dropout or padding mask is needed.

Output

Return the weighted values as a tensor of shape (seq_len, d_v), leaving inputs unchanged. Return values, not scores or attention weights.

The grid below shows permitted query/key pairs rather than numerical weights. Select Causal to see the pattern used here.

Loading visualization…

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

Attention combines values using learned weights

Attention lets each output position combine information from source positions. It uses three sets of vectors:

A query for each output position.
A key for each source position, compared with the query to get a score.
A value for each source position, containing the information to combine.

Scaled dot-product attention has three steps: compute query-key scores, turn those scores into weights with softmax, then take a weighted sum of the values.

$$A=\operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}}+M\right),\qquad O=AV.$$

Here, d_k is the width of each query and key vector. The mask M is zero for allowed positions and negative infinity for blocked positions.

For a single head, Q has shape (query_length, key_width), K has shape (source_length, key_width), and V has shape (source_length, value_width). The weights A have shape (query_length, source_length). Softmax runs over source positions, so each query gets its own set of weights that sum to one. The output has shape (query_length, value_width).

For weights [0.25, 0.75] and scalar values [2, 6], the output is 0.25 * 2 + 0.75 * 6 = 5. Both values contribute; attention does not simply choose the value with the largest score.

Mask before softmax

A causal mask blocks future positions. A padding mask blocks positions added to make sequences the same length. Setting a blocked score to negative infinity makes its exponential zero, so it receives zero weight. Setting it to zero does not work: exp(0) = 1.

Each row needs at least one allowed, finite score. Applying ordinary softmax to an all-negative-infinity row produces undefined values. If you instead zero weights after softmax, the remaining weights no longer sum to one unless you normalise them again.

Why divide by the square root of the key width?

If query and key components are independent, with mean zero and variance one, the variance of their dot product is d_k. Dividing by sqrt(d_k) brings it back to one. This helps keep softmax from concentrating nearly all the weight on one source position, where its derivatives become small. Scaling helps under these assumptions; it does not prevent extreme weights for every possible input.

In self-attention, queries, keys, and values come from the same sequence. In cross-attention, queries come from one sequence and keys and values come from another.

Detecting runtime…