Attention playground

Attention is the operation at the centre of every transformer. Each token asks a question (its query), compares it with what every token offers (its key), and returns a weighted average of what the tokens carry (their values). On this page you build it yourself, one step at a time, with numbers small enough to check by hand.

You will see three figures:

The attention lab. One head on a sentence of up to 8 words. Every vector has 2 numbers, so every vector is a point in a plane.
Two heads side by side. Two heads read the same sentence in different ways, and their outputs are combined.
Sinusoidal positional encodings. How a transformer gets the order of the tokens.

1. The attention lab #

The lab computes, for every token at the same time:

Step	What happens	Formula
1. Embed	look up the vector of each word, and add a position code	$x_t = e(\text{word}_t) + s\,p_t$
2. Project	three matrices give a query, a key and a value	$q_t = x_tW_Q,\ k_t = x_tW_K,\ v_t = x_tW_V$
3. Score	compare the query of token $i$ with the key of token $j$	$s_{ij} = q_i^\top k_j / \sqrt{d_k}$
4. Softmax	turn each row of scores into weights that sum to 1	$a_{ij} = \dfrac{e^{s_{ij}/\tau}}{\sum_l e^{s_{il}/\tau}}$
5. Mix	average the values with these weights	$o_i = \sum_j a_{ij}\,v_j$

Try this #

Does the same word always get the same output? Choose "the cat sat on the mat". Click the first "the" (position 0), then the second "the" (position 4), and compare the two rows of the weight matrix. Then turn on the causal mask. Then turn it off again and move the slider $s$ of the position code to 1.
What you should see
Without positions and without a mask, the two rows are equal and the two outputs are equal: $o = (0.44, -0.28)$ for both. The two tokens have the same embedding, so they have the same query, and they see the same set of keys. Attention alone treats the input as a set of tokens. The causal mask breaks the tie (the first "the" sees only itself, the second sees five tokens), and so does the position code: with $s = 1$ the two tokens get different vectors $x_0 \ne x_4$ .
What does the temperature do? Choose "it → cat" and click "it". Move $\tau$ down to 0.1, then up to 3. Then set $\tau = 0.5$ and compare with $\tau = 1$ and the preset $W_Q = 2I$ .
What you should see
With $\tau = 0.1$ the row of "it" is almost one-hot: weight 1.00 on "cat", and the output jumps onto the value of "cat". With $\tau = 3$ the weights are flat (largest 0.28) and the output moves towards the mean of all values. Doubling $W_Q$ doubles every score, which gives exactly the same weights as halving $\tau$ . This is why the scores are divided by $\sqrt{d_k}$ : the scale of the scores controls how sharp the softmax is.
What if the query carries no information? Choose the preset $W_Q = 0$ . Look at the weight matrix. Then turn on the causal mask.
What you should see
All queries are zero, so all scores are 0 and every row of $A$ is uniform: $a_{ij} = 1/T$ . Every token gets the same output, the mean of all values. With the causal mask, row $i$ is uniform over the first $i + 1$ tokens, so $o_i$ is the running mean $\frac{1}{i+1}\sum_{j \le i} v_j$ (for example $1/3$ on each of the first three tokens in row 2).
Can context change the meaning of a word? Choose "river bank" and click "bank". Note its output $o$ . Then choose "money bank" and click "bank" again.
What you should see
"bank" has the same embedding in both sentences. In the first sentence its largest weight (0.47) is on "river" and $o = (-1.37, 0.27)$ . In the second its largest weight (0.39) is on "money" and $o = (-0.90, -1.09)$ . The output of attention is a contextual vector: it depends on the other words in the sentence.

Quick check

In the lab, you set $W_Q = 2I$ and keep $\tau = 1$ . A friend keeps $W_Q = I$ and sets $\tau = 0.5$ . Everything else is the same. How do the two weight matrices compare?

Quick check

The lab uses $W_Q = 0$ and the causal mask. The values of the first three tokens are $v_0 = (3, 0)$ , $v_1 = (0, 3)$ and $v_2 = (0, 0)$ . What is the output $o_2$ ?

2. Two heads side by side #

Multi-head attention runs $H$ heads in parallel. Each head $h$ has its own matrices $W_Q^h, W_K^h, W_V^h$ , so it computes its own weights $A^h$ and its own outputs $o^h_i$ . The outputs are put side by side (concatenated) and multiplied by an output matrix $W_O$ . Write $W_O$ as a stack of blocks $W_O^1, \dots, W_O^H$ , one block for each head. Then the result is a sum over heads:

\begin{aligned} \operatorname{MultiHead}(X)_i &= [\,o^1_i, \dots, o^H_i\,]\,W_O \\ &= \sum_{h=1}^{H} o^h_i\,W_O^h . \end{aligned}

In the figure, $W_V^h = I$ and $W_O^1 = W_O^2 = \tfrac12 I$ , so the result is the midpoint of the two head outputs.

Try this #

Why use more than one head? In "it → cat", click "hungry". Which token does each head attend to most? Then give head 2 the same preset as head 1.
What you should see
Head 1 (axis 1) puts 0.65 on "cat"; head 2 (axis 2) puts 0.49 on "ate". One head can find the subject while the other finds the verb, and the combined output mixes both. When both heads have the same matrices, they compute the same weights and the same outputs, and the second head adds nothing. Different heads are useful only because they learn different projections.

Quick check

Both heads use the same $W_Q$ , $W_K$ and $W_V$ , and $W_O^1 = W_O^2 = \tfrac12 I$ as in the figure. How does the multi-head output compare with the output $o_i$ of one head with these matrices?

3. Sinusoidal positional encodings #

Attention alone does not know the order of the tokens (experiment 1). The original transformer adds a fixed vector $PE(t)$ to the embedding of the token at position $t$ . For $i = 0, \dots, D/2 - 1$ :

\begin{aligned} PE(t)_{2i} &= \sin(\theta_i t), \\ PE(t)_{2i+1} &= \cos(\theta_i t), \\ \theta_i &= b^{-2i/D}, \end{aligned}

with base $b = 10000$ . Each pair $(PE(t)_{2i}, PE(t)_{2i+1})$ is a point on the unit circle at the angle $\theta_i t$ . A step of $k$ positions turns every pair by the angle $\theta_i k$ . This turn does not depend on $t$ , so $PE(t + k) = M_k\,PE(t)$ with a fixed block rotation $M_k$ , and

PE(t)^\top PE(t + k) = \sum_{i=0}^{D/2 - 1} \cos(\theta_i k),

a function of the offset $k$ only. The lab above uses one such pair, $p_t = (\sin 0.6t, \cos 0.6t)$ , multiplied by the slider value $s$ .

Try this #

What does the base $b$ change? Keep $D = 32$ . Read the similarity at offset $k = 8$ under the chart, first with $b = 10000$ , then with $b = 100$ . Press Play with each base.
What you should see
With $b = 10000$ the similarity at $k = 8$ is 0.66; with $b = 100$ it is 0.40. A smaller base makes the slow clocks turn faster, so nearby positions become less similar and the encoding separates positions better inside a short window. A larger base keeps slow pairs that are still different after thousands of positions, which matters for long sequences.

4. The math in one place #

For a sequence of $T$ tokens with vectors $x_t \in \mathbb{R}^{D}$ stacked as the rows of $X \in \mathbb{R}^{T \times D}$ :

\begin{aligned} &Q = XW_Q, \quad K = XW_K, \quad V = XW_V, \\ &A = \operatorname{softmax}\!\Big(\frac{QK^\top}{\sqrt{d_k}} + M\Big), \\ &O = AV . \end{aligned}

The softmax acts on each row of the $T \times T$ score matrix separately, and row $i$ of $O$ is the output $o_i$ of token $i$ .

The mask $M$ has $M_{ij} = -\infty$ where query $i$ may not see key $j$ , and 0 elsewhere. The causal mask hides the later tokens: $M_{ij} = -\infty$ for $j > i$ . Since $e^{-\infty} = 0$ , masked keys get weight exactly 0, and the other weights still sum to 1.
Why $\sqrt{d_k}$ . If the entries of $q$ and $k$ are independent with mean 0 and variance 1, then $q^\top k$ has variance $d_k$ . Dividing by $\sqrt{d_k}$ keeps the scores of order 1, so the softmax is not too sharp at the start of training.
The output is a weighted average. The weights in a row are positive and sum to 1, so $o_i$ lies in the convex hull of the values $v_j$ (the grey region in the lab).
Temperature. The lab divides the scores by $\tau$ before the softmax. A standard transformer has $\tau = 1$ . We add $\tau$ only to show the effect of the scale of the scores.
Cost. The score matrix has $T \times T$ entries, so attention costs $O(T^2 d_k)$ time and $O(T^2)$ memory for each head.

Key takeaways

Attention computes, for each token $i$ , the output $o_i = \sum_j a_{ij} v_j$ : a weighted average of the values of all visible tokens.
The weights are a softmax of the scores $s_{ij} = q_i^\top k_j/\sqrt{d_k}$ with $q = xW_Q$ and $k = xW_K$ ; each row of weights is positive and sums to 1.
Larger scores give sharper weights: multiplying all scores by $c$ is the same as dividing the temperature by $c$ ; the factor $1/\sqrt{d_k}$ keeps the scores of order 1.
A causal mask sets $s_{ij} = -\infty$ for $j > i$ before the softmax, so each token reads only itself and earlier tokens.
Without positions or masks, attention is permutation-equivariant: equal tokens get equal outputs, whatever their place.
Multi-head attention runs several heads with their own projections and combines them: $\operatorname{MultiHead}(X)_i = \sum_h o^h_i W_O^h$ .
Sinusoidal encodings give each pair of dimensions a rotation of angle $\theta_i t$ with $\theta_i = b^{-2i/D}$ ; then $PE(t)^\top PE(t+k) = \sum_i \cos(\theta_i k)$ depends only on the offset $k$ .

5. Learn more in the tracks #

Deep Learning, lesson 9: Attention: queries, keys, values and masks: the full derivation, the softmax Jacobian, padding masks, multi-head shapes and cross-attention.
Deep Learning, lesson 10: Positional encodings: sinusoidal and learned encodings, RoPE and ALiBi.
Deep Learning, lesson 11: The transformer: how attention, the MLP, normalization and residual connections form a block.
Linear Algebra, lesson 15: an attention head as two low-rank matrices $W_QW_K^\top$ and $W_VW_O$ around a softmax.
Linear Algebra, lesson 4: rotations and orthogonal matrices, used by the positional encodings.
Calculus, lesson 6: backpropagation through the softmax.
Large Language Models: how attention is used to build and train language models.