0. Why Transformers Replaced RNNs
Before transformers, sequence models (RNNs, LSTMs) processed tokens one at a time, carrying a hidden state forward. This made them inherently sequential: you couldn't compute the representation for token 100 without first computing tokens 1–99, which meant no parallelism across the sequence dimension during training, and made it hard to preserve information over long distances (gradients had to flow through every intermediate step, reintroducing the vanishing gradient problem from the Foundations page). Transformers replace recurrence entirely with self-attention: every position computes its new representation by directly looking at every other position in the sequence, in parallel, in a single matrix operation. This removes the sequential bottleneck (huge GPU-training speedup) and gives every token a direct, one-hop path to every other token, which drastically eases long-range dependency learning.
1. Core Idea: Attention Over Sequences
1.1 Intuition first
Imagine reading the sentence "The trophy didn't fit in the suitcase because it was too big." To understand what "it" refers to, you (a human reader) implicitly look back at the other words and weigh them by relevance: "trophy" and "suitcase" are candidates, and world knowledge plus sentence structure tells you "it" more likely means "trophy." Self-attention formalizes exactly this process mathematically: for every token, compute a relevance score against every other token, turn those scores into weights that sum to 1, and build a new representation as the weighted combination of all tokens' information, weighted by relevance.
1.2 Queries, keys, and values
Each token's hidden vector \(x_i \in \mathbb{R}^d\) is linearly projected into three different vectors using learned weight matrices \(W_Q, W_K, W_V\):
\[ q_i = x_i W_Q, \qquad k_i = x_i W_K, \qquad v_i = x_i W_V \]
Intuitively: the query represents "what this token is looking for," the key represents "what this token has to offer / how it advertises itself," and the value represents "the actual content this token will contribute if attended to." Relevance between token \(i\) (asking) and token \(j\) (being looked at) is computed as the dot product \(q_i \cdot k_j\) — a similarity score between what \(i\) wants and what \(j\) offers.
1.3 Scaled dot-product attention, step by step
Stacking every token's queries, keys, and values into matrices \(Q, K, V \in \mathbb{R}^{n \times d_k}\) (for a sequence of length \(n\)), attention is computed for the whole sequence at once:
\[ \text{Attention}(Q,K,V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right) V \]
Breaking this into steps:
- Similarity scores: \(QK^\top\) computes the dot product between every query and every key, producing an \(n \times n\) matrix of raw relevance scores.
- Scaling: divide by \(\sqrt{d_k}\), where \(d_k\) is the dimensionality of each key vector. Without scaling, dot products grow large in magnitude as \(d_k\) grows (since they're a sum of \(d_k\) terms), which pushes the softmax into regions with extremely small gradients (saturation). Scaling by \(\sqrt{d_k}\) keeps the variance of the scores roughly constant regardless of dimensionality.
- Softmax: convert each row of scores into a probability distribution that sums to 1: \(\text{softmax}(s)_j = \dfrac{e^{s_j}}{\sum_{l} e^{s_l}}\). Each token now has a weight for every other token, representing "how much attention to pay."
- Weighted sum: multiply these attention weights by \(V\), producing, for every token, a new representation that is a weighted blend of all tokens' value vectors — dominated by whichever tokens had high relevance.
Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) * VRefer: The Illustrated Transformer · Attention Is All You Need
import torch
import torch.nn.functional as F
def scaled_dot_product_attention(Q, K, V, mask=None):
d_k = Q.size(-1)
scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
if mask is not None:
scores = scores.masked_fill(mask == 0, float("-inf"))
weights = F.softmax(scores, dim=-1)
return torch.matmul(weights, V)1.4 Multi-head attention
Instead of computing attention once with the full hidden dimension \(d\), the model splits \(Q, K, V\) into \(h\) smaller "heads," each of dimension \(d_k = d/h\), runs scaled dot-product attention independently and in parallel within each head, and concatenates the results back together before a final linear projection:
\[ \text{MultiHead}(Q,K,V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)W_O, \qquad \text{head}_i = \text{Attention}(QW_Q^{(i)}, KW_K^{(i)}, VW_V^{(i)}) \]
Each head can specialize in a different type of relationship — one head might learn to track subject-verb agreement, another might track coreference (like the "it" example above), another might track positional adjacency. This gives the model multiple independent "views" of the sequence's relationships at every layer, at roughly the same total compute cost as one large attention operation.
2. Transformer Block Anatomy
A standard decoder block (used in GPT-style LLMs) contains two sublayers, each wrapped in a residual connection and preceded by normalization (pre-norm, see Foundations §3):
- Pre-norm self-attention with residual connection: \(x \leftarrow x + \text{MultiHeadAttention}(\text{RMSNorm}(x))\)
- Pre-norm feed-forward network (FFN) with residual connection: \(x \leftarrow x + \text{FFN}(\text{RMSNorm}(x))\)
The FFN is typically two linear layers with a non-linearity (or gated activation, like SwiGLU) in between, applied independently and identically to every token position:
\[ \text{FFN}(x) = W_2\big(\text{SiLU}(xW_1) \odot (xW_3)\big) \]
Conceptually, attention is the only place information mixes across token positions; the FFN operates per-token, acting as a memory/transformation module that processes whatever information attention has just gathered for that position. Stacking many blocks (attention, then FFN, repeated \(L\) times) lets the model build increasingly abstract representations layer by layer. Encoder-decoder models (the original Transformer, T5) add a third sublayer, cross-attention, where decoder queries attend over encoder keys/values; most LLMs today are decoder-only and skip this.
3. Positional Information
3.1 Why position must be injected explicitly
Attention as defined above is permutation-invariant: if you shuffled the input tokens, the set of pairwise dot products would be the same, just rearranged — the mechanism has no innate sense of word order. But word order carries meaning ("dog bites man" vs "man bites dog"), so models must inject positional information explicitly.
3.2 Approaches
- Absolute positional embeddings: learn (or fix via sinusoids) a separate vector for each position \(1, 2, 3, \dots\) and add it to the token embedding before the first layer (early GPT, BERT). Simple, but generalizes poorly beyond the maximum length seen during training, and treats position as an absolute property rather than encoding useful relative distance information directly inside attention.
- RoPE (Rotary Position Embedding): rotates the query and key vectors by an angle proportional to their position before computing the dot product. For a 2D pair of dimensions at position \(m\), this looks like multiplying by a rotation matrix \(R_{\theta m} = \begin{pmatrix}\cos(m\theta) & -\sin(m\theta)\\ \sin(m\theta) & \cos(m\theta)\end{pmatrix}\), applied across paired dimensions with different frequencies \(\theta\) (analogous to sinusoidal frequencies). The key mathematical trick: the dot product between a rotated query at position \(m\) and a rotated key at position \(n\) depends only on their relative distance \(m-n\), not their absolute positions. This means the model automatically learns relative positional relationships, which composes cleanly with long-context extension techniques and generalizes better to sequence lengths beyond training. RoPE is the default in LLaMA, Mistral, and most modern LLMs.
- ALiBi (Attention with Linear Biases): instead of modifying \(Q\) or \(K\), adds a fixed linear penalty directly to the attention scores based on the distance between positions: \(\text{score}_{ij} = q_i \cdot k_j - m \cdot |i-j|\), where \(m\) is a head-specific slope. Distant tokens are penalized more, which naturally biases attention toward local context while still allowing distant attention if the dot-product similarity is strong enough. This lets models extrapolate to longer contexts than seen during training without retraining position tables, though its extrapolation quality and the raw quality ceiling can differ from RoPE-based approaches in practice.
RoPE has become the default across most modern open LLMs because it composes well with long-context extensions (like position interpolation or NTK-aware scaling) and keeps relative position information embedded directly inside the attention mechanics rather than as a separate additive term.
4. Efficiency Variants
4.1 Grouped-query and multi-query attention
Standard multi-head attention (MHA) uses a completely separate set of K and V projections for every one of the \(h\) heads. At inference time, autoregressive generation caches all past keys and values (the KV cache, see Section 5) so they don't need to be recomputed at every new token step — but storing a separate K/V per head for every layer and every token consumes a large amount of GPU memory, which directly limits batch size and context length.
- Multi-query attention (MQA) shares a single K/V head across all query heads, cutting KV cache memory by roughly a factor of \(h\) (the number of heads), at some cost to model quality since all heads now attend using the same keys/values.
- Grouped-query attention (GQA) is a middle ground: heads are split into a small number of groups, and each group shares one K/V head. This recovers most of MQA's memory savings while retaining more of MHA's quality, and is the choice used in production LLMs like LLaMA 2/3 and Mistral.
4.2 Mixture of Experts (MoE)
Instead of a single dense FFN per layer that every token passes through, an MoE layer maintains many separate "expert" FFNs and a small learned gating network that routes each token to only a small subset of experts (commonly the top-1 or top-2 by gate score):
\[ y = \sum_{e \in \text{TopK}(g(x))} g_e(x) \cdot \text{FFN}_e(x) \]
where \(g(x)\) is typically a softmax over a linear projection of \(x\) producing per-expert scores. Total parameter count grows large (many experts), but the compute cost per token stays roughly constant (only a small number of experts activate), meaning an MoE model can have far more total capacity than a dense model at the same inference FLOP cost. Challenges include load balancing (without extra loss terms, the gate tends to collapse onto favoring a few experts, wasting the rest of the model's capacity) and substantially more complex training/serving infrastructure (experts often need to be sharded across many devices).
5. Causal Masking and Autoregressive Decoding
5.1 Causal masking
Decoder-only LLMs must respect the fact that during generation, position \(i\) can only ever have seen positions \(1, \dots, i\) — it hasn't been told the future yet. To enforce this even during parallel training (where the whole sequence is available at once), a causal mask sets attention scores for any \(j > i\) to \(-\infty\) before the softmax, so those positions receive exactly zero attention weight:
\[ \text{scores}_{ij} = \begin{cases} q_i \cdot k_j / \sqrt{d_k} & j \le i \\ -\infty & j > i \end{cases} \]
This forces the model to learn genuinely left-to-right, autoregressive next-token prediction, matching how it will actually be used at inference time.
5.2 The KV cache
At inference, tokens are generated one at a time: predict the next token, append it to the sequence, repeat. Naively, this would mean recomputing attention over the entire growing sequence from scratch at every single step — wasteful, since the keys and values for all previously-generated tokens never change once computed. Instead, systems cache the \(K\) and \(V\) vectors computed for every past token (the KV cache) and, at each new step, only compute \(Q, K, V\) for the newest token, attending it against the cached keys/values from all previous positions. This turns each generation step from \(O(n^2)\) (recomputing full attention) into roughly \(O(n)\) work relative to sequence length, at the cost of memory proportional to sequence length \(\times\) number of layers \(\times\) number of KV heads \(\times\) head dimension — which is exactly why MQA/GQA (Section 4.1) matter so much for serving long contexts efficiently.
6. Scaling and Design Tradeoffs
Transformer quality improves predictably with more parameters, data, and compute, following empirical scaling laws (see the LLMs page for the Chinchilla relationship in detail). Within a fixed compute and memory budget, engineers make several architecture tradeoffs:
- More layers vs wider hidden dimensions — depth tends to help compositional/hierarchical reasoning, width increases per-layer representational capacity; empirically there are compute-optimal ratios that most model families converge toward.
- Dense FFN vs MoE routing — MoE trades training/serving complexity for a better capacity-to-inference-FLOPs ratio.
- Full MHA vs GQA — trades a small amount of quality for large inference memory and latency savings, essential for long-context serving at scale.
- Context length vs quadratic attention cost — naive attention is \(O(n^2)\) in sequence length \(n\) (every token attends to every other token), so doubling context roughly quadruples raw attention compute (though FFN cost only scales linearly), making long-context support an active systems engineering problem, not just an architecture choice.
FlashAttention and other IO-aware attention kernels reduce memory-bandwidth bottlenecks (they avoid materializing the full \(n \times n\) attention matrix in slow GPU memory, computing softmax in fused, tiled blocks instead) without changing the mathematical definition of attention itself — they are a systems-level optimization, not an algorithmic approximation, which is why they've become close to universal in production training and inference stacks.
7. Putting It Together: One Full Block, End to End
For a single input token stream \(x \in \mathbb{R}^{n \times d}\), one full decoder block computes, in order:
\[ h = x + \text{MultiHeadAttention}\big(\text{RMSNorm}(x)\big) \]
\[ y = h + \text{FFN}\big(\text{RMSNorm}(h)\big) \]
with RoPE applied inside the attention step to \(Q\) and \(K\) before the dot product, a causal mask applied to the scores, and GQA sharing K/V heads across groups of query heads for memory efficiency. Stack this block \(L\) times (32, 80, or more in production LLMs), and finish with a final normalization layer plus a linear projection to vocabulary-sized logits, and you have the complete architecture underlying essentially every modern LLM.