Attention Visualizer

About

Attention is how a transformer lets each word gather information from the other words in its context. For every token it computes a set of weights over all tokens, then takes a weighted average of what they offer. This lab computes that for real, on embeddings small enough to read.

The weights in the presets are hand-built, not learned, so each pattern has an explanation you can check in the maths panel. Try this order: follow the previous token head along the sentence with the arrow keys; switch the causal mask on and notice the upper triangle disappear; drag temperature to 0.1 and to 3 and watch rows go from one-hot to uniform; then add a random head to see what attention looks like before training.

Related: the Transformers lesson explains why attention replaced recurrence, and the Transformer Explorer places this step inside a whole model.

How it works

Scaled dot-product attention, as defined by Vaswani et al. (2017):

1. Project. Each embedding row x is multiplied by three matrices. The query q = x·W_Q is what a token is looking for; the key k = x·W_K is what it advertises; the value v = x·W_V is what it hands over if chosen.
2. Score. The dot product q_i·k_j measures the match. Dividing by √d_k keeps the variance of the scores roughly constant as d_k grows, otherwise softmax saturates and gradients vanish.
3. Normalise. Softmax turns each row of scores into weights that are positive and sum to 1.
4. Mix. The output for token i is Σ_j a_ij · v_j. It has the same shape as a value vector, so it can be added back to the token’s representation.

How the previous-token head works. Position p is stored as (sin pω, cos pω). A rotation by −ω is a linear map, so W_Q can turn position i into position i − 1. The dot product with key j is then proportional to cos((i − 1 − j)ω), which peaks when j = i − 1. Two frequencies (ω = 1 and 0.45) remove the ambiguity a single sinusoid would have. This is exactly the property Vaswani et al. cite for choosing sinusoids: PE(p + k) is a linear function of PE(p).

Heads in real models

In trained models, many heads do recognisable jobs. Clark et al. (2019) found BERT heads that attend to the previous or next token, to the direct objects of verbs, and to the determiners of nouns, as well as heads that put most of their weight on the special [SEP] token.

Anthropic’s interpretability work (Elhage et al., 2021; Olsson et al., 2022) showed that a previous-token head can combine with a second head to form an induction head: “find where the current token appeared before, and copy what came next”. Induction heads appear abruptly during training and coincide with a jump in in-context learning.

The first-token preset mirrors attention sinks: Xiao et al. (2023) observed that large language models put a lot of attention on the first token regardless of its meaning, because softmax weights must sum to 1 and a head with nothing to do needs somewhere to put them. Keeping those first tokens in the cache is what lets StreamingLLM run on very long streams.

What is simplified

Everything you see is computed exactly, but the setting is deliberately small:

• Embeddings are 10 hand-designed features. A trained model’s embeddings have hundreds to thousands of learned dimensions and encode meaning, not just word type (see the Embedding lab).
• Word types come from a small lookup table and suffix rules, which is why you can correct them by clicking.
• Preset weights are chosen by hand to produce clean patterns. Trained heads are messier and often do several things at once.
• Real models tokenise into sub-word pieces, stack dozens of layers, and add the attention output to a residual stream.

A caution when reading attention maps from real models: high attention weight is not the same as importance for the output. Jain and Wallace (2019) showed that very different attention distributions can yield the same predictions, so treat heatmaps as a clue, not an explanation.

References

• Vaswani, A. et al. (2017). Attention is all you need. NeurIPS.
• Bahdanau, D., Cho, K., Bengio, Y. (2014). Neural machine translation by jointly learning to align and translate.
• Clark, K. et al. (2019). What does BERT look at? An analysis of BERT’s attention. BlackboxNLP.
• Elhage, N. et al. (2021). A mathematical framework for transformer circuits. Anthropic.
• Olsson, C. et al. (2022). In-context learning and induction heads. Anthropic.
• Xiao, G. et al. (2023). Efficient streaming language models with attention sinks.
• Jain, S., Wallace, B. C. (2019). Attention is not explanation. NAACL.