Multi-Head Self-Attention

A classic question you might encounter when discussing modern architectures is to explain multi-head self-attention in the context of Transformers and why we want multiple heads. It’s essential not only to recall the definition but also to convey a clear understanding of why multi-head attention is so powerful.

In the conversation that follows, your ability to articulate both the mechanics—queries, keys, values—and the rationale behind using multiple attention heads demonstrates theoretical depth and practical insight. Let’s walk through everything you need about multi-head self-attention, from a concise conceptual overview to implementation details, testing approaches, and common follow-up questions.

Why self-attention?

Imagine every word in a sentence having the power to converse directly with every other word. This is the essence of self-attention—a breakthrough that transforms how models understand context.

In traditional attention models, particularly those used in encoder-decoder architectures for machine translation, there are two distinct sequences: one produces queries (from the decoder), and the other provides keys and values (from the encoder). Essentially, one sequence asks questions while the other offers the answers.

Self-attention revolutionizes this approach by deriving queries, keys, and values all from the same sequence. Every token in the input becomes both the inquirer and the provider. This means each token can “look at” every other token, enabling the model to directly capture relationships and long-range dependencies that were difficult for older models like RNNs.

Imagine the sentence: “The cat sat on the mat.” Each word (or token) needs to understand which other words in the sentence contribute to its meaning. Self-attention allows every word to “look at” every other word to gather context. Here’s how it works:

• Each word generates three vectors:

  • Query (Q): Think of this as what the word is looking for in the sentence.

  • Key (K): This indicates what the word can offer or its characteristics.

  • Value (V): This carries the actual content or information of the word.

• Each word computes a dot product between its query vector and the key vectors of all words in the sentence, including itself, to measure their contextual relevance. This operation produces scores that measure how much each word should focus on the others. For example, if “cat” and “mat” share related features, the score between “cat’s query” and “mat’s key” might be high, suggesting a strong contextual relationship.

• These raw scores can be uneven or not directly comparable. We apply softmax to convert them into probabilities that add up to 1. Softmax scales the numbers so that higher scores become closer to 1 (high importance), and lower scores approach 0. This normalization helps determine the relative weight each token should receive.

• Each word then uses these normalized scores to weigh the value vectors of all other words. By summing these weighted values, the word gets a new representation that captures all the relevant context from the entire sentence.

As a result, single-head self-attention scales as $O\bigl(n^2 \times d_{\text{model}}\bigr)$, because we compute an $(n\times n)$ ...