Distill
How LLMs Work

Self-attention mechanism

6 sources · updated 1 week ago

Self-attention (also called intra-attention) is the mechanism that lets a Transformer model determine how much each element of a sequence should influence every other element when building a contextual representation. Unlike recurrent networks, which process tokens one at a time and pass a hidden state forward, self-attention examines the entire sequence simultaneously — reducing the path length between any two positions to a constant O(1) sequential operations, compared with O(n) for recurrent layers. This is what enables Transformer architecture to capture long-range dependencies that RNNs and LSTMs struggled with.

Intuition: context changes meaning

A clean illustration: in the sentence "on Friday, the judge issued a sentence", the word "sentence" is ambiguous. Self-attention resolves this by having the representation of "sentence" attend strongly to "judge" and "issued" — both of which push its meaning toward the legal concept rather than the grammatical one — while largely ignoring peripheral words. The mechanism formalizes exactly this kind of context-sensitive re-weighting.

Scaled dot-product attention

For each token, the model generates three vectors by multiplying the token's embedding by three learned weight matrices: a query vector Q (what this token is seeking), a key vector K (what each token offers), and a value vector V (the information to be aggregated). A useful analogy: Q is like the search text you type into a search engine — the thing you want to find more about; K is like the title of each web page in the results — what each candidate token represents; and V is like the actual content of those pages — the information retrieved once the right match is found.

Alignment between a query and all keys is measured by their dot product; the result is scaled by 1/√d_k to prevent the dot products from growing so large they push the softmax into regions of near-zero gradient. A softmax then normalizes these scores into attention weights summing to 1, and the output is a weighted sum of value vectors:

Attention(Q, K, V) = softmax(QKᵀ / √d_k) · V

In the original paper, d_k = d_v = 64 for the base model. Dividing by √d_k is not merely a stability trick — without it, for large d_k the dot products grow proportionally to the variance of the inputs (mean 0, variance d_k), pushing softmax into saturation.

Multi-head attention

Rather than computing one attention function over d_model-dimensional inputs, the Transformer projects queries, keys, and values into h parallel lower-dimensional subspaces, runs scaled dot-product attention in each, then concatenates and linearly projects the results:

MultiHead(Q, K, V) = Concat(head₁, …, headₕ) · Wᴼ where headᵢ = Attention(Q·WᵢQ, K·WᵢK, V·WᵢV)

The original paper uses h = 8 heads with d_k = d_v = 64 each (d_model/h = 512/8), keeping total computation similar to single-head attention at full dimensionality. GPT-2 (small) uses h = 12 heads over 768-dimensional embeddings, giving 64 dimensions per head — the same ratio. The Q, K, and V matrices for all heads are computed simultaneously and then split, enabling efficient parallel computation on modern hardware.

Each head learns to attend to different aspects of meaning simultaneously — in practice, attention visualizations show individual heads specializing in tasks like anaphora resolution (tracking what "it" refers to), syntactic dependencies (completing phrases across long distances), or short-range versus long-range context. Ablation experiments confirm that both too few and too many heads hurt quality, with the 8-head base configuration in the original paper performing best.

Three uses of attention in the Transformer

The original Transformer deploys attention in three distinct configurations: (1) encoder self-attention, where every position can attend to all positions in the previous encoder layer; (2) decoder self-attention, masked so each position can only attend to positions up to and including itself; and (3) encoder-decoder cross-attention, where decoder queries attend over all encoder key-value pairs, enabling the decoder to focus on relevant parts of the input.

From attention weights to updated embeddings

After attention weights are computed, each value vector is scaled by its weight and the results are summed to produce the attention output for that position. This output is passed through a residual connection (added back to the original position-encoded embedding) and then layer-normalized, preventing information loss as representations travel deeper into the network. The resulting updated embedding carries both the token's original meaning and the contextual information contributed by other positions. See LLM training and pretraining for how these weights are learned.