Causal Self-Attention in NumPy: Stop Future-Token Leakage

Imagine that a text model is learning to continue this status message:

greenhouse sensor reports soil is dry

When the model processes reports, it may use greenhouse, sensor, and reports. It must not read soil, is, or dry, because those words are later in the sequence. If it can see them during training, information from the targets it is meant to predict can leak into the calculation.

A causal mask prevents that leak. It blocks every attention link from a token to a later token. This tutorial builds the mask with NumPy, applies it before softmax, and tests the result by changing a future token. The goal is not to train a language model. It is to make one information-flow rule inside a next-token model visible and testable.

This lesson starts after the query-key score calculation. If queries, keys, values, and scaled dot-product attention are new to you, first read scaled dot-product attention in NumPy. Here, we focus on how token order changes which scores are allowed.

Setup and the Small Teaching Dataset

You need Python 3.11 or newer and basic experience with NumPy arrays. Create a virtual environment and install the packages used by the executed example. On macOS or Linux, run:

python -m venv .venv
source .venv/bin/activate
python -m pip install numpy pandas matplotlib

On Windows PowerShell, replace the activation command with .venv\Scripts\Activate.ps1.

The published run used Python 3.13.2, NumPy 2.5.1, pandas 3.0.3, and Matplotlib 3.11.0. Download greenhouse_token_vectors.csv and save it beside your Python file.

The CSV is an original synthetic dataset released under CC0-1.0. Its sentence and numeric features were created only for this lesson. They do not come from a trained model or a real greenhouse. That makes the example reproducible and small enough to inspect.

Load the six tokens and their four features:

import numpy as np
import pandas as pd

tokens = pd.read_csv("greenhouse_token_vectors.csv")
vectors = tokens.drop(columns="token").to_numpy(float)

print(tokens.to_string(index=False))
print("vector shape:", vectors.shape)

     token  entity  action  condition  urgency
greenhouse     1.0     0.1        0.0      0.2
    sensor     0.8     0.5        0.0      0.1
   reports     0.1     0.8        0.2      0.0
      soil     0.7     0.1        0.9      0.0
        is     0.0     0.7        0.1      0.1
       dry     0.2     0.1        1.0      0.8
vector shape: (6, 4)

Each row represents one token. Each column is a fictional teaching feature. In a trained Transformer, learned projection matrices create separate query, key, and value vectors. Queries and keys determine the attention weights; values carry the information that those weights combine. We use the same small vector for all three roles so that extra parameters do not hide the masking rule.

Mental Model: Read Left, Never Right

Picture the sequence as six covered cards placed from left to right. At the reports position, only the first three cards have been revealed. The model can compare those visible cards and combine their information. The remaining three cards must stay covered. A causal mask is the numeric record of which cards are visible at each position.

This visibility window grows by one token per row. It never skips ahead. The pattern explains the triangle: the first query has one visible key, the second has two, and the sixth has all six. The mask does not decide which visible token is most useful. It only defines which comparisons are allowed; softmax decides how to divide weight among those positions.

Why Unmasked Attention Leaks the Future

Self-attention compares each token with every token. Without a mask, matrix multiplication produces a complete 6 × 6 grid of scores. Python uses zero-based indexing, so row 2 belongs to reports; columns 3, 4, and 5 belong to later words.

The function below converts each score row into attention weights. Softmax turns arbitrary scores into non-negative shares that sum to 1. Subtracting the largest score first avoids numeric overflow without changing the result.

def softmax(scores):
    shifted = scores - scores.max(axis=-1, keepdims=True)
    exp_scores = np.exp(shifted)
    return exp_scores / exp_scores.sum(axis=-1, keepdims=True)

scores = vectors @ vectors.T / np.sqrt(vectors.shape[-1])
unmasked_weights = softmax(scores)

print(np.round(unmasked_weights[2], 3))

[0.147 0.171 0.189 0.158 0.179 0.156]

Read this row from left to right. At the reports position, the calculation gives 15.8% weight to soil, 17.9% to is, and 15.6% to dry. Together, those later positions receive 49.3% of the attention weight. The arithmetic is valid, but the information flow is wrong for next-token prediction.

This restriction is called causal because an output may depend only on information available at that point in the sequence. Changing a later token must not change an earlier representation. Some documentation also calls this an autoregressive mask. Autoregressive generation builds a sequence one new token at a time from the prefix already available.

Build a Lower-Triangular Mask

The allowed links form a lower triangle. Row 0 can read only column 0. Row 1 can read columns 0 and 1. The final row can read every column because every token is available by then.

NumPy’s tril keeps the lower triangle, including the main diagonal by default. Create a Boolean matrix where True means “this query may read this key”:

sequence_length = len(tokens)
causal_mask = np.tril(
    np.ones((sequence_length, sequence_length), dtype=bool)
)

print(causal_mask.astype(int))

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

Rows are query positions and columns are key positions. A 1 allows a link. A 0 blocks it. The main diagonal stays open because a token is allowed to use its own representation.

Apply this mask to the scores before softmax. Invalid positions receive negative infinity. Since exp(-inf) is zero, their final attention weight is exactly zero:

masked_scores = np.where(causal_mask, scores, -np.inf)
causal_weights = softmax(masked_scores)

print("reports weights:", np.round(causal_weights[2], 3))
print("row sums:", causal_weights.sum(axis=1))

reports weights: [0.29  0.337 0.374 0.    0.    0.   ]
row sums: [1. 1. 1. 1. 1. 1.]

The reports row now reads only greenhouse, sensor, and itself. Its three allowed weights are larger than before because softmax redistributes the full total of 1 across only those positions. The three future weights are exactly zero, and every row still sums to 1.

Current PyTorch documentation exposes the same rule through the is_causal=True option in scaled_dot_product_attention. That API can use optimized implementations. Our NumPy version is intentionally direct so that you can inspect the mask.

In the left heatmap, every row spreads weight across all six columns. In the right heatmap, the empty upper-right triangle shows the blocked future links. The first row has one weight of 1 because the first token can read only itself. Each later row has one more available key.

Create Context Without Reading Ahead

Attention weights are not the final output. The calculation uses them to make a weighted mixture of value vectors, called a context vector. Because our teaching example uses the input vectors as values, one matrix multiplication creates the result:

causal_context = causal_weights @ vectors

print("context shape:", causal_context.shape)
print("reports context:", np.round(causal_context[2], 3))

context shape: (6, 4)
reports context: [0.596 0.496 0.075 0.092]

There is one output row per input token and one output feature per value feature. The reports context is a weighted combination of only the first three rows. It contains no contribution from soil, is, or dry.

Here is a reusable version that returns both outputs and weights. Returning the weights makes the function easier to test:

def causal_self_attention(vectors):
    width = vectors.shape[-1]
    scores = vectors @ vectors.T / np.sqrt(width)
    allowed = np.tril(np.ones(scores.shape, dtype=bool))
    masked_scores = np.where(allowed, scores, -np.inf)
    weights = softmax(masked_scores)
    context = weights @ vectors
    return context, weights

This function handles one sequence stored in a two-dimensional array. Production code commonly adds batch and attention-head dimensions. The central rule remains the same: the mask must align its final two dimensions with query rows and key columns.

Test Causality, Not Just the Mask Shape

A triangular picture is useful, but a behavior test gives stronger evidence. If we replace the final token vector with very large numbers, the earlier outputs must remain unchanged. The last token is in their future.

The next code compares that change with and without the causal mask:

changed_vectors = vectors.copy()
changed_vectors[-1] = [9.0, 9.0, 9.0, 9.0]

changed_scores = changed_vectors @ changed_vectors.T / np.sqrt(4)
changed_open = softmax(changed_scores) @ changed_vectors
changed_causal, _ = causal_self_attention(changed_vectors)

open_context = unmasked_weights @ vectors
earlier_rows = slice(0, -1)

open_change = np.abs(changed_open[earlier_rows] - open_context[earlier_rows]).max()
causal_change = np.abs(changed_causal[earlier_rows] - causal_context[earlier_rows]).max()

print("maximum earlier change without mask:", round(open_change, 6))
print("maximum earlier change with mask:", round(causal_change, 6))

maximum earlier change without mask: 8.753238
maximum earlier change with mask: 0.0

Without a mask, changing dry changes at least one value in the earlier context vectors by as much as 8.753238. With the causal mask, the maximum change across all earlier rows is exactly zero. The final row may change because it is allowed to read the final token. That is the behavior we want.

Add two compact assertions to make this property part of an automated test:

assert np.count_nonzero(np.triu(causal_weights, k=1)) == 0
assert np.allclose(changed_causal[:-1], causal_context[:-1])

The first assertion checks that every future-position weight is zero. The second checks the more important behavior: editing a future token cannot affect any earlier context vector.

Common Mistakes and Troubleshooting

The triangle points the wrong way. With query positions in rows and key positions in columns, allowed positions are on and below the main diagonal. If your code swaps those axes, verify the result with the first row: the first token should read only itself.

The mask is applied after softmax. Zeroing weights after softmax makes a row total less than 1. Apply negative infinity to forbidden scores first, then run softmax so the remaining weights are normalized correctly.

The diagonal is blocked. np.tril(..., k=-1) removes the main diagonal. Standard next-token self-attention normally keeps the current position available. Use the default k=0 for the mask shown here.

Boolean mask meaning changes between libraries. In this tutorial, True means allowed. Library APIs do not all use the same convention. Check the current documentation before moving a mask between NumPy, PyTorch, TensorFlow, or another framework.

A row becomes all nan values. Softmax cannot form a distribution when every score in a row is negative infinity. Make sure every query has at least one valid key. Keeping the diagonal open guarantees that for this square self-attention mask.

Padding and causality are treated as the same mask. They solve different problems. A causal mask blocks later positions. A padding mask blocks placeholder tokens. A batch of unequal-length sequences may need both rules combined before softmax.

The test changes an earlier token. To test causal isolation, change a token strictly after the outputs you compare. Earlier tokens are allowed to affect later outputs.

Recap

Causal self-attention adds an order rule to ordinary self-attention:

1. Calculate query-key scores.
2. Build a lower-triangular matrix of allowed links.
3. Replace future-position scores with negative infinity.
4. Apply stable softmax so allowed weights sum to 1.
5. Mix the values into context vectors.
6. Test that changing a future token leaves earlier outputs unchanged.

The durable mental model is read left, never right. At any sequence position, a next-token model may use the current token and the prefix to its left, but not information to its right. The triangular mask encodes that rule, and the future-token mutation test confirms that the rule works in the output—not only in a diagram.