Vision Transformers

Week 3 of the residency. What attention does, why the architecture is called a transformer, and how you feed it an image.

Week 3 opened with a question I had never thought to ask. Why is it called a transformer?

Leaves fall every fall season.

Going in, every word is ambiguous. Leaves could be foliage or departing. Fall could be the season or the verb, and it appears twice meaning different things.

Week 1’s embedding gives each word one vector. That single vector has to cover every sense of the word. fall gets the same numbers whether it means autumn or dropping.

Coming out, each word has one meaning. The block has rewritten each word’s vector using the other words present. That’s its job. It transforms a representation of a word into a representation of that word in this context.

Once that clicked, the rest of the week followed from it.

Attention

Each token produces three vectors. A query, which is what it is looking for. A key, which is what it offers. A value, which is what it passes on if chosen.

Every token’s query is compared against every token’s key by dot product, giving a score for each pair. Those scores go through a softmax, which exponentiates them and divides by the total so they become weights adding up to one. Each token’s new representation is then the weighted sum of everyone’s values.

If that sounds like a lookup table with fuzzy keys, that is roughly right. The difference is that nothing is looked up exactly; every token contributes a little, in proportion to how well its key matched.

fall asks “am I near a season word or a motion word”, season answers, and fall’s vector moves accordingly.

The Square Root

That denominator, \(\sqrt{d_k}\), looks like a detail. It is not, and the workshop showed why rather than asserting it: the same attention computed at three embedding sizes, with and without the scaling.

At one dimension there is nothing to choose between them.

At four, the unscaled scores are spreading out.

At 512 the unscaled version has fallen over. A dot product sums one term per dimension, so with 512 dimensions the scores are simply bigger numbers. Exponentiate bigger numbers and the largest one dominates completely. Almost all the weight lands on a single token and everything else gets close to nothing.

Attention stops being a weighted average and becomes a hard lookup. Worse, a weight pinned at zero or one barely responds to small changes in the inputs, so the model stops being able to learn from it.

Dividing by \(\sqrt{d_k}\) keeps the scores in the range where softmax is still soft. It is one symbol in the formula and it is the difference between the mechanism working and not.

The Task

Week 3’s project is both halves of that diagram. Recognise one digit with an encoder, then read a multi-digit number with an encoder and a decoder. This post is the left-hand side; the next one is the right.

The awkward part is that a transformer eats a sequence of tokens. An image is not one. It is a grid of pixels with no natural order, and there are far too many of them.

Attention compares every position with every other. Cost grows with the square of the sequence length. A 28×28 digit is 784 positions. A photograph is hopeless.

Patches

The Vision Transformer’s answer is to stop treating pixels as the unit. Cut the image into fixed-size squares and treat each square as a word.

Each patch is flattened and pushed through one Linear layer. That is the entire modification.

A 196-pixel patch becomes a 64-number vector. A word became a 64-number vector in week 1. The transformer reading them cannot tell the difference.

One Encoder, Any Input

This is the slide that makes week 3 worth more than a digit classifier.

Text goes through a tokenizer and an embedding table. An image goes through a linear projection of flattened patches. Audio goes through a 1-D convolution over frequency. They all produce the same thing: a sequence of vectors.

Everything after that is identical. The encoder does not know what modality it is reading. Which means the work of supporting a new input type is writing a new front-end, not designing a new model, and that is what makes weeks 4 and 5 possible.

Building One

What follows is that architecture shrunk until it trains while this page builds: 7×7 patches, two encoder layers, 64 dimensions, six thousand digits. The real service uses 4×4 patches into 128 dimensions, six layers and eight heads, with the attention and encoder blocks written out by hand rather than taken from PyTorch.

I use learned position embeddings here rather than the sinusoidal ones the service uses, so I can delete them and see what they were holding up.

import sys

sys.path.insert(0, "..")
import numpy as np
import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download

REVISION = "f705fed08827ff6c36e3b5329495c943a5e544e8"
data = np.load(hf_hub_download("roshbeed/ai-residency-blog-data", "mnist/mnist-small.npz",
                               repo_type="dataset", revision=REVISION))

x_train = torch.from_numpy(data["x_train"]).float() / 255.0
y_train = torch.from_numpy(data["y_train"]).long()
x_test = torch.from_numpy(data["x_test"]).float() / 255.0
y_test = torch.from_numpy(data["y_test"]).long()

PATCH = 7
PATCHES = (28 // PATCH) ** 2


def patchify(images):
    """28x28 -> 16 patches of 7x7, flattened. This is the whole trick."""
    batch = images.shape[0]
    tiles = images.unfold(1, PATCH, PATCH).unfold(2, PATCH, PATCH)
    return tiles.reshape(batch, PATCHES, PATCH * PATCH)


baseline = torch.bincount(y_test).max().item() / len(y_test)
print(f"{len(x_train):,} training digits, each becoming {PATCHES} tokens of {PATCH * PATCH}")
print(f"always guessing the most common digit: {baseline:.4f}")

6,000 training digits, each becoming 16 tokens of 49
always guessing the most common digit: 0.1173

The model below is a patch projection, two encoder blocks and a classification head reading position zero. train takes a flag for whether the position embeddings are added at all, which is what the next section uses.

DIM, HEADS, LAYERS = 64, 4, 2


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.attention = nn.MultiheadAttention(DIM, HEADS, batch_first=True)
        self.norm1, self.norm2 = nn.LayerNorm(DIM), nn.LayerNorm(DIM)
        self.feedforward = nn.Sequential(nn.Linear(DIM, 4 * DIM), nn.GELU(),
                                         nn.Linear(4 * DIM, DIM))

    def forward(self, x, want_weights=False):
        attended, weights = self.attention(x, x, x, need_weights=want_weights,
                                           average_attn_weights=True)
        x = self.norm1(x + attended)
        return self.norm2(x + self.feedforward(x)), weights


class VisionTransformer(nn.Module):
    def __init__(self, use_positions=True):
        super().__init__()
        self.use_positions = use_positions
        self.project = nn.Linear(PATCH * PATCH, DIM)
        self.cls = nn.Parameter(torch.zeros(1, 1, DIM))
        self.positions = nn.Parameter(torch.randn(1, PATCHES + 1, DIM) * 0.02)
        self.blocks = nn.ModuleList([Block() for _ in range(LAYERS)])
        self.head = nn.Linear(DIM, 10)

    def forward(self, images, want_weights=False):
        tokens = self.project(patchify(images))
        tokens = torch.cat([self.cls.expand(len(images), -1, -1), tokens], dim=1)
        if self.use_positions:
            tokens = tokens + self.positions

        weights = None
        for i, block in enumerate(self.blocks):
            tokens, w = block(tokens, want_weights and i == LAYERS - 1)
            if w is not None:
                weights = w
        return self.head(tokens[:, 0]), weights   # position 0 is the [CLS] token


def train(use_positions, epochs=15, batch=128, seed=0):
    torch.manual_seed(seed)
    model = VisionTransformer(use_positions)
    optimiser = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
    generator = torch.Generator().manual_seed(seed)

    accuracies = []
    for _ in range(epochs):
        perm = torch.randperm(len(x_train), generator=generator)
        for i in range(0, len(perm) - batch, batch):
            b = perm[i:i + batch]
            loss = nn.functional.cross_entropy(model(x_train[b])[0], y_train[b])
            optimiser.zero_grad()
            loss.backward()
            optimiser.step()
        with torch.no_grad():
            accuracies.append((model(x_test)[0].argmax(1) == y_test).float().mean().item())
    return model, accuracies


model, with_positions = train(use_positions=True)
print(f"{sum(p.numel() for p in model.parameters()):,} parameters")
print(f"best test accuracy: {max(with_positions):.4f}   (baseline {baseline:.4f})")

104,970 parameters
best test accuracy: 0.9387   (baseline 0.1173)

Drawn out, that’s the whole model:

from _arch import diagram

diagram(VisionTransformer(), input_shape=(1, 28, 28), style="flow")

Figure 1: The model end to end. Each patch is projected to 64 numbers, two identical encoder blocks run over the 17 tokens, and the head reads position 0. The tall blue and green pairs are the feed-forward expansion inside each block; the outlines arching over are the residual connections.

Position Embeddings

Attention compares every token with every other and takes a weighted sum. Nothing in that calculation refers to where a token is, so shuffling the tokens gives the same outputs in a different order. Without something to break that symmetry a Vision Transformer sees a bag of patches rather than a picture.

Deleting them measures what they were worth.

_, without_positions = train(use_positions=False)
print(f"with position embeddings:    {max(with_positions):.4f}")
print(f"without position embeddings: {max(without_positions):.4f}")
print(f"difference:                  {max(with_positions) - max(without_positions):+.4f}")

with position embeddings:    0.9387
without position embeddings: 0.8707
difference:                  +0.0680

Plotted against each other, with the majority-class baseline underneath:

from _style import COLOURS, MUTED, figure, style_axes

epochs = range(1, len(with_positions) + 1)
fig, ax = figure(height=3.8)
ax.axhline(baseline, color=MUTED, linewidth=1.2, linestyle="--")
ax.plot(epochs, with_positions, color=COLOURS[0], linewidth=2)
ax.plot(epochs, without_positions, color=COLOURS[1], linewidth=2)
ax.annotate("with position embeddings", xy=(epochs[-1], with_positions[-1]), xytext=(-6, 6),
            textcoords="offset points", ha="right", fontsize=9, color=COLOURS[0])
ax.annotate("without", xy=(epochs[-1], without_positions[-1]), xytext=(-6, -16),
            textcoords="offset points", ha="right", fontsize=9, color=COLOURS[1])
ax.annotate("always guess the most common digit", xy=(1, baseline), xytext=(2, 8),
            textcoords="offset points", fontsize=9, color=MUTED)
ax.set_ylim(0, 1.0)
style_axes(ax, "Epoch", "Test accuracy")
fig.tight_layout()

Figure 2: The same model with and without position embeddings. A bag of patches still says a lot about which digit it is; the layout is worth the rest.

It still works without them, which surprised me until I thought about what a bag of 7×7 patches contains. A 0 and a 1 are made of visibly different pieces however you shuffle them. Position buys the distinctions that depend on layout.

Since the classifier reads only the [CLS] token, the attention out of that position says which patches the answer actually rests on.

import matplotlib.pyplot as plt

with torch.no_grad():
    logits, weights = model(x_test[:4], want_weights=True)

cls_attention = weights[:, 0, 1:].reshape(-1, 4, 4)

fig, axes = plt.subplots(1, 4, figsize=(8.2, 2.4))
for ax, image, attention, predicted in zip(axes, x_test[:4], cls_attention, logits.argmax(1)):
    ax.imshow(image, cmap="gray_r")
    ax.imshow(attention, cmap="inferno", alpha=0.45, extent=(0, 28, 28, 0),
              interpolation="bilinear")
    ax.set_title(f"predicted {predicted.item()}", fontsize=9, color=MUTED, pad=6)
    ax.axis("off")
fig.tight_layout()

Figure 3: Attention from the [CLS] token in the final block. Nothing in the loss says to look at the strokes.

Conclusion

The accuracy is not the point. A small convolutional network beats this at MNIST scale, and the ViT paper says transformers only overtake convolutions with a lot of data behind them.

What week 3 gives you:

• Why it is called a transformer: it rewrites each token’s meaning using context
• Why attention divides by the square root of the width
• An image is a sequence of vectors, the same as a sentence
• The front end is swappable, and the model behind it does not change

Week 4 takes that last point literally and bolts a vision model to a language model.

The full project, with the blocks written from scratch, is on GitHub.