An Image Is Worth 16×16 Words — ViT, Patch by Patch

Everything in Attention Meets Pixels — Windows, Axes and the First Fully Attentional Networks shrank the key set so that attention could survive on a feature map. [DOSO20] (October 2020) noticed that all of that machinery "has not yet been scaled effectively on modern hardware accelerators", and that in vision "attention is either applied in conjunction with convolutional networks, or used to replace certain components of convolutional networks while keeping their overall structure in place" (§1). Their problem statement is the plain one: apply "a standard Transformer directly to images, with the fewest possible modifications" (§3), so that NLP-scale code works "almost out of the box". The pipeline is one line: image → N flattened patches → one linear map → N + 1 tokens → the encoder of Attention Is All You Need — The Transformer, Block by Block → read the first token.

Brush up before you start ;)

• Q/K/V, softmax, √d, multi-head: Query, Key, Value — Dot-Product Attention from Scratch. The paper's Appendix A is a restatement of it.
• The block, residuals, LayerNorm, the sinusoidal positions and the encoder: Attention Is All You Need — The Transformer, Block by Block. I will only say what ViT changes.
• Why pixels are too many tokens: Attention Meets Pixels — Windows, Axes and the First Fully Attentional Networks. I will reuse one number from there, 50,176.

1. Patches as Words

My picture: a 224×224 photo is a 50,176-word essay if every pixel is a word. At 16×16 it is a 196-word paragraph. The transformer of Attention Is All You Need — The Transformer, Block by Block reads paragraphs fine; essays it cannot afford. The paper's own sentence is the title: "image patches are treated the same way as tokens (words) in an NLP application".

So the plan has three parts, and the rest of the post is one section per part:

• turn the picture into a sequence (§2),
• run the standard encoder on it with as few changes as possible (§3),
• and, since the model has given up the convolution's built-in assumptions, feed it enough data to learn them (§6).

2. Patch Embedding, the Class Token and Positions

2.1 The Equation

Let the image be $x\in {ℝ}^{H\times W\times C}$ and pick a patch side $P$. Eq. 1 of the paper, in two lines:

Let's read it piece by piece.

• $N$: the number of patches, which is the sequence length. $x_p^i\in {ℝ}^{P^{2}C}$: the $i$-th patch, flattened into one long vector of pixel values.
• $E$: one shared linear map from $P^{2}C$ to the model width $D$. It is the same matrix for every patch, which is exactly a $P\times P$ convolution with stride $P$. This is the patch embedding.
• $x_{\text{class}}\in {ℝ}^D$: a learned class token put at index 0, "similar to BERT's [class] token" (§3.1). It is a blank first line; after L rounds of reading everybody else's notes, that line is the summary you hand in.
• $E_{\text{pos}}$: one learned vector per index 0..N. Nothing two-dimensional about it: the model is told "this is token 37", not "this is row 2, column 9". Table 8 of the paper tests 1-D, 2-D and relative variants on ViT-B/16 (5-shot linear): none 0.614, 1-D 0.642, 2-D 0.640, relative 0.640. Having positions matters; the kind does not.
• $z_0\in {ℝ}^{(N+1)\times D}$: the sequence the encoder reads.

So: cut the picture into a grid of squares, flatten each square, multiply by the same matrix, line the results up as a sentence, put a blank summary word at the front, and tag each word with a learned position. Shrink $P$ in the figure below and watch the sequence grow as $1/P^2$; raise the image side with $P$ fixed, as fine-tuning does, and $N$ grows while the patch vector does not.

patch side P (snaps to 4, 8, 14, 16, 28, 32, 56)

16

image side (snaps to 224, 384, 512)

224

2.2 The Numbers

For 224×224, $D=768$ (ViT-Base), per layer, with N excluding the class token:

Table. 1 Tokens and per-layer cost against patch size, computed from Eq. 1 and Eq. 3 for a 224×224 image and D = 768; attention counts both the QKᵀ and the AV product, as Query, Key, Value — Dot-Product Attention from Scratch §6 does. The pixel row is the one Attention Meets Pixels — Windows, Axes and the First Fully Attentional Networks could not afford.

• The sequence length is "inversely proportional to the square of the patch size" (§4.1): halving P multiplies N by 4 and the attention term by 16.
• At P = 16 the attention matrix of one layer with 12 heads is 1.8 MB in fp32.

And the parameter count of ViT-B/16, $D=768$, $L=12$, MLP width 3072, a 1000-class head:

Table. 2 Where the 86 M parameters of ViT-B/16 (Table 1 of [DOSO20]) live, computed from the shapes in §2.1 and §3.

• 98 % of the weights are in the twelve blocks, two thirds of each block in the MLP. The "vision-specific" parts, E, the class token and the positions, are 0.86 % of the model. That is what "fewest possible modifications" looks like in parameters.

3. The Encoder

Eq. 2, 3 and 4 of the paper, for $\ell =1\dots L$:

• MSA is the multi-head self-attention of Query, Key, Value — Dot-Product Attention from Scratch; LN is LayerNorm; the MLP has two layers with a GELU between. All of it works on an $(N+1)\times D$ matrix.
• $z_L^0$: row 0 after the last block, the class token. $y\in {ℝ}^D$ is the image representation; the head is an MLP with one hidden layer in pre-training and a single zero-initialised $D\times K$ linear layer in fine-tuning.

Exactly three things differ from the encoder of Attention Is All You Need — The Transformer, Block by Block:

1. The input: patch vectors through E instead of word pieces (§2).
2. The positions: learned, 1-D, added once at the input, instead of the sinusoids.
3. The read-out: one class row instead of a sequence. (Appendix D.3 shows global average pooling over the patch tokens works as well, with a different learning rate; the class token is a convention, not a requirement.)

NOTE: pre-LN, as the NOTE in §2.3 of Attention Is All You Need — The Transformer, Block by Block warned: the LayerNorm sits before MSA and before the MLP, inside the residual branch.

One more detail that matters at fine-tuning time. ViT is fine-tuned at a higher resolution than it was pre-trained at (224 → 384, up to 512/518 for the headline numbers) with the same P, so N grows and the pre-trained position embeddings no longer fit; they are 2-D-interpolated to the new grid. The paper calls this interpolation and the patch cut "the only points at which an inductive bias about the 2D structure of the images is manually injected" (§3.2).

4. Why 16×16 Is Affordable

Let's derive the cost of one layer as a function of P, from Eq. 1 and the $O(n^{2}d)$ of Query, Key, Value — Dot-Product Attention from Scratch:

• Line 2 is the $Q{K}^{\top }$ and $AV$ products: quadratic in N, so quartic in 1/P.
• Line 3 is everything linear: the Q, K, V and output projections (4 matrices of $D\times D$ per token) and the MLP ($D\times 4D$ twice, 8 $D^2$ per token). Linear in N.
• Set the two equal, $2N^{2}D=12N{D}^2$, and they tie at $N=6D$: 4,608 tokens for $D=768$, a patch side of about 3.3 pixels on a 224² image. At every patch size ViT actually uses, the linear layers dominate and attention is the cheap part: 40 % of the layer at P = 4, 4 % at P = 16 (Table. 1). The quadratic term that killed pixels is a rounding error here.

NOTE: this is also why ViT does not help with the dense-prediction problem. At detection or segmentation resolution N climbs back into the thousands and the quadratic term returns; that is the problem Swin Transformer — Windows, Shifts and a Pyramid for Dense Prediction solves.

5. Models and the Recipe

Table. 3 Table 1 of [DOSO20]. "ViT-L/16" is Large with 16×16 patches; for hybrids the number after the slash is the ResNet downsampling ratio, not a patch size (§4.4).

• Pre-training (§4.1 and Table 3 of the paper): Adam with β = (0.9, 0.999), batch 4096, 10k warm-up steps, resolution 224. JFT-300M for 7 or 14 epochs with weight decay 0.1 and no dropout; ImageNet-21k for 30 or 90 epochs with weight decay 0.03 and dropout 0.1; ImageNet for 300 epochs with weight decay 0.3 and dropout 0.1.
• Fine-tuning: SGD with momentum 0.9, batch 512, resolution 384 (512 or 518 for the Table 2 ImageNet numbers).
• Three pre-training sets of three sizes: ImageNet (1.3 M images), ImageNet-21k (14 M), JFT-300M (300 M). The whole result of the paper is what happens as you move along that axis.

6. Data Scale Against Inductive Bias

6.1 The Headline

Table. 4 Table 2 of [DOSO20], means over three fine-tuning runs; the BiT and Noisy Student rows are as reported there (CIFAR-10, Pets and Flowers columns dropped). ViT-L/16 on JFT ties BiT-L on ReaL and VTAB and beats it everywhere else.

• ViT-L/16 on JFT beats BiT-L by 0.22 points with 14.6× less compute (0.68k vs 9.9k core-days). ViT-L/16 on ImageNet-21k "could be trained using a standard cloud TPUv3 with 8 cores in approximately 30 days" (§4.2).

6.2 The Axis That Matters

One figure for the two scaling questions. Move the first slider across the three pre-training sets and read the gap between ViT-L/16 and ViT-B/16 off the tile; it changes sign. Move the second slider, a compute budget, and see what a ResNet needs to match the best ViT under it.

pre-training images, log10 (6.1 = ImageNet, 7.1 = 21k, 8.5 = JFT)

6.1

pre-training budget, log10 exaFLOPs

2.9

The left panel is Table 5 of the paper (the points of its Figure 3), fine-tuned to ImageNet at 384:

Table. 5 ImageNet top-1 after fine-tuning at 384, Table 5 of [DOSO20]. The BiT band drawn in the paper's Figure 3 is not tabulated and is not reproduced.

• On ImageNet alone, bigger is worse. L/16 at 76.53 is below B/16 at 77.91, and both are, in the abstract's words, "a few percentage points below ResNets of comparable size". With no locality built in, a 307 M-parameter model has nothing to stop it overfitting 1.3 M images.
• On ImageNet-21k the order flips; on JFT-300M the Large and Huge models pull away.
• Figure 4 of the paper makes the same point with linear 5-shot probes (a linear probe trains only a linear classifier on the frozen features, here from 5 examples per class; nothing in the backbone moves) on JFT subsets of 9 M, 30 M, 90 M and 300 M images: ResNets are ahead at 9 M and ViT is ahead from 90 M on, for both B/32 against ResNet50 and L/16 against ResNet152x2. There is no table behind that figure, so I state the claim and do not plot it.
• The paper's reading (§4.3): "the convolutional inductive bias is useful for smaller datasets, but for larger ones, learning the relevant patterns directly from data is sufficient, even beneficial." Or, from §1: "large scale training trumps inductive bias."

NOTE: "ViT has no inductive bias" is too strong, and the paper says much less, not none (§3.1). The MLPs are local and translation-equivariant, the patch cut and the position interpolation inject 2-D structure, and Figure 7 (centre) shows the learned position embeddings organising themselves into rows and columns. The precise sentence is the one §7 of Attention Is All You Need — The Transformer, Block by Block ends on: no locality or translation equivariance in the attention layers.

6.3 Compute

The right panel is Table 6 (the points of Figure 5): every model pre-trained on JFT-300M, ImageNet top-1 against the pre-training compute in exaFLOPs.

Table. 6 ImageNet top-1 against pre-training compute, Table 6 of [DOSO20] (JFT-300M, 7 epochs unless marked). The hybrid number after the slash is the ResNet downsampling ratio.

• At matched compute ViT is ahead of every ResNet, "approximately 2−4× less compute to attain the same performance" (their average over five datasets). On ImageNet: ViT-B/32 reaches 80.73 for 55 exaFLOPs, the cheapest ResNet above that is R152x1 at 141 (2.6×); ViT-L/16 reaches 86.30 for 783, the cheapest ResNet above that is R200x3 at 3,306 (4.2×). Slide the budget in the figure and read the ratio off the tile.
• The hybrid wins at small budgets, 84.90 for 106 exaFLOPs against ViT-B/16's 84.15 for 224, and meets ViT at 87.12 near 1,600; only ViT-H/14 goes further.
• A tangent the paper leaves open (§4.6): masked-patch self-supervision gives ViT-B/16 79.9 % on ImageNet, +2 over training from scratch and −4 against supervised pre-training. That gap is a later post, Masked Autoencoders — Hide 75 %, Encode the Rest.

7. What the Heads Look At

How far does a token look? Appendix D.7 defines mean attention distance: for one head, the distance in image space between a query and each key, weighted by the attention weight, averaged over images. The paper calls it "analogous to receptive field size in CNNs"; it is a weighted average, not a hard field.

• Figure 7 (right), mean attention distance per head and layer: in the lowest layers some heads already span most of the image while others stay local. The local heads "may serve a similar function as early convolutional layers in CNNs"; the global ones are something a CNN cannot do in layer 1.
• "The attention distance increases with network depth": the deeper the layer, the farther its heads look on average.
• Together with the row-and-column structure of the learned positions (Figure 7, centre) and Table 8, the message of §2.1 holds: ViT is told nothing about the grid and learns it anyway, given enough images.

8. What Came Next

ViT left two doors open, and the next two posts walk through them. The ImageNet-only column of Table. 5, where ViT loses to ResNets, is the gap DeiT — Training a Vision Transformer on ImageNet Alone closes with a training recipe and a distillation token, no extra data. And the quadratic term that §4 waved away at P = 16 comes back at dense-prediction resolution; Swin Transformer — Windows, Shifts and a Pyramid for Dense Prediction brings back the windows of Attention Meets Pixels — Windows, Axes and the First Fully Attentional Networks, adds a shift and a hierarchy, and makes a ViT-shaped backbone work for detection and segmentation. Stay tuned.

9. References

1. [DOSO20] Dosovitskiy, Beyer, Kolesnikov, Weissenborn, Zhai, Unterthiner, Dehghani, Minderer, Heigold, Gelly, Uszkoreit, Houlsby. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ICLR 2021, arXiv 2010.11929 (v2, the camera-ready; its exaFLOPs figures are the corrected ones).
2. [VASW17] Vaswani et al. Attention Is All You Need. 2017, arXiv 1706.03762.