II · THE IDEA · ARTIFICIAL INTELLIGENCE
Self-Attention, Step by Step
▶ Listen · narrated
A transformer does not read left to right. It compares every word to every other word in parallel, decides what to attend to, then updates each position accordingly.
At a glance
- Input
- A sequence of token embeddings, each a vector of fixed dimension
- Output
- A new sequence of the same length, each vector now informed by the others
- Core operation
- Scaled dot-product attention: query · key / √d, then softmax, then weighted sum
- Learned parameters
- Three weight matrices (Q, K, V) that project embeddings into query, key and value spaces
Imagine you are reading a sentence and trying to understand each word. To understand the word "it", you need to remember what noun it refers to earlier in the sentence. Self-attention is a mechanism that lets every word look back at every other word and decide which ones are relevant. It does this by scoring each pair of words — how much does this word relate to that word? — then using those scores to mix information together. The word "it" ends up carrying a weighted blend of the words it paid attention to, which might be the noun it refers to plus a bit of context from nearby words. Every word does this simultaneously, each building its own mixture, and the result is a new version of the sentence where each word now knows something about the others.
Self-attention operates on a sequence of vectors, typically token embeddings. Three learned weight matrices — Wq, Wk, Wv — project each input vector into query, key, and value spaces. For a sequence of length n, you obtain n queries, n keys, and n values. The attention score for query i and key j is the dot product qᵢ · kⱼ, scaled by 1/√dₖ where dₖ is the key dimension. Scaling prevents the dot products from growing large enough to push softmax into regions of vanishing gradient. You compute scores for all n² pairs, then apply softmax row-wise to convert each query's scores into a probability distribution over keys. The output for position i is the weighted sum of all value vectors, using the probabilities from query i as weights. This happens in parallel for all positions, producing an output sequence of the same length. Multi-head attention runs multiple independent attention operations with different projection matrices, concatenates their outputs, and applies a final linear transformation. The mechanism is permutation-equivariant: without positional encodings, shuffling the input shuffles the output in the same way, because attention has no inherent notion of order. Computational cost is O(n²d) for sequence length n and model dimension d, dominated by the n² pairwise comparisons. Memory cost for storing the attention matrix is O(n²), which becomes prohibitive for very long sequences.
Look closer
The dot product measures alignment
Given two vectors of the same dimension, their dot product is the sum of element-wise multiplications. If the vectors point in similar directions, the result is large and positive. If they point in opposite directions, it is large and negative. If they are perpendicular, it is near zero. Self-attention uses this property to measure how much one token's query aligns with another token's key. A high dot product means strong relevance; a low one means little connection. The softmax then converts these scores into probabilities that sum to one.
The scale factor prevents saturation
Dot products grow with the dimension of the vectors. If each query and key has dimension 64, the dot product is the sum of 64 multiplications, and its magnitude can become large enough that the softmax concentrates almost all probability on a single position, leaving gradients near zero everywhere else. Dividing by the square root of the dimension — √64 = 8 in this case — keeps the scores in a range where softmax remains sensitive. Vaswani et al. state this explicitly: without the scaling, large dot products push softmax into regions of extremely small gradient.
The output is a weighted average of values
After softmax, each position has a probability distribution over all positions in the sequence. These probabilities are then used as weights to combine the value vectors. If position 3 attends strongly to positions 1 and 5, its output will be dominated by the values at those two positions, with smaller contributions from the rest. The result is a new vector for position 3 that encodes what the model decided was relevant context. This happens in parallel for every position, producing a new sequence of the same length.
The story
Start with a sequence of three tokens, already embedded. Call them x₁, x₂, x₃, each a vector of dimension four. The numbers might be [1.0, 0.5, 0.2, 0.8], [0.3, 1.2, 0.6, 0.4], and [0.9, 0.1, 1.1, 0.7]. These are the inputs.
Self-attention begins by projecting each embedding into three different spaces using learned weight matrices Wq, Wk, and Wv. Multiply x₁ by Wq to get a query vector q₁. Multiply x₁ by Wk to get a key vector k₁. Multiply x₁ by Wv to get a value vector v₁. Do the same for x₂ and x₃. The dimensions of these matrices determine the dimensions of the resulting vectors; often the query and key dimensions match, and the value dimension can differ.
Now compute attention scores. For the first position, take its query q₁ and compute the dot product with every key: q₁ · k₁, q₁ · k₂, q₁ · k₃. Suppose these come out to 2.4, 1.8, and 3.1. Divide each by the square root of the key dimension — if keys are dimension 4, divide by 2 — giving 1.2, 0.9, and 1.55. These are the scaled scores.
Apply softmax to convert scores into probabilities. Softmax exponentiates each score, then divides by the sum of all exponentials. For the scores [1.2, 0.9, 1.55], compute e^1.2 ≈ 3.32, e^0.9 ≈ 2.46, e^1.55 ≈ 4.71. The sum is 10.49. Divide each exponential by this sum: 3.32/10.49 ≈ 0.32, 2.46/10.49 ≈ 0.23, 4.71/10.49 ≈ 0.45. These are the attention weights for position 1. They sum to one, and they tell you how much position 1 should attend to each position in the sequence.
Finally, compute the output for position 1 as a weighted sum of the value vectors. Multiply v₁ by 0.32, v₂ by 0.23, and v₃ by 0.45, then add them element-wise. If v₁ = [0.5, 0.8, 0.3, 0.6], v₂ = [0.7, 0.4, 0.9, 0.2], and v₃ = [0.6, 0.5, 0.7, 0.8], the weighted sum is [0.5×0.32 + 0.7×0.23 + 0.6×0.45, ...] for each dimension. The result is a new vector for position 1, incorporating information from all three positions according to the attention weights.
Repeat this process for positions 2 and 3. Each position gets its own query, computes its own attention weights, and produces its own output. The three outputs together form a new sequence of the same length as the input. This new sequence is the result of one attention head. A transformer layer typically uses multiple heads in parallel, each with its own Q, K, V matrices, and concatenates their outputs before a final linear projection.
The mechanism is the same regardless of sequence length. A sequence of 512 tokens means each query computes 512 dot products, softmax over 512 scores, and a weighted sum of 512 value vectors. The computation scales quadratically with sequence length, which is why long contexts are expensive.
Why it mattered then
Before attention, sequence models were dominated by recurrence. An LSTM or GRU processed tokens one at a time, carrying a hidden state forward. This made parallelisation difficult: you could not compute position 100 until you had finished position 99. It also made long-range dependencies hard to learn, because information had to survive many sequential updates without degrading. Attention changed the architecture. By computing all positions in parallel and allowing each position to look directly at any other position, transformers removed the sequential bottleneck. Training became faster because every position could be computed simultaneously on the same batch. Long-range dependencies became easier because a token at position 500 could attend directly to a token at position 3, without routing the signal through 497 intermediate steps. The Vaswani et al. paper demonstrated that attention alone, without recurrence or convolution, was sufficient for state-of-the-art machine translation. The model they called the Transformer outperformed existing architectures on English-to-German and English-to-French benchmarks while requiring substantially less training time. The result was not incremental. It suggested that the inductive biases built into recurrent and convolutional architectures — the assumption that nearby tokens matter more, or that order must be processed sequentially — were less important than the ability to compare everything to everything and let the model learn the structure itself.
Why it matters now
Self-attention is the mechanism behind every large language model in wide use. GPT, BERT, LLaMA, Claude, and their descendants are all transformers, and transformers are built from stacked self-attention layers. The dot-product-softmax-weighted-sum pattern appears dozens of times in a single forward pass. The quadratic scaling with sequence length remains the primary architectural constraint. A model with a context window of 128,000 tokens must compute roughly 16 billion attention scores per layer per forward pass, and a deep model has many layers. This is why context length is a headline feature and why alternatives to full self-attention — sparse attention, linear attention, state-space models — continue to attract research effort. The mechanism works, but it is expensive, and the cost grows quickly. Understanding self-attention also clarifies what a model can and cannot do. Attention is a comparison operation: it measures similarity, retrieves relevant context, and blends information. It does not, by itself, perform symbolic reasoning or execute algorithms. Those behaviours emerge from stacking many attention layers with nonlinear transformations in between, and from training on data where such behaviours are rewarded. The mechanism is surprisingly simple. The complexity comes from scale and composition.
The surprising detail
The softmax operation, which converts scores into probabilities, is differentiable but not invertible. Once you have the attention weights, you cannot reconstruct the original scores uniquely. This means attention is lossy: information about the relative magnitudes of the scores is compressed into a probability distribution. A score of [2.0, 1.0, 0.5] and a score of [20.0, 10.0, 5.0] produce different probability distributions, but both concentrate most weight on the first position. The scaling factor √d helps keep scores in a range where softmax is sensitive, but the operation itself flattens differences. This is one reason why very deep transformers can struggle: repeated softmax operations, layer after layer, compress information in ways that gradients must work against during training.
Remember this
Self-attention is three projections, a scaled dot product, softmax, and a weighted sum. Every position compares itself to every other position in parallel.
Test yourself
You have a sequence of 1,000 tokens and you apply self-attention with queries and keys of dimension 64. You then apply self-attention again to a different sequence of 2,000 tokens, using queries and keys of dimension 128. By what factor does the total number of dot products increase?
The number of dot products scales with the square of the sequence length, not with the dimension of the queries and keys. For 1,000 tokens, you compute 1,000 × 1,000 = 1,000,000 dot products (each of the 1,000 queries is compared to each of the 1,000 keys). For 2,000 tokens, you compute 2,000 × 2,000 = 4,000,000 dot products. The factor is 4. The dimension affects the cost of each individual dot product — a dimension-128 dot product requires twice as many multiplications as a dimension-64 one — but the question asks about the count of dot products, not their individual cost. This distinction matters when reasoning about context-length scaling: doubling the sequence length quadruples the number of comparisons, regardless of embedding dimension.
Go deeper
- Attention Is All You Need · arXiv · Ashish Vaswani et al. · 2017-06-12
- The Annotated Transformer · nlp.seas.harvard.edu
Image: Original diagram, The Daily Triptych. Licence: Original work. Source.