Causal Attention Mask

Causal attention lets each sequence position use its own value and earlier values, while excluding later positions. This is the masking rule used for autoregressive self-attention. This task computes the full attention output, not just the mask.

Read the score matrix by row

For each sequence in the batch, first calculate scores = q @ transpose_last_two_axes(k) / sqrt(d). Its shape is (T, T): row i is the query position, and column j is the key position. Keep a pair when j <= i; mask it when j > i.

For three positions, the allowed pairs are:

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

The diagonal stays allowed. With only one position, the output therefore equals its value vector. If all allowed scores are equal, each query returns the mean of the value vectors in its allowed prefix.

Mask before softmax

Replace future-position scores with -1e9, then apply softmax over the last axis and multiply by v. Merely setting those scores to zero would still give them positive softmax weights. Masking after softmax without renormalizing would leave the retained weights summing to less than one.

This exercise uses float32 scores and moderate input magnitudes, for which -1e9 makes masked weights underflow to zero. The sentinel is not universal: it lies outside float16's finite range and must be sufficiently below the allowed scores. An ideal additive mask uses negative infinity, provided each row has at least one finite allowed score. An entirely masked row needs separate handling; ordinary softmax does not define a useful distribution there.

Your task

Implement causal_attention(q, k, v) for three float32 tensors of the same shape (N, T, d), with positive dimensions. Use batched matrix products, a lower-triangular mask including the diagonal, and a stable row-wise softmax. For this exercise, use the finite sentinel -1e9. Do not call a built-in attention function or add padding masks, dropout or projection layers.

In PyTorch, transpose the last two axes of k, not the batch axis, and create the mask on the input device. In JAX, use jnp.swapaxes(k, -2, -1). One (T, T) mask broadcasts across all N sequences.

Output

Return the attended values with shape (N, T, d). Leave the inputs unchanged. For example, zero queries and keys with values [2, 4, 9] in a single width-one sequence produce [2, 3, 5], the successive prefix means.

The grid below shows allowed pairs rather than softmax weights. Select Causal to see this exercise's lower-triangular pattern.