Attention in transformers explained with numbers

Attention is the mechanism that lets each token in a transformer look back at other tokens and decide which ones matter for what comes next. This page serves the futurism goal of this wiki, along with learning and research. It belongs to a short series on the mathematics of language models, with Next-token prediction and cross-entropy loss, Token embeddings and the geometry of meaning, and How central gradient descent is to modern AI. The aim is to make the famous formula concrete enough to compute by hand, then explain why it replaced older designs, what it costs as context grows, and what researchers have found inside it.

The formula

The 2017 paper "Attention Is All You Need" introduced the transformer, a network "based solely on attention mechanisms, dispensing with recurrence and convolutions entirely." Its core operation is scaled dot-product attention:

Attention(Q, K, V) = softmax( Q K^T / sqrt(d_k) ) V

Each token's embedding is multiplied by three learned matrices to produce three vectors:

• a query (Q), which describes what this token is looking for;
• a key (K), which describes what this token offers to others;
• a value (V), which is the information it hands over if chosen.

The dot product of a query with each key gives a relevance score. Softmax, the same function that turns logits into next-token probabilities, turns those scores into weights that sum to 1. The output is the weighted average of the values. The number d_k is the length of the key vectors.

A worked example

The numbers below are illustrative and chosen to be easy to follow; real models learn their vectors and use far more dimensions. Take three tokens, "the cat sat," with 2-dimensional keys and values that have already been projected:

token   key K     value V
the     (0, 1)    (1, 0)
cat     (2, 0)    (0, 1)
sat     (1, 1)    (1, 1)

Suppose the query for "sat" is (2, 0), loosely meaning "who is doing the sitting?"

Step 1, scores. Dot the query with each key: the gives 2*0 + 0*1 = 0, cat gives 2*2 + 0*0 = 4, sat gives 2*1 + 0*1 = 2.

Step 2, scale. Divide by sqrt(d_k) = sqrt(2), about 1.414. The scores become 0, 2.828, and 1.414.

Step 3, softmax. Exponentiate and normalize. The weights come out to about 0.045 for the, 0.768 for cat, and 0.187 for sat.

Step 4, mix values. The output is 0.045*(1, 0) + 0.768*(0, 1) + 0.187*(1, 1), which is about (0.232, 0.955). The new representation of "sat" is mostly made of what "cat" offered.

Without the scaling step, the softmax of (0, 4, 2) gives weights of about 0.016, 0.867, and 0.117. The choice is sharper, which seems harmless in two dimensions but becomes a problem in many.

Why divide by sqrt(d_k)

The paper explains that if the parts of a query and a key are independent random numbers with mean 0 and variance 1, their dot product has variance d_k. For large d_k, "the dot products grow large in magnitude, pushing the softmax function into regions where it has extremely small gradients." A quick simulation makes the point: with random vectors, the typical size (standard deviation) of the dot product is about 1.4 for 2 dimensions, 8 for 64 dimensions, and 23 for 512 dimensions. Dividing by the square root of d_k brings all of these back to about 1, so softmax stays in a range where gradient descent can still learn.

Causal masking and multiple heads

A language model must not peek at the answer. When "cat" is processed, "sat" has not been written yet. The transformer enforces this by setting the scores for future positions to minus infinity before softmax, which the paper describes as "masking out" illegal connections. After softmax those positions get a weight of exactly 0. In the example, if the query for "cat" were (0, 1), its scores against "the" and "cat" would be 1 and 0, scaled to 0.707 and 0. The masked weights are about 0.67 and 0.33, and "sat" receives none.

One attention calculation can only express one kind of relevance at a time. Multi-head attention runs several in parallel, each with its own projections; it "allows the model to jointly attend to information from different representation subspaces at different positions." The original model used 8 heads of 64 dimensions each. One head might track grammar while another tracks which noun a pronoun refers to, although heads do not always split so neatly.

Why attention replaced recurrent networks

Attention began as an add-on to recurrent networks. In 2014 Bahdanau, Cho, and Bengio argued that squeezing a whole source sentence into "a fixed-length vector is a bottleneck" for translation, and let the decoder "(soft-)search for parts of a source sentence that are relevant to predicting a target word." The transformer then removed the recurrence. A recurrent network must process tokens one after another, and the 2017 paper notes that this "sequential nature precludes parallelization within training examples." Self-attention connects every pair of positions in a constant number of steps, which suits GPUs and made the model "more parallelizable and requiring significantly less time to train." That trade is a large part of why scaling language models became practical.

The cost of long context

Every token scores against every other token, so the work grows with the square of the context length n. A 1,000-token context needs about one million scores per head per layer; a 100,000-token context needs about ten billion. Researchers have answered this in several ways.

• Sparse or local attention. Longformer replaced full attention with "a local windowed attention" plus a few global tokens, so its cost "scales linearly with sequence length."
• Better memory use. FlashAttention computes exact attention but is "IO-aware," splitting the work into tiles so fewer reads and writes go to slow GPU memory. It reported a 3 times speedup on GPT-2 training.
• The KV cache. During generation, keys and values for earlier tokens are stored so they are not recomputed. Pope and colleagues found that at large batch sizes and long sequences, the time spent loading this cache from memory dominates. As illustrative arithmetic, a model with 32 layers, 8 key-value heads of 128 dimensions, and 2-byte numbers stores 131,072 bytes per token, or about 12 GiB for 100,000 tokens. Multi-query attention shares one set of keys and values across all heads to shrink this cache, and the vLLM system manages it with paging borrowed from operating systems, reporting 2 to 4 times higher throughput.

For anyone Running AI models locally, this cache is often what limits how long a conversation can get on a single GPU.

What researchers have found inside attention

Anthropic researchers identified induction heads, attention heads that complete patterns of the form [A][B] ... [A] -> [B] by finding an earlier copy of the current token and predicting what followed it. They report that these heads "develop at precisely the same point as a sudden sharp increase in in-context learning ability," visible as a bump in the training loss, and present evidence that they may drive much of in-context learning. A caution goes with this. Attention weights are tempting to read as explanations, but Jain and Wallace found that they are "frequently uncorrelated with gradient-based measures of feature importance" and that very different attention patterns can give the same prediction.

Teaching and learning with attention

By hand first. Have students compute the worked example above on paper, then change the query to (0, 1) and predict which token wins before calculating. The habit of predicting, then checking, builds intuition fast.

Then in code. This NumPy snippet reproduces the example.

import numpy as np
K = np.array([[0, 1], [2, 0], [1, 1]], float)
V = np.array([[1, 0], [0, 1], [1, 1]], float)
q = np.array([2, 0], float)
s = K @ q / np.sqrt(2)
w = np.exp(s) / np.exp(s).sum()
print(w.round(3), (w @ V).round(3))

Learners who want a visual tour can read Jay Alammar's "The Illustrated Transformer," and those ready to build a full model can study the short nanoGPT code base.

Discussion questions, essay ideas, and learning related AI prompt ideas

• In the worked example, what query vector would make "sat" attend mostly to "the," and what does that reveal about how queries and keys divide the work?
• Why does dividing by the square root of the key length keep training stable, and what would happen in a model with 4,096-dimensional heads if the step were removed?
• Attention cost grows with the square of context length. Will future models keep full attention and buy more hardware, or will a different mechanism win?
• If induction heads drive much of in-context learning, what does that suggest about how people learn from examples in a conversation?
• Can attention weights ever serve as an honest explanation of a model's decision, and what tests would show it?
• How might the KV cache shape who can run long-context AI at home versus in a data center, and what would change that?
• Plan a classroom activity where students act as tokens holding cards with keys and values, and a volunteer with a query decides where to direct attention.

See also

• Next-token prediction and cross-entropy loss
• Token embeddings and the geometry of meaning
• How central gradient descent is to modern AI
• Artificial intelligence
• Natural language processing
• Open design GPU
• Running AI models locally
• Mathematics
• Computer science
• Learning