Multi-Head Attention
by marimo team
A primer on multi-head attention, the standard attention block used in modern Transformers.
About this notebook
This notebook builds the multi-head attention operator from the single-head definition. It uses a small synthetic sequence of 12 tokens in 8 dimensions. NumPy draws the tokens as three clusters of four points around random centers, with a fixed seed. A custom NumPy function computes the queries, keys, and values, splits them into per-head blocks, and applies a row-wise softmax to each head. Matplotlib draws the per-head attention matrices as heatmaps, and anywidget with traitlets drives a 3D view.
A dropdown sets the number of heads. You can select 1, 2, 4, or 8 heads, and the default is 4. When you change the value, the notebook recomputes the attention matrices for that number of heads. A Matplotlib grid then shows one heatmap per head, where each cell is the weight one query row places on one key row. A 3D widget stacks the same matrices as translucent planes, one plane per head. Buttons let you emphasize a single head or show all heads together, and you can rotate the stack with the mouse.
The math runs on the CPU in NumPy, and the weight matrices come from a fixed random seed. The notebook does not train the model or load external weights. The 3D widget renders with Three.js, which the widget loads from the esm.sh CDN, and it adds orbit controls for rotation. Text cells also note that multi-head attention stays permutation equivariant, so positional information must still be added to the input when order matters.
