NNX Causal MHA

Restrict each query to its prefix

Causal self-attention lets position i attend to positions 0 through i, including itself. It blocks keys at positions greater than i. This is useful for next-token prediction: the representation of a prefix must not depend on later input tokens. The prediction target is shifted relative to the input, so allowing the current input position is correct.

Start with four biased linear projections for queries, keys, values and the output. Split the projected features into heads, as in the previous exercise. The new step is a mask on the query-key scores before softmax. There is no dropout, cache, position encoding or residual connection here.

Read the mask by row

Scores have shape (H, T, T): head, query position, key position. Create a Boolean lower triangle with jnp.tril(jnp.ones((T, T), dtype=bool)). True means the key is allowed. The (T, T) mask broadcasts over all heads.

For three positions, its rows are:

query 0:  [True,  False, False]
query 1:  [True,  True,  False]
query 2:  [True,  True,  True ]

Replace forbidden scores with negative infinity, then apply softmax along the key axis. With finite allowed scores, forbidden keys receive zero weight and each row sums to one up to roundoff. The diagonal ensures that every row has at least one allowed key.

Why zero or a finite sentinel can fail

A score of zero still contributes exp(0) = 1 to softmax. It is not a mask. A finite value such as -1e9 works only when it is sufficiently below the allowed scores. If an allowed score is even more negative, the supposedly blocked key can receive most of the probability mass. Use -jnp.inf here so the mask is independent of score magnitude.

Mask before normalization. Multiplying already-normalized weights by a mask without renormalizing leaves rows whose weights sum to less than one. If another mask removes every key in a row, plain softmax over all negative infinities gives NaNs. That case needs a separate policy; it cannot occur with this nonempty, diagonal-inclusive causal mask alone.

A complete example with prefix averages

import jax
import jax.numpy as jnp
from flax import nnx

class CausalMHA(nnx.Module):
    def __init__(self, d_model, num_heads, rngs):
        assert num_heads > 0 and d_model % num_heads == 0
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads
        self.q_proj = nnx.Linear(d_model, d_model, rngs=rngs)
        self.k_proj = nnx.Linear(d_model, d_model, rngs=rngs)
        self.v_proj = nnx.Linear(d_model, d_model, rngs=rngs)
        self.out_proj = nnx.Linear(d_model, d_model, rngs=rngs)

    def __call__(self, x):
        T, D = x.shape
        H, Dh = self.num_heads, self.head_dim
        q = self.q_proj(x).reshape(T, H, Dh).transpose(1, 0, 2)
        k = self.k_proj(x).reshape(T, H, Dh).transpose(1, 0, 2)
        v = self.v_proj(x).reshape(T, H, Dh).transpose(1, 0, 2)
        scores = (q @ k.transpose(0, 2, 1)) / jnp.sqrt(Dh)
        allowed = jnp.tril(jnp.ones((T, T), dtype=bool))
        weights = jax.nn.softmax(jnp.where(allowed, scores, -jnp.inf), axis=-1)
        per_head = weights @ v
        combined = per_head.transpose(1, 0, 2).reshape(T, D)
        return self.out_proj(combined)

model = CausalMHA(2, 2, nnx.Rngs(0))
model.q_proj.kernel[...] = jnp.zeros((2, 2))
model.k_proj.kernel[...] = jnp.zeros((2, 2))
model.v_proj.kernel[...] = jnp.eye(2)
model.out_proj.kernel[...] = jnp.eye(2)
x = jnp.array([[2., 4.], [6., 8.], [10., 12.]], dtype=jnp.float32)
print(model(x))  # [[2., 4.], [4., 6.], [6., 8.]]
changed = x.at[2].set(jnp.array([100., 200.]))
assert jnp.allclose(model(changed)[:2], model(x)[:2])

All biases start at zero. The chosen Q and K kernels make allowed scores zero, so each row averages only its prefix. Changing the final input row cannot affect the first two outputs. The example replaces weights for clarity; use the initialized projections for the graded function.

Training and decoding

A full sequence can be processed in parallel while each query is restricted to its prefix. Single-token decoding needs no additional future mask if its keys contain only valid past and current positions. A preallocated cache with unused slots still needs a validity mask, and prompt prefill or multi-token chunks still need causal ordering. This exercise implements only the full-sequence calculation, without a cache.

In the visualization, choose Causal to see the lower triangle. Filled cells show allowed connections, not their numerical attention weights.