Transformers: A Deep Dive

Why the transformer matters

The transformer (Vaswani et al., 2017) is the architecture behind essentially every modern large language model, most speech and translation systems, many vision backbones, and the denoisers inside leading image and video generators. If you work in AI engineering you will read transformer code, size transformer deployments, and debug transformer training. This lesson builds the whole thing from tokens to a working attention layer, with every equation defined and a numeric example you can check by hand.

Intuition: a room full of people passing notes

Picture each token in a sentence as a person in a room holding a card with their word on it. In every round, each person writes a query ("what am I looking for?"), a key ("what do I offer?") and a value ("what I will tell you if you pick me"). Each person compares their query against everyone's keys, decides how much to listen to each person, and updates their card with a weighted blend of the values they heard. After a few dozen rounds, the word "bank" in "river bank" has absorbed information from "river", and the card for "it" knows which noun it refers to. That round is self-attention; the room is one layer; stacking rounds gives the deep network. Unlike an RNN, everyone talks at once, which is what makes the transformer fast to train.

Step 1: tokenization with byte-pair encoding

Models operate on integer tokens, not characters or words. Word-level vocabularies cannot handle unseen words; character-level sequences are long. Byte-pair encoding (BPE), adapted for neural translation by Sennrich, Haddow and Birch (2015), finds a middle ground: start from characters (or raw bytes) and repeatedly merge the most frequent adjacent pair into a new symbol.

Worked example. Corpus word counts: hug ×10, pug ×5, pun ×12, bun ×4, hugs ×5. Start with characters.

1. Pair counts: u g = 10 + 5 + 5 = 20, p u = 5 + 12 = 17, u n = 12 + 4 = 16, h u = 15. Merge u g → ug.
2. Recount: u n = 16, h ug = 15, p u = 12 (only pun still has a separate u). Merge u n → un.
3. Recount: h ug = 15, p un = 12. Merge h ug → hug.
4. Merge p un → pun.

The learned merge list is applied in order to new text: bugs → b u g s → b ug s, three tokens; hugs → hug s. Production tokenizers run tens of thousands of merges, operate on UTF-8 bytes so that any string is representable (byte-level BPE, used since GPT-2), and pre-split text on spaces and punctuation. Related schemes include WordPiece (BERT) and the Unigram model in SentencePiece. Vocabulary sizes of roughly 32k to a few hundred thousand tokens are typical. Because merges are learned from the training corpus, common English words often become single tokens, while code, numbers, rare words and many non-English languages are split into more pieces per word, which raises their cost and eats context length.

Step 2: embeddings and positional information

Each token id indexes a row of an embedding matrix $E \in \mathbb{R}^{V \times d}$, where $V$ is vocabulary size and $d$ is the model width. Many models tie the output projection to $E^\top$ to save parameters.

Self-attention, as defined below, is permutation-equivariant: shuffle the input tokens and the outputs shuffle identically. Word order must therefore be injected explicitly.

Sinusoidal encodings (the original transformer) add a fixed vector to each embedding:

for position $p$ and dimension index $i$. Each pair of dimensions oscillates at a different frequency, from fast to very slow, giving every position a unique signature. BERT and GPT-2 instead learned a position-embedding table, which cannot extrapolate past its trained length.

Rotary position embedding (RoPE) (Su et al., 2021) rotates the query and key vectors instead of adding to the input. Split a $d$-dimensional vector into $d/2$ pairs; at position $m$, rotate pair $i$ by angle $m\theta_i$ with $\theta_i = 10000^{-2i/d}$:

Because rotations compose, $\langle R_m q, R_n k \rangle = \langle q, R_{n-m} k \rangle$: the attention score depends only on the relative offset $n - m$. Mini-check in two dimensions with $q = k = (1, 0)$ and $\theta = \pi/6$: positions $(1, 3)$ give a rotation difference of $2\theta = \pi/3$ and a score of $\cos(\pi/3) = 0.5$; positions $(5, 7)$ give exactly the same 0.5. RoPE was adopted by GPT-J and GPT-NeoX and then by LLaMA, and is now widely used in open-weight decoder models; many context-extension methods work by rescaling its frequencies.

Step 3: scaled dot-product attention

Stack the $n$ token vectors of a sequence as rows of $X \in \mathbb{R}^{n \times d}$. Learned matrices project them to queries, keys and values:

Then

where the softmax is applied to each row, and $M$ is a mask matrix (all zeros for unmasked attention). Row $i$ of $QK^\top$ holds the similarity of query $i$ to every key; softmax turns it into weights that sum to 1; multiplying by $V$ returns a weighted average of value vectors.

Why divide by $\sqrt{d_k}$? If the components of $q$ and $k$ are independent with mean 0 and variance 1, then $q \cdot k = \sum_{j=1}^{d_k} q_j k_j$ has variance $d_k$. For $d_k = 128$ the raw scores have a standard deviation around 11, pushing the softmax into saturation where one weight is near 1 and gradients for the others are near 0. Dividing by $\sqrt{d_k}$ restores unit variance.

Worked example: three tokens, $d_k = 2$

Scores. $QK^\top = \begin{bmatrix} 1 & 0 & 1 \\ 0 & 1 & 1 \\ 1 & 1 & 2 \end{bmatrix}$; dividing by $\sqrt{2} = 1.4142$ gives $\begin{bmatrix} 0.707 & 0 & 0.707 \\ 0 & 0.707 & 0.707 \\ 0.707 & 0.707 & 1.414 \end{bmatrix}$.

Softmax, row 1. $e^{0.707} = 2.028$, $e^{0} = 1$, so the row sums to $5.056$ and the weights are $(0.401, 0.198, 0.401)$. Row 2 by symmetry: $(0.198, 0.401, 0.401)$. Row 3: $e^{0.707} = 2.028$ twice and $e^{1.414} = 4.113$, sum $8.169$, weights $(0.248, 0.248, 0.503)$.

Outputs (weights times $V$):

• Token 1: $0.401(1,2) + 0.198(3,4) + 0.401(5,6) = (3.000, 4.000)$
• Token 2: $0.198(1,2) + 0.401(3,4) + 0.401(5,6) = (3.407, 4.407)$
• Token 3: $0.248(1,2) + 0.248(3,4) + 0.503(5,6) = (3.510, 4.510)$

Each output is a blend of all values, tilted toward the tokens whose keys best match that token's query.

Causal masking

A language model predicting token $i+1$ must not see tokens after $i$. The causal mask sets $M_{ij} = -\infty$ for $j > i$ and 0 otherwise, so those entries become $e^{-\infty} = 0$ after the softmax. In our example: token 1 can see only itself, so its output is $v_1 = (1, 2)$. Token 2 sees scores $(0, 0.707)$, giving weights $(0.330, 0.670)$ and output $0.330(1,2) + 0.670(3,4) = (2.340, 3.340)$. Token 3 sees everything and is unchanged at $(3.510, 4.510)$. Masking is what lets a decoder train on all positions of a sequence in parallel while each position still only uses its past.

Step 4: multi-head attention

One attention pattern cannot track syntax, coreference and position all at once. Multi-head attention runs $h$ attentions in parallel, each with its own projections of width $d_k = d/h$, concatenates the results and mixes them:

with $W_O \in \mathbb{R}^{d \times d}$. In code, all heads come from one $d \times 3d$ matrix multiply reshaped into heads, so the cost is the same as one wide attention.

import math, torch, torch.nn as nn, torch.nn.functional as F

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.h, self.dk = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)       # W_O

    def forward(self, x):                                         # x: (B, T, D)
        B, T, D = x.shape
        q, k, v = self.qkv(x).split(D, dim=-1)
        q, k, v = (t.view(B, T, self.h, self.dk).transpose(1, 2) for t in (q, k, v))
        scores = q @ k.transpose(-2, -1) / math.sqrt(self.dk)     # (B, h, T, T)
        future = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)
        scores = scores.masked_fill(future, float("-inf"))        # causal mask
        out = scores.softmax(dim=-1) @ v                          # (B, h, T, dk)
        return self.proj(out.transpose(1, 2).reshape(B, T, D))    # merge heads

torch.manual_seed(0)
layer = CausalSelfAttention(d_model=64, n_heads=4)
x = torch.randn(2, 10, 64)
q, k, v = (t.view(2, 10, 4, 16).transpose(1, 2) for t in layer.qkv(x).split(64, dim=-1))
fused = F.scaled_dot_product_attention(q, k, v, is_causal=True)   # PyTorch fused kernel
ref = layer.proj(fused.transpose(1, 2).reshape(2, 10, 64))
print(torch.allclose(layer(x), ref, atol=1e-5))                    # True

torch.triu(..., 1) marks the strictly upper triangle — the future. F.scaled_dot_product_attention computes the same thing and, on supported GPUs, dispatches to memory-efficient or FlashAttention kernels; use it in real code.

Step 5: the transformer block

A block wraps attention and a position-wise MLP in residual connections and normalization. Modern models use pre-norm, normalizing the input to each sub-layer:

The original paper used post-norm, $x \leftarrow \operatorname{LN}(x + \operatorname{Sublayer}(x))$, which needs careful learning-rate warmup; pre-norm keeps an unnormalized residual "stream" from input to output and trains stably at depth (analysed by Xiong et al., 2020). LayerNorm normalizes each token vector across its features: $\operatorname{LN}(x) = \gamma \odot (x - \mu)/\sqrt{\sigma^2 + \epsilon} + \beta$, with $\mu$, $\sigma^2$ the mean and variance of $x$'s components and $\gamma, \beta$ learned. Many LLMs use RMSNorm, which skips the mean subtraction and $\beta$.

The MLP acts on each token independently: classically $W_2\, \phi(W_1 x)$ with hidden width $d_{ff} = 4d$ and GELU activation $\phi$. Many LLMs use SwiGLU, $W_2\big(\operatorname{SiLU}(W_1 x) \odot W_3 x\big)$, with three matrices and $d_{ff} \approx \tfrac{8}{3} d$ to keep the parameter count similar. Roughly: attention moves information between tokens; the MLP processes information within each token and holds much of the model's stored knowledge.

class Block(nn.Module):
    def __init__(self, d_model, n_heads, d_ff):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(d_model), nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_heads)
        self.mlp = nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))
    def forward(self, x):
        x = x + self.attn(self.ln1(x))     # pre-norm residual around attention
        return x + self.mlp(self.ln2(x))   # pre-norm residual around the MLP

A full decoder-only model is: embedding → $L$ blocks → final norm → linear projection to $V$ logits, trained with cross-entropy on next-token prediction.

Three families

In cross-attention the queries come from the decoder and the keys and values from the encoder — exactly Bahdanau attention from the previous lesson, in matrix form.

Counting parameters in a decoder layer

Take $d = 4096$, 32 heads of size 128, standard multi-head attention without biases, a SwiGLU MLP with $d_{ff} = 11{,}008$, and two RMSNorms.

• Attention: $W_Q, W_K, W_V, W_O$ are each $d \times d$: $4d^2 = 4 \times 16{,}777{,}216 = 67{,}108{,}864$.
• MLP: three matrices of $d \times d_{ff}$: $3 \times 4096 \times 11{,}008 = 135{,}266{,}304$.
• Norms: $2d = 8{,}192$.
• Per layer: 202,383,360 ≈ 202M, two-thirds of it in the MLP.

Thirty-two layers give $6.476$B. Add an input embedding and an untied output head with $V = 32{,}000$: $2Vd = 262{,}144{,}000$, plus a final norm. Total ≈ 6.74B — the shape of the familiar 2023-era "7B" models. Rule of thumb: a block has about $12d^2$ parameters with a $4d$ GELU MLP ($4d^2$ attention + $8d^2$ MLP).

Cost: O(n²) attention and the KV cache

For a sequence of $n$ tokens, $QK^\top$ is an $n \times n$ matrix per head. Computing it and multiplying by $V$ costs $O(n^2 d)$ per layer, and storing it naively costs $O(n^2)$ memory per head. Doubling context quadruples the attention cost, while the projections and MLP grow only linearly. At short contexts the MLP dominates; at long contexts attention does.

During generation, a decoder caches each past token's keys and values so it never recomputes them. Per token, the KV cache takes $2 \times \text{layers} \times \text{KV heads} \times \text{head dim} \times \text{bytes per element}$. For the model above in 16-bit: $2 \times 32 \times 32 \times 128 \times 2 = 524{,}288$ bytes = 512 KiB per token, or 16 GiB at a 32k-token context. See kv-cache-and-context-cost and the AI Index methodology for how this feeds memory estimates.

Efficient attention variants

Multi-query attention (MQA) (Shazeer, 2019) shares a single key head and value head across all query heads. Grouped-query attention (GQA) (Ainslie et al., 2023) is the middle ground: $h$ query heads share $g$ KV heads in groups. With 32 query heads and 8 KV heads, the KV cache above drops 4× to 128 KiB per token (4 GiB at 32k), and $W_K, W_V$ shrink from $d \times d$ to $d \times 1024$, cutting attention parameters from 67.1M to 41.9M per layer. Decoding is usually limited by memory bandwidth, so a smaller cache directly means faster generation and larger batches, with little quality loss for GQA. GQA and other KV-compression schemes are now common in open-weight LLMs; check a model's config for num_key_value_heads.

FlashAttention (Dao et al., 2022) does not change the math — it computes exact attention. It changes where data lives. GPUs have large but comparatively slow high-bandwidth memory (HBM) and tiny, fast on-chip SRAM. Naive attention writes the full $n \times n$ score matrix to HBM, reads it back for the softmax, and again for the multiply by $V$; for long sequences this memory traffic, not arithmetic, is the bottleneck. FlashAttention splits $Q$, $K$, $V$ into tiles that fit in SRAM and computes the softmax incrementally with the online softmax trick: for each new block of scores $s$ it keeps a running maximum $m$, normalizer $\ell$ and output accumulator $o$, updating

and dividing $o$ by $\ell$ at the end. The $n \times n$ matrix is never materialized, memory becomes linear in $n$, and the backward pass recomputes attention tiles instead of storing them. The result is the same output, much less memory, and substantially faster wall-clock time. Later versions improved GPU work partitioning further.

Common pitfalls

• Mask errors: an off-by-one causal mask (masking the diagonal, or not masking the future) gives a model that looks brilliant in training and fails at generation because it was peeking.
• Padding masks forgotten in batched encoders, or combined incorrectly with causal masks; also left- versus right-padding mismatches in batched generation.
• Softmax in low precision: compute attention softmax and LayerNorm statistics in float32 or with fused kernels; fp16 overflow produces NaNs.
• Missing the $\sqrt{d_k}$ scale when hand-writing attention, which saturates the softmax.
• Tokenizer mismatch: using a different tokenizer, or different special tokens or chat template, from the one the model was trained with.
• Assuming linear cost in context: prefill attention is quadratic, and KV-cache memory grows linearly per sequence and per concurrent user.

In practice

• Use maintained implementations (PyTorch scaled_dot_product_attention, Hugging Face transformers, vLLM and similar serving engines) rather than hand-rolled attention in production; hand-roll only to learn and test.
• Pick architecture by task: encoders for embeddings and classification, decoders for generation, encoder–decoders where input and output are clearly separate (speech recognition, translation).
• Budget memory from parameters plus KV cache, and prefer GQA models for long-context or high-concurrency serving.
• When fine-tuning, keep the original tokenizer and template; when extending context, use the model's documented RoPE scaling rather than inventing your own.
• The next lesson, beyond-transformers, covers architectures that attack the $O(n^2)$ cost.