Chapter 27 · Section 27.2 · 60 min

The ViT architecture

A Vision Transformer, or ViT, is a Transformer encoder that reads an image as a sequence of patches. Nothing inside it knows that it is looking at pixels. This section builds the whole model one tensor shape at a time, and then counts what it costs in parameters, in arithmetic, and in memory. You need those counts to budget a training run and to choose an image resolution. You also need them to load someone else's weights without silent errors, and to answer the estimation questions that labs ask in interviews.

The input was built in §27.1. A 224 by 224 color image is cut into a 14 by 14 grid of 16 by 16 squares, called patches, 196 in all. Each patch is flattened into its $16\cdot16\cdot3 = 768$ pixel values. One learned matrix turns each such list into a vector of width 768, the patch's token. A learned vector called the class token goes in front of the patches. A learned position vector is added to every row, so the model can tell where each patch came from. For the Base model with 16-pixel patches, written ViT-B/16, the result is a matrix of 197 rows and 768 columns.

What reads that matrix is a stack of Transformer blocks. §6.2 built GPT-2 small, a 124-million-parameter language model, from twelve of them. Each block is an attention layer, in which every token reads from the others, followed by a small two-layer network applied to each token on its own. A shortcut connection runs around each of the two. Take those twelve blocks and delete the causal mask that stops each token from seeing the tokens after it. Run them over the 197 rows. That is very nearly a ViT.

So the interesting parts are the two ends and the cost. The input is a projection of pixels instead of a table lookup, so the vocabulary table that held almost a third of GPT-2's parameters is gone. The output is one vector per image, and a causal mask left in by mistake makes that vector blind to the image. The cost depends on the number of tokens, which grows with the square of the image side. Attention is 4% of the arithmetic at 224 pixels and almost half of it at 1,024 pixels. By the end of the section you will derive both numbers and check them against published training runs.

A ViT is a Transformer encoder over patches

An encoder is a Transformer in which every token may attend to every other token, in both directions. BERT, the text encoder of §7.1 that is trained to fill in hidden words, is one. In a ViT this means every patch reads from every other patch in every layer. Information can travel from the top-left corner of the image to the bottom-right in a single layer. A convolutional network looks at a small neighborhood in each layer, so it needs many layers before one of its outputs depends on the whole image (§0.8). A ViT has that view from its first layer.

At the end, a classifier reads one vector, the final state of the class token. That token has no pixels of its own. It spends every layer gathering information from the patches, so it ends as a summary of the image.

Here is the whole model for ViT-B/16 on a batch of $B$ images, with the shape of every tensor. Three abbreviations appear in it. LN is a LayerNorm, which rescales each token's vector to mean zero and variance one and then applies a learned gain and shift (§5.2). MHA is multi-head attention, which runs several smaller attentions side by side and joins their outputs (§4.5). MLP is the two-layer network applied to each token on its own, with the GELU, a smooth relative of $\max(0, x)$ , between its layers (§5.3). Read the diagram from top to bottom.

image            (B, 224, 224, 3)      pixels, channels last
  │  cut into 16 x 16 patches, flatten each (§27.1)
patches X_p      (B, 196, 768)         768 = 16·16·3 numbers per patch
  │  X_p W_E + b_E, prepend x_cls, add P_pos
X⁽⁰⁾              (B, 197, 768)         T = 196 patches + 1 class token
  │  12 times:  H = X + MHA(LN₁(X))      12 heads of 64, no mask
  │             X = H + MLP(LN₂(H))      768 → 3072 → 768, GELU
X⁽¹²⁾             (B, 197, 768)
  │  take row 0 (the class token), LN_f
z                (B, 768)
  │  z W_head + b_head
logits           (B, 1000)

Each block adds the output of attention to its input, then adds the output of the MLP. The running sum that every block reads from and writes back to is called the residual stream (§5.1). Each sublayer reads a normalized copy of the stream rather than the stream itself. That placement of the norm is called pre-LN.

Three models built from the same block make the differences concrete. The ViT authors took their Base and Large sizes straight from BERT (Dosovitskiy et al. 2021). So all three columns below share one block shape, with width $d = 768$ , 12 heads, and an MLP of width $d_{\text{ff}} = 3{,}072$ .

	GPT-2 small (§6.2)	BERT-base (§7.1)	ViT-B/16
input	token table $\mathbf{E}\in\R^{50257\times768}$	token table, 30,522 rows	linear map of pixels, $\mathbf{W}_E\in\R^{768\times768}$
attention mask	causal	none	none
norm placement	pre-LN, plus a final LN	post-LN (norm after each residual add)	pre-LN, plus a final LN
positions	learned, 1,024	learned, 512, plus segments	learned, 197 (a 14 by 14 grid plus one)
output	a distribution over the next token at each position	a vector per position, and the [CLS] vector for sentence-level tasks	one vector per image
GELU	tanh approximation	depends on the code base	exact in timm and torchvision
LayerNorm $\epsilon$	$10^{-5}$	$10^{-12}$	$10^{-6}$
parameters	124,439,808	about 110M	86,567,656

The twelve blocks hold the same tensors in all three models, $12d^2 + 13d = 7{,}087{,}872$ parameters each with $d = 768$ (§5.6 counts them). BERT only places its norms differently. Everything else that differs sits at the two ends of the network.

The small rows of the table matter whenever you load weights trained by someone else. Most ViT weights are loaded through one of two PyTorch libraries, timm (PyTorch Image Models) and torchvision. The $\epsilon$ of a LayerNorm is the small number added to the variance before the square root, so that a constant vector does not divide by zero. A LayerNorm built with PyTorch's default $\epsilon = 10^{-5}$ instead of ViT's $10^{-6}$ will not crash. Its outputs will be slightly wrong, and nothing will tell you. When an interviewer asks how you would port a checkpoint, list these constants, the $\epsilon$ , the GELU form, and the norm placement. Then say you would compare your outputs with the original code's on one fixed input.

The forward pass, written exactly

Here is the model as a definition, with every symbol introduced. Weights are stored as (input, output) and applied to tokens stored as rows, as in $\mathbf{X}\mathbf{W}$ . That has been the book's convention since Chapter 4.

📐 Vision Transformer (ViT)

Let $\mathbf{X}_{\text{img}} \in \R^{H\times W\times C}$ be an image with height $H$ , width $W$ , and $C$ channels (3 for color). $H$ and $W$ are multiples of the patch size $P$ . Let $\mathbf{X}_p \in \R^{N_p\times P^2C}$ hold its $N_p = HW/P^2$ flattened patches, one per row. The rows are in raster order, left to right along each row of patches and then top to bottom (§27.1). The model has width $d$ , $h$ heads of size $d_k = d/h$ , MLP width $d_{\text{ff}}$ , $L$ blocks, and $K$ classes. It computes the following.

Input. The first block reads

\mathbf{X}^{(0)} = [\,\mathbf{x}_{\text{cls}}^\top;\ \mathbf{X}_p\mathbf{W}_E + \mathbf{1}\mathbf{b}_E^\top\,] + \mathbf{P}_{\text{pos}} \in \R^{T\times d}, \qquad T = N_p + 1 .

The patch embedding is $\mathbf{W}_E \in \R^{P^2C\times d}$ with bias $\mathbf{b}_E\in\R^d$ . The class token $\mathbf{x}_{\text{cls}}\in\R^d$ and the position table $\mathbf{P}_{\text{pos}}\in\R^{T\times d}$ are learned. Here $\mathbf{1}$ is a column of ones, so $\mathbf{1}\mathbf{b}^\top$ adds the bias $\mathbf{b}$ to every row. The bracket $[\mathbf{a}^\top; \mathbf{M}]$ stacks the row $\mathbf{a}^\top$ on top of the matrix $\mathbf{M}$ .

Blocks. For $\ell = 1, \dots, L$ ,

\mathbf{H}^{(\ell)} = \mathbf{X}^{(\ell-1)} + \operatorname{MHA}\big(\operatorname{LN}_1(\mathbf{X}^{(\ell-1)})\big), \qquad \mathbf{X}^{(\ell)} = \mathbf{H}^{(\ell)} + \operatorname{MLP}\big(\operatorname{LN}_2(\mathbf{H}^{(\ell)})\big).

Inside MHA, a normalized input $\mathbf{Y}\in\R^{T\times d}$ is projected once, $[\mathbf{Q}\ \mathbf{K}\ \mathbf{V}] = \mathbf{Y}\mathbf{W}_{QKV} + \mathbf{1}\mathbf{b}_{QKV}^\top$ , with $\mathbf{W}_{QKV}\in\R^{d\times3d}$ and $\mathbf{b}_{QKV}\in\R^{3d}$ . Head $j$ uses columns $jd_k$ to $(j+1)d_k - 1$ of each of $\mathbf{Q}, \mathbf{K}, \mathbf{V}$ , written $\mathbf{Q}_j, \mathbf{K}_j, \mathbf{V}_j$ . Then

\operatorname{MHA}(\mathbf{Y}) = \big[\operatorname{softmax}\big(\mathbf{Q}_j\mathbf{K}_j^\top/\sqrt{d_k}\big)\mathbf{V}_j\big]_{j=0}^{h-1}\mathbf{W}_O + \mathbf{1}\mathbf{b}_O^\top, \qquad \operatorname{MLP}(\mathbf{Y}) = \operatorname{GELU}\big(\mathbf{Y}\mathbf{W}_1 + \mathbf{1}\mathbf{b}_1^\top\big)\mathbf{W}_2 + \mathbf{1}\mathbf{b}_2^\top,

with $\mathbf{W}_O\in\R^{d\times d}$ , $\mathbf{W}_1\in\R^{d\times d_{\text{ff}}}$ , and $\mathbf{W}_2\in\R^{d_{\text{ff}}\times d}$ . The softmax runs over each row, and there is no mask. The outer brackets place the $h$ head outputs side by side. Each LN has its own gain $\boldsymbol\gamma$ and shift $\boldsymbol\beta$ in $\R^d$ , and $\epsilon = 10^{-6}$ . GELU is the exact $x\Phi(x)$ , with $\Phi$ the standard normal cumulative distribution function.

Output. The image representation is $\mathbf{z} = \operatorname{LN}_f\big(\mathbf{X}^{(L)}_{0}\big) \in \R^d$ , the final state of the class token, which is row 0. The logits, the unnormalized class scores, are $\mathbf{z}^\top\mathbf{W}_{\text{head}} + \mathbf{b}_{\text{head}}^\top \in \R^{1\times K}$ , with $\mathbf{W}_{\text{head}}\in\R^{d\times K}$ .

In plain words, embed each patch, put the class token in front, and add positions. Then, in every block, let each token read from all the others (attention) and then transform itself on its own (the MLP). Each sublayer reads a normalized copy of the stream and adds its answer back. At the end, normalize the class token's vector and read the class scores off it with one matrix.

Every piece has been derived earlier in the book. Pre-LN, and why it trains without a careful warmup, is in §5.2. The GELU MLP is in §5.3. The split into heads and its reshape are in §4.5, and the division by $\sqrt{d_k}$ , which keeps the scores from growing with the head size, is in §4.3.

🔬 The paper's notation, and variants that change only the ends

The definition is Equations 1 to 4 of the ViT paper in the book's notation. The paper writes $D$ for our $d$ and $N$ for our $N_p$ , since we keep $N$ for a parameter count. It writes $\mathbf{E}$ for our $\mathbf{W}_E$ , since we keep $\mathbf{E}$ for a token table, and $\mathbf{E}_{pos}$ for $\mathbf{P}_{\text{pos}}$ . It writes $\mathbf{z}_\ell$ for our $\mathbf{X}^{(\ell)}$ and $\mathbf{y}$ for our $\mathbf{z}$ , so its $\mathbf{z}$ is not ours. Its Equation 1 has no patch bias $\mathbf{b}_E$ , though the released code does have one.

Some models change the two ends and leave the blocks alone. The class token can be replaced by the mean of the patch outputs, or by a learned query that attends over them, and both are taught in §27.1. A few extra learned tokens called registers can be inserted after the class token. They take part in attention but the head never reads them, and §27.4 explains why large ViTs need them. With $R$ registers, $T = N_p + 1 + R$ .

The GELU also varies. The timm and torchvision libraries use the exact $x\Phi(x)$ . The original JAX code calls the nn.gelu of the Flax library, whose default is the tanh approximation. The two differ by less than $10^{-3}$ per activation. That matters only when you compare outputs digit by digit, as §6.2 found for GPT-2.

The tracer below follows one batch, tensor by tensor, through five choices of model and resolution. Before you step through it, predict the shape of the attention scores in one block of ViT-L/14 at 336 pixels for a batch of 32 images, from $P = 14$ , the resolution, and $h = 16$ heads.

At 336 pixels a 14-pixel patch gives a 24 by 24 grid, so $N_p = 576$ and $T = 577$ . The scores have shape $(32, 16, 577, 577)$ , about 170 million numbers in every block. Every other tensor inside a block has $T$ rows per image. The scores are the only tensor whose size grows with the square of the token count, and that is where the cost of high resolution will come from.

Two heads, one for pretraining and one for transfer

The head is the small network that turns $\mathbf{z}$ into class scores, and the ViT paper uses two different ones. During pretraining, $\mathbf{z}$ goes through one hidden layer with a tanh, $\tanh(\mathbf{z}^\top\mathbf{W}_{\text{pre}} + \mathbf{b}_{\text{pre}}^\top)$ with $\mathbf{W}_{\text{pre}}\in\R^{d\times d}$ , and then a linear layer to the pretraining classes. There are about 18,000 classes for JFT-300M, Google's internal set of 303 million labeled images, and about 21,800 for ImageNet-21k, a public set of 14 million. The paper says this design is inherited from the Transformer for text. It matches the small network that BERT puts on its [CLS] token.

Fine-tuning means continuing to train a pretrained model on a new task. To fine-tune, the authors throw the whole two-layer head away. They replace it with a single $d\times K$ linear layer whose weights start at zero (Appendix B.1.1 of the paper). Zero is a deliberate choice, and it has an exact consequence for the backbone, which is everything before the head.

∴ A zero-initialized head passes no gradient to the backbone on the first step

Let the logits be $\mathbf{s} = \mathbf{W}_{\text{head}}^\top\mathbf{z} + \mathbf{b}_{\text{head}}$ and the loss the cross-entropy $\mathcal{L} = -\log\operatorname{softmax}(\mathbf{s})_y$ for the true class $y$ . Write $\mathbf{p} = \operatorname{softmax}(\mathbf{s})$ and $\mathbf{e}_y$ for the one-hot target. If $\mathbf{W}_{\text{head}} = \mathbf{0}$ , then $\partial\mathcal{L}/\partial\mathbf{z} = \mathbf{0}$ , so every parameter of the backbone gets zero gradient. The head's own gradient, $\partial\mathcal{L}/\partial\mathbf{W}_{\text{head}} = \mathbf{z}(\mathbf{p} - \mathbf{e}_y)^\top$ , is not zero whenever $\mathbf{z} \ne \mathbf{0}$ .

Proof. The gradient of softmax cross-entropy with respect to the logits is $\partial\mathcal{L}/\partial\mathbf{s} = \mathbf{p} - \mathbf{e}_y$ (§0.2). By the chain rule through $\mathbf{s} = \mathbf{W}_{\text{head}}^\top\mathbf{z} + \mathbf{b}_{\text{head}}$ , we get $\partial\mathcal{L}/\partial\mathbf{z} = \mathbf{W}_{\text{head}}(\mathbf{p} - \mathbf{e}_y) = \mathbf{0}$ and $\partial\mathcal{L}/\partial\mathbf{W}_{\text{head}} = \mathbf{z}(\mathbf{p} - \mathbf{e}_y)^\top$ . Every backbone gradient is a product with $\partial\mathcal{L}/\partial\mathbf{z}$ , so it vanishes.

✏️ The zero head, by hand

Take $d = 3$ , $K = 4$ classes, the feature $\mathbf{z} = (1, -2, 0.5)$ , and true class $y = 2$ , counting from 0. The head and its bias are zero. The logits are all 0, so $\mathbf{p} = (0.25, 0.25, 0.25, 0.25)$ and $\mathbf{p} - \mathbf{e}_y = (0.25, 0.25, -0.75, 0.25)$ .

The gradient reaching the backbone is $\mathbf{W}_{\text{head}}(\mathbf{p} - \mathbf{e}_y) = (0, 0, 0)$ . The head's gradient is the outer product $\mathbf{z}(\mathbf{p} - \mathbf{e}_y)^\top$ . Its first row is $(0.25, 0.25, -0.75, 0.25)$ and its second row is $(-0.5, -0.5, 1.5, -0.5)$ . So the first step moves the head toward class 2 for this feature and leaves the backbone where it was.

In plain words, the first update trains only the head. At that moment the model is a linear probe, a linear classifier trained on top of features that stay fixed. A randomly initialized head would instead push random gradients into features that took hundreds to thousands of accelerator-days to learn. The paper reports that the zero head was "a little more robust" than re-initializing only the last layer. That is an empirical finding, not a theorem. In an interview, the expected answer is that the zero head protects the pretrained features while the new head learns.

🔬 Weight decay, and why the protection lasts longer than one step

The theorem is about gradients, and decoupled weight decay moves weights without one. With AdamW (§0.6), every step also applies $\boldsymbol\theta \leftarrow \boldsymbol\theta - \eta\lambda\boldsymbol\theta$ , with learning rate $\eta$ and weight decay $\lambda$ . That shrinks the backbone weights on step 1 even though their gradient is zero. The ViT paper fine-tunes with SGD, momentum 0.9, and no weight decay (its Appendix B.1.1 and Table 4). So in its recipe the backbone does not move at all on the first step.

The protection also lasts longer than one step. From step 2 on, the gradient into the backbone is $\partial\mathcal{L}/\partial\mathbf{z} = \mathbf{W}_{\text{head}}(\mathbf{p} - \mathbf{e}_y)$ . The head is only a few small steps away from zero, so the backbone receives small gradients until the head has grown and points somewhere sensible. Layer-wise learning-rate decay (§27.4) gives the head the largest learning rate and the earliest layers the smallest. It comes from the same wish to protect what pretraining learned.

A ViT from scratch in NumPy

Here is our version, which the rest of the section uses. It is a tiny ViT with width $d = 32$ , $L = 2$ blocks, $h = 4$ heads, and MLP width $d_{\text{ff}} = 128$ . It reads 32 by 32 images with $P = 8$ , so there are $N_p = 16$ patches and $T = 17$ tokens. Each line of the forward pass is one line of the definition. The first cell holds the pieces and the second the model.

Python · cell 1

Python · cell 2

The shapes match the diagram at toy size. The scores are $(B, h, T, T) = (2, 4, 17, 17)$ , one $T\times T$ matrix per head per image. That matrix is the object whose size will dominate the cost at high resolution.

Python · cell 3

With the mask the difference is exactly zero. A causal mask lets row $i$ of each attention matrix read only positions 0 to $i$ . Row 0 may read only position 0, so the class token reads nothing but itself, in every layer. Its final state depends on $\mathbf{x}_{\text{cls}}$ , its position row, and the weights, and never on a single pixel.

What does this change in practice? A ViT with a causal mask copied from a GPT gives the same output for every image, so all it can learn is how often each class occurs. On ImageNet-1k, whose 1,000 classes are close to balanced, the loss stalls near $\ln 1000 \approx 6.9$ and accuracy stays at chance. That is easy to mistake for a data bug. BERT's [CLS] token also sits at position 0, and it works for the same reason a ViT's class token does, because the encoder has no mask.

The cells above build the plain model, with a class token and nothing else. Checkpoints people actually load often differ in two ways. Many add registers, the extra learned tokens that take part in attention but are never read by the head (§27.4 explains why large models need them). Many also pool by averaging the patch outputs instead of reading a class token, which is called global average pooling (GAP). Both change the bookkeeping in ways that fail silently, so the interview version of this problem asks for the general model.

Interview

A general ViT forward pass, with registers and either pooling, no loops over tokens or heads

Implement vit_forward(params, img, pool="cls"). It takes a batch of images img of shape (B, H, W, C), channels last. $H$ and $W$ are multiples of $P$ but need not be equal. It returns logits of shape (B, K). The blocks are those of the definition above. What changes is the sequence the blocks read and the vector the head reads.

The parameters are a dict with these entries.

params["P"] is the patch size and params["h"] the number of heads.
W_E is (P*P*C, d) and b_E is (d,). Patches are in raster order, left to right and then top to bottom. Inside a patch the pixels are also flattened row by row, with the $C$ channels of each pixel next to each other. So pixel $(u, v)$ and channel $c'$ sit at index $(uP + v)C + c'$ .
reg is (R, d), the register tokens, with $R \ge 0$ . It may have zero rows.
With pool="cls", cls is (d,) and pos is (1 + N_p, d), with row 0 for the class token and the rest for the patches. With pool="gap" there is no cls, and pos is (N_p, d).
blocks is a list of dicts with ln1_g, ln1_b, W_qkv (d, 3d), b_qkv, W_o (d, d), b_o, ln2_g, ln2_b, W_1 (d, d_ff), b_1, W_2 (d_ff, d), b_2. The first d columns of the QKV output are the queries, then come the keys, then the values. Head j takes the columns j*d_k:(j+1)*d_k of each, a Python slice that excludes (j+1)*d_k itself.
lnf_g, lnf_b, W_head (d, K), b_head (K,) form the output.

Build the sequence in this order. First add pos to the class token (if there is one) and the patches. Then insert the registers right after the class token, with no position of their own. The sequence is the class token, the $R$ registers, then the $N_p$ patches. It is the input sequence defined in §27.1, and the layout of timm and of DINOv2, a ViT from Meta trained without labels (§28.4).

Pool the way timm does. With pool="cls", apply the final LayerNorm to row 0. With pool="gap", average the patch rows only, never the registers, and apply the final LayerNorm to the mean, which timm calls fc_norm.

Use pre-LN, LayerNorm with the population variance and $\epsilon = 10^{-6}$ , the exact GELU $x\Phi(x)$ (use scipy.special.erf), and no mask. You may loop over the blocks, but not over tokens or heads. The setup gives you make_vit_params(rng, H, W, R=..., pool=...), which makes random parameters to test with while you work.

import numpy as np
from scipy.special import erf

def vit_forward(params, img, pool="cls"):
    """img: (B, H, W, C) -> logits (B, K). pool is "cls" or "gap"."""
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Almost every parameter sits in the blocks

Now the parameter count, done term by term the way §6.2 counted GPT-2. The parameters live in five places. They are the patch embedding, the prefix vectors (the class token and any registers), the position table, the blocks, and the final LayerNorm with the head. Here is ViT-B/16 first, by hand.

One block with $d_{\text{ff}} = 4d$ holds $12d^2 + 13d$ parameters (§5.6). The four attention projections with their biases give $4d^2 + 4d$ . The two MLP matrices with their biases give $8d^2 + 5d$ . The two LayerNorms give $4d$ .

✏️ ViT-B/16 at 224, counted by hand

With $P = 16$ , $C = 3$ , $d = 768$ , $L = 12$ , $d_{\text{ff}} = 3072$ , $N_p = 196$ , one class token, no registers, and $K = 1000$ classes,

piece	arithmetic	parameters
patch embedding	$(768 + 1)\cdot768$	590,592
class token	$768$	768
positions	$197\cdot768$	151,296
one block	$12\cdot768^2 + 13\cdot768 = 7{,}077{,}888 + 9{,}984$	7,087,872
twelve blocks	$12\cdot7{,}087{,}872$	85,054,464
final LayerNorm	$2\cdot768$	1,536
head	$769\cdot1000$	769,000
total		86,567,656

The blocks are 98.3% of the model, the head 0.9%, the patch embedding 0.7%, and the position table 0.2%. Without the head, which is what timm builds with num_classes=0, the backbone has 85,798,656 parameters.

Fine-tuning at 384 pixels changes only the position table, which grows from 197 to 577 rows. That adds $380\cdot768 = 291{,}840$ parameters, so ViT-B/16 at 384 has 86,859,496. Compare GPT-2 small, which has the same blocks and whose token table alone holds 38.6 million.

The same sum, written with symbols, counts any ViT. Write $n_{\text{cls}} \in \{0, 1\}$ for whether the model has a class token and $R$ for the number of registers.

∴ Parameter count of a ViT

A ViT with patch size $P$ , $C$ channels, width $d$ , $L$ blocks, MLP width $d_{\text{ff}}$ , $N_p$ patches, $n_{\text{cls}}$ class tokens, $R$ registers (with no position embedding), and a $K$ -class linear head has

N = \underbrace{(P^2C + 1)\,d}_{\mathbf{W}_E,\ \mathbf{b}_E} + \underbrace{(n_{\text{cls}} + R)\,d}_{\text{prefix vectors}} + \underbrace{(N_p + n_{\text{cls}})\,d}_{\mathbf{P}_{\text{pos}}} + L\underbrace{\big(4d^2 + 2dd_{\text{ff}} + d_{\text{ff}} + 9d\big)}_{N_{\text{block}}} + \underbrace{2d}_{\operatorname{LN}_f} + \underbrace{(d + 1)K}_{\text{head}}

parameters. With $d_{\text{ff}} = 4d$ the block term is $12d^2 + 13d$ . The count does not depend on the number of heads $h$ . It depends on the resolution only through $N_p$ , in the position table.

Proof. Count each tensor of the definition. The patch embedding is a $P^2C\times d$ matrix plus a bias of $d$ . The class token and each register are one $d$ -vector. The position table has one row per patch and one for the class token. In a block, $\mathbf{W}_{QKV}$ and $\mathbf{b}_{QKV}$ give $3d^2 + 3d$ , and $\mathbf{W}_O$ and $\mathbf{b}_O$ give $d^2 + d$ . $\mathbf{W}_1$ and $\mathbf{b}_1$ give $dd_{\text{ff}} + d_{\text{ff}}$ , $\mathbf{W}_2$ and $\mathbf{b}_2$ give $d_{\text{ff}}d + d$ , and the two LayerNorms give $2\cdot2d$ . These sum to $4d^2 + 2dd_{\text{ff}} + d_{\text{ff}} + 9d$ . Splitting $d$ into $h$ heads only regroups the columns of $\mathbf{W}_{QKV}$ , so $h$ does not appear. The final LayerNorm has a gain and a shift, $2d$ , and the head has a $d\times K$ matrix and $K$ biases.

In plain words, almost everything is in the blocks, and the blocks know nothing about images. The image-specific parts, the patch projection and the position table, are under one percent of a Base model.

The cell below checks the formula against the tiny model we built, by adding up the sizes of its arrays. Then it counts the five standard sizes, Tiny, Small, Base, Large, and Huge.

Python · cell 4

The B, L, and H counts are exactly those of torchvision's vit_b_16, vit_l_16, and vit_h_14, and all five match the corresponding timm models. Published tables often disagree with them, for reasons the box below traces. The lesson for practice is short. Whenever you quote or compare a parameter count, say which head, resolution, and patch size it includes.

🔬 Why published parameter counts disagree

The ViT paper's Table 1 lists ViT-Base at 86M, ViT-Large at 307M, and ViT-Huge at 632M. It does not say which patch size, head, or resolution it counts. Base and Huge match B/16 and H/14 above. L/16 has 304.3M, though, and 307M is the count of L/32. Its patch embedding is four times larger ( $3072\cdot1024$ weights instead of $768\cdot1024$ ), which gives 306,535,400 with a 1,000-class head. Later papers, among them the Swin Transformer paper (Liu et al. 2021) and the DeiT paper (Touvron et al. 2021a), print "307M" next to L/16.

Zhai et al. (2022) report ViT-g/14 (giant) at 1,011M and ViT-G/14 (gigantic) at 1,843M. The formula reproduces both with no class token and no head. Those models pool with attention, a learned query that attends over the patch outputs, and the published counts leave out that pooling head and the classifier. ViT-22B's 21,743M (Dehghani et al. 2023) is reproduced to the million by the block weight matrices alone, $48\,(4d^2 + 2dd_{\text{ff}})$ with $d = 6144$ and $d_{\text{ff}} = 24{,}576$ .

The 7-billion-parameter ViT of DINOv3 uses a gated MLP called SwiGLU, which has three $d\times d_{\text{ff}}$ matrices instead of two (§5.3). Its blocks hold $40\,(4\cdot4096^2 + 3\cdot4096\cdot8192) = 6.71$ billion weights, the published "6.7B" (Siméoni et al. 2025, Table 2).

Core

Count any ViT exactly

Write vit_param_count(cfg). It returns the exact number of parameters, as a Python int, for a config dict with these keys. They are img (the side of the square image, a multiple of P), P, C, d, L, h, d_ff, n_classes, cls_token (a bool), n_registers, and pool ("cls" or "gap").

The model is the definition above, with these rules.

The patch embedding has a bias.
A class token exists only if cls_token is true.
Each register is one $d$ -vector with no position embedding.
The position table has one row per patch, plus one for the class token if there is one.
Each block has biases on all four attention projections and on both MLP layers, and two LayerNorms with a gain and a shift each.
The final LayerNorm is always there, whatever the pooling. Pooling itself has no parameters.
The head is a $d\times K$ matrix plus $K$ biases, and there is no head when n_classes is 0.

def vit_param_count(cfg):
    # placeholder: the 12 L d^2 rule of thumb, which ignores everything outside the blocks
    return 12 * cfg["L"] * cfg["d"] ** 2

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Vision papers count multiply-adds and call them FLOPs

Parameters say how big a model is. Compute says what it costs to run, and for a ViT the two are only loosely related. The same weights are applied to every token, and the number of tokens is set by the image, not by the model.

Start with one matrix product. In ViT-B/16, the query, key, and value projection multiplies the 197 by 768 matrix of normalized tokens by the 768 by 2,304 matrix $\mathbf{W}_{QKV}$ . Each of the $197\cdot2{,}304$ outputs is a dot product of length 768, with one multiplication and one addition per term. That is $197\cdot768\cdot2{,}304 = 348{,}585{,}984$ multiply-adds, for one layer of one image. Whether you report it as 349 million operations or 697 million depends on which field you work in.

📐 Multiply-adds (MACs) and FLOPs

A multiply-add (MAC) is one multiplication and one addition, $a \leftarrow a + xy$ . A FLOP is one floating-point operation, so one MAC is two FLOPs. The product of an $m\times k$ and a $k\times n$ matrix costs $mnk$ MACs, or $2mnk$ FLOPs (§5.6).

Language-model papers count true FLOPs. Most vision papers count MACs and call them FLOPs. The fvcore library, Meta's operation counter that many vision code bases use, says so in the docstring of its FlopCountAnalysis class ("We count one fused multiply-add as one flop", fvcore). The cost tables of Swin and ConvNeXt, two widely used vision backbones (§27.5), follow that convention. This book writes MACs for multiply-adds and FLOPs for true floating-point operations.

The whole forward pass is a sum of such products, one term per kind of layer.

∴ Forward multiply-adds of a ViT, and where attention overtakes the linear layers

For one image with $N_p$ patches and $T$ tokens (patches plus prefix tokens), the forward pass of the ViT above costs

\text{MACs} = \underbrace{N_p P^2C\,d}_{\text{patch embedding}} + L\Big(\underbrace{T\,(4d^2 + 2dd_{\text{ff}})}_{\text{linear layers}} + \underbrace{2T^2d}_{\text{attention}}\Big) + \underbrace{dK}_{\text{head}} ,

ignoring LayerNorms, biases, softmax, GELU, and additions. Forward FLOPs are twice this. Training costs about three times the forward pass, so one training image costs about $6\times\text{MACs}$ FLOPs. In each block, attention costs as much as the linear layers when

2T^2d = T\,(4d^2 + 2dd_{\text{ff}}) \iff T = 2d + d_{\text{ff}},

which is $T = 6d$ when $d_{\text{ff}} = 4d$ . For a square image of side $H$ , $T = H^2/P^2 + n_{\text{pre}}$ , with $n_{\text{pre}}$ the number of prefix tokens. So the linear terms grow like $H^2$ and attention like $H^4$ .

Proof. The patch embedding multiplies $\mathbf{X}_p\in\R^{N_p\times P^2C}$ by $\mathbf{W}_E\in\R^{P^2C\times d}$ , which is $N_pP^2Cd$ MACs (the class token is not projected). In a block, $\mathbf{Y}\mathbf{W}_{QKV}$ is $T\cdot d\cdot3d$ and the output projection is $T\cdot d\cdot d$ . The MLP's two matrices are $T\cdot d\cdot d_{\text{ff}}$ each. Together they make $T(4d^2 + 2dd_{\text{ff}})$ . For each head, $\mathbf{Q}_j\mathbf{K}_j^\top$ is $T\cdot d_k\cdot T$ , and multiplying the head's attention weights $\mathbf{A}_j$ by $\mathbf{V}_j$ is $T\cdot T\cdot d_k$ . So the $h$ heads cost $2T^2hd_k = 2T^2d$ . The head reads one vector, $dK$ . Doubling turns MACs into FLOPs. The backward pass of every matrix product costs two products of the same size, one for each input's gradient (§5.6), so training is about three forward passes. Setting the per-block attention term equal to the linear term and dividing by $2Td$ gives $T = 2d + d_{\text{ff}}$ .

In plain words, the linear layers cost a fixed amount per token, so their cost grows with the area of the image. Attention costs an amount per pair of tokens, so its cost grows with the square of the area. The ratio of the two is $T/(2d + d_{\text{ff}})$ , which is $T/(6d)$ for a standard MLP. §5.6 found the same ratio for language models.

✏️ ViT-B/16 at 224, 384, and 1,024 pixels

At 224 pixels, $N_p = 196$ and $T = 197$ .

Patch embedding, $196\cdot768\cdot768 = 115{,}605{,}504$ MACs.
Linear layers, $12\cdot197\cdot7{,}077{,}888 = 16{,}732{,}127{,}232$ .
Attention, $12\cdot2\cdot197^2\cdot768 = 715{,}327{,}488$ .
Head, $768{,}000$ .

The total is 17,563,828,224 MACs, or 17.56 G, and attention is 4.1% of it. The Swin paper's Table 1 lists DeiT-B, the same network, at "17.5G FLOPs", which is the MAC count. The true forward cost is 35.1 GFLOPs. That is exactly what Zhai et al. (2022) print for ViT-B/16 at 224 in their Table 2, because they count real FLOPs. Two tables, the same model, a factor of two.

At 384 pixels, $N_p = 576$ and $T = 577$ . The linear layers grow by $577/197 = 2.93$ and attention by $(577/197)^2 = 8.58$ . The total becomes 55.48 G MACs (Swin's table says 55.4G). That is 3.16 times the cost at 224, while the pixel count grew only $(384/224)^2 = 2.94$ times. Attention is now 11.1% of the total. At 1,024 pixels, $T = 4{,}097$ , and attention is 46.9% of 659.8 G MACs.

Where is the crossover for ViT-B/16? It is at $T = 6\cdot768 = 4{,}608$ tokens, which is $4{,}607$ patches, a square of side $\sqrt{4607}\cdot16 \approx 1{,}086$ pixels. The first resolution that is a multiple of 16 and reaches it is 1,088 pixels, with 68 by 68 patches and 4,625 tokens. Narrow models cross earlier, because the crossover is proportional to $d$ . ViT-Ti/16 has $d = 192$ , so it crosses at 1,152 tokens, or 544 pixels, and attention is already 14% of its cost at 224.

The cell below sweeps the resolution for ViT-B/16 and draws the three parts of the cost.

Python · cell 5

The linear layers' curve grows with the square of the image side and attention's with the fourth power, so past the crossover attention runs away. At 1,536 pixels attention costs twice as much as the linear layers, two thirds of the total. The explorer below shows the same three parts for every model size, as shares of the total, together with the memory the attention scores would take.

Here is the whole family, counted with the formulas of this section. The first five are the standard sizes, where Ti and S were added by DeiT, a recipe for training ViTs on ImageNet alone (Touvron et al. 2021a). The sixth, SoViT-400m/14, has a shape chosen by a scaling study. It is the "so400m" image tower of SigLIP and SigLIP 2, models that learn to match images with their captions using a sigmoid loss (§30.4). It is also the vision encoder of PaliGemma, Google's open vision-language model, and of many others. The last four are among the largest published ViTs, listed with their own reported parameter counts, which leave out the head.

✏️ The ViT family at 224 pixels

model	$L$	$d$	$h$	$d_{\text{ff}}$	$P$	parameters	tokens	forward MACs
ViT-Ti/16	12	192	3	768	16	5,717,416	197	1.25 G
ViT-S/16	12	384	6	1,536	16	22,050,664	197	4.60 G
ViT-B/16	12	768	12	3,072	16	86,567,656	197	17.56 G
ViT-L/16	24	1,024	16	4,096	16	304,326,632	197	61.55 G
ViT-H/14	32	1,280	16	5,120	14	632,045,800	257	167.3 G
SoViT-400m/14	27	1,152	16	4,304	14	412M	256	109.5 G
ViT-g/14	40	1,408	16	6,144	14	1,011M	256	266.0 G
ViT-G/14	48	1,664	16	8,192	14	1,843M	256	481.8 G
ViT-22B	48	6,144	48	24,576	14	21,743M	256	5.61 T
DINOv3 ViT-7B	40	4,096	32	8,192 (SwiGLU)	16	6.7B	201	1.36 T

The first five counts include a class token and a 1,000-class head. SoViT, g, G, and 22B pool with attention instead of a class token, hence 256 tokens. SoViT's 412M is the formula with no head, and its paper's 428M adds the attention-pooling head. Twice our 109.5 G is within 1% of the 221 GFLOPs it reports at 224 (Alabdulmohsin et al. 2023, Table 1). DINOv3's 201 tokens are 196 patches, a class token, and 4 registers. Our MACs for g/14 and G/14 are within 0.25% of half the GFLOPs that Zhai et al. report (533.1 and 965.3). The SwiGLU MLP of the 7B model costs $3dd_{\text{ff}}$ per token instead of $2dd_{\text{ff}}$ .

Read the parameter and multiply-add columns side by side and two patterns appear. Both are worth memorizing, because they let you estimate the cost of any ViT in your head.

The standard sizes grow depth, width, and MLP width together, roughly in step. Whether that is the best use of a compute budget is an empirical question, and two scaling studies answered it.

🔬 How the family should grow with compute

Zhai et al. (2022) trained ViTs from about 5 million to 2 billion parameters (s/28 to G/14) on up to 3 billion images. They found that the best error rate at each compute budget follows a power law that flattens at both ends, which they call a double-saturating law,

\operatorname{err}(C_{\text{train}}) \approx a\,(C_{\text{train}} + c_0)^{-b} + \operatorname{err}_\infty ,

with $C_{\text{train}}$ the training compute and positive constants $a$ , $b$ , $c_0$ , and $\operatorname{err}_\infty$ fitted for each task. The floor $\operatorname{err}_\infty$ says that even unlimited compute leaves some error, the irreducible term of §9.1. The shift $c_0$ says that a model with almost no compute still does better than the pure power law predicts, since guessing the most common class already earns some accuracy.

Alabdulmohsin et al. (2023) fit such laws to each shape dimension separately. As compute grows, the MLP width should grow fastest, then the depth, and the model width slowest. The standard family grows all three together, so it is not compute-optimal. Their SoViT-400m/14, given the pretraining compute of ViT-g/14, matches it on ImageNet after fine-tuning, with 90.3% against 90.2% top-1 accuracy at 518 pixels. It does so at less than half the inference cost, 109.5 G against 266.0 G multiply-adds at 224 in the table above.

Core

Multiply-adds by term, and the crossover resolution

Write vit_macs(cfg, res) for a square image of side res, a multiple of P. The config dict is the one of Exercise 27.2b, with keys P, C, d, L, h, d_ff, n_classes, cls_token, n_registers, and pool. The function returns a dict with Python int values for the keys "patch", "linear", "attention", "head", and "total". It counts forward multiply-adds, not FLOPs, as the theorem does.

The token count $T$ includes the class token (if any) and the registers, since both take part in attention. The head reads one vector, so it costs $dK$ , or zero if n_classes is 0.

Then write attention_crossover_res(cfg). It returns the smallest resolution that is a multiple of P at which the attention MACs are at least the linear MACs. Use res, not cfg["img"], for the image side. Use Python ints, not NumPy integers, because the browser's NumPy uses 32-bit integers by default and these products overflow them.

def vit_macs(cfg, res):
    return {"patch": 0, "linear": 0, "attention": 0, "head": 0, "total": 0}

def attention_crossover_res(cfg):
    return 224

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Six FLOPs per multiply-add reproduces the paper's training compute

A formula is only as good as its agreement with real runs, and the ViT paper gives a rare chance to check. Its Table 6 lists the total pretraining compute of each model in exaFLOPs, where one exaFLOP is $10^{18}$ FLOPs. Every model was pretrained at 224 pixels on JFT-300M, which has 303 million images, for 7 or 14 epochs, full passes over the data. If our accounting is right, the compute should be

C_{\text{train}} \approx 6 \times \text{MACs per image} \times \text{epochs} \times 303\times10^6 .

In plain words, each image seen costs six FLOPs for every forward multiply-add, two for turning MACs into FLOPs and three for the forward and backward passes. We write $C_{\text{train}}$ with a subscript because plain $C$ is the channel count in this chapter. The cell checks the formula against the paper.

Python · cell 6

Five of the six rows agree to within 0.15%, and B/32 is 2.3% high. So the paper's exaFLOPs are six FLOPs per forward multiply-add, for every image seen. That tells us two things at once. Their numbers are true FLOPs, two per multiply-add. And they count training as three forward passes, the same rule the language-model chapters use.

The last line of the output is worth a moment. The $6ND$ rule of §5.6 says training costs about six FLOPs per parameter per token, with $N$ parameters and $D$ training tokens. It works for ViTs too, provided $D$ counts tokens, not images. ViT-B/16 sees $197\times7\times303$ million $\approx 4.2\times10^{11}$ tokens in pretraining. Six times the blocks' 84.9M matrix weights times that gives 213 exaFLOPs. That is 5% short of the exact count, because it leaves out attention (4.1% of the cost at 224) and the patch embedding (0.7%). Applying $6N$ per image instead of per token underestimates the bill by a factor of 197.

The same paper reports how long the runs took, which lets us measure how well they used the hardware. Model FLOPs utilization (MFU) is the FLOPs the model needs divided by the FLOPs the hardware could have done at its peak in the same time (§6.7). An MFU of 50% means that half of the chips' peak arithmetic went into the model's own math. The rest was lost to memory traffic, communication, and idle time.

About a third of peak was common for large training runs of that time. Later runs did better, and the same arithmetic lets you check a published utilization figure instead of taking it on trust.

🔬 Checking ViT-22B's 54.9% utilization

In 2023, ViT-22B reported 54.9% MFU on TPU v4, processing 1,150 tokens per second per core (Dehghani et al. 2023). We can check it. Each token costs $6N$ FLOPs for the weights plus $12LTd$ for attention, six times the $2Td$ attention MACs per token per layer. That is $6\cdot21.74\times10^9 + 12\cdot48\cdot256\cdot6144 = 1.314\times10^{11}$ FLOPs per token. At 1,150 tokens per second the device does $1.51\times10^{14}$ FLOP/s. A TPU v4 chip peaks at 275 TFLOP/s in bfloat16, which gives 54.9%.

The paper says "core", and the arithmetic works only if a core means one v4 chip, which JAX exposes as a single device. A v4 chip has two TensorCores, and at half a chip the utilization would be over 100%.

Run the arithmetic the other way and a published compute figure becomes a budget on today's hardware. That is the form the interview question usually takes.

Interview

Pretraining compute, GPU-days, and the utilization a paper implies

This is the estimate interviewers ask for in one form or another. The setup gives you forward_macs(cfg, res), the forward multiply-adds of one image by the theorem's formula, head included. Write three functions.

pretraining_exaflops(cfg, res, epochs, n_images) returns the training compute in exaFLOPs ( $10^{18}$ true FLOPs) for epochs passes over a dataset of n_images images at resolution res.
gpu_days(exaflops, peak_tflops, mfu) returns the number of device-days a run of that size takes on accelerators with dense peak peak_tflops (in $10^{12}$ FLOP/s) at model FLOPs utilization mfu.
implied_mfu(exaflops, device_days, peak_tflops) returns the utilization a published run achieved.

def pretraining_exaflops(cfg, res, epochs, n_images):
    return forward_macs(cfg, res) * epochs * n_images / 1e18     # placeholder: forward multiply-adds only

def gpu_days(exaflops, peak_tflops, mfu):
    raise NotImplementedError

def implied_mfu(exaflops, device_days, peak_tflops):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

Attention memory forces windows at high resolution

FLOPs are not the only cost. During training, a straightforward attention implementation keeps every head's $T\times T$ weight matrix $\mathbf{A}$ for the backward pass (§4.3). That is $hT^2$ numbers per layer per image (§4.7). For ViT-B/16 at 224 it is $12\cdot197^2 = 465{,}708$ numbers, about 0.93 MB per layer in bfloat16, the 2-byte format used for training. That is trivial.

High resolution changes this completely. The cleanest example is the image encoder of Segment Anything (SAM), a model that outputs a mask for any object you point at in an image (§27.6).

FlashAttention removes the memory problem without reducing the arithmetic (§12.3). It never stores $\mathbf{A}$ . It keeps only the output and one log-sum-exp per query row per head, which is $O(Td)$ memory per layer, and it recomputes $\mathbf{A}$ tile by tile in the backward pass. It does nothing for the $2T^2d$ MACs, though, and those are what SAM avoids.

In windowed attention, each token attends only to the tokens in its own 14 by 14 window. SAM's encoder uses it in 28 of its 32 blocks. The other four blocks, equally spaced, use global attention, so information still crosses the whole image. The 64 by 64 grid is padded to 70 by 70 and cut into 25 windows of 196 tokens. A windowed layer's attention costs $2\cdot25\cdot196^2\cdot1280$ MACs instead of $2\cdot4096^2\cdot1280$ , about 17 times less. Over the whole encoder, attention falls from 1.37 T MACs to 0.24 T. ViTDet, a detector built on a plain ViT, uses the same plan (Li et al. 2022). §27.5 builds the window machinery properly, with the shifted windows of Swin.

Python · cell 7

Core

Attention memory, materialized and recomputed

Write attention_memory_bytes(cfg, res, batch, bytes_per=2, flash=False) for a square image of side res and the config dict of Exercise 27.2b.

With flash=False, return the bytes needed to keep every head's $T\times T$ attention weights, in every layer, for the whole batch, at bytes_per bytes per number.
With flash=True, return what a FlashAttention-style kernel keeps instead of those weights, which is one float32 log-sum-exp per query row, per head, per layer, per image. The inputs and the output are stored in both cases, so leave them out.

Then write max_res_materialized(cfg, budget_bytes, batch=1, bytes_per=2). It returns the largest resolution that is a multiple of P whose stored weights (the flash=False number) fit in the budget. As before, $T$ counts the class token and the registers.

def attention_memory_bytes(cfg, res, batch, bytes_per=2, flash=False):
    N_p = (res // cfg["P"]) ** 2
    return cfg["L"] * batch * N_p * N_p * bytes_per          # placeholder

def max_res_materialized(cfg, budget_bytes, batch=1, bytes_per=2):
    raise NotImplementedError

Tests, hints, verified solutions, and AI help, with Pro at $7.99 a month.

FLOPs do not predict speed

Arithmetic counts are the right tool for budgets and the wrong tool for latency. The roofline model of §11.2 explains why. Large matrix products are limited by arithmetic. Softmax, LayerNorm, small matrices, and data movement are limited by memory bandwidth, so they take longer than their FLOPs suggest.

The DeiT paper's Table 5 measured images per second on one V100 GPU (Touvron et al. 2021a), and its numbers show the effect in three ways.

DeiT-B runs at 292 images per second at 224 and 85.9 at 384. That is 3.40 times slower for 3.16 times the multiply-adds. Much of the gap is the softmax and other memory-bound work on the $T\times T$ scores, which grows faster than the large matrix products.
DeiT-Ti has 14 times fewer multiply-adds than DeiT-B but runs only 8.7 times faster. Its small matrices ( $d = 192$ ) cannot keep the GPU busy.
Swin-B, with 15.4 G multiply-adds, is slower than DeiT-B with 17.5 G, at 278 against 292 images per second in Swin's own Table 1. A likely reason is that it has twice as many, smaller blocks and spends time moving tokens into and out of windows, work that does no arithmetic.

Hardware also pads the work. A TPU pads 257 tokens up to 384, which §27.1 gave as one reason to drop the class token. The ViT paper's own timing study (its Appendix D.5) found that the fourth-power growth of cost with image side "only barely starts happening" for its largest models at its largest resolutions. That fits the counts above, since at moderate sizes attention is a small share of the work. In an interview, the answer to "will halving the FLOPs halve the latency?" is usually no, and the reason is the roofline.

In PyTorch

The PyTorch version is about twenty lines. The patch embedding is an nn.Conv2d with kernel and stride $P$ , the equivalence §27.1 proved. Its weight is stored in PyTorch's $(d, C, P, P)$ layout and applied to images in PyTorch's $(B, C, H, W)$ layout. Note also that nn.Linear stores its weight as (output, input), the transpose of our convention.

The block below runs three checks. First, our block equals PyTorch's own nn.TransformerEncoderLayer with norm_first=True, which is a ViT block. Second, our model at ViT-B/16 size has the 86,567,656 parameters we counted by hand. Third, PyTorch's FLOP counter agrees with the multiply-add formula exactly. That counter counts two FLOPs per multiply-add in every matrix product and convolution, and ignores elementwise work.

🔥 In PyTorchfor a machine with PyTorch

The counter reports 35,127,656,448 FLOPs. Half of that is 17,563,828,224, our multiply-add count exactly. The model computes the scores with explicit matrix products so that the counter can see them. In production you would call F.scaled_dot_product_attention, as the block comparison does, which dispatches to a fused FlashAttention kernel on GPUs.

The libraries most people load pretrained ViTs from give the same counts. The block below is shown, not run, because creating the models with pretrained weights downloads them.

The largest ViTs change the block

The 2020 block trains well up to a few billion parameters. ViT-G/14, with 1.8 billion, still uses it unchanged (Zhai et al. 2022). Past that, the published large models change the block, and each change has a reason this book has already taught.

The main block change in ViT-22B is parallel layers, which §5.6 built and fused. Both sublayers read one normalized copy of the stream, and their outputs are added together,

\mathbf{X}^{(\ell)} = \mathbf{X}^{(\ell-1)} + \operatorname{MHA}(\mathbf{Y}) + \operatorname{MLP}(\mathbf{Y}), \qquad \mathbf{Y} = \operatorname{LN}(\mathbf{X}^{(\ell-1)}).

In plain words, attention and the MLP no longer wait for each other. Because both read the same $\mathbf{Y}$ , their input matrices fuse into one matrix product and their output matrices into another. A block loses one LayerNorm and one output bias and keeps every multiply-add. The gain is utilization, since a few large products keep the hardware busier than many smaller ones.

change	what it is	why	used by
QK-norm	LayerNorm on queries and keys before the dot product	attention logits grew until the weights were one-hot and training diverged, around 8B parameters (§4.3, §5.7)	ViT-22B
parallel layers	attention and MLP read the same LN and are summed	fused matrix products, as in §5.6; PaLM measured about 15% faster training at large scale	ViT-22B
no biases in QKV and LayerNorms	bias-free QKV projections, and norms with a gain only, no centering and no shift	3% better accelerator utilization with no loss in quality, as reported	ViT-22B (it keeps the MLP biases)
attention pooling (MAP, multihead attention pooling)	one learned query attends over the outputs	a learned summary with no extra token in the sequence (§27.1)	ViT-g, ViT-G, ViT-22B, SoViT, SigLIP
SwiGLU MLP	gated MLP with three matrices (§5.3)	better quality per parameter in language models, carried over	DINOv2 ViT-g, DINOv3 ViT-7B
axial RoPE	rotary positions over the row and the column (§27.1)	one model for many resolutions and aspect ratios	DINOv3 ViT-7B
LayerScale	a learned per-channel scale on each branch, initialized near zero	lets deep stacks start close to the identity (§27.4)	CaiT, DINOv2, DINOv3
registers	a few extra tokens with no position and no output	remove high-norm artifact tokens in large ViTs (§27.4)	DINOv2 with registers, DINOv3

None of these changes how the cost scales with resolution, which is still quadratic in the token count. The ViT-22B numbers are from Dehghani et al. (2023). The 15% speedup from parallel layers was measured on PaLM, Google's 540-billion-parameter language model (Chowdhery et al. 2022). LayerScale comes from CaiT, a deeper ViT (Touvron et al. 2021b). ViT-22B also reached a shape bias of 87%, the highest recorded for a vision model at the time. Shape bias is how often a model labels an image by its shape rather than its texture when the two disagree, and humans score 96%. §27.3 explains that test.

The same encoder appears all through the rest of the book. Masked autoencoders train it without labels by hiding most patches and reconstructing them (masked image modeling). DINO trains it by matching a student network to a slowly averaged teacher (DINO). CLIP and SigLIP train it to match images with their captions (CLIP's image tower). Video models extend its patches into space-time blocks (tubelets). Diffusion transformers cut the latents of an autoencoder into patches and run the same blocks, with conditioning added (DiT).

⚠️ Where ViT accounting goes wrong

Quoting multiply-adds as FLOPs, or mixing a vision paper's "FLOPs" with a language model's MFU, is off by two.
Using $N_p$ where $T$ belongs forgets the class token and the registers, which attend and are attended to.
Giving registers a row in the position table, or forgetting the patch bias or the final LayerNorm, shifts a count by hundreds to a few thousand. The count then disagrees with every library.
Writing $12d^2$ per block when $d_{\text{ff}} \ne 4d$ , or $2dd_{\text{ff}}$ for a SwiGLU MLP (it has three matrices, $3dd_{\text{ff}}$ ), is wrong by millions.
Applying $6N$ per image instead of per token underestimates training compute by a factor of $T$ .
Comparing parameter counts without naming the head, the resolution (which sets the position table), and the patch size produces the 304M against 307M confusion for ViT-L/16.
In code, a causal mask copied from a GPT makes the class token blind to the image. A GELU or LayerNorm $\epsilon$ that differs from the checkpoint's makes ported weights subtly wrong. Reshaping heads straight from $(B, h, T, d_k)$ to $(B, T, d)$ without the transpose scrambles them silently.

🧭 What to remember

A ViT is a pre-LN Transformer encoder over patch tokens. It differs from GPT-2 only at the ends, with a pixel projection instead of a token table, no mask, and one pooled vector read by a linear head. A causal mask left in makes the class token blind to the image.
The exact count is $N = (P^2C + 1)d + (n_{\text{cls}} + R)d + (N_p + n_{\text{cls}})d + L(4d^2 + 2dd_{\text{ff}} + d_{\text{ff}} + 9d) + 2d + (d+1)K$ , which gives 86,567,656 for ViT-B/16. Published counts differ because they include or omit heads, use other patch sizes, or count only block matrices.
Forward cost is $N_pP^2Cd + L(T(4d^2 + 2dd_{\text{ff}}) + 2T^2d) + dK$ multiply-adds, 17.56 G for ViT-B/16 at 224. Vision papers call multiply-adds FLOPs, so double them before comparing with language models, and multiply by six (not three) for training FLOPs.
Linear layers grow with the image area and attention with its square. Attention equals the linear layers at $T = 2d + d_{\text{ff}}$ , which is $6d$ for a standard MLP and about 1,088 pixels for ViT-B/16. Stored attention weights reach half a gigabyte per layer per image at SAM's 1,024 pixels, which is why high-resolution encoders use windows.
$6\times\text{MACs}\times\text{images seen}$ reproduces the ViT paper's pretraining compute to within 0.15% for five of its six models (2.3% for B/32). Published run times then give utilizations, about 32% for ViT-H/14 on TPU v3 and 55% for ViT-22B on TPU v4.
Error falls with compute as a power law that flattens at both ends, and shape matters as well as size (SoViT-400m/14 matches ViT-g/14 at less than half the cost). The largest ViTs add QK-norm, parallel layers, bias-free projections, attention pooling, SwiGLU, axial RoPE, LayerScale, and registers, and none of these changes the quadratic cost in tokens.
Next, §27.3 asks whether this architecture, which builds in almost nothing about images, is a good idea at all. It answers with symmetry, a theorem, and the data scale at which a ViT beats a CNN.

Dosovitskiy et al. (2021). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. Equations 1 to 4 define the model, Table 1 its sizes, Table 6 the pretraining compute we reproduced, and Appendix D.5 the measured speed against resolution.
Touvron et al. (2021a). Training data-efficient image transformers & distillation through attention. Introduces the Ti and S sizes and gives the V100 throughput table used here to show that FLOPs do not predict speed.
Zhai et al. (2022). Scaling Vision Transformers. The double-saturating power law of error against compute, the g and G models, a table of true GFLOPs at 224 and 384, and the reasons for replacing the class token with attention pooling.
Alabdulmohsin et al. (2023). Getting ViT in Shape: Scaling Laws for Compute-Optimal Model Design. Scaling laws for width, depth, and MLP size separately, and the SoViT-400m/14 shape that SigLIP's so400m tower uses.
Dehghani et al. (2023). Scaling Vision Transformers to 22 Billion Parameters. Parallel layers, QK-norm, omitted biases, and the 54.9% MFU we checked by hand.
Siméoni et al. (2025). DINOv3. Table 2 gives the 7B teacher's shape, with SwiGLU, axial RoPE, and four registers.
Kirillov et al. (2023). Segment Anything. Appendix A describes the ViT-H/16 encoder at 1,024 pixels with windowed attention and four global blocks.
Liu et al. (2021). Swin Transformer: Hierarchical Vision Transformer using Shifted Windows. Table 1 lists ViT and DeiT at 224 and 384 in multiply-adds, next to Swin's own costs and throughputs.

Made for larger screens

The ViT architecture

A ViT is a Transformer encoder over patches

The forward pass, written exactly

Two heads, one for pretraining and one for transfer

A ViT from scratch in NumPy

A general ViT forward pass, with registers and either pooling, no loops over tokens or heads

Almost every parameter sits in the blocks

Count any ViT exactly

Vision papers count multiply-adds and call them FLOPs

Multiply-adds by term, and the crossover resolution

Six FLOPs per multiply-add reproduces the paper's training compute

Pretraining compute, GPU-days, and the utilization a paper implies

Attention memory forces windows at high resolution

Attention memory, materialized and recomputed

FLOPs do not predict speed

In PyTorch

The largest ViTs change the block

Read more

Watch next