Transformers

Why attention?

Before transformers, recurrent neural networks processed sequences one token at a time, threading a hidden state forward through the chain: $h_t = f(h_{t-1}, x_t)$. Consider what this means concretely. To compute the representation of token 500, the network must first compute the representations of tokens 1 through 499 — sequentially, one after another. This creates three problems.

Problem 1: No parallelism

A modern GPU has thousands of cores that can multiply matrices simultaneously. But an RNN forces sequential computation — $h_2$ depends on $h_1$, $h_3$ depends on $h_2$, and so on. Those thousands of cores sit idle, waiting. Training a 500-token sequence takes 500 sequential steps regardless of how many GPUs you have. This is why RNN training was painfully slow at scale, and why scaling to billions of parameters was impractical.

Problem 2: Vanishing gradients over distance

Backpropagation through $T$ timesteps involves a product of $T$ Jacobian matrices (recall Chapter 01 §6). When each Jacobian has spectral norm slightly less than 1, the product shrinks exponentially. After 100 steps, the gradient from token 100 that reaches token 1 is essentially zero. The network cannot learn long-range dependencies — it forgets the beginning of the sentence by the time it reaches the end. LSTMs and GRUs mitigated this with gating mechanisms, but they couldn't fully solve it for very long sequences.

Problem 3: Information bottleneck

The hidden state $h_t$ has a fixed dimension — say, 512. By token 500, the entire history of the sequence (all 500 tokens, their interactions, their context) must be compressed into 512 numbers. This is an information-theoretic impossibility for complex sequences. Information about early tokens gets progressively overwritten as new tokens arrive. In machine translation, this means the model "forgets" the beginning of a long source sentence by the time it finishes encoding it.

The attention solution

The key insight: let every token look at every other token directly, without routing information through a chain of intermediate states.

With attention, token 500 can directly attend to token 1 in a single computation step. There is no chain of hidden states to route through — the connection is direct. This simultaneously solves all three problems: computation becomes fully parallel (every token-to-token interaction can be computed simultaneously), gradient paths are short (one step from any token to any other), and there is no bottleneck (each token can selectively access the full context).

The cost is quadratic — computing all pairwise interactions between $n$ tokens requires $O(n^2)$ operations. For a sequence of 4,096 tokens, that's about 16 million interactions. But on modern GPUs, this is a single matrix multiply — and matrix multiplies are exactly what GPUs are built for.

Scaled dot-product attention

Think of attention as a soft dictionary lookup. You have a question (query) and a collection of entries (keys), each associated with a piece of information (value). The attention mechanism scores how well each key matches the query, then returns a weighted average of the values — with weights proportional to the match quality.

Concretely, each token in the sequence produces three vectors by linear projection:

• Query ($q$) — "what am I looking for?" This encodes what information the current token needs from other tokens. When the word "it" appears in a sentence, its query encodes the question "what noun do I refer to?"
• Key ($k$) — "what do I contain?" This advertises what information the token offers. The word "cat" produces a key that encodes "I am an animal, a noun, a potential antecedent for pronouns."
• Value ($v$) — "what information do I provide if selected?" This is the actual content that gets aggregated. While the key advertises, the value delivers.

The projections are simple linear transformations: $Q = XW_Q$, $K = XW_K$, $V = XW_V$, where $X \in \mathbb{R}^{n \times d}$ is the input (a stack of $n$ token embeddings) and the $W$ matrices are learned parameters.

The attention computation proceeds in four steps:

1. Compute scores: $S = QK^T \in \mathbb{R}^{n \times n}$. Each entry $S_{ij}$ is the dot product $q_i \cdot k_j$ — how similar token $i$'s query is to token $j$'s key. Recall from Chapter 01 §3 that the dot product measures similarity: parallel vectors give large values, orthogonal vectors give zero.
2. Scale: $S \leftarrow S / \sqrt{d_k}$. This prevents large values from pushing softmax into flat regions where gradients vanish. Without scaling, the variance of the dot products grows with $d_k$, which means softmax sees increasingly extreme inputs as the model dimension grows. Dividing by $\sqrt{d_k}$ keeps the variance near 1 regardless of dimension.
3. Normalize: $A = \text{softmax}(S)$, applied row-wise. Each row of $A$ is now a probability distribution over the sequence — it sums to 1 and tells you how much token $i$ "attends to" each other token. High values mean "pay attention here"; near-zero values mean "ignore this."
4. Aggregate: $\text{output} = AV$. Each output token is a weighted average of all value vectors, with weights given by the attention distribution. Token $i$'s output is $\sum_j A_{ij} v_j$ — a mixture of information gathered from across the sequence, with the mixture weights learned end-to-end.

Attention with actual numbers

Let's trace the full computation for a tiny example. Suppose we have 3 tokens with $d_k = 2$ (absurdly small, but it makes the math visible). After projection, suppose:

Step 1: Scores. $QK^T$ computes all pairwise dot products:

Reading row 1: token 1's query is most similar to token 1's key (score 1.0) and least similar to token 2's key (score 0). Token 1 "wants" information from itself and from token 3.

Reading row 3: token 3's query is equally similar to all keys (scores 1, 1, 1). Token 3 "wants" information from everyone equally.

Step 2: Scale. Divide by $\sqrt{d_k} = \sqrt{2} \approx 1.41$: scores become $[0.71, 0, 0.35; \ 0, 0.71, 0.35; \ 0.71, 0.71, 0.71]$.

Step 3: Softmax. Apply softmax to each row. For row 1: $e^{0.71} \approx 2.03$, $e^0 = 1$, $e^{0.35} \approx 1.42$. Sum $= 4.45$. So the attention weights are approximately $[0.46, 0.22, 0.32]$. Token 1 attends 46% to itself, 22% to token 2, and 32% to token 3.

For row 3, since all scores are equal, softmax produces $[0.33, 0.33, 0.33]$ — uniform attention. Token 3 gathers equally from everyone.

Step 4: Aggregate. Multiply attention weights by values: $\text{output}_1 = 0.46 \cdot [1,0] + 0.22 \cdot [0,1] + 0.32 \cdot [0.5, 0.5] = [0.62, 0.38]$. Token 1's output is a weighted blend — mostly its own value, with some information mixed in from the others.

This is the entire mechanism. Every operation is a matrix multiply or elementwise function. There are no sequential dependencies between tokens — the entire $n \times n$ score matrix is computed in one parallel operation. This is why transformers train so fast on GPUs.

Preventing the model from seeing the future

For language generation, we need a critical constraint: when predicting the next token at position $t$, the model should only see tokens at positions $\leq t$. It cannot look at future tokens — that would be cheating. You can't use the answer to predict the answer.

Causal masking enforces this by setting the upper triangle of the score matrix to $-\infty$ before softmax. Since $e^{-\infty} = 0$, these entries vanish after softmax. The result: each token can only attend to itself and earlier tokens.

A crucial subtlety: the causal mask is applied only once during training, but it enables generation. During training, the model processes all positions in parallel — the mask ensures each position only uses valid context, so every position simultaneously receives the right training signal. During generation, tokens are produced one at a time, and the mask naturally extends as the sequence grows.

Multi-head attention — why one head isn't enough

Consider the sentence: "The cat sat on the mat because it was comfortable." The word "it" needs to simultaneously attend to:

• "mat" (because that's what "it" refers to — the mat is comfortable)
• "sat" (because that's the action that led to the state)
• "because" (the causal connector that links the two clauses)

A single attention distribution is a probability vector — it sums to 1. If it puts high weight on "mat," it must put lower weight on "sat" and "because." But all three relationships are important simultaneously. One distribution cannot capture multiple independent relationships at once.

Multi-head attention solves this by running $h$ independent attention operations in parallel, each with its own learned projections $W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}$:

Each head operates in a lower-dimensional subspace: with $d_\text{model} = 768$ and $h = 12$ heads, each head has $d_k = 64$. The total parameter count is $4 d_\text{model}^2$ — identical to a single-head attention with the same total dimension. You get $h$ different attention patterns at no additional parameter cost.

Empirically, trained heads specialize. Some heads learn positional patterns — always attending to the previous token (useful for local syntax). Some learn syntactic patterns — attending from a verb to its subject, or from a pronoun to its antecedent. Some learn semantic patterns — attending to topically related words. And some are nearly degenerate — doing very little and potentially prunable. This specialization emerges from training; it is not engineered.

Positional encoding — because attention is order-blind

Attention computes pairwise similarities between tokens regardless of their position. The attention score between "cat" and "sat" is the same whether they are adjacent or 500 tokens apart. But position matters enormously: "The dog bit the man" and "The man bit the dog" contain exactly the same tokens.

Since attention has no built-in notion of order, position must be injected explicitly. The original transformer adds positional encodings to the token embeddings: $x_i = \text{embed}(w_i) + \text{PE}(i)$.

Sinusoidal encoding (Vaswani et al., 2017)

The original approach uses fixed sine and cosine functions at different frequencies:

Why sine and cosine? Because any fixed offset in position corresponds to a linear transformation of the encoding. Specifically, $\text{PE}(\text{pos} + k)$ can be written as a fixed matrix times $\text{PE}(\text{pos})$. This means the model can learn to compute relative positions using a simple linear operation — it doesn't need to memorize every possible distance.

Why different frequencies? The low-frequency dimensions (large $i$) change slowly across positions — they encode "coarse" position (roughly where in the sequence). The high-frequency dimensions (small $i$) oscillate rapidly — they encode "fine" position (exactly where). Together, they give a unique fingerprint for each position, analogous to how a clock uses hours, minutes, and seconds at different frequencies to uniquely identify a time.

Learned positional embeddings (GPT-2)

A trainable lookup table $P \in \mathbb{R}^{n_\text{max} \times d}$, where position $i$ gets embedding $P[i]$. More flexible than sinusoidal — the model can learn arbitrary positional patterns. But it cannot extrapolate beyond $n_\text{max}$, because positions beyond the table have no embedding.

RoPE — Rotary Position Embedding (LLaMA, Mistral)

The dominant modern approach. Instead of adding a positional embedding to the token embedding, RoPE rotates the query and key vectors in 2D subspaces by a position-dependent angle. The dot product between rotated query $\tilde{q}_m$ and rotated key $\tilde{k}_n$ depends only on the relative position $m - n$, not the absolute positions. This gives the best extrapolation behavior and is used in virtually all modern LLMs.

The transformer block — attention + MLP + residuals

A single transformer block combines multi-head attention with a position-wise feed-forward network (MLP), connected by residual connections and layer normalization. The standard "pre-norm" layout used in most modern LLMs:

Let's unpack each component's role in the information flow:

• Layer norm normalizes the input before each sub-layer. It stabilizes activations, preventing the internal values from drifting to extreme magnitudes over many layers. Pre-norm (before the sub-layer) is more stable than post-norm (after); this is why modern transformers use pre-norm.
• Multi-head attention enables tokens to exchange information — it is the only operation where different positions interact. Without attention, each token would be processed independently.
• Residual connection ($+ h$) adds the input directly to the output, so the sub-layer only needs to learn the residual — how much to change the representation. This provides a gradient highway (recall Chapter 01 §9) and is the reason 100+ layer transformers can train at all.
• MLP processes each token independently (no inter-token interaction). It applies a two-layer network that expands the dimension by 4×, applies GELU activation, then compresses back: $d \to 4d \to d$.

A transformer is just $N$ of these blocks stacked. GPT-2 small: $N = 12$. LLaMA-7B: $N = 32$. LLaMA-70B: $N = 80$. The input enters the first block as token embeddings plus positional encodings; the output of the last block feeds into the language modeling head.

What the MLP actually does

The MLP is often treated as an afterthought — "just a feed-forward network." But it contains two-thirds of the transformer's parameters and plays a distinct, essential role that attention cannot fill.

Attention is a mixing operation — it combines information across tokens but does not transform individual token representations in expressive ways. The MLP is a processing operation — it transforms each token's representation independently through a nonlinear function. Together, they form a complete computational unit: attention gathers relevant context, then the MLP processes it.

The 4× expansion is crucial: $\text{MLP}(x) = W_2\, \text{GELU}(W_1 x + b_1) + b_2$ where $W_1 \in \mathbb{R}^{4d \times d}$ and $W_2 \in \mathbb{R}^{d \times 4d}$. The expanded intermediate representation has 4 times the dimension of the input. This gives the network a high-dimensional space in which to compute nonlinear functions before compressing back. Without this expansion, the MLP would be too constrained to learn useful transformations.

Encoder, decoder, or both

The encoder-decoder architecture (T5, BART, the original transformer, Whisper) uses bidirectional attention in the encoder and causal attention + cross-attention in the decoder. The encoder processes the input with full bidirectional context; the decoder generates the output autoregressively, attending to the encoder's representations via cross-attention. This is covered in depth in Chapter 08.

Since 2022, decoder-only architectures have dominated. The key discovery: a sufficiently large decoder-only model, trained on enough data, can match or exceed encoder-only and encoder-decoder models on virtually every task — understanding, generation, translation, summarization, reasoning — using appropriate prompting.

The GPT architecture

Weight tying: a common trick where the output projection matrix $W_\text{output} \in \mathbb{R}^{V \times d}$ shares weights with the embedding matrix $W_\text{embed} \in \mathbb{R}^{V \times d}$. The intuition: generating a token (projecting the hidden state onto vocabulary space) should use the same representation as reading a token (embedding it). This reduces parameters and often improves performance.

Next-token prediction — the only objective

This is cross-entropy loss at every position, averaged over the sequence. The model sees a batch of text, predicts the next token at every position in parallel (thanks to causal masking), and receives a gradient for every prediction.

No labels needed beyond the text itself — the next token IS the label. This is self-supervised learning. The entire internet becomes a training set. This unlimited data supply, combined with the clean loss function and full parallelism, is why autoregressive transformers have scaled to trillions of parameters while other architectures hit data and compute walls.

Teacher forcing

During training, the model receives the ground-truth tokens at every position (not its own predictions). This is called teacher forcing — the "teacher" provides the correct prefix at each step. Without teacher forcing, early mistakes would compound: a wrong prediction at position 5 would mean positions 6, 7, 8, etc. are all conditioned on an incorrect prefix, making learning very slow.

The training recipe

The standard transformer training recipe (established by GPT-2/3 and refined by Chinchilla and LLaMA):

• Optimizer: AdamW with $\beta_1 = 0.9$, $\beta_2 = 0.95$, weight decay $0.1$.
• Learning rate schedule: linear warmup over the first ~2000 steps, then cosine decay to 10% of the peak LR. The warmup prevents early training instability; the cosine decay provides a smooth annealing.
• Gradient clipping: clip the global gradient norm to 1.0 to prevent explosive updates.
• Batch size: as large as possible — modern LLM training uses effective batch sizes of millions of tokens per step, achieved via gradient accumulation across many GPUs.

Training a transformer on sequence reversal

Why sequence reversal?

It's a task that requires attention to solve. The model must learn that output position 1 should copy from input position $n$, output position 2 from position $n-1$, and so on. A model that merely copies or uses local context will fail. Success requires long-range, position-dependent attention — and when it works, we can verify the learned algorithm by examining the attention maps.

Attention maps reveal the learned algorithm

After training, extracting attention weights from the output positions reveals a strikingly clean anti-diagonal pattern: output token 1 attends almost exclusively to input token $n$, output token 2 to $n-1$, and so on. The model has literally learned to "look at the right place." The attention mechanism is implementing a reversal algorithm — and we can see it directly in the attention weights.

Head specialization and entropy

Measuring attention entropy per head reveals functional specialization. The reversal-specialized heads have very low entropy — they place nearly all weight on a single position (the correct source position for each output). Other heads have higher entropy, attending more broadly — these handle contextual tasks like separator detection. This is the multi-head principle in action: different heads learn different patterns, and the network routes each task to the appropriate head.

Temperature effects on generation

At $\tau = 0.1$ (near-deterministic), the model always picks the highest-probability token and gets reversals correct. At $\tau = 2.0$ (high randomness), sampling introduces errors — the model "knows" the right answer but the soft distribution allows incorrect tokens to be sampled occasionally. This demonstrates the fundamental tradeoff between deterministic accuracy and stochastic diversity that governs all autoregressive generation (covered in depth in Chapter 07).

Interactive: Transformer Explainer

The Transformer Explainer, built by Cho et al. at Georgia Tech's Polo Club, runs GPT-2 (124M parameters) live in your browser. Type a prompt, watch tokens flow through embeddings, attention layers, and MLPs, observe attention patterns, and experiment with sampling parameters.

Try these experiments:

1. Type "The capital of France is" and generate. Examine the attention pattern for "is" — which earlier tokens does it attend to? You should see strong attention to "France" and "capital."
2. Set temperature to 0.1, generate several tokens. Then set it to 2.0. At low temperature, the model picks the most likely next token deterministically. At high temperature, rare tokens get sampled, producing surprising (often nonsensical) text.
3. Expand the attention view and look for a "previous-token head" — a head showing a diagonal stripe in the attention matrix (each token attending to the token immediately before it). This positional pattern is one of the most common head specializations.
4. Watch the final softmax distribution. How peaked is it? For factual continuations ("The capital of France is ..."), the distribution is usually very peaked (one token has >90% probability). For open-ended prompts ("Once upon a ..."), it's more spread out.

The tool runs a real GPT-2 model via ONNX Runtime. It may take a few seconds to load. If the embed doesn't work, open it directly ↗.

Takeaways & exercises

The seven things to carry forward

1. Attention is a soft dictionary lookup: Q asks the question, K advertises the content, V provides the answer. The softmax(QKᵀ/√d_k) formula computes match scores, normalizes them, and aggregates values.
2. The √d_k scaling prevents softmax saturation — without it, large dot products push softmax into near-zero gradient regions, preventing learning.
3. Multi-head attention gives multiple simultaneous views. Each head can track a different relationship (syntax, semantics, position). The total parameter count is the same as single-head attention with the same dimension.
4. Causal masking enables parallel training of sequential generation. The lower-triangular mask ensures each position only sees past tokens, so every position receives the correct training signal simultaneously.
5. Position must be injected explicitly because attention is a set operation. Sinusoidal, learned, and RoPE are three approaches; RoPE dominates modern LLMs because it encodes relative position and extrapolates well.
6. The MLP contains two-thirds of the parameters and processes each token independently. It acts as a nonlinear processing step and may function as an associative key-value memory for storing factual knowledge.
7. The training objective is next-token prediction — cross-entropy at every position, self-supervised, embarrassingly parallel. To predict well, the model must learn everything about language.

Further reading

• Vaswani et al. (2017) — Attention Is All You Need. The paper. Read it.
• Alammar (2018) — The Illustrated Transformer. The best visual walkthrough.
• Radford et al. (2018, 2019) — GPT and GPT-2. Decoder-only for language modeling.
• Devlin et al. (2019) — BERT. Encoder-only for understanding.
• Su et al. (2021) — RoFormer: Rotary Position Embedding.
• Olsson et al. (2022) — In-Context Learning and Induction Heads. Mechanistic interpretability.
• Geva et al. (2021) — Transformer Feed-Forward Layers Are Key-Value Memories.
• Cho et al. (2024) — Transformer Explainer. The interactive visualization above.
• P. Kulkarni — Module 02 notebook. Transformer from scratch, sequence reversal.