Chapter 2.1 β€” Attention Is All You Need

Contents

  1. Picking up where recurrence left off
  2. Scaled dot-product attention: the mechanism itself
  3. Why the scaling factor exists: a derivation
  4. Multi-head attention: why one attention function isn't enough
  5. The encoder and decoder stacks
  6. Three uses of one mechanism: self-attention, masked self-attention, and cross-attention
  7. The real payoff: parallelism
  8. Interview angle
  9. Self-check questions
  10. Sources

1. Picking up where recurrence left off

Chapter 1.5 ended on a pointed question. Bahdanau attention and Luong attention had already shown that a decoder could look back at every encoder state and learn where to focus, rather than being forced to compress an entire source sentence into one fixed-size vector. That was a genuine breakthrough, but in every seq2seq system built on it, attention was still a helper mechanism bolted onto an RNN. The recurrent backbone was still doing the heavy lifting of representing the sequence, and attention was just a smarter way of reading from it. The question that closed the chapter was blunt: if attention is good enough to let a decoder retrieve exactly the information it needs from anywhere in the source, why keep the recurrence at all? What if attention were not a helper mechanism but the entire mechanism β€” the thing that builds representations, not just the thing that reads them?

This is exactly the question Vaswani et al. answered in 2017 with the Transformer. Their argument was not primarily that attention produces better representations than recurrence, though it often does. Their argument was architectural and computational: recurrence has an inherent sequential dependency β€” you cannot compute the hidden state at position t until you've computed it at position t-1 β€” and that dependency is a training-time bottleneck that gets worse, not better, as sequences and datasets grow. If you could replace recurrence with a mechanism that lets every position attend to every other position directly, with no step-by-step dependency chain, you could compute an entire sequence's representations in parallel. That single design decision is arguably the most consequential architectural choice in the history of deep learning for language, because it is what made it computationally feasible to train models on the scale of data and parameters that define the LLM era. Everything else in this chapter is in service of that one idea.

2. Scaled dot-product attention: the mechanism itself

Strip away the RNN entirely and ask: how would you let a token look at every other token in a sequence and decide, on its own, which ones matter? The Transformer's answer is scaled dot-product attention, and it's worth building it up from scratch rather than just stating the formula, because every piece of it is there for a reason.

Every token's embedding is projected through three separate learned linear maps into three vectors: a query vector $q$, a key vector $k$, and a value vector $v$. The query represents "what this token is looking for." The key represents "what this token has to offer, as a label for retrieval." The value represents "what this token actually contributes, as content, if it's attended to." This query/key/value split is a direct borrowing from the retrieval systems metaphor: think of it like a soft lookup table, where a query is compared against a set of keys to produce a relevance score for each entry, and those scores determine how much of each entry's value gets pulled out.

For a given query $q_i$ at position $i$, its similarity to every key $k_j$ in the sequence is computed as a dot product, $q_i \cdot k_j$, for all $j$. These raw scores are then divided by $\sqrt{d_k}$, where $d_k$ is the dimensionality of the key vectors, and passed through a softmax to turn them into a probability distribution over positions. Each output is then the weighted sum of all the value vectors, weighted by that distribution: $$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V$$ where $Q$, $K$, and $V$ are matrices stacking the query, key, and value vectors for every position in the sequence. The output for position $i$ is a weighted average of every value vector in the sequence, where the weights are determined entirely by how well position $i$'s query matched every position's key. Nothing here depends on position $i-1$'s output the way an RNN would β€” every position's attention weights can be computed independently and in parallel, given the full set of $Q$, $K$, $V$ matrices.

3. Why the scaling factor exists: a derivation

The $\sqrt{d_k}$ term looks like a minor implementation detail, but it's actually solving a specific, quantifiable numerical problem, and it's the kind of thing that separates people who've derived the Transformer from people who've only used it. Assume, as is roughly true at initialization, that the components of $q$ and $k$ are independent random variables with mean 0 and variance 1. The dot product $q \cdot k = \sum_{l=1}^{d_k} q_l k_l$ is then a sum of $d_k$ independent products, each with mean 0 and variance 1 (since $\text{Var}(q_l k_l) = E[q_l^2]E[k_l^2] = 1$ under these assumptions). Summing $d_k$ independent, identically distributed terms means the variance of the sum is $d_k$ times the variance of one term, so $\text{Var}(q \cdot k) = d_k$.

That matters because of what happens next: the dot products get fed into a softmax. As $d_k$ grows, the raw dot products grow in magnitude proportionally to $\sqrt{d_k}$ (the standard deviation), which means for large $d_k$ some of these scores can be large in absolute value. A softmax applied to inputs with large variance becomes extremely peaked β€” it pushes almost all of its probability mass onto whichever input is largest, and squashes the gradient everywhere else toward zero, because the softmax's Jacobian shrinks as its output approaches a one-hot vector. In practice this means that for large $d_k$, the attention mechanism would tend to collapse into a near-deterministic hard selection at initialization, and the gradients flowing back through the softmax into the query and key projections would vanish. Dividing the dot products by $\sqrt{d_k}$ exactly cancels the growth in variance β€” $\text{Var}(q\cdot k / \sqrt{d_k}) = 1$, independent of $d_k$ β€” which keeps the softmax's input distribution well-behaved regardless of how large the key dimension is, and keeps useful gradient flowing through attention during training. It's a small piece of arithmetic with an outsized effect: without it, scaling up model dimension would actively work against training stability, exactly backwards from what you want in an architecture whose entire purpose is to scale.

4. Multi-head attention: why one attention function isn't enough

A single attention function, however well-scaled, forces every token to combine all of its "reasons for attending" into one shared weighting over the sequence. But there are usually many distinct kinds of relationships worth tracking in a sentence at once β€” one might be about syntactic dependency (which verb does this noun belong to), another about coreference (which earlier noun does this pronoun refer to), another about simple positional proximity. A single softmax distribution over the sequence can only express one weighting at a time, which means a single-head attention mechanism is forced to average all of these signals together, potentially blurring distinctions that would be useful to keep separate.

Multi-head attention solves this by running several independent attention operations in parallel, each in its own lower-dimensional subspace, and letting the model learn what each one specializes in. Concretely, the model learns separate projection matrices for each of $h$ heads, projecting the same input into $h$ different sets of queries, keys, and values, each of dimension $d_k = d_{\text{model}}/h$. Each head computes its own scaled dot-product attention independently, producing $h$ output vectors per position, which are concatenated back together and passed through one final learned linear projection to return to dimension $d_{\text{model}}$. Because each head operates in its own subspace with its own learned projections, different heads are free to specialize β€” empirically, once trained, individual heads in a Transformer often show clearly interpretable behavior, such as one head consistently attending to the previous token, another tracking subject-verb agreement, another attending to the first token of the sequence as a kind of no-op or default. The point isn't that any single head is individually more powerful than the one big attention function it replaces; the point is that giving the model several independent, lower-dimensional attention subspaces gives it more representational flexibility to track multiple kinds of relationships simultaneously, at roughly the same total compute cost as one full-dimensional attention operation.

5. The encoder and decoder stacks

The original Transformer, as introduced for machine translation, keeps the encoder-decoder structure that seq2seq models had already established, but rebuilds both halves entirely out of attention and simple feed-forward layers, with no recurrence anywhere. The encoder is a stack of identical layers, each containing two sublayers: a multi-head self-attention sublayer, where every position attends to every position in the same (source) sequence, and a position-wise feed-forward network, which is applied identically and independently to each position's representation β€” essentially a small two-layer MLP with a nonlinearity in between, run token by token. Each sublayer is wrapped in a residual connection followed by layer normalization, a design choice whose full implications are the subject of Chapter 2.4.

The decoder mirrors this structure but with an extra sublayer and an important restriction. Each decoder layer contains masked self-attention over the target sequence generated so far, then cross-attention where the decoder's positions serve as queries against the encoder's output as keys and values, and finally the same kind of feed-forward sublayer. Stacking six of these encoder layers and six decoder layers (in the original paper's base configuration) gives a model where information from the source flows into every decoder position through cross-attention at every layer, and the target sequence builds its own internal representation of what's been generated so far through masked self-attention. Nothing in this stack has a sequential dependency across the encoder's input positions; the sequential dependency only appears in the decoder, and only because generation is inherently sequential β€” you can't attend to a token you haven't generated yet.

6. Three uses of one mechanism: self-attention, masked self-attention, and cross-attention

It's worth being precise about the fact that the Transformer doesn't actually contain three different mechanisms β€” it contains one mechanism, scaled dot-product attention, used in three different configurations, and conflating these is a common source of confusion. Self-attention, as used in the encoder, lets every position attend to every position in the same sequence, including positions after it β€” there is no restriction, because the entire input sequence is available at once and there's no notion of causality to preserve. This is what makes encoder representations bidirectional: a word's representation is built from context on both sides.

Masked (causal) self-attention, used in the decoder, is the same operation with one modification: before the softmax, positions are masked so that position $i$ can only attend to positions $j \le i$. This is typically implemented by adding $-\infty$ to the attention scores for all $j > i$ before the softmax, which drives their post-softmax probability to zero. This masking is what preserves the autoregressive property required for generation β€” at inference time, you generate one token at a time and genuinely don't have future tokens to attend to, so training has to respect that same constraint or the model would learn to rely on information it won't have at inference time. Cross-attention, used in the decoder's second sublayer, is again the identical mechanism, but with queries coming from the decoder's own sequence and keys and values coming from the encoder's output β€” this is the direct descendant of Bahdanau and Luong attention from Chapter 1.5, now expressed as a special case of the same general operation rather than a bespoke mechanism. Recognizing that all three are one operation with different inputs, rather than three separate ideas, is exactly the kind of unifying insight that a strong grasp of this architecture should produce.

7. The real payoff: parallelism

It's tempting to focus entirely on representational quality when explaining why the Transformer won, but the historically decisive advantage was computational, and it's worth stating plainly. In an RNN, computing the hidden state at every position requires having already computed the hidden state at the previous position β€” the computation graph has a sequential dependency of length equal to the sequence length, and no amount of extra hardware parallelism can shorten that critical path, because each step genuinely needs the output of the step before it. Training on a sequence of length $n$ therefore requires $O(n)$ sequential operations that cannot be parallelized across the sequence dimension, even though you might have many GPUs sitting idle.

Self-attention has no such dependency across sequence positions during a single layer's computation. For a given layer, every position's query, key, and value can be computed independently, all pairwise dot products can be computed as one matrix multiplication ($QK^\top$), and the resulting weighted sums can likewise be computed as a single matrix multiplication. This means the entire self-attention computation for a sequence of length $n$ can be executed as a small number of large matrix multiplications, which is exactly the kind of workload that modern GPU and TPU hardware is built to devour in parallel. The path length between any two positions in the network is also now $O(1)$ instead of $O(n)$ β€” information from position 1 can influence position $n$'s representation directly through a single attention operation, rather than having to propagate through $n$ sequential recurrent steps, which also tends to make optimization easier since gradients don't have to survive a long chain of sequential transformations. This is the concrete, measurable reason the Transformer could be trained on datasets and at scales that were impractical for RNNs: not that attention is a cleverer idea in the abstract, but that it converts a fundamentally sequential computation into a fundamentally parallel one, and parallel computation is what modern accelerators are good at.

None of this comes for granted, though β€” a Transformer's self-attention has no built-in notion of order at all. Because attention treats its input as a set of vectors, and the weighted sum in the attention formula is invariant to the order in which those vectors are listed, a Transformer with no other modification literally cannot tell "the dog bit the man" from "the man bit the dog" from their word representations alone, since permuting the input positions leaves every attention computation unchanged. Solving that problem β€” giving a fundamentally order-blind, fully parallel mechanism a way to know where each token sits in the sequence β€” is exactly where Chapter 2.2 picks up.

8. Interview angle

"Derive why the attention scores are divided by the square root of the key dimension. What would go wrong without it?" A strong answer walks through the variance argument directly: assuming unit-variance, zero-mean, independent query and key components, the dot product of two $d_k$-dimensional vectors has variance $d_k$, since it's a sum of $d_k$ independent terms each of variance 1. Without correction, that variance grows with model dimension, pushing softmax inputs to extremes and causing the softmax to saturate β€” producing near one-hot outputs and vanishing gradients through the attention weights. Dividing by $\sqrt{d_k}$ normalizes the variance back to 1 regardless of dimension, keeping the softmax in a well-conditioned regime. The strongest answers note this is a variance-normalization argument, not an ad hoc constant, and connects it to why this matters more, not less, as models scale to larger $d_k$.

"Why use multiple attention heads instead of one larger attention operation?" A strong answer avoids the vague "so it can look at different things" non-answer and instead explains that multi-head attention partitions the representation into $h$ lower-dimensional subspaces, each with its own learned query/key/value projections, so different heads can specialize in different types of relationships (syntactic, positional, coreference-like) without being forced to average them into a single weighting. It should also note the compute-cost framing: splitting into $h$ heads of dimension $d_{\text{model}}/h$ keeps total computation roughly comparable to one head of the full dimension, so this is closer to "reallocating" capacity than "adding" it.

"Is self-attention permutation-invariant? Why does that matter?" Yes β€” the answer should show, not just assert, this: since attention output at position $i$ is a weighted sum over value vectors with weights from a softmax over dot products, permuting the order of the input tokens permutes which value corresponds to which token but produces the exact same set of dot products and weights, so the mechanism has no way to distinguish sequence order on its own. This is why positional information has to be injected explicitly β€” it isn't a minor omission, it's a direct structural consequence of treating the sequence as an unordered set of vectors for attention purposes.

"Why did removing recurrence matter so much in practice, beyond representational quality?" A strong answer centers on parallelism: RNNs have an inherent $O(n)$ sequential dependency chain in their computation graph that no amount of parallel hardware can shorten, whereas self-attention's core computation is a handful of large matrix multiplications with no cross-position sequential dependency within a layer, which maps efficiently onto GPU/TPU hardware. The candidate should also mention the $O(1)$ path length between any two positions, versus $O(n)$ for RNNs, which also tends to ease optimization.

"Walk through the difference between self-attention, masked self-attention, and cross-attention." The answer should state clearly that these are the same scaled dot-product attention operation with different sources for queries, keys, and values, and different masking. Encoder self-attention: Q, K, V all come from the same (source) sequence, no masking, fully bidirectional. Decoder masked self-attention: Q, K, V come from the same (target) sequence, with future positions masked out (added $-\infty$ before softmax) to preserve autoregressive generation. Cross-attention: Q comes from the decoder, K and V come from the encoder's output, letting every decoder position read from the full source representation.

9. Self-check questions

  1. Starting from the chain-rule decomposition established in Chapter 1.1, and the encoder-decoder framing from Chapter 1.4, explain in your own words why a mechanism that lets a token attend to any other token directly is a natural way to compute $P(x_i \mid x_1, \ldots, x_{i-1})$ without recurrence.
  2. Derive the variance of $q \cdot k$ under the standard initialization assumptions, and explain precisely how dividing by $\sqrt{d_k}$ fixes the softmax saturation problem.
  3. If you increased $d_k$ without changing the scaling factor, what would you expect to observe in the attention weight distributions early in training, and why?
  4. Explain the query/key/value split using the retrieval-table analogy, and explain what role each of the three learned projection matrices plays.
  5. What computational property of self-attention, specifically, is responsible for enabling parallel training across sequence positions, and what is the corresponding property of RNNs that this replaces?
  6. Why is self-attention permutation-invariant, and what does that imply must be true of any complete Transformer architecture (i.e., what's missing)?
  7. Describe how masked self-attention differs mechanically from unmasked self-attention, and explain why that specific difference is necessary for autoregressive generation.

10. Sources

  • Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention Is All You Need. NeurIPS 2017. arXiv:1706.03762. https://arxiv.org/abs/1706.03762

Self-check quiz

Test your understanding of this chapter with a short quiz β€” a mix of concept questions and quick calculations.

What specific numerical problem does dividing attention scores by $\sqrt{d_k}$ solve?

Explanation: Var(q·k) grows with d_k, which pushes softmax inputs to extremes and saturates it. Dividing by √d_k normalizes that variance back to 1 regardless of d_k, keeping gradients flowing.

In the original Transformer, why is masking applied to the decoder's self-attention but not the encoder's self-attention?

Explanation: Training has to respect the same constraint the model will face at inference time β€” it can't see future tokens then, so it shouldn't be allowed to during training either.

In cross-attention (the decoder's second sublayer), where do the queries, keys, and values come from?

Explanation: Cross-attention lets every decoder position read from the full source representation: queries are the decoder's own positions, keys and values come from the encoder's output.

Why does self-attention allow far more parallelism during training than an RNN?

Explanation: An RNN has an O(n) sequential dependency chain that no extra hardware can shorten. Self-attention's core computation is a handful of large matrix multiplications with no cross-position dependency within a layer.

A Transformer uses $d_{\text{model}} = 512$ and $h = 8$ attention heads. What is the dimension $d_k$ of each head's query/key vectors?

Explanation: d_k = d_model / h = 512 / 8 = 64.

Suppose a model's query and key vector components are independent, zero-mean, unit-variance random variables, and $d_k = 64$. What is $\text{Var}(q \cdot k)$ before any scaling is applied?

Explanation: The dot product is a sum of d_k independent terms each of variance 1, so Var(qΒ·k) = d_k = 64.

Two 2-dimensional vectors are given: $q = (3, 4)$ and $k = (1, 2)$, with $d_k = 2$. Compute the scaled attention score $\dfrac{q \cdot k}{\sqrt{d_k}}$ (round to two decimal places).

Explanation: qΒ·k = 3(1) + 4(2) = 11. Scaled: 11 / √2 β‰ˆ 11 / 1.41421 β‰ˆ 7.78.

The original Transformer's base configuration stacks 6 encoder layers and 6 decoder layers, each contributing exactly one self-attention sublayer (not counting cross-attention or feed-forward sublayers). How many self-attention sublayers are there in total across the full encoder-decoder stack?

Explanation: 6 self-attention sublayers in the encoder + 6 masked self-attention sublayers in the decoder = 12.