AI Attention and Transformers
Calculate one masked attention head and identify what is missing.
How it works
Scaled dot-product attention scores query/key compatibility, applies a mask and softmax, then combines values. A transformer includes additional layers, position information, residual paths, normalisation and trained parameters. This two-token calculation illustrates one mechanism only. It is not a complete transformer or a test of language understanding.
Run the code
Try in playground loads this complete standard-library Python example without executing it. Choose Run when ready. Local review uses Python 3.12.14; the existing browser editor uses Pyodide 0.27.7. These small fictional examples need no account, external data, model-provider key or database service.
Example
import math
queries, keys = [[1.,0.],[0.,1.]], [[1.,0.],[0.,1.]]
values = [[10.,0.],[0.,20.]]
weights, outputs = [], []
for i, query in enumerate(queries):
scores = [sum(a*b for a,b in zip(query,key,strict=True))/math.sqrt(2) if j<=i else -math.inf for j,key in enumerate(keys)]
exps = [math.exp(score-max(scores)) for score in scores]
row = [v/sum(exps) for v in exps]
weights.append(row)
outputs.append([sum(row[j]*values[j][k] for j in range(2)) for k in range(2)])
print("first row:", weights[0])
print("rows normalised:", all(abs(sum(row)-1)<1e-12 for row in weights))
print("finite outputs:", all(math.isfinite(v) for row in outputs for v in row))
Expected result
first row: [1.0, 0.0] rows normalised: True finite outputs: True
The causal mask prevents the first token from attending to the later token. The second can use both values. The test checks normalisation, finite outputs and mask behaviour. Learned attention weights alone are not a faithful general explanation of a model’s decisions.