Self-Attention
by marimo team
A primer on self-attention, a building block for modern Transformers.
About this notebook
This notebook builds the self-attention operator from scratch and explains it step by step. It uses marimo, NumPy, and Matplotlib. There is no external dataset. The notebook generates its own tokens with a NumPy random generator. It creates three clusters of four tokens each in an eight- dimensional space. It sets the query, key, and value weights to the identity, so the attention pattern reflects the raw geometry of the input. The math cells cover softmax, queries, keys, values, permutation equivariance, and two interpretations of the operator.
The notebook gives you three sliders. The first slider changes a single logit and shows how softmax reshapes the output distribution across two bar charts. The second slider picks a permutation seed and reports the numerical gap between permuting the input and permuting the output. The gap stays at the level of floating-point rounding. The third slider sets the temperature. A Matplotlib heatmap shows the full attention matrix, and a bar chart shows one query row as a distribution over keys. At temperature one you may see a block-diagonal pattern where each query attends to its own cluster. A small temperature sharpens the pattern toward a hard lookup. A large temperature smooths it toward uniform mixing.
The attention function computes scaled dot-product logits, subtracts the row maximum for stability, and normalizes each row with softmax. All computation runs on the CPU with NumPy. The token data comes from a seeded random generator, so the clusters stay the same across runs.
