Causal Attention

Compute one attention head where each token can read only itself and earlier tokens.

Start with the concrete trace below. It shows the state the algorithm must carry as it runs.

Problem trace

Causal Attention: Form scaled query-key scores, mask the strict upper triangle before softmax, then use prefix-only weights to mix values.

Input and goalCompute one attention head where each token can read only itself and earlier tokens.

Multiply queries by keysUse Q=K=[[1,0],[0,1]] and V=[[10,0],[0,20]]. QK^T is the 2 by 2 identity score matrix.

1001

axesquery rows x key rows

Recognize it

Use it when each sequence position may aggregate only its own and earlier value rows, requiring future key positions to receive exactly zero normalized weight.

Keep true

Attention row i contains scores for query i against every key; after masking and softmax, entries j>i are zero and allowed prefix weights j<=i sum to one.

Reuse it

For constrained attention, encode allowed information flow as a score-space mask before normalization so forbidden positions receive zero mass, then mix values with the resulting row-stochastic weights.

Pattern: Matrix multiplication, scaling, mask, softmax, matrix multiplication.

Simple idea: Build query-key scores. Divide by the square root of key width. Set future scores to negative infinity before softmax. Use the probabilities to mix value rows.

import numpy as np

def causal_attention(
   query: np.ndarray, key: np.ndarray, value: np.ndarray
) -> np.ndarray:
   scores = query @ key.T / np.sqrt(query.shape[-1])
   future = np.triu(np.ones(scores.shape, dtype=bool), k=1)
   scores[future] = -np.inf
   exponentials = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
   weights = exponentials / np.sum(exponentials, axis=-1, keepdims=True)
   return weights @ value

Cost: $O(sequenc{e}^2\times width)$ time and $O(sequenc{e}^2)$ score space.