Vision Transformers: Attention for Images

Third Foundations post in ML & AI, and the one that opens up the Vision side of this section. The pitch behind the Vision Transformer (Dosovitskiy et al., ICLR 2021) is almost aggressively simple: take the same discretization trick from the tokenization post — turn raw input into a sequence of discrete units — and apply it to pixels instead of characters, then feed that sequence into the exact same architecture from How the Transformer Actually Works. No new architecture required.

Patches are the image's tokens

A 224×224 image gets split into fixed-size patches — 16×16 pixels each in the original paper — flattened, and linearly projected into an embedding, exactly the way a token gets embedded in an NLP model. For a 224×224 image with 16×16 patches, that's (224×224)/(16×16) = 196 patches: a sequence length in the same ballpark as a paragraph of text.

The [class] token, borrowed directly from BERT

Classification needs a single output representing "the whole image," not 196 separate patch representations. ViT's answer is lifted straight from BERT: prepend a learnable embedding — the [class] token — to the sequence, and after the sequence passes through the full encoder, use that one token's final-layer output as the image's representation. Nothing about attention treats it specially; it's just another position in the sequence that happens to get used differently afterward.

Position embeddings get added too, for the same reason as text: attention has no built-in sense of order, only content, so without them the model couldn't tell a patch from the top-left corner from one in the bottom-right.

After that, it's the same encoder

This is the part worth sitting with: once you have a sequence of patch embeddings plus a [class] token plus position embeddings, everything downstream — multi-head self-attention, the feed-forward blocks, residual connections and layer norm — is unmodified from the original Transformer encoder. ViT doesn't introduce a vision-specific attention mechanism. The only vision-specific part of the whole model is the patch embedding step that turns pixels into a sequence in the first place.

Why this doesn't just work immediately

Convolutional networks have two properties baked directly into the convolution operation: locality (a filter only looks at a small neighborhood of pixels at a time) and translation equivariance (a learned filter for detecting an edge fires the same way no matter where in the image that edge appears). Those are strong assumptions about how images work, and they turn out to be correct often enough that CNNs get a real head start from architecture alone.

A Transformer has no such assumptions built in — self-attention treats the sequence of patches the same way regardless of position, so any notion of "nearby patches are more related" has to be learned from data rather than assumed. The paper's own finding: trained on a mid-sized dataset like ImageNet alone without heavy regularization, ViT scores a few points below comparable ResNets. Trained on enough data (JFT-300M, 300 million images), it matches or exceeds them. The paper's summary of this tradeoff: "large scale training trumps inductive bias."

Model sizes

Model	Layers	Hidden size	Heads	Parameters
ViT-Base	12	768	12	86M
ViT-Large	24	1024	16	307M
ViT-Huge	32	1280	16	632M

Note how closely this mirrors the original Transformer's own dimensions from the NLP side (d_model, number of heads, number of layers) — same knobs, same architecture, just scaled up and fed patches instead of word pieces.

The pattern, again

This is the same throughline as the tokenization post: attention itself doesn't care what the tokens represent. Text tokens, image patches — the mechanism from softmax(QKᵀ/√d_k)V doesn't change. What changes each time is the discretization step in front of it, and how much data it takes to make up for whatever structural assumptions that step doesn't encode.

---

References

• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale — Dosovitskiy et al. (2021)

Patch-embedding diagram above is original artwork made for this post.