← All Posts
Deep Learning · Transformers· Foundations

Causal Attention: The Information Boundary

Causality is a statement about what a prediction can depend on. A next-token model may use its prefix, including the token at the current input position, but may not read the target token or later tokens. The attention mask enforces this boundary inside each layer.

Inside each head → mask scores before softmax. View the model diagram →

Predict the next token from what is already known

Suppose the training text is “the cat sat on mat”, with one token per word in this example. At the input position “sat”, the model should predict “on”. It can use “the”, “cat”, and “sat”. It must not read “on” yet: that would reveal the answer.

Each row reads its prefix

Rows are query positions; columns are key positions. ✓ = may read, — = blocked.

Query ↓
Key →
thecatsatonPredict
theallowedblockedblockedblocked→ cat
catallowedallowedblockedblocked→ sat
satallowedallowedallowedblocked→ on
onallowedallowedallowedallowed→ mat
The “sat” row reads the first three inputs and predicts “on”. The diagonal is allowed because a position predicts the following token, not its own input token.

The triangular shape is the causal mask. It keeps each prediction within its available prefix, even when all the training tokens are processed together.

In symbols, the row containing $x_i$ predicts $x_{i+1}$. The input sequence $(x_1,\ldots,x_{n-1})$ is paired with targets $(x_2,\ldots,x_n)$. A full sequence probability is a product of next-token probabilities:

$$p(x_1,\ldots,x_n)=\prod_{i=1}^{n}p(x_i\mid x_{<i}).$$

$x_{<i}$ means all tokens before $x_i$. A beginning-of-sequence token can supply the input used to predict the first ordinary token.

Turn that boundary into a score mask

Add zero to allowed scores and $-\infty$ to blocked scores, then apply softmax. With query position $i$ and key position $j$, the rule is:

$$M_{ij}=\begin{cases}0,&j\le i,\\-\infty,&j>i.\end{cases}$$

Keep the original positions when using a cache

After generating a prefix, the model can reuse its saved keys and values. A new query still has its position in the full sequence, even if it is the only row in the new query tensor. If there are $P$ cached tokens and $r$ new query tokens, local query row $a$ has absolute position $P+a$, while the keys span positions $0$ through $P+r-1$.

query_positions = prefix_length + np.arange(new_length)
key_positions = np.arange(prefix_length + new_length)
allowed = key_positions[None, :] <= query_positions[:, None]

With $P=3$ and two new queries, the allowed rows are [1,1,1,1,0] and [1,1,1,1,1]. A triangular mask built directly from a $2\times5$ matrix without offset handling would hide most of the prefix. A single new query can read every existing key, including its newly appended key.

Why training is parallel but generation is sequential

During teacher-forced training, the entire correct sequence is available. The mask lets the model compute all prefix-conditioned predictions in one forward pass, while preventing future information from influencing each row. During generation, the next sampled token does not exist until the current prediction has been made, so subsequent computation depends on that sample.

Parallel training does not remove the network's depth dependence, and it does not make token-by-token generation parallel. Techniques such as speculative decoding introduce a separate verification procedure to reduce sequential target-model calls.

Padding and document boundaries add more restrictions

A padding mask excludes keys outside the example. A document-boundary mask prevents one packed example from attending to another. A causal mask excludes future positions within the same example. A valid allowed matrix is the intersection of all required conditions.

Masking padding in the loss alone does not stop valid query tokens from attending to padding representations. Conversely, masking padding keys does not stop the loss from scoring invalid target positions. Attention validity and target validity are related but distinct arrays.

A functional test for leakage

Turn off dropout, hold a prefix fixed, and change only future input tokens. Earlier logits should remain unchanged within numerical tolerance. Test several layers: a mistake in any layer can create a future-to-past path that later layers amplify.

Then compare a full forward pass with a cached pass using the same weights, positions, and input sequence. They should agree in evaluation mode for an exactly cached causal model. A mismatch suggests mask offsets, positional IDs, cache append order, or numerical implementation differences. The lab demonstrates the attention-level equivalence.

The mask is not the entire causal guarantee

Any operation that mixes positions must respect the boundary. A convolution needs a causal receptive field; normalization over token positions can leak; an external feature computed from the full target can leak. Position-wise LayerNorm and MLPs do not themselves mix sequence positions.

Try it: A cached prefix contains ten tokens and the decoder processes one new input token. How many keys may that query read?

Eleven: the ten cached keys plus the new token’s own key, assuming all are valid and no window restriction applies. Its output predicts the following token.

From reading context to transforming features

The mask determines which positions attention can read. After that attention update is added, the feed-forward network transforms each position’s features. For more on reusing past work, see KV caching. Causal self-attention is part of the original Transformer decoder.