The transformer: the block, model families and vision transformers

Before this lesson: lesson 05 (LayerNorm, residual connections, pre-LN and post-LN), lesson 07 (strided convolutions), lesson 08 (tokens, embeddings, next-token prediction), lesson 09 (attention, masks, multi-head attention) and lesson 10 (positional encodings). Time: about 95 minutes, plus the notebook.

---

Why this matters #

Almost every large model of today is a transformer: GPT, LLaMA and most chat models (decoder-only), BERT (encoder-only), T5 and Whisper (encoder–decoder), and ViT, CLIP and the diffusion transformer DiT for images. The same block, with small changes, runs in AlphaFold-style protein models and in weather models such as ClimaX. If you know the block, its shapes and its costs, you can read most architecture papers, estimate how much memory and compute a model needs, and debug an implementation. Interviews ask exactly this: draw the block, count the parameters of GPT-2, estimate the training FLOPs of a 7B model, and explain the KV cache.

1. Predict first #

2. The picture #

Two kinds of mixing #

Write the tokens of a sequence as the rows of a matrix $X \in \mathbb{R}^{T \times D}$: $T$ tokens, each a vector of width $D$. (As in the whole track, we use the row convention of PyTorch.) A network can combine numbers of this matrix in two directions:

• across tokens (down a column): token 3 uses information from token 7. Only attention does this, with weights that depend on the input (lesson 09).
• across features (along a row): the $D$ numbers of one token are combined into new numbers. The MLP (also called the feed-forward network, FFN) does this, with the same weights for every token.

A transformer block does both, one after the other. Around each of the two steps it puts a residual connection (the step's output is added to its input) and a LayerNorm (each token vector is normalized to mean 0 and variance 1, then scaled; lesson 05). The modern order, called pre-LN, is

Here $\operatorname{MHA}$ is multi-head attention. The input $X$, the middle value $X'$ and the output $Y$ all have the same shape $T \times D$. So blocks can be stacked: the output of one block is the input of the next. The sequence of vectors that flows through the stack, and that each block adds to, is called the residual stream.

Interactive figure: loading…

The whole model: text in, probabilities out #

A decoder-only language model such as GPT-2 is a stack of $L$ such blocks between an input part and an output part:

1. Token embedding. Each token id selects one row of a learned matrix $E \in \mathbb{R}^{|\mathcal V| \times D}$, where $|\mathcal V|$ is the vocabulary size (lesson 08).
2. Positions. Attention alone cannot see the order of the tokens (lesson 09), so position information is added. GPT-2 adds a learned vector for each position; LLaMA rotates the queries and keys instead (RoPE). Lesson 10 compares the methods.
3. $L$ blocks, each with a causal mask: token $t$ attends only to tokens $1, \dots, t$.
4. A final LayerNorm, then the unembedding: a matrix product that maps each vector of width $D$ to $|\mathcal V|$ logits, one score per possible next token.
5. Softmax over the vocabulary gives the probabilities $p_t$ of the next token. Training minimizes the cross-entropy $-\log p_t[s_{t+1}]$, summed over all positions at once.

Because of the causal mask, one forward pass over a text of $T$ tokens gives $T$ predictions in parallel: position $t$ predicts token $t + 1$ from tokens $1, \dots, t$ only.

An image as a sequence of patches #

A transformer needs a sequence of vectors. An image is a grid of pixels. The vision transformer (ViT, Dosovitskiy et al., 2021) makes a sequence like this:

1. Cut the image into non-overlapping squares of $P \times P$ pixels, called patches. A $224 \times 224$ image with $P = 16$ gives a $14 \times 14$ grid: $196$ patches.
2. Read the patches row by row, like the words of a page. Patch $(r, c)$ of the grid becomes token number $r \cdot 14 + c$ (counting from 0).
3. Flatten each patch to one vector ($3 \cdot 16 \cdot 16 = 768$ numbers for a color image) and map it with one shared linear layer to width $D$.
4. Put one extra learned token, [CLS], in front, and add a learned position vector to every token. Then run ordinary encoder blocks (no causal mask).

Interactive figure: loading…

3. Definitions and main results #

3.1 The transformer block #

Definition (MLP of a block). For a token matrix $Z \in \mathbb{R}^{T \times D}$,

with $W_1 \in \mathbb{R}^{D \times d_{ff}}$, $W_2 \in \mathbb{R}^{d_{ff} \times D}$, an activation $\sigma$ (GELU in GPT-2, BERT and ViT; ReLU in the original transformer) and usually $d_{ff} = 4D$. LLaMA uses a SwiGLU MLP instead (lesson 04). The MLP acts on each row of $Z$ separately.

Definition (transformer block). With multi-head attention $\operatorname{MHA}$ (with or without a causal mask) and two LayerNorms with their own scale and shift vectors:

The original transformer (2017) and BERT use post-LN. GPT-2, LLaMA, ViT and most modern models use pre-LN. Lesson 05 explains the reason. In pre-LN, the direct path from $X$ to $Y$ is the identity, so the gradient reaches early blocks at any depth. In post-LN, every LayerNorm sits on the main path, the gradients near the output are large at initialization, and training needs a learning-rate warmup.

What each part is for:

• Attention is the only part that moves information between tokens.
• The MLP is the computation on each token. It holds two thirds of the block's parameters (section 3.3). One common reading is that the MLPs store much of the model's knowledge, while attention routes information.
• Residual connections: each sublayer reads the residual stream and adds its result. A block can learn to change the stream only a little.
• LayerNorm keeps the input of each sublayer at a fixed scale. A pre-LN stack never normalizes the stream itself, so it needs a final LayerNorm after the last block.

Shapes through one pre-LN block (batch $B$, $H$ heads of width $d_k = D/H$):

Step	Shape
input $X$, $\operatorname{LN}_1(X)$	$(B, T, D)$
$Q, K, V$ after the split into heads	$(B, H, T, d_k)$
attention weights $A$ of all heads	$(B, H, T, T)$
heads merged, after $W_O$; $X'$ after the residual	$(B, T, D)$
MLP hidden layer $\sigma(ZW_1 + \mathbf{1}b_1)$	$(B, T, 4D)$
output $Y$	$(B, T, D)$

3.2 The decoder-only language model #

Definition. Let the input be token ids $s_1, \dots, s_T$. A decoder-only model with token embedding $E \in \mathbb{R}^{|\mathcal V| \times D}$, learned position embedding $E_{\text{pos}} \in \mathbb{R}^{T_{\max} \times D}$ and $L$ causal pre-LN blocks computes

with logits $Z \in \mathbb{R}^{T \times |\mathcal V|}$ and the unembedding $W_U \in \mathbb{R}^{D \times |\mathcal V|}$. Row $t$ of $\operatorname{softmax}(Z)$ is $p_t$, the predicted distribution of token $t + 1$. $T_{\max}$ is the context length, the largest number of positions the model has embeddings for.

Weight tying. Many models set $W_U = E^\top$ (Press and Wolf, 2017). Then the logit of token $v$ is the dot product of the final vector with the embedding of $v$: tokens with similar embeddings get similar scores. Tying saves $|\mathcal V| D$ parameters; for GPT-2 small that is 38.6 million, about a third of the model. GPT-2 ties; LLaMA does not.

3.3 Counting parameters and FLOPs #

A FLOP is one floating-point operation. A multiply followed by an add counts as 2 FLOPs.

Theorem (parameters of one block). A block with width $D$, $d_{ff} = 4D$, biases and two LayerNorms has $12D^2 + 13D$ parameters, for every number of heads.

Proof. Count each part.

1. Multi-head attention: $W_Q, W_K, W_V, W_O \in \mathbb{R}^{D \times D}$ and four bias vectors of length $D$: $4D^2 + 4D$ (lesson 09).
2. MLP: $W_1$ has $D \cdot 4D$ entries and $b_1$ has $4D$; $W_2$ has $4D \cdot D$ and $b_2$ has $D$: $8D^2 + 5D$.
3. Two LayerNorms, each with a scale and a shift of length $D$: $4D$.
4. Sum: $4D^2 + 8D^2 = 12D^2$, and $4D + 5D + 4D = 13D$. $\blacksquare$

For $D = 768$ this is $7{,}077{,}888 + 9{,}984 = 7{,}087{,}872$. The $13D$ term is about 0.1%, so a block has about $12D^2$ parameters: $4D^2$ in attention and $8D^2$ in the MLP.

Theorem (FLOPs of a linear layer). Let $Y = XW$ with $X \in \mathbb{R}^{B \times n}$ (a batch of $B$ token vectors) and $W \in \mathbb{R}^{n \times m}$. Then

1. the forward pass costs $2Bnm$ FLOPs: $2$ FLOPs per token per weight;
2. the backward pass costs $4Bnm$ FLOPs: $\partial\mathcal{L}/\partial X = (\partial\mathcal{L}/\partial Y)\,W^\top$ and $\partial\mathcal{L}/\partial W = X^\top(\partial\mathcal{L}/\partial Y)$ each cost $2Bnm$ (lesson 03);
3. so one training step costs $6Bnm$ FLOPs: 6 FLOPs per token per weight.

Proof. Each entry $Y_{ij} = \sum_{k=1}^{n} X_{ik}W_{kj}$ needs $n$ multiplications and $n - 1$ additions, about $2n$ FLOPs, and there are $Bm$ entries. The two backward products have the same three sizes $B$, $n$, $m$ in a different order, so each also costs $2Bnm$. $\blacksquare$

Corollary (FLOPs per token of a transformer). Let $N$ be the number of weights that are multiplied for every token: $12LD^2$ in the blocks, plus $D|\mathcal V|$ for the unembedding (the embedding lookup reads one row and costs no FLOPs). For a context of $T$ tokens,

The term $4LTD$ comes from attention: for one query, $QK^\top$ costs $2Td_k$ per head and $AV$ another $2Td_k$, so $4TD$ for all heads, in each of the $L$ blocks. (This counts all $T$ keys; a kernel that skips the masked half of a causal matrix does about half of it.) The rule "training costs about $6N$ FLOPs per token" (Kaplan et al., 2020) holds when $T$ is small compared with $6D$, as we now show.

When does attention dominate? In one block, for one token, the weights cost $2 \cdot 12D^2 = 24D^2$ FLOPs and the attention matrix costs $4TD$. They are equal when $4TD = 24D^2$, that is, at $\mathbf{T = 6D}$. For GPT-2 small ($D = 768$) that is $T = 4608$, far beyond its context of 1024: at $T = 1024$ attention is 18% of the per-layer FLOPs. For long contexts, the $T^2$ term wins.

Interactive figure: loading…

3.4 Three model families #

The same block makes three kinds of model. They differ in the masks and in what they are trained to do.

Family	Self-attention mask	Cross-attention	Trained to	Examples
encoder-only	padding only: every token sees every token (bidirectional)	none	fill in masked tokens; classify	BERT, ViT
decoder-only	causal + padding	none	predict the next token	GPT, LLaMA
encoder–decoder	encoder: padding; decoder: causal	decoder queries, encoder keys and values	map an input sequence to an output sequence	original transformer, T5, Whisper

• Encoder-only. BERT (Devlin et al., 2019) selects about 15% of the input tokens, hides most of them behind a [MASK] token, and learns to predict the original tokens from both sides. This is masked language modelling. The output vectors are good features for classification and search. An encoder cannot generate text left to right in a natural way, because every token sees the future.
• Decoder-only. GPT is trained on next-token prediction with teacher forcing (lesson 08): the input is the true text, the target is the same text shifted by one. It generates text by feeding its own output back in (section 3.5).
• Encoder–decoder. The encoder reads the whole input (a sentence to translate, an audio clip). The decoder writes the output. A decoder block has three sublayers: causal self-attention, cross-attention (queries from the decoder, keys and values from the final encoder output, lesson 09), and the MLP. So it has about $16D^2$ parameters instead of $12D^2$. T5 (Raffel et al., 2020) is trained to fill in missing spans of text; the original transformer was trained to translate.

3.5 Generating text: decoding and the KV cache #

Autoregressive generation. Start with a prompt $s_1, \dots, s_t$. Repeat: run the model, take the logits $z \in \mathbb{R}^{|\mathcal V|}$ of the last position, choose a token $s_{t+1}$, append it, and set $t \leftarrow t + 1$. Stop at an end token or at a length limit.

Decoding rules (how to choose the token from $z$):

• Greedy: $s_{t+1} = \arg\max_v z_v$. Deterministic, but often repetitive.
• Temperature sampling: draw $s_{t+1}$ from

$\tau = 1$ is the model's own distribution. $\tau < 1$ makes it sharper; as $\tau \to 0$ it becomes greedy (when the largest logit is unique). $\tau > 1$ makes it flatter; as $\tau \to \infty$ it becomes uniform.

• Top-$k$ and top-$p$ (nucleus) sampling (Holtzman et al., 2020): set all probabilities to 0 except the $k$ largest, or except the smallest set of tokens whose probabilities sum to at least $p$; renormalize and sample. This removes the long tail of unlikely tokens.

The KV cache. Running the whole model on all $t$ tokens at every step repeats almost all the work. With a causal mask, it is not needed:

Theorem (cached keys and values do not change). In a decoder with causal self-attention, the hidden vector of token $j$ in every layer depends only on the tokens $s_1, \dots, s_j$. So the keys and values of token $j$, in every layer, do not change when later tokens are appended.

Proof (induction over the layers).

1. Layer 0: row $j$ of $X_0$ is $E[s_j] + E_{\text{pos}}[j]$. It depends only on $s_j$.
2. Suppose that in layer $l - 1$, row $i$ depends only on $s_1, \dots, s_i$, for every $i$.
3. In block $l$, the LayerNorms, the MLP, the residual additions and the projections $W_Q, W_K, W_V$ act on each row separately. The causal mask lets query $j$ read only the keys and values of rows $i \le j$, which by step 2 depend only on tokens with index $\le i \le j$.
4. So row $j$ of layer $l$ depends only on $s_1, \dots, s_j$. Its key and value are computed from it, so they do too. $\blacksquare$

So we store the keys and values of all past tokens in every layer: the KV cache. At each new step only the newest token goes through the model. Its query attends to the cached keys, and its new key and value are appended to the cache. The result is exactly the same as running the full model. For an encoder (no causal mask) the theorem fails: a new token changes the vectors of all earlier tokens, so nothing can be reused.

Work. Without a cache, step $t$ runs the dense layers on $t$ tokens, so generating $T$ tokens costs about $1 + 2 + \dots + T = T(T + 1)/2$ token passes. With a cache, it costs $T$ token passes. The attention term still grows: the new query compares with $t$ cached keys.

Size of the KV cache. Each of the $L$ layers stores a key and a value of width $D$ for each of the $T$ tokens. So the cache of one sequence holds

(with standard multi-head attention, $D = H d_k$). Multiply by the batch size and by the bytes per number (2 for fp16 or bf16). Multi-query attention (MQA; Shazeer, 2019) uses one key head and one value head shared by all query heads; grouped-query attention (GQA; Ainslie et al., 2023) uses $H_{kv}$ groups. The cache is then $2 \cdot L \cdot T \cdot H_{kv} d_k$ numbers, $H/H_{kv}$ times smaller.

Interactive figure: loading…

3.6 The vision transformer #

Patch tokens. An image $x \in \mathbb{R}^{C \times H_{\text{in}} \times W_{\text{in}}}$ is cut into $T = (H_{\text{in}}/P)(W_{\text{in}}/P)$ non-overlapping $P \times P$ patches ($P$ must divide both sides). Each patch is flattened to a vector of $CP^2$ numbers and mapped by one shared matrix $W_E \in \mathbb{R}^{CP^2 \times D}$ plus a bias $b_E \in \mathbb{R}^D$: $CP^2D + D$ parameters.

Theorem (patch embedding is a strided convolution). Let a convolution have $D$ filters of size $C \times P \times P$, stride $P$ and no padding, and let filter $c$, flattened in the order (channel, row, column), be column $c$ of $W_E$, with the same bias. Then output channel $c$ at grid position $(r, s)$ equals entry $c$ of $\operatorname{patch}_{rs}\,W_E + b_E$.

Proof.

1. A convolution computes, at each output position, the dot product of each filter with the input window at that position, plus the filter's bias (lesson 07).
2. With kernel size $P$ and stride $P$, the window of output position $(r, s)$ covers rows $rP, \dots, rP + P - 1$ and columns $sP, \dots, sP + P - 1$. These windows do not overlap and they tile the image: window $(r, s)$ is exactly patch $(r, s)$.
3. The dot product of filter $c$ with the window equals the dot product of the two flattened vectors, if both are flattened in the same order. That is entry $c$ of $\operatorname{patch}_{rs}\,W_E$. Adding the bias finishes the proof. $\blacksquare$

The output size formula of lesson 07 agrees: with $n = H_{\text{in}}$, $k = s = P$ and $p = 0$, $o = \lfloor (H_{\text{in}} - P)/P \rfloor + 1 = H_{\text{in}}/P$. Reading the $\frac{H_{\text{in}}}{P} \times \frac{W_{\text{in}}}{P}$ output grid row by row gives the $T \times D$ token matrix. In PyTorch: nn.Conv2d(C, D, kernel_size=P, stride=P), then .flatten(2).transpose(1, 2).

The model (Dosovitskiy et al., 2021). With a learned class token $x_{\text{class}} \in \mathbb{R}^{1 \times D}$ and learned position embeddings $E_{\text{pos}} \in \mathbb{R}^{(T+1) \times D}$:

then $L$ pre-LN encoder blocks (no causal mask) and a final LayerNorm. The final vector of the [CLS] token is the representation of the image; a linear layer on it gives the class logits. (Many later models take the mean of the patch vectors instead of a [CLS] token.) The position embeddings are 1-D: one learned vector per index in reading order. In the ablation of the ViT paper, 2-D and relative position embeddings were not better, but no position embedding at all was clearly worse.

Counting. $T + 1$ tokens. Parameters: patch embedding $CP^2D + D$, class token $D$, position embeddings $(T + 1)D$, about $12D^2$ per block, and the head. ViT-B/16 ($D = 768$, $L = 12$, $H = 12$, $P = 16$, $224 \times 224$ images) has 86.6 million parameters with a 1000-class head (section 4). The attention matrices cost $\propto (T + 1)^2$ and $T \propto P^{-2}$, so halving $P$ gives about 16 times the attention cost.

Inductive bias and data. An inductive bias is an assumption that is built into the architecture. A CNN assumes locality (a pixel interacts first with its neighbors) and translation equivariance in every layer (lesson 07). A ViT assumes much less: only the cut into patches and the per-token MLPs are local; self-attention connects every patch with every other patch from the first layer; the learned position embeddings start with no 2-D structure. The ViT must learn these structures from data. So:

• trained on ImageNet-1k alone (1.3 million images) with the usual regularization, a ViT ends a few points below a ResNet of similar size;
• pre-trained on ImageNet-21k (14 million images) or JFT-300M (300 million), it matches or beats the best CNNs;
• DeiT (Touvron et al., 2021) trains ViTs well on ImageNet-1k alone, with strong data augmentation (RandAugment, Mixup, CutMix, random erasing), stochastic depth, label smoothing and AdamW (lesson 06). DeiT-B (86 million parameters) reaches 81.8% top-1 accuracy at $224 \times 224$.

After training, nearby patches have similar position embeddings: the model has learned the 2-D layout.

A new resolution. To fine-tune at a higher resolution, keep $P$ (the patch-embedding weights fix it). Then $T$ grows, and the position embeddings no longer fit. Reshape the $T$ patch rows of $E_{\text{pos}}$ into their 2-D grid, resize the grid by 2-D interpolation (DeiT uses bicubic), flatten it again, and keep the [CLS] row unchanged. The block weights need no change: they work for any number of tokens.

Why the blocks accept any number of tokens

Every weight matrix in a block acts along the feature axis: $W_Q, W_K, W_V, W_O, W_1, W_2$ have shapes $D \times \cdot$ or $\cdot \times D$, never $T \times \cdot$. Attention forms a $T \times T$ matrix from the data, not from a parameter. So the same weights apply to 197 or 577 tokens. Only the position embeddings (a $(T+1) \times D$ parameter) and the causal mask (a $T \times T$ pattern) depend on $T$. This is also why a decoder with learned absolute positions cannot read beyond $T_{\max}$: it has no embedding for those positions, while the blocks themselves would run (lesson 10).

Answers to section 1. (1) The MLP. It has $8D^2 + 5D$ parameters and the attention $4D^2 + 4D$: about twice as many, for every number of heads. (2) Step $t$ processes $t$ tokens, so the total is $1 + 2 + \dots + T \approx T^2/2$ token passes: quadratic in $T$. With a causal mask, the keys and values of past tokens never change (theorem of section 3.5), so the model stores them in a KV cache, and each step processes only the new token. (3) The patch size stays 16 and the block weights stay the same. The number of tokens grows from $14^2 = 196$ to $24^2 = 576$ (plus [CLS]). The position embeddings must be resized from a $14 \times 14$ grid to a $24 \times 24$ grid by interpolation.

4. Worked example #

(a) One pre-LN block by hand #

Take $T = 2$ tokens of width $D = 2$, one head ($d_k = 2$), LayerNorm scale 1 and shift 0 (and we ignore its small $\epsilon$), no biases, $W_Q = W_K = W_V = W_O = I$, an MLP with $d_{ff} = 2$, $W_1 = I$, ReLU and $W_2 = \frac12 I$. (ReLU keeps the numbers simple; GPT-2 uses GELU.) The input is

Step 1, LayerNorm of each row. Row 1: mean $1$, variance $\frac{(2-1)^2 + (0-1)^2}{2} = 1$, so $(2 - 1, 0 - 1)/1 = (1, -1)$. Row 2: mean $0.5$, variance $0.25$, standard deviation $0.5$, so $(0 - 0.5, 1 - 0.5)/0.5 = (-1, 1)$. (With $D = 2$, a LayerNorm can only output $\pm(1, -1)$; real models have $D$ in the hundreds or thousands.)

Step 2, attention. $Q = K = V = U$, and $UU^\top = \begin{pmatrix} 2 & -2 \\ -2 & 2 \end{pmatrix}$. The scores are

Row 1: $e^{1.414} = 4.113$ and $e^{-1.414} = 0.243$, sum $4.356$, weights $(0.944, 0.056)$. Row 2 is the mirror image: $(0.056, 0.944)$. Row 1 of the output is $0.944\,(1, -1) + 0.056\,(-1, 1) = (0.888, -0.888)$, and row 2 is its mirror image:

With $W_O = I$ this is the attention output.

Step 3, first residual. $X' = X + AV = \begin{pmatrix} 2.888 & -0.888 \\ -0.888 & 1.888 \end{pmatrix}$.

Step 4, LayerNorm of $X'$. Row 1 has mean $1$ and the deviations $\pm 1.888$, so it becomes $(1, -1)$. Row 2 has mean $0.5$ and the deviations $\mp 1.388$, so it becomes $(-1, 1)$.

Step 5, MLP. Row 1: $\operatorname{ReLU}((1, -1)\,I) = (1, 0)$, times $\frac12 I$ gives $(0.5, 0)$. Row 2: $\operatorname{ReLU}(-1, 1) = (0, 1)$, then $(0, 0.5)$.

Step 6, second residual. Add the MLP output to $X'$, row by row: $(2.888 + 0.5, -0.888 + 0)$ and $(-0.888 + 0, 1.888 + 0.5)$.

With a causal mask, token 1 sees only itself: its weights are $(1, 0)$, its attention output is $(1, -1)$, $X'_1 = (3, -1)$, the MLP adds $(0.5, 0)$, and $Y_1 = (3.5, -1)$. Token 2 sees both tokens and does not change: $Y_2 = (-0.888, 2.388)$.

(b) GPT-2 small: parameters and FLOPs #

GPT-2 small has $|\mathcal V| = 50{,}257$, context $T_{\max} = 1024$, $D = 768$, $L = 12$, $H = 12$, tied weights.

Part	Count	Parameters
token embedding (= unembedding)	$50{,}257 \cdot 768$	$38{,}597{,}376$
position embedding	$1024 \cdot 768$	$786{,}432$
12 blocks	$12 \cdot 7{,}087{,}872$	$85{,}054{,}464$
final LayerNorm	$2 \cdot 768$	$1{,}536$
total		$\mathbf{124{,}439{,}808}$

FLOPs per token: the weights multiplied per token are $N = 12 \cdot 12 \cdot 768^2 + 768 \cdot 50{,}257 = 84{,}934{,}656 + 38{,}597{,}376 \approx 123.5$ million. Forward: $2N \approx 247$ MFLOP, plus attention $4LTD = 4 \cdot 12 \cdot 1024 \cdot 768 \approx 37.7$ MFLOP at $T = 1024$: about 285 MFLOP per token. Training: about $3 \cdot 285 \approx 854$ MFLOP per token. On 10 billion training tokens that is about $8.5 \cdot 10^{18}$ FLOPs.

(c) Generation: temperature and the KV cache #

The model gives the logits $z = (2, 1, 0)$ for three candidate tokens.

• $\tau = 1$: $e^z = (7.389, 2.718, 1)$, sum $11.107$, so $p = (0.665, 0.245, 0.090)$.
• $\tau = 0.5$: $z/\tau = (4, 2, 0)$, $e^{z/\tau} = (54.60, 7.389, 1)$, sum $62.99$, so $p = (0.867, 0.117, 0.016)$.
• $\tau = 2$: $z/\tau = (1, 0.5, 0)$, $e^{z/\tau} = (2.718, 1.649, 1)$, sum $5.367$, so $p = (0.506, 0.307, 0.186)$.
• $\tau \to 0$: greedy, always token 1.

KV cache of GPT-2 small for one sequence of $T = 1024$ tokens in fp16: $2 \cdot 12 \cdot 1024 \cdot 768 = 18{,}874{,}368$ numbers, times 2 bytes $\approx 37.7$ MB. The weights in fp16 take $124.4 \cdot 10^6 \cdot 2 \approx 249$ MB. At a batch of 32 sequences the cache needs about 1.2 GB, five times the weights.

A larger model with grouped-query attention: LLaMA-2 70B has $L = 80$, $H = 64$ query heads of width $d_k = 128$ and $H_{kv} = 8$ key-value heads. At $T = 4096$ in fp16: $2 \cdot 80 \cdot 4096 \cdot 8 \cdot 128 \cdot 2$ bytes $\approx 1.34$ GB per sequence. With 64 key-value heads it would be 8 times as much, about 10.7 GB.

(d) ViT-B/16 #

$C = 3$, $224 \times 224$ images, $P = 16$, $D = 768$, $L = 12$, 1000 classes.

Part	Count	Parameters
patch embedding	$3 \cdot 16 \cdot 16 \cdot 768 + 768$	$590{,}592$
class token	$768$	$768$
position embeddings	$197 \cdot 768$	$151{,}296$
12 blocks	$12 \cdot 7{,}087{,}872$	$85{,}054{,}464$
final LayerNorm	$2 \cdot 768$	$1{,}536$
head	$768 \cdot 1000 + 1000$	$769{,}000$
total		$\mathbf{86{,}567{,}656}$

Fine-tuning at $384 \times 384$. Keep $P = 16$. The grid becomes $24 \times 24 = 576$ patches, 577 tokens. Resize the $14 \times 14$ grid of position embeddings to $24 \times 24$ and keep the [CLS] row. The per-token costs grow by $577/197 \approx 2.9$, and the attention matrices by $(577/197)^2 \approx 8.6$.

5. In code #

PyTorch checks the worked examples. The notebook builds the block, a tiny GPT with a KV cache and the patch embedding yourself.

import torch
import torch.nn as nn
import torch.nn.functional as F

torch.manual_seed(0)

# (a) the tiny pre-LN block of the worked example
blk = nn.TransformerEncoderLayer(d_model=2, nhead=1, dim_feedforward=2,
                                 dropout=0.0, activation="relu", bias=False,
                                 batch_first=True, norm_first=True)
with torch.no_grad():
    blk.self_attn.in_proj_weight.copy_(torch.eye(2).repeat(3, 1))  # W_Q=W_K=W_V=I
    blk.self_attn.out_proj.weight.copy_(torch.eye(2))               # W_O = I
    blk.linear1.weight.copy_(torch.eye(2))                          # W_1 = I
    blk.linear2.weight.copy_(0.5 * torch.eye(2))                    # W_2 = I/2
X = torch.tensor([[[2., 0.], [0., 1.]]])
print(blk(X).detach().numpy().round(3))
# [[[ 3.388 -0.888]
#   [-0.888  2.388]]]
causal = nn.Transformer.generate_square_subsequent_mask(2)
print(blk(X, src_mask=causal).detach().numpy().round(3))
# [[[ 3.5   -1.   ]
#   [-0.888  2.388]]]

# (b) 12 D^2 + 13 D parameters per block, and GPT-2 small
layer = nn.TransformerEncoderLayer(768, 12, 3072, batch_first=True, norm_first=True)
print(sum(p.numel() for p in layer.parameters()), 12 * 768**2 + 13 * 768)
# 7087872 7087872
V, Tmax, D, L = 50257, 1024, 768, 12
print(V * D + Tmax * D + L * (12 * D**2 + 13 * D) + 2 * D)
# 124439808

# (c) temperature and the KV cache
z = torch.tensor([2., 1., 0.])
for tau in (1.0, 0.5, 2.0):
    print(tau, torch.softmax(z / tau, -1).numpy().round(3))
# 1.0 [0.665 0.245 0.09 ]
# 0.5 [0.867 0.117 0.016]
# 2.0 [0.506 0.307 0.186]
print(2 * L * Tmax * D * 2 / 1e6, "MB")   # 37.748736 MB

# (d) ViT-B/16: the patch embedding is a convolution with kernel = stride = 16
conv = nn.Conv2d(3, 768, kernel_size=16, stride=16)
x = torch.randn(1, 3, 224, 224)
tokens = conv(x).flatten(2).transpose(1, 2)              # (1, 196, 768)
patches = x.reshape(3, 14, 16, 14, 16).permute(1, 3, 0, 2, 4).reshape(196, 768)
by_hand = patches @ conv.weight.reshape(768, -1).T + conv.bias
print(tokens.shape, torch.allclose(tokens[0], by_hand, atol=1e-5))
# torch.Size([1, 196, 768]) True
pos = torch.randn(1, 197, 768)                           # [CLS] row + 14 x 14 grid
grid = pos[:, 1:].reshape(1, 14, 14, 768).permute(0, 3, 1, 2)
grid = F.interpolate(grid, size=(24, 24), mode="bicubic", align_corners=False)
grid = grid.permute(0, 2, 3, 1).reshape(1, 576, 768)
new_pos = torch.cat([pos[:, :1], grid], dim=1)           # keep the [CLS] row
print(new_pos.shape)
# torch.Size([1, 577, 768])

The bias=False option (PyTorch 2.1 or later) removes the biases of the linear layers and the shifts of the LayerNorms. The $\epsilon = 10^{-5}$ inside LayerNorm changes the result only after the third decimal.

6. Where it shows up #

• Large language models. GPT-2, GPT-3, LLaMA, Mistral and most chat models are decoder-only stacks. LLaMA changes three parts of the GPT-2 block: RMSNorm instead of LayerNorm (lesson 05), a SwiGLU MLP (lesson 04) and RoPE instead of learned positions (lesson 10). The plan of this lesson (embedding, $L$ pre-norm blocks, final norm, unembedding) stays the same.
• Serving models. Inference has two phases. In prefill the model reads the prompt in one parallel pass and fills the KV cache. In decode it makes one token per step; each step reads all weights and the whole cache from memory, so speed is limited by memory bandwidth, not by FLOPs. GQA, cache quantization and paged cache memory (as in vLLM) all make the cache cheaper.
• Compute budgets. The rule "6 FLOPs per parameter per training token" gives quick estimates: a 7-billion-parameter model trained on 2 trillion tokens needs about $6 \cdot 7 \cdot 10^9 \cdot 2 \cdot 10^{12} = 8.4 \cdot 10^{22}$ FLOPs. Scaling laws are written in these units.
• Encoders. BERT-style encoders produce text embeddings for search and classification. Encoder–decoders run speech recognition (Whisper: an audio encoder and a text decoder) and translation.
• Vision. ViTs are the image encoders of most CLIP models, of DINO (lesson 14) and of the Segment Anything model. The diffusion transformer (DiT) applies the same patch tokens to the latent images of a diffusion model. Swin transformers restrict attention to local windows, which brings back a CNN-like locality bias and a cost linear in the number of patches.
• Science. ClimaX is a ViT-style weather and climate model that makes patch tokens for each physical variable separately and then merges them per patch. Pangu-Weather uses a 3-D transformer over latitude, longitude and pressure levels. Seismic and other time series are cut into 1-D patches (a strided Conv1d), as in PatchTST for forecasting.

7. Common mistakes #

• No final LayerNorm in a pre-LN stack. The residual stream is never normalized inside the blocks; without the final norm its scale grows with depth and reaches the logits.
• "Attention is the expensive part." At short and medium lengths, the MLP and the projections dominate: attention matrices cost more only when $T > 6D$.
• FLOP counts off by 2 or 3. A multiply-add is 2 FLOPs; a training step is about 3 times the forward pass (the backward pass computes two products per linear layer).
• A wrong position for the new token in cached generation. The new token at step $t$ must get position $t$ (its learned embedding or its RoPE angle), not position 1. Without a cache this bug cannot happen, so test cached against uncached generation.
• A KV cache for an encoder. Without a causal mask, a new token changes the keys and values of all earlier tokens, so a cache gives wrong results.
• Temperature 0 as a division. $z/\tau$ with $\tau = 0$ gives infinities and NaN. Implement $\tau = 0$ as greedy decoding.
• Changing the patch size for a new resolution. Keep $P$, let $T$ grow, interpolate the position grid. Do not interpolate the [CLS] row, and do not forget it in the token count ($T + 1$).
• A patch flattening order that does not match the convolution weights. A filter of nn.Conv2d is flattened in the order (channel, row, column); your patches must use the same order.
• "A ViT has no inductive bias." It has less image-specific bias than a CNN, not none: the patches, the per-token MLPs and the permutation symmetry of attention are all built-in choices.

8. Check yourself #

1. A pre-LN block has $B = 2$, $T = 5$, $D = 64$, $H = 4$ and $d_{ff} = 4D$. Give the shapes of $\operatorname{LN}_1(X)$, of the queries after the split into heads, of the attention weights, of the MLP hidden layer and of the block output.
   Answer

   $d_k = 64/4 = 16$. $\operatorname{LN}_1(X)$: $(2, 5, 64)$. Queries: $(2, 4, 5, 16)$. Attention weights: $(2, 4, 5, 5)$. MLP hidden layer: $(2, 5, 256)$. Output: $(2, 5, 64)$, the same as the input.
2. A decoder-only model has $|\mathcal V| = 32{,}000$, $T_{\max} = 1024$ learned positions, $D = 512$, $L = 8$, $d_{ff} = 4D$, biases and a final LayerNorm. Count its parameters with tied weights, and without tying.
   Answer

   Embedding $32{,}000 \cdot 512 = 16{,}384{,}000$; positions $1024 \cdot 512 = 524{,}288$; blocks $8 \cdot (12 \cdot 512^2 + 13 \cdot 512) = 8 \cdot 3{,}152{,}384 = 25{,}219{,}072$; final LayerNorm $1{,}024$. Tied: $42{,}128{,}384$. Untied: add $16{,}384{,}000$ for $W_U$ (no bias in GPT-style models), $58{,}512{,}384$. Here the embeddings are about 40% of the tied model.
3. Why can a decoder reuse the keys and values of earlier tokens when it appends a token, while BERT cannot?
   Answer

   With a causal mask, the vector of token $j$ in every layer depends only on tokens $1, \dots, j$ (proof by induction over the layers), so its key and value never change. In BERT every token attends to every token, so a new token changes the vectors, keys and values of all earlier tokens in every layer after the first.
4. A ViT for CIFAR-10 images ($3 \times 32 \times 32$) uses $P = 4$, $D = 192$ and a [CLS] token. Count the tokens, the patch-embedding parameters (with bias) and the position-embedding parameters.
   Answer

   $T = (32/4)^2 = 64$ patches, 65 tokens. Patch embedding: $3 \cdot 4 \cdot 4 \cdot 192 + 192 = 9{,}408$. Position embeddings: $65 \cdot 192 = 12{,}480$.
5. At $224 \times 224$ you halve the patch size from 16 to 8. By which factors do the number of tokens and the size of each attention matrix grow?
   Answer

   Patches: $196 \to 784$, a factor 4 (197 to 785 tokens with [CLS]). The attention matrix has $(T + 1)^2$ entries: a factor of about 16 ($785^2/197^2 \approx 15.9$).
6. A model has $D = 1024$. From which context length $T$ on does the attention matrix cost more FLOPs per token than all the weights of a block?
   Answer

   Per token and per block, the weights cost $24D^2$ FLOPs and the attention matrix $4TD$. So $4TD > 24D^2 \iff T > 6D = 6144$.

Recall cardsAnswer in your head first, then reveal. Add the cards to your review deck to see them again just before you would forget.

1. Write the two equations of a pre-LN transformer block.

   Show answer

2. Which part of a transformer block mixes information between tokens?

   Show answer

3. How many parameters does a transformer block with width $D$ and $d_{ff} = 4D$ have, and how are they split?

   Show answer

4. Why does training cost about 6 FLOPs per parameter per token?

   Show answer

5. At which context length does the attention matrix cost as much as the weights of a block?

   Show answer

6. What masks do encoder-only, decoder-only and encoder–decoder models use?

   Show answer

7. Why does a KV cache give exactly the same output as running the full model again?

   Show answer

8. How many numbers does the KV cache of one sequence hold?

   Show answer

9. How many tokens does a ViT make from a $224 \times 224$ image with $P = 16$ and a [CLS] token?

   Show answer

10. How do you fine-tune a ViT at a higher resolution?

    Show answer

Add these cards to my reviewOpen the review →

9. Interview angle #

• "Draw a transformer block and give all shapes." — Pre-LN: LN, multi-head attention, add; LN, MLP ($D \to 4D \to D$), add. Every arrow is $(B, T, D)$; inside attention $(B, H, T, d_k)$ and $(B, H, T, T)$. Name the final LayerNorm.
• "How many parameters does GPT-2 small have?" — $12D^2$ per block gives $85$ M for 12 blocks; the token embedding adds $38.6$ M and positions $0.8$ M; tied unembedding; total $124.4$ M.
• "How many FLOPs to train a 7B model on 2T tokens?" — $6 \times$ parameters $\times$ tokens $\approx 6 \cdot 7 \cdot 10^9 \cdot 2 \cdot 10^{12} = 8.4 \cdot 10^{22}$, plus the attention term if the context is long.
• "What is the KV cache, how big is it, and how do you make it smaller?" — Stored keys and values of past tokens, exact because of the causal mask; $2LTD$ numbers per sequence; shrink it with MQA/GQA, a lower-precision cache, or a sliding window.
• "BERT, GPT, T5: what is the difference?" — Masks (bidirectional, causal, both), training task (masked tokens, next token, span infilling), cross-attention in the decoder of T5.
• "How does a ViT work, and why does it need more data than a CNN?" — Patches as tokens (a strided convolution), [CLS] and learned positions, a standard encoder. It has fewer built-in assumptions about images, so it learns them from data; DeiT compensates with augmentation and regularization.

10. Practice #

The notebook for this lesson runs in Colab (PyTorch on a CPU, under a minute). It has five exercises, each with a check cell and a folded solution:

1. Write a pre-LN transformer block from tensor operations and match PyTorch's nn.TransformerEncoderLayer (with norm_first=True) with copied weights.
2. Write a function that counts the parameters of a GPT-style model and check it against sum(p.numel() for p in model.parameters()) and GPT-2 small.
3. Write temperature sampling (with $\tau = 0$ as greedy) and use it on a tiny character-level GPT that the notebook trains on a short text in a few seconds.
4. Add a KV cache to the tiny GPT and check that cached and uncached generation give the same tokens.
5. Write the ViT patch embedding as a reshape and a matrix product, and check it against nn.Conv2d(kernel_size=P, stride=P).

11. Watch and read #

Transformers, the tech behind LLMs | Deep Learning Chapter 53Blue1Brown · 27 min

The whole model, from embeddings to the unembedding and softmax, drawn as data flowing through matrices.

How might LLMs store facts | Deep Learning Chapter 73Blue1Brown · 23 min

What the MLP of a block does, and how its parameters are counted.

Let's build GPT: from scratch, in code, spelled out.Andrej Karpathy · 116 min

From 1:19:11 he turns one attention head into full blocks (multi-head, MLP, residual, LayerNorm) and compares encoder and decoder blocks.

LLaMA explained: KV-Cache, Rotary Positional Embedding, RMS Norm, Grouped Query Attention, SwiGLUUmar Jamil · 71 min

From 40:22, the KV cache and then grouped-query attention, with pictures of the matrices at each step.

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale (Paper Explained)Yannic Kilcher · 30 min

The ViT architecture, its data requirements and what the position embeddings learn.

• Vaswani et al., Attention Is All You Need (2017): §3 (the encoder–decoder architecture, the FFN) and Table 3 (model sizes).
• Prince, Understanding Deep Learning, free PDF: chapter 12 (§12.4 transformer layers, §12.6–12.8 encoders, decoders, encoder–decoders, §12.10 transformers for images).
• Rush et al., The Annotated Transformer: the whole model in commented PyTorch code.
• Kaplan et al., Scaling Laws for Neural Language Models (2020): §2.1, Table 1 (parameters and FLOPs per token of a transformer).
• Kipply, Transformer Inference Arithmetic: KV cache size, FLOPs, and why decoding is limited by memory bandwidth.
• Dosovitskiy et al., An Image is Worth 16x16 Words (ViT, 2021): §3.1–3.2 (model, inductive bias, higher resolution), §4.3 (data size), Appendix D.4 (position embeddings).
• Touvron et al., Training data-efficient image transformers (DeiT, 2021): §4–6 (distillation token, training recipe).
• Devlin et al., BERT (2019) and Raffel et al., T5 (2020): the training tasks of the encoder-only and encoder–decoder families.
• Ainslie et al., GQA (2023): grouped-query attention, and how to convert a multi-head model to it.