Causal Multi-Head Self-Attention

Let each position use only its prefix

In next-token training, an input token at position i is used to predict token i + 1. Its attention layer can read positions 0 through i, including the current input token. Reading later input positions would reveal information unavailable when generating the next token.

A causal mask enforces that boundary. Each row represents a query and each column a key. Entry (i, j) is allowed exactly when j <= i. This exercise initializes one self-attention layer and applies that mask. It does not train a language model or generate tokens.

Keep the diagonal

For four positions, the allowed entries are:

[[1, 0, 0, 0],
 [1, 1, 0, 0],
 [1, 1, 1, 0],
 [1, 1, 1, 1]]

The first query can use only key zero. The second can use keys zero and one. Including the diagonal ensures that every query has at least one allowed key. Masking every key in a row is not a valid way to request a zero attention output.

Use a boolean lower triangle, or an equivalent array with ones for allowed entries and zeros for blocked entries. This is a permission mask, not a matrix of additive score penalties.

Pass the mask to Linen

nn.MultiHeadDotProductAttention takes mask=mask during initialization and application. Its default attention implementation replaces blocked scores with the most negative finite value of the score dtype before softmax. It does not multiply those scores by zero: a score of zero would still receive positive softmax weight.

For this unbatched input, attention weights have shape (H, T, T). A (T, T) mask broadcasts over heads. A (1, T, T) mask works too; three-dimensional masks are not inherently invalid. For batched attention with weight shape (B, H, T, T), a (B, 1, T, T) mask can provide a different mask per batch item while sharing it across heads.

With ordinary finite scores and at least one allowed key, the causal mask gives the first query weights [1, 0, 0, 0]. Later queries split their weight among allowed keys according to their query/key scores.

Check that future tokens cannot change the prefix

This complete example changes the final input row and checks that the earlier outputs stay the same, using the same initialized parameters:

import jax
import jax.numpy as jnp
import flax.linen as nn

x = jnp.arange(24, dtype=jnp.float32).reshape(4, 6) / 24.0
T = x.shape[0]
mask = jnp.tril(jnp.ones((T, T), dtype=jnp.bool_))
attn = nn.MultiHeadDotProductAttention(num_heads=2, qkv_features=8)
variables = attn.init(jax.random.PRNGKey(0), x, mask=mask)
out = attn.apply(variables, x, mask=mask)
changed = x.at[-1].add(1.0)
changed_out = attn.apply(variables, changed, mask=mask)

assert out.shape == (4, 6)
assert jnp.allclose(out[:-1], changed_out[:-1])
assert jnp.allclose(
    out, attn.apply(variables, x, mask=mask[None, :, :])
)

For this layer, evaluating a shorter prefix with the same parameters also gives the corresponding prefix of the full causal output, up to numerical rounding. The mask does not create a key/value cache. Cached decoding needs its own cache state and position handling.

Dimensions and defaults

x has shape (T, D_in). The module separately projects it to queries, keys and values with total width D = qkv_features. Each of the H heads has width D // H. The output projection returns to D_in. D must be divisible by H; D_in need not be.

Keep the module's default biases and initializers, zero dropout and disabled decoding. No position encoding, residual connection or surrounding normalization is added. Initialization creates fresh weights; application reads them without training them.