Causal Attention Lab

Your project

Build the contextualisation core of a tiny sequence model. Project token vectors, apply causal self-attention, and inspect the resulting attention map. The given projections are fixed: this project implements a forward pass, not language-model training or a full transformer. No tokenizer, residual connection or layer norm is needed.

1. Project the sequence

tokens has shape (T, D). wq and wk have shape (D, K), and wv has shape (D, V). Compute Q = tokens @ wq, Kmat = tokens @ wk, and values = tokens @ wv. All dimensions are positive; inputs are float32 tensors.

2. Attend only to the past and present

Compute scores = Q @ Kmat.T / sqrt(K). Row i may use columns j only when j <= i. Set all other scores to negative infinity before stable softmax. Subtract each row’s maximum, exponentiate and divide by that row’s sum. The diagonal is allowed, so every row has at least one valid key.

3. Return contextual vectors

Return a dictionary of ordinary Python nested lists: attention of shape (T, T) and contextual = attention @ values of shape (T, V). Do not modify input tensors. No dropout, biases, or positional encoding.

Worked example

With two identical scalar tokens [[1],[1]] and scalar identity projections, attention is [[1,0],[0.5,0.5]]; contextual output is [[1],[1]]. Masking after softmax without renormalising would give an incorrect first row.

Experiment after you pass

Export the script, change only the final token and compare earlier contextual vectors: they should stay unchanged. Then remove the mask locally and observe future-token leakage. Larger attention weight means a larger mixture coefficient, not necessarily a causal explanation of a trained model’s prediction.