2.2 — Vision Transformer (ViT)

Learning objectives

After this lesson, you should be able to:

• Explain how an image becomes a sequence of patch tokens.
• Compute the number of patches for a given image and patch size.
• Explain the roles of the class token and position embeddings.
• Explain why doubling image size makes attention about 16 times costlier.

The one-sentence idea

ViT turns an image into a sequence of patch vectors so a transformer can read it like text.

Patches in, tokens out. Same transformer as for text.

An image as a sentence

Transformers read sequences of tokens. An image is a grid of pixels, not a sequence. The Vision Transformer (ViT) fixes this in a simple way: cut the image into small square patches and treat each patch as one “word”. The paper applies a standard Transformer directly to a sequence of image patches, and its title, “An image is worth 16x16 words”, gives the patch size it uses.^[1]

The recipe has four steps: split the image into patches, linearly embed each patch, add position embeddings, and feed the sequence to a standard Transformer encoder. A learnable extra token is added at the front, and its output is used for classification.^[1] In a VLA, the patch outputs themselves are used (lesson 1.4), but the front half is the same.

Step 1 and 2: patches and embeddings

With 16×16 patches the base model turns a 224×224 image into 196 patches, each embedded to a vector of width 768.^[1] The embedding is a single matrix multiplication shared by all patches:

import numpy as np
rng = np.random.default_rng(0)

image = rng.random((224, 224, 3))
P, D = 16, 768
n = 224 // P
patches = image.reshape(n, P, n, P, 3).transpose(0, 2, 1, 3, 4).reshape(n * n, P * P * 3)
W = rng.normal(size=(P * P * 3, D)) * 0.02              # the learned "linear embedding" (random here)
tokens = patches @ W

print("patches:", patches.shape)
print("embedded tokens:", tokens.shape)

patches: (196, 768)
embedded tokens: (196, 768)

Toy model. The weights are random. In a trained ViT the matrix is learned, so that patches that look alike get similar vectors.

The patch matrix has 768 columns because 16 × 16 × 3 = 768 (height × width × colour channels). It is a coincidence that this equals the embedding width in the base model.

Step 3: a class token and positions

The model also needs to know where each patch came from, since attention ignores order (lesson 2.1). ViT adds position embeddings to the patch embeddings, and prepends a learnable classification token.^[1] Here are the shapes, and a check that positions really change the result:

import numpy as np
rng = np.random.default_rng(0)

n_patches, D = 196, 768
tokens = rng.normal(size=(n_patches, D))
cls = rng.normal(size=(1, D))                           # learnable in a real model
pos = rng.normal(size=(n_patches + 1, D)) * 0.02        # learnable in a real model

seq = np.concatenate([cls, tokens]) + pos
print("sequence into the Transformer:", seq.shape)

def softmax(x):
    e = np.exp(x - x.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)
def attend(X):
    return softmax(X @ X.T / np.sqrt(X.shape[1])) @ X

base = np.concatenate([cls, tokens])
perm = rng.permutation(n_patches + 1)
no_pos = np.allclose(attend(base[perm]), attend(base)[perm])
with_pos = np.allclose(attend(base[perm] + pos), attend(base + pos)[perm])
print("without positions, shuffling the tokens only shuffles the output:", no_pos)
print("with positions, shuffling the tokens only shuffles the output:", with_pos)

sequence into the Transformer: (197, 768)
without positions, shuffling the tokens only shuffles the output: True
with positions, shuffling the tokens only shuffles the output: False

With positions added, a shuffled image is a different input, not a rearranged copy of the same one: the position vectors stay in their slots while the tokens move, so each token is now combined with a different position. That is exactly the order-awareness we wanted.

Why resolution gets expensive

Attention compares every token with every other token, so the work grows with the square of the number of tokens. Higher resolution means more patches, which means a much bigger attention table:

P = 16
for side in (224, 336, 448, 896):
    n = (side // P) ** 2
    print(side, "px ->", n, "patches ->", n * n, "attention entries per head")

224 px -> 196 patches -> 38416 attention entries per head
336 px -> 441 patches -> 194481 attention entries per head
448 px -> 784 patches -> 614656 attention entries per head
896 px -> 3136 patches -> 9834496 attention entries per head

Doubling the side of the image from 224 to 448 pixels gives 4 times as many patches and 16 times as many attention entries. This is one reason image size is a design choice in robot models, and why speed comes up in lesson 1.5. The ViT paper’s central finding, besides the architecture, is that large-scale pre-training matters: at sufficient scale the Transformer matches or beats convolutional networks on image classification.^[1] VLA encoders are in practice built from pre-trained ones: OpenVLA, for example, fuses pretrained DINOv2 and SigLIP features.^[2]

Common misconceptions

Engineering takeaways

1. 196 patches for a 224 px image with 16 px patches.
2. One shared matrix embeds every patch.
3. Positions and a class token are added.
4. Cost grows with the square of the patch count.
5. Pre-training at scale matters.

Cut the picture into words and the transformer can read it.

Knowledge check

Try each question first, then reveal the explanation.

What’s next

Next: CLIP, which teaches a model to match images with the words that describe them.

Back to the VLA roadmap