ViT: An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

Cut an image into 16×16 patches, feed them to a standard Transformer, and with enough pre-training data it beats convolutional networks.

The Vision Transformer runs the text Transformer almost unchanged: it slices the picture into a grid of patches and treats each patch as a word. Trained on ImageNet's 1.3 million images it trails a ResNet; pre-trained on 303 million it passes the best convolutional models.

Posted October 2020. The largest model reaches 88.55% ImageNet top-1 for about a quarter of the pre-training compute of the ResNet it beats.

By 2020 natural language processing had one recipe: pre-train a large Transformer on a very large text corpus, then fine-tune it for each task, as BERT did. Image recognition still ran on convolutional networks, and at large scale the state of the art was a ResNet variant. Attempts to use attention in vision either combined it with convolutions or used restricted attention patterns that needed custom kernels to run fast on accelerators.

Dosovitskiy and colleagues at Google Brain tried the direct version: split the image into patches, treat each patch as a token, and run the text Transformer's encoder on the sequence with as few changes as possible. The result, in the paper's words: "large scale training trumps inductive bias". The sentence has a condition in it. On ImageNet alone the model lands a few points below a ResNet of similar size, which the authors say they expected; the lead appears only after pre-training on 14 to 303 million images.

This page covers, in order:

1. why pixels are too many tokens, and how 16×16 patches fix the count;
2. the input sequence, Eq. (1), and the encoder, Eqs. (2) to (8), with every tensor shape for ViT-B/16;
3. which assumptions about images a convolution builds in and ViT leaves out;
4. the data and compute experiments, with the paper's tables in the figures;
5. fine-tuning at higher resolution, and what the trained model's position embeddings and attention heads show.

Why vision had not used the Transformer

A Transformer layer's main operation is self-attention. Each element of the input sequence (a token) computes a score against every token, turns the scores into weights that sum to 1, and replaces itself with a weighted average of all tokens. For a sequence of $N$ tokens that is $N^2$ scores per attention head per layer.

A sentence has a few hundred tokens, so $N^2$ is at most a few hundred thousand. The obvious tokens for an image are its pixels, and a 224×224 image has 50,176 of them. Squared, that is about 2.5 billion scores per head per layer, and a 12-layer model with 12 heads would compute 144 such matrices for a single image.

Earlier work avoided the square by restricting it: attention over local neighbourhoods only, sparse or axial patterns that attend along rows and columns, or attention applied on top of a convolutional network that had already shrunk the image to a small feature map. The closest prior model, by Cordonnier, Loukas and Jaggi (2020), cut images into 2×2 patches and ran full self-attention on them, which kept $N$ manageable only for small-resolution images. ViT uses the same idea with much larger patches and much more pre-training data.

From an image to a token sequence

Take a 224×224 RGB image and cut it into non-overlapping 16×16 squares. That gives a 14×14 grid, 196 patches. Each patch holds $16\times16\times3 = 768$ numbers; flatten them into one vector of length 768. The image is now a sequence of 196 vectors, the same kind of input a text Transformer takes. In general:

$(H,W)$ is the image resolution, $C$ the number of channels, $P$ the patch side, and $N$ the number of patches, which is also the sequence length the encoder sees. Because $N$ scales as $1/P^2$, the patch size sets the cost. Halving $P$ from 16 to 8 quadruples $N$ from 196 to 784 and multiplies the attention matrix by 16. Model names carry the patch size for this reason: ViT-L/16 is the Large model on 16×16 patches, and ViT-L/32 is the same network on 32×32 patches, 49 tokens instead of 196.

The 768 numbers in a patch are raw pixel intensities. The first learned step multiplies each flattened patch by a trainable matrix $\mathbf{E}$ of shape $(P^2\cdot C)\times D$, where $D$ is the width the Transformer keeps through all its layers (768 for ViT-Base). Each patch comes out as a $D$-dimensional vector, its patch embedding. In ViT-B/16 the input and output widths are both 768; that is a coincidence of $16\cdot16\cdot3$ and BERT-Base's width, and for ViT-L/16 the same 768 pixel values map to $D=1024$. The official code implements the cut, flatten and multiply as one convolution whose kernel size and stride both equal $P$, so each output position sees exactly one patch; the arithmetic is the same.

Two more pieces complete the input. The first is the [class] token: as in BERT, a learnable vector $\mathbf{x}_{\text{class}}$ is prepended to the sequence. It belongs to no patch. Through the layers it attends to the patch tokens like any other token, and its state after the last layer is the representation the classifier reads.

The second is a set of position embeddings. Self-attention has no notion of order. Swap two patches in the input and, because every token is compared with every token using the same weights, the two corresponding outputs swap and nothing else changes. The operation is permutation-equivariant, so without extra information the model sees a bag of patches with no layout. ViT adds a learned vector to each slot: slot 0 for the [class] token, slots 1 to 196 for the patches in raster order (left to right, then top to bottom). The table is 1D, with no built-in knowledge that the slots form a 14×14 grid. In the official code the [class] vector starts at zeros and the position table at small random values (normal, standard deviation 0.02); both are trained with the rest of the model.

Together these give the paper's first equation, the input to the encoder:

The $N+1$ rows of $\mathbf{E}_{pos}$ cover the 196 patches and the [class] token, so the encoder's input for ViT-B/16 is a 197×768 matrix. Figure 2 steps through the pipeline; the slider changes the patch size.

Nothing after $\mathbf{z}_0$ refers to images. The encoder receives a 197×768 matrix and processes it exactly as it would process 197 embedded words.

The encoder: a standard Transformer

The encoder is a stack of $L$ identical blocks. Each block has two sublayers, multi-head self-attention (MSA) and a per-token MLP, and each sublayer is wrapped in a LayerNorm (LN) and a residual connection:

Read (2) left to right: normalize the block's input, run attention on it, add the unnormalized input back. Equation (3) does the same with the MLP. Equation (4) takes row 0 of the last layer's output, the [class] token's state $\mathbf{z}_L^0$, normalizes it, and calls it the image representation $\mathbf{y}$.

LayerNorm rescales each token's vector to zero mean and unit variance across its $D$ features, then applies a learned per-feature gain and bias. Its position differs from the 2017 Transformer, which ViT otherwise says it follows "as closely as possible". The 2017 model is post-norm: it computes $\operatorname{LN}(\mathbf{x}+\operatorname{Sublayer}(\mathbf{x}))$, normalizing the sum after each residual add. ViT is pre-norm: the LN sits inside the branch, and the running sum $\mathbf{z}$ is never normalized between blocks. From $\mathbf{z}_0$ to $\mathbf{z}_L$ the residual path is a chain of additions, so the gradient of the loss with respect to an early block's output includes an identity term from every later block. The paper gives no reason for the switch; it cites Wang et al. (2019), who used pre-norm to train 30-layer translation encoders, and Baevski and Auli (2019).

Inside MSA is scaled dot-product attention. The paper's appendix writes it for a generic input $\mathbf{z}$ of $N$ rows; in ViT the sequence has $N+1$ rows, 197 for B/16. One matrix produces a query, a key and a value for every token:

Row $i$ of $\mathbf{q}\mathbf{k}^{\top}$ holds token $i$'s dot product with every token's key. The softmax runs along each row, so row $i$ of $A$ is a set of 197 non-negative weights summing to 1, and row $i$ of $A\mathbf{v}$ is the weighted average of all 197 value vectors. $A$ is 197×197: every token against every token, the [class] token included. Dividing by $\sqrt{D_h}$ keeps the dot products at unit scale when q and k have unit-variance entries, so the softmax does not saturate into a one-hot row as $D_h$ grows.

Multi-head attention runs $k$ of these in parallel, each head with its own slice of $\mathbf{U}_{qkv}$, concatenates the $k$ outputs, and mixes them with $\mathbf{U}_{msa}$. Each head is $D_h=D/k$ wide, which keeps compute and parameter count the same whatever $k$ is. For ViT-B/16, $k=12$ and $D_h=64$: each head builds its own 197×197 attention matrix from 64-dimensional queries and keys. This is the 2017 Transformer's attention unchanged; ViT packs q, k and v into one matrix with a single head width, where Vaswani et al. wrote separate $d_k$ and $d_v$, which gives the same computation when they are equal. In code, for one image:

def msa(z):                              # z: [197, 768], 12 heads, D_h = 64
    q, k, v = split(z @ U_qkv, 3)        # each [197, 768]            eq (5)
    q, k, v = [t.reshape(197, 12, 64).transpose(1, 0, 2) for t in (q, k, v)]
    A = softmax(q @ k.transpose(0, 2, 1) / sqrt(64))  # [12, 197, 197]  eq (6)
    out = A @ v                          # [12, 197, 64]              eq (7)
    out = out.transpose(1, 0, 2).reshape(197, 768)    # concat the heads
    return out @ U_msa                   # [197, 768]                 eq (8)

The MLP sublayer is two linear layers with a GELU nonlinearity between them (a smooth variant of ReLU), applied to each token independently. It widens each token from $D$ to $4D$ and back: 768 to 3072 to 768 in ViT-Base. The paper uses three sizes; Base and Large are BERT's, Huge is new:

• ViT-Base: 12 layers, $D=768$, MLP 3072, 12 heads, 86M parameters.
• ViT-Large: 24 layers, $D=1024$, MLP 4096, 16 heads, 307M parameters.
• ViT-Huge: 32 layers, $D=1280$, MLP 5120, 16 heads, 632M parameters.

On top of $\mathbf{y}$ sits a classification head: an MLP with one tanh hidden layer during pre-training, and a single linear layer for fine-tuning. The full forward pass:

# ViT-B/16 forward pass for one 224x224 RGB image; shapes in the comments
x = image.reshape(14, 16, 14, 16, 3)     # split H and W into 14 blocks of 16
x = x.transpose(0, 2, 1, 3, 4)           # [14, 14, 16, 16, 3]
x = x.reshape(196, 768)                  # N = 196 patches of P*P*C = 768 values
z = x @ E + b_E                          # [196, 768]  patch embeddings, D = 768
z = concat([x_class, z]) + E_pos         # [197, 768]  eq (1)
for blk in blocks:                       # L = 12
    z = z + blk.msa(blk.ln1(z))          # eq (2), LayerNorm inside the branch
    z = z + blk.mlp(blk.ln2(z))          # eq (3)
y = ln_final(z[0])                       # [768]   eq (4): the [class] token
logits = y @ W_head + b_head             # [1000]  linear head when fine-tuning
# trained by backprop: E, b_E, x_class, E_pos, every block, ln_final, head

Shapes and parameter counts

The shapes above determine the parameter count exactly. For ViT-B/16 the patch embedding $\mathbf{E}$ and its bias hold $768\cdot768+768=590{,}592$ numbers, the [class] token 768, and $\mathbf{E}_{pos}$ $197\cdot768=151{,}296$. One block holds 7,087,872: 1,771,776 in $\mathbf{U}_{qkv}$, 590,592 in $\mathbf{U}_{msa}$, 4,722,432 in the MLP and 3,072 in its two LayerNorms. Twelve blocks make 85.1M, and with the final LayerNorm and a 1000-class linear head the total is 86.6M, the 86M of Table 1. The MLPs hold two thirds of every block.

Figure 3 recomputes every row for any model and patch size. Two things to check in it: how little the patch size changes the parameter count, and how the bar at the bottom splits one block's multiply-adds between the parts linear in the token count $n=N+1$ (the four projections and the MLP, $12\,nD^2$) and the attention matrix itself ($\mathbf{q}\mathbf{k}^{\top}$ and $A\mathbf{v}$, $2\,n^2 D$).

The patch size moves the parameter count only through $\mathbf{E}$ and $\mathbf{E}_{pos}$: ViT-B/32 has 88.2M against B/16's 86.6M, because its larger patches need a larger $\mathbf{E}$. Compute moves a lot, since every block runs on 4 times as many tokens at /16 as at /32. The count matches Table 1 for Base (86.6M) and Huge (632.0M); for Large/16 it gives 304.3M against the table's 307M, and Large/32 counts 306.5M.

The attention matrix is 4.1% of a ViT-B/16 block's multiply-adds, because 197 tokens is small next to $6D=4608$; its share is $n/(6D+n)$. At pixel level, with $n=50{,}177$, the same formula gives 91.6%, and every block would do about 2,900 times the work of the patch model.

What ViT gives up

An inductive bias is an assumption built into a model's architecture rather than learned. The paper names three that a convolutional network builds into every layer:

• Locality. A 3×3 filter reads only a pixel and its eight neighbours, so every feature in an early layer is a function of a small neighbourhood.
• Two-dimensional neighbourhood structure. The filter slides over a grid, so above, below, left and right are part of the computation.
• Translation equivariance. The same filter weights are applied at every position, so shifting an object across the image shifts its feature map by the same amount. (Invariance, the same output wherever the object is, comes later from pooling layers.)

These assumptions are mostly true of natural images, and they mean a convolutional network does not have to learn them from examples. ViT keeps almost none of them. In ViT only the MLP layers are local and translation-equivariant: the MLP applies the same weights to each token separately, so moving a patch to another slot moves its MLP output with it. Self-attention is global from the first layer. Any patch can take weight from any other patch in one step, and the weights are computed from the tokens' contents, with no term for how far apart they are.

Figure 4 compares one layer's reach. Hover over the grid to move the query patch: attention reads all 49 patches, and the 3×3 convolution reads the query and its immediate neighbours. A convolutional network extends its reach by stacking layers; each 3×3 layer grows the receptive field by one patch in each direction.

Global reach matters when related content is far apart: an object and its reflection, two hands on opposite sides of a frame. In Figure 5 the amber object patch sits in one corner and a related teal patch in the other. Switch between the two models and click to move the query. The 3×3 window cannot contain both; attention can connect them in the first layer if training finds that useful.

ViT has much less image-specific inductive bias than a CNN, not none. Two pieces of 2D knowledge are put in by hand: cutting the image into a 2D grid of patches, and, when fine-tuning at a new resolution, 2D-interpolating the position embeddings (covered in the fine-tuning section). At initialization the position embeddings carry no information about the 2D positions of the patches, so every spatial relation between patches is learned from the training data.

The paper also tests a hybrid, in which the patches come from a ResNet's intermediate feature map instead of raw pixels. $\mathbf{E}$ is applied to patches of the feature map, with 1×1 patches as a special case, and the rest of the model is unchanged.

How much data it takes

To see what the missing assumptions cost, the paper pre-trains the same ViT models on three datasets of increasing size and fine-tunes each on ImageNet:

• ImageNet (ILSVRC-2012): 1.3M images, 1k classes.
• ImageNet-21k, its superset: 14M images, 21k classes.
• JFT-300M, Google's in-house set: 303M images, 18k classes, not public.

On the two smaller sets the authors tuned weight decay, dropout and label smoothing (softened one-hot targets) to help the models along. Figure 6 plots the resulting ImageNet top-1 for each model from the paper's Table 5. The convolutional reference is BiT (Big Transfer, Kolesnikov et al. 2020), a family of large ResNets pre-trained on the same three datasets and fine-tuned the same way; the figure uses its largest model, ResNet-152x4. Start at ImageNet and read the ranking, then move the slider.

Pre-trained on ImageNet, ViT-L/16 reaches 76.53% and ViT-B/16 77.91%: the larger model is worse, and both trail the ResNet at 81.30%. With ImageNet-21k the gap nearly closes (L/16 85.15%, H/14 85.13%, ResNet 85.39%), and the Large and Base models are in the expected order. On JFT-300M, H/14 reaches 88.04% and L/16 87.12%, against 87.54% for the ResNet. The larger models gain the most from more data: from ImageNet to JFT-300M, L/16 gains 10.6 points and B/16 6.2.

A second experiment removes the regularization question. The authors pre-train on random subsets of JFT-300M (9M, 30M, 90M images and the full set) with the same hyperparameters everywhere and early stopping, and measure a linear probe: a linear classifier fit on the frozen features from a few labelled images per class, cheaper than full fine-tuning. ViT overfits more than ResNets of similar compute on the small subsets: ViT-B/32, slightly faster than ResNet-50, is much worse on 9M images and better from 90M up, and the same holds for ViT-L/16 against ResNet-152x2.

The paper offers its reading of these results as an intuition: "the convolutional inductive bias is useful for smaller datasets, but for larger ones, learning the relevant patterns directly from data is sufficient, even beneficial." With few examples, a network told in advance that nearby pixels go together does not have to spend data discovering it. With many, a network allowed to relate any patch to any other can fit structure a stack of small filters represents only indirectly.

Accuracy per unit of compute

The headline model, ViT-H/14 pre-trained on JFT-300M and fine-tuned, reaches these numbers (Table 2):

• 88.55% ImageNet top-1;
• 90.72% on ImageNet-ReaL, the same validation images with cleaned-up labels;
• 94.55% on CIFAR-100;
• 77.63% on VTAB, a suite of 19 varied tasks with 1,000 training examples each.

The two strongest convolutional baselines score lower or equal on ImageNet: BiT-L (a ResNet-152x4, also on JFT-300M) 87.54%, and Noisy Student (an EfficientNet-L2 trained with pseudo-labels on 300M unlabeled images) 88.4% to 88.5%. The 88.55% is not a from-scratch ImageNet number. It comes after pre-training on 303M images, and the fine-tuning ran at 518×518 px with Polyak averaging (a running average of the weights), which the next section explains.

Pre-training cost is reported in TPUv3-core-days: TPU v3 cores (two per chip) times days of training. ViT-H/14 took 2.5k, ViT-L/16 on JFT-300M 0.68k, BiT-L 9.9k and Noisy Student 12.3k, so ViT-H/14 used about a quarter of BiT-L's compute (2.5/9.9 ≈ 0.25). ViT-L/16 on the public ImageNet-21k took 0.23k core-days: a standard cloud TPUv3 with 8 cores for about 30 days. The authors add a hedge: pre-training efficiency also depends on the schedule, optimizer and weight decay, not only the architecture.

To control for that, Section 4.4 trains 6 ViTs, 7 ResNets and 5 hybrids on JFT-300M for 7 or 14 epochs and compares transfer accuracy with total pre-training compute. Figure 7 plots all 18 rows of the paper's Table 6. Select a model and the readout finds the cheapest ResNet that matches or beats it on the chosen dataset (the cheapest ViT, if you select a ResNet).

On ImageNet, ViT-B/16 reaches 84.15% for 224 exaFLOPs; the cheapest ResNet at or above it is ResNet-152x2 at 563 exaFLOPs, 2.5 times the compute. ViT-L/16 at 14 epochs (87.12%, 1,567 exaFLOPs) is matched only by ResNet-200x3 (87.22%, 3,306 exaFLOPs). Averaged over five datasets, the paper finds that ViT needs about 2 to 4 times less compute than a ResNet for the same accuracy. Hybrids beat plain ViT at small budgets, and the difference vanishes for the larger models. None of the ViT curves had flattened at the largest size tried, which the authors present as a reason to keep scaling.

The training recipe behind these runs: Adam with $\beta_1=0.9$, $\beta_2=0.999$, batch size 4096, 10k steps of learning-rate warmup, 224 px images, and on JFT-300M a weight decay of 0.1 with linear learning-rate decay. The ResNets were trained with the same optimizer, which the authors note is unusual for ResNets but worked slightly better than SGD in their setting.

Fine-tuning at a higher resolution

To fine-tune, the pre-training head is removed and replaced by a zero-initialized $D\times K$ linear layer, where $K$ is the number of classes in the downstream task. Fine-tuning uses SGD with momentum at batch size 512, and it usually runs at a higher resolution than pre-training, which improves accuracy. The paper fine-tunes at 384 px for most results, and at 512 px (ViT-L/16) and 518 px (ViT-H/14) for Table 2.

The patch size stays fixed, so more pixels means more patches. ViT-H/14 pre-trains at 224 px on a 16×16 grid, 257 tokens; at 518 px the grid is 37×37, 1,369 patches, 1,370 tokens, and each attention matrix has 28.4 times as many entries. The encoder handles any sequence length; its weights do not depend on $N$. The one tensor that does is $\mathbf{E}_{pos}$, which has one row per pre-training slot: 257 rows for a model that now needs 1,370.

ViT resizes the table as a 2D image. The [class] row is set aside, the other 256 rows are reshaped into a 16×16 grid of $D$-dimensional vectors, and the grid is interpolated to 37×37, so each new patch gets the embedding found at its location in the original image. The official code does this with bilinear interpolation (scipy.ndimage.zoom, order 1). This is the second of the two places where ViT is told the image is 2D. Figure 8 shows the resize on one channel of the table, next to the naive alternative of reusing rows in their 1D order.

With interpolation the pattern stays where it was in image coordinates, only sampled more finely. Reusing the rows in raster order fails two ways. Take a /16 model fine-tuned at 384 px, where the 14×14 grid becomes 24×24: only the first 196 of 576 patches get a row at all, and those land at the wrong places, because slot 16 (patch 15, counting patches from 0) was row 1, column 1 on the 14-wide grid and is row 0, column 15 on the 24-wide one.

Resolution accounts for part of the headline number. Table 5 fine-tunes the same JFT-300M ViT-H/14 at 384 px without Polyak averaging and reports 88.04%; at 518 px with averaging, Table 2 reports 88.55%.

What the trained model learned

Section 4.5 inspects a trained model in three places: the patch embedding, the position embeddings, and the attention heads.

The columns of $\mathbf{E}$, each a 16×16×3 pattern, act as filters applied to each patch. The top principal components of the learned filters resemble, in the paper's words, plausible basis functions for a low-dimensional representation of the fine structure within each patch.

For the position embeddings, the paper computes the cosine similarity between each patch's learned position vector and every other patch's, for ViT-L/32 on its 7×7 grid. Figure 9 models that plot. Hover a patch to select it.

Nearby patches have the most similar embeddings, and patches in the same row or column stay similar at any distance. The model was given 49 unrelated vectors and learned from data that they form a 7×7 grid. For larger grids a sinusoid-like structure sometimes appears in the trained vectors; the 2017 Transformer built sinusoids in by hand, and here they emerge. This explains a negative result in Appendix D.4: on a 5-shot ImageNet linear probe with ViT-B/16, learned 1D embeddings score 0.642, hand-built 2D embeddings 0.640 and relative embeddings 0.640, while removing position embeddings drops the score to 0.614.

For the attention heads, the paper defines attention distance: for one head, the distance in pixels between a query position and every other position, weighted by the attention weight between them and averaged over 128 example images. It plays the role a receptive field plays in a CNN. Figure 10 plots it for the 16 heads of each layer of ViT-L/16 on 224 px images.

In the lowest layers the heads split. Some already attend across most of the image, so the global reach is used from layer 1. Others keep small distances, which resembles the local filters of a CNN's early layers. Attention distance grows with depth, and in the second half of the network most heads attend widely. In the hybrid, whose input already went through ResNet layers, the highly localized heads are less pronounced, which the authors take as a hint that those heads do the job of early convolutions. Maps of the attention from the output token back to the input pixels land on the image regions that are semantically relevant for the label (the paper's Figure 6).

Self-supervision, pooling, limits

BERT's pre-training needs no labels: hide some words and predict them. Section 4.6 tries the image version, a preliminary experiment the authors label as such. They corrupt 50% of patch embeddings (80% of those replaced by a learned [mask] embedding, 10% by another patch's embedding, 10% left as is), and predict for each corrupted patch its mean colour quantized to 3 bits per channel, one of 512 classes. Pre-trained this way on JFT-300M, ViT-B/16 reaches 79.9% on ImageNet, 2 points above training from scratch and 4 points below supervised pre-training. Contrastive pre-training and closing that gap are left to future work.

Appendix D.3 answers whether the [class] token is needed. The alternative, global average pooling (GAP) of the patch outputs followed by a linear classifier, as a ResNet does, first performed very poorly. The authors traced the difference to the learning rate: with its own learning rate, GAP works as well as the class token. They kept the token to stay close to the text Transformer.

The limits follow from the results. On mid-sized datasets without large-scale pre-training, a ResNet of similar size is still more accurate. The strongest numbers depend on JFT-300M, which is not public; the ImageNet-21k models are the reproducible ones. The paper covers classification only and names detection and segmentation as the next applications.

Later work removed several of these limits while keeping the patch tokenizer and the encoder. DeiT (Touvron et al., 2021) trains ViT on ImageNet alone using distillation from a CNN teacher and strong augmentation. MAE returns to masked-patch prediction, reconstructing pixels of 75% hidden patches, and DINO trains ViTs without labels by self-distillation. CLIP trains a ViT image encoder against text captions, and that encoder is the vision input of LLaVA. SAM segments with an MAE-pre-trained ViT-H, and DiT applies the same patch-and-Transformer recipe to the latents of a diffusion model.