Multi-head attention is a core idea behind transformer models used in language processing, vision, and many other AI systems. It helps a model decide which parts of an input sequence matter most when processing each token. Instead of looking at the sequence in only one way, the model uses several attention heads in parallel.
This makes it better at capturing different relationships such as nearby context, long-range dependencies, and grammatical structure.
The mechanism starts by turning each token embedding into three learned vectors called queries, keys, and values. Each head computes attention scores with score = QK^T / sqrt(d_k), converts them to weights with softmax, and uses those weights to combine the value vectors. The outputs from all heads are concatenated and passed through another learned projection.
Because different heads can focus on different patterns, the model builds a richer representation than a single attention operation alone.
Understanding Multi-Head Attention Visualized
A useful way to picture a head is as a small pattern detector with its own learned viewpoint. Every token begins as a long list of numbers. Before attention is calculated, each head maps that list into a smaller internal space.
This gives the head freedom to treat the same word differently from another head. One head may learn that a pronoun often depends on a noun earlier in a sentence. Another may notice punctuation, word order, or repeated terms.
These roles are not assigned by programmers. They emerge gradually if they help reduce training error. Some heads become easy to interpret, while others spread their weight across many tokens or appear less useful.
The attention weights form a row for each token being processed. The entries in that row describe how strongly that token uses information from each permitted position. A large weight does not mean a token is important in every sense.
It only means that this particular head found the other token useful for its current calculation. The final output is a weighted blend of value vectors, not a copy of the selected token.
This matters because several smaller weights can combine meaningful details from many positions. The scaling step keeps the score values in a range where softmax can still distinguish options without becoming nearly all or nothing too early in training.
After the heads finish, their outputs are placed side by side. The output projection then mixes these pieces into one vector with the original model width. This mixing is important.
Without it, each head would remain in a separate compartment. The projection lets later layers use a combination such as local word order from one head plus a distant reference from another. In a transformer, this process is repeated across many layers.
Early layers may use simple surface patterns. Later layers can build representations that reflect meaning in a wider context. Position information is needed too, since attention by itself does not know whether a token came first or last.
Students often encounter this idea in text prediction, translation, chat systems, image transformers, and code completion. In text generation, a causal mask blocks a token from attending to later tokens. This prevents a model from seeing the answer before predicting it.
In tasks that read a whole input at once, such as sentence classification, tokens can often attend in both directions. When studying diagrams, track the shapes of the data. Notice the number of tokens, the width within one head, the number of heads, and the width after joining them.
Keep separate the attention weight matrix, which describes token to token links, from the value output, which carries the information being blended. That distinction prevents many common mistakes.
Key Facts
- Each attention head uses learned projections: Q = XW_Q, K = XW_K, V = XW_V
- Scaled dot-product attention is Attention(Q, K, V) = softmax(QK^T / sqrt(d_k))V
- The factor 1 / sqrt(d_k) keeps large dot products from making softmax too extreme
- Multi-head attention is MultiHead(Q, K, V) = Concat(head_1, ..., head_h)W_O
- For head i, head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)
- Attention weights for one query sum to 1 after softmax, so they act like a distribution over tokens
Vocabulary
- Token embedding
- A numerical vector that represents a token such as a word or subword in a form a model can process.
- Query
- A learned vector that represents what information a token is looking for from other tokens.
- Key
- A learned vector that represents what kind of information a token offers to other tokens.
- Value
- A learned vector containing the information that gets combined using the attention weights.
- Attention head
- One parallel attention computation that learns to focus on a particular type of relationship in the input.
Common Mistakes to Avoid
- Thinking each head sees different input tokens, which is wrong because all heads usually receive the same input embeddings and differ by their learned projection matrices.
- Forgetting the scaling factor sqrt(d_k), which is wrong because unscaled dot products can become too large and make softmax produce unstable or overly sharp weights.
- Assuming attention weights are applied to keys, which is wrong because the weights are computed from queries and keys but are used to combine the value vectors.
- Believing multi-head attention is just repeated identical attention, which is wrong because each head has different learned parameters and can capture different patterns.
Practice Questions
- 1 A head has d_k = 64. If the raw dot product Q · K = 16 for one query-key pair, what scaled score is sent into softmax?
- 2 A query attends to three tokens with scaled scores [1, 0, -1]. Using e approximately [2.72, 1.00, 0.37], compute the softmax attention weights to two decimal places.
- 3 Explain why using several attention heads can represent sentence structure better than using only one head.