An Image Is Worth 16×16 Words

Dosovitskiy et al., 2020 — the Vision Transformer

The whole paper is one reshape. I kept looking for the clever part, and it was .reshape the entire time.

Download vit_from_scratch.py — everything below runs from it. NumPy only.

I'd shipped image classifiers: fine-tuned ResNets, exported to ONNX, argued about input normalisation in code review. Then transformers ate vision and I assumed there was a new trick for pixels. There isn't. Cut the image into tiles, flatten each into a row, hand the rows to the exact encoder from the attention page. The interesting engineering is what that reshape throws away, and why it wins anyway.

Tokenising an image

You know this pipeline. Every language model starts with it:

ids    = tokenizer("the cat ate food")   # [464, 3797, 8268, 2057]
tokens = embedding_table[ids]            # (4, D)   one row per word
out    = transformer(tokens)             # (4, D)   same shape, now in context

The paper is that pipeline with a different tokeniser. Words come from a vocabulary; tiles come from a grid:

patches = patchify(image, p=16)          # (196, 768)  196 tiles, 16*16*3 each
tokens  = patches @ E                    # (196, D)    a matmul, not a lookup
tokens  = concat([cls, tokens]) + pos    # (197, D)    one extra row, plus order
out     = transformer(tokens)            # (197, D)    unchanged since 2017

Two differences from text, and they're the only image-specific ideas in the paper. No vocabulary: a tile is 768 continuous numbers, not an id, so the embedding is a matmul, not a lookup. And the tokeniser has no rules. It's a grid.

TermMeansProgrammer's analogy
PatchOne p×p×3 tile of the imageA word
FlattenThat tile as one row of p·p·3 numbersThe word's bytes
EA (p·p·3, D) matrix every row is multiplied byThe embedding table, as a matmul because there's no vocab
CLS tokenOne learned row prepended to the sequenceA sentinel. The row you read at the end
Position embeddingA learned (N+1, D) table added to the rowsThe index in the list. Without it the list is a set

Why it beat the CNN — and when it didn't

A convolution is a for loop over a sliding window:

for y, x in every_pixel:
    out[y, x] = kernel · img[y-1:y+2, x-1:x+2]   # only the 3x3 neighbourhood,
                                                 # same kernel at every (y, x)

Two assumptions are baked into that loop: nearby pixels matter more than far ones, and a cat in the corner is the same cat as one in the middle. A CNN gets both for free from its structure. ViT gets neither: tile 0 and tile 195 are related only through attention weights it must learn, and "nearby pixels matter" is something it must discover from data. The paper's phrase is much less image-specific inductive bias. Fewer assumptions is a trade: worse when data is scarce, better once there's enough data to learn the assumptions and then some. On ImageNet-1k alone, ViT lands below a ResNet of comparable size. Pretrained on ImageNet-21k (14M images) or JFT-300M it wins; the biggest, ViT-H/14, tops the paper's headline table at 88.55% on ImageNet. The interview question is "why does ViT need more data than a CNN?", and the answer is that sentence about inductive bias.

The one formula

z0 = [xclass; xp1E; xp2E; …; xpNE] + Epos

Equation 1. Read it left to right as shapes, for ViT-Base on a 224×224 image with patch 16:

StepShapeWhat it does
x(224, 224, 3)The image. Height, width, channel
xp(196, 768)The reshape. N = 224·224 / 16² = 196 tiles, each 16·16·3 = 768 numbers
E(768, 768)Patch embedding. Square only because ViT-Base picked D = 768
xpE(196, 768)196 tokens. From here on, this is a sentence
xclass(1, 768)The learned CLS row, prepended → (197, 768)
+ Epos(197, 768)Learned, 1D, one vector per slot. Added, not concatenated

Equations 2 and 3 are the encoder block, stacked 12 deep for Base, 24 for Large, 32 for Huge; Equation 4 says read row 0 after a final LayerNorm. Both are on the attention page already. The one change since 2017: LayerNorm moves to before each sublayer (pre-norm), and the MLP uses GELU. At fine-tune time the paper feeds a bigger image, keeps the patch size so there are more tokens, and 2D-interpolates the learned position table to the new grid; they note that's the only place a 2D assumption goes in by hand.

The reshape, with every shape printed

The script uses a 32×32×3 toy image and 8×8 patches so the numbers stay readable: 16 tiles of 192 numbers. Three steps. Cut each spatial axis into (how many tiles, pixels per tile). Group the two "which tile" axes together and the two "which pixel" axes together, which is a transpose. Flatten to one row per tile. Demo 1 does it one line at a time so you can breakpoint each step.

Don't take my word for it. Here's the actual output:

  image    (32, 32, 3)       (H, W, C)
  cut      (4, 8, 4, 8, 3)   (H/p, p, W/p, p, C)
  grouped  (4, 4, 8, 8, 3)   (H/p, W/p, p, p, C)
  patches  (16, 192)         (N, p*p*C)          <- the sentence
  tokens   (16, 64)          (N, D)              patches @ E
  + cls    (17, 64)          (N+1, D)            BERT-style
  + pos    (17, 64)          (N+1, D)            added
  block    (17, 64)          (N+1, D)            shape unchanged
  attn     (17, 17)          (N+1, N+1)          all patches see all
  cls out  (64,)             (D,)                the row we read
  logits   (10,)             (n_classes,)

unpatchify(patchify(img)) == img? True

3072 numbers in, 3072 out, regrouped into 16 rows. Nothing is computed until patches @ E. From (16, 64) onward the shapes are the attention page's (n_tokens, d_model), and the 17×17 attention matrix means every tile sees every tile, plus CLS. The round-trip at the bottom is the test I'd write: if it fails, the reshape is wrong, whatever the shapes say.

Position: a bag of tiles

Attention is a set operation. The attention page proved it with "cat ate food" versus "food ate cat"; here it's a jigsaw with the pieces swapped. Shuffle the 16 tiles into a new image, run both, and compare the CLS row with the position table off and on: 

tile permutation: [ 1 12  7 10 14  4  5  8  0  9  2 13 11  6  3 15]
shuffled image == original? False

position embeddings OFF: CLS max abs diff = 1.44e-15  allclose? True   ← blind
position embeddings ON : CLS max abs diff = 0.0495    allclose? False  ← sees it

That 1.44e-15 is not zero, and I won't pretend it is. The shuffle changes the order the attention sum runs in, and float64 rounds differently in a different order. Fifteen decimal places is "identical" by any test you'd write; the 0.0495 with positions on is thirteen orders of magnitude bigger. That + self.pos is the only thing telling the model where a tile was. The paper keeps the table deliberately dumb: learned, 1D, one vector per slot, like a list index. They tried 2D-aware versions and report no gain, because the plain table learns the grid on its own; the paper's inspection section has the picture, where the similarity between learned position vectors forms rows and columns.

The same shuffle settles which row you read at the end. BERT-style it's row 0, but there are three obvious candidates: 

max abs diff under the shuffle, position embeddings OFF:
  z[0]        CLS token       1.44e-15
  z[1:].mean  mean-pool       6.66e-16
  z[6]        patch token 6   1.47
  (token 6 now holds tile 4, not tile 5)

CLS and the mean of the patch rows are both reads over the whole set, so both survive. A single patch row is a read at a position, so it moves. The paper's appendix tried mean-pooling instead of CLS and found it works just as well once the learning rate is re-tuned; the CLS token is a BERT habit, not a requirement.

Why 16×16: the n² bill

Patch size is the one dial the paper adds, and it sets everything. Tokens n = (224/p)², and every attention layer builds an n×n score matrix per head. The script prints the bill for a 224×224 image: 

   patch  per token  tokens n          n x n   vs 16x16
    1x1           3     50176  2,517,630,976  65536.00x
    4x4          48      3136      9,834,496    256.00x
    8x8         192       784        614,656     16.00x
   16x16        768       196         38,416      1.00x
   32x32       3072        49          2,401      0.06x

Halve the patch, quadruple the tokens, sixteen times the attention matrix. One pixel per token would be 2.5 billion scores per head per layer, which is why nobody does that. 16×16 gives 196 tokens: a paragraph, not a book. Two things fall out of the table. 16·16·3 = 768 is exactly ViT-Base's D, so its embedding matrix happens to be square. And the model names: ViT-L/16 is Large with patch 16; ViT-H/14 uses patch 14 for more tokens at more cost.

The title is a shape. "16×16 words" is the tile, and 196 of them is the sentence. Everything a transformer costs scales with that sentence length squared.

The wrong reshape is silent

This is the bug you will actually ship. Skip the cut and the transpose and go straight to img.reshape(16, 192). NumPy reads the image row by row, so each "patch" is two full-width strips, not an 8×8 square. The shape is identical. Nothing throws. 

shapes: (16, 192) (16, 192)  equal? True

patch 5, first 8 values (R G B R G B R G):
  right: [0.258 0.258 0.    0.29  0.258 0.    0.323 0.258]
  wrong: [0.    0.323 1.    0.032 0.323 1.    0.065 0.323]

what is 'patch 5', really?
  right == img[8:16, 8:16]  an 8x8 square?  True
  wrong == img[10:12, :]    a 2x32 strip?   True

model runs on the wrong patches anyway: logits (10,)

Red ramps left to right in the toy image, green top to bottom. The right row starts at pixel (8, 8), so red and green both read 0.258. The wrong row starts at pixel (10, 0): red 0, green 0.323, blue on the wrong checkerboard square. Same shape, different picture, ten plausible logits. It will even train; the accuracy is just quietly bad and you'll blame the learning rate. The PyTorch flavour is running this recipe on a (C, H, W) tensor without moving the channel axis first. The failure you want is the loud one: 32 isn't divisible by 5, the cut can't happen, and NumPy refuses. Real pipelines resize to 224 before patchifying for exactly this reason. 

  p=8: patches (16, 192)
  p=5: ValueError: cannot reshape array of size 3072 into shape (6,5,6,5,3)

The code

Pure NumPy. No torch, no autograd, no training — just the forward pass, so you can watch the reshape happen. Every function under 20 lines. softmax, attention and layernorm are the attention page's, unchanged: in the script, not repeated here.

Download the full script — 160 lines plus six runnable demos. Then: 

python3 vit_from_scratch.py
import numpy as np


def patchify(img, p):
    H, W, C = img.shape
    # (1) CUT: split each spatial axis into (how many patches, pixels each).
    #     (H, W, C) -> (H/p, p, W/p, p, C). Throws if p does not divide H, W.
    cut = img.reshape(H // p, p, W // p, p, C)
    # (2) GROUP: bring the two "which patch" axes together, and the two
    #     "which pixel inside it" axes together. This transpose is the line
    #     everyone skips. Skipping it is silent garbage (demo 5).
    grouped = cut.transpose(0, 2, 1, 3, 4)       # (H/p, W/p, p, p, C)
    # (3) FLATTEN: one row per patch. p*p pixels, R G B R G B ... in order.
    return grouped.reshape((H // p) * (W // p), p * p * C)


class EncoderBlock:
    def __init__(self, D, seed):
        rng = np.random.default_rng(seed)
        s = 1 / np.sqrt(D)
        self.W_q, self.W_k, self.W_v, self.W_o = [
            rng.normal(0, s, (D, D)) for _ in range(4)]
        self.W_1 = rng.normal(0, s, (D, 4 * D))             # MLP: D -> 4D -> D
        self.W_2 = rng.normal(0, 1 / np.sqrt(4 * D), (4 * D, D))

    def __call__(self, z):
        h = layernorm(z)                                     # LN first
        a, weights = attention(h @ self.W_q, h @ self.W_k, h @ self.W_v)
        z = z + a @ self.W_o                                 # residual: a diff
        z = z + gelu(layernorm(z) @ self.W_1) @ self.W_2     # MLP, residual
        return z, weights


class ViT:
    def __init__(self, H, W, C, p, D, n_blocks, n_classes, seed=0):
        rng = np.random.default_rng(seed)
        self.p, self.N = p, (H // p) * (W // p)
        self.E = rng.normal(0, 1 / np.sqrt(p * p * C), (p * p * C, D))
        self.cls = rng.normal(0, 0.02, (1, D))
        self.pos = rng.normal(0, 0.02, (self.N + 1, D))
        self.blocks = [EncoderBlock(D, seed=i + 1) for i in range(n_blocks)]
        self.W_head = rng.normal(0, 1 / np.sqrt(D), (D, n_classes))

    def __call__(self, img, use_pos=True):
        return self.forward(patchify(img, self.p), use_pos)  # (N, p*p*C) in

    def forward(self, patches, use_pos=True):
        tokens = patches @ self.E                            # (N, D)   x_p E
        z = np.concatenate([self.cls, tokens], axis=0)       # (N+1, D) prepend
        if use_pos:
            z = z + self.pos                                 # (N+1, D) + E_pos
        for block in self.blocks:
            z, weights = block(z)                            # (N+1, D) unchanged
        z = layernorm(z)                                     # Eq. 4: LN(z_L^0)
        return z[0] @ self.W_head, z, weights                # logits, rows, attn

Now break it

This is the part that actually makes it stick. Reading won't do it. You debug for a living — reverse-engineering from broken behaviour is your native mode.

Break thisWatch for
Call patchify(img, 5)ValueError: cannot reshape array of size 3072 into shape (6,5,6,5,3). Loud, immediate, at the first line. This is the good failure.
Delete z = z + self.posShuffle the tiles and the CLS row is unchanged: max abs diff 1.44e-15, allclose? True. The same answer for any tile order. It can't see layout at all.
In ViT.__call__, swap patchify for patchify_wrongNothing. Shapes match, logits (10,), no error. "Patch 5" is now img[10:12, :], a 2×32 strip. Silent, and it'll train.
Read z[6] instead of z[0] for the headUnder the shuffle it moves by 1.47 while CLS and z[1:].mean(0) stay at 1e-15. Mean-pooling is a legitimate head (the paper's appendix, with a re-tuned learning rate). One patch row is not.
Train it on a small datasetYou can't, here: there is no training loop. The paper's finding is that on ImageNet-1k alone ViT trails a comparable ResNet, and only overtakes it after ImageNet-21k or JFT-300M pretraining, because it has less image-specific inductive bias and must learn locality from data.
Then read the paper — but only §3.1, which is the script with Greek letters, §4.3 for the "how much data does it need" curves, and the two appendix ablations on head type and position embeddings. Skip the compute tables and the self-supervision section. You'll find the paper is just the comments you've already read.

On the maths anxiety

There is no new maths in this paper. If you cleared the attention page you have cleared this one: a reshape and a transpose are array indexing, the patch embedding is one matmul, and the rest is the same softmax and weighted average as before.

The one thing worth being fluent in is the arithmetic in the n² table: tokens from patch size, attention cost from tokens, so memory and latency from patch size. That's the sum you'll do at work when someone asks whether the model fits on the GPU at 384×384, and it's three multiplications. You don't need to know why 1D position embeddings learn a 2D grid. Nobody fully does.

Where to go next

This page gave you one thing: an image as a sentence, in code you can break. These fill in what I deliberately skipped.

Order I'd take them in: §3.1 of the paper → vit-pytorch → the Hugging Face fine-tune. How, then real code, then your own dataset. If you only read one, read vit-pytorch with the script open beside it.