Suppose you want to recognise an object in a photograph, wherever it appears in the frame. A CNN starts with a useful assumption for that job, applying the same filters to local groups of pixels. A Vision Transformer reads the image differently, turning patches into a sequence it can attend to. Building vision transformers from scratch lets us inspect what changes when the model has to learn more of the image's structure from examples.
We compare patch attention and local filters on small data. The CNN scores 85.2% against the ViT's 70.8%, with six times fewer parameters. Google's 2020 ViT matched leading CNNs after pretraining on 300 million images. Augmentation tests check how the model benefits from varied examples.
By the end of this post you will have built a Vision Transformer from scratch in PyTorch, trained it on CIFAR-10, and measured a from-scratch ViT trailing a far smaller CNN and the improvement from data augmentation.
Let's Build It
First, see the whole thing run. Click the badge to open the notebook in Google Colab:
The single idea behind a ViT is that an image is a sequence. Slice a 32x32 image into 4x4 patches and you get an 8x8 grid, 64 patches, which you unroll into a sequence of 64 tokens, exactly the shape a Transformer expects.
Here is a complete Vision Transformer. We will unpack every block afterwards.
import torch
import torch.nn as nn
class PatchEmbedding(nn.Module):
"""Split the image into patches and linearly embed each one.
A strided conv with kernel = stride = patch_size IS that linear projection."""
def __init__(self, img_size=32, patch=4, in_chans=3, dim=192):
super().__init__()
self.n_patches = (img_size // patch) ** 2
self.proj = nn.Conv2d(in_chans, dim, kernel_size=patch, stride=patch)
def forward(self, x):
x = self.proj(x) # (B, dim, 8, 8)
return x.flatten(2).transpose(1, 2) # (B, 64, dim)
class MultiHeadSelfAttention(nn.Module):
def __init__(self, dim, n_heads):
super().__init__()
self.n_heads, self.head_dim = n_heads, dim // n_heads
self.scale = self.head_dim ** -0.5
self.qkv = nn.Linear(dim, dim * 3)
self.proj = nn.Linear(dim, dim)
def forward(self, x):
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2]
attn = (q @ k.transpose(-2, -1)) * self.scale # how much each token looks at every other
attn = attn.softmax(dim=-1)
out = (attn @ v).transpose(1, 2).reshape(B, N, C)
return self.proj(out)
class EncoderBlock(nn.Module):
"""Pre-norm Transformer block: attention, then MLP, each wrapped in a residual."""
def __init__(self, dim, n_heads, mlp_ratio=2.0):
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = MultiHeadSelfAttention(dim, n_heads)
self.norm2 = nn.LayerNorm(dim)
hidden = int(dim * mlp_ratio)
self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(), nn.Linear(hidden, dim))
def forward(self, x):
x = x + self.attn(self.norm1(x))
x = x + self.mlp(self.norm2(x))
return x
class VisionTransformer(nn.Module):
def __init__(self, img_size=32, patch=4, n_classes=10, dim=192, depth=6, n_heads=6):
super().__init__()
self.patch_embed = PatchEmbedding(img_size, patch, 3, dim)
n = self.patch_embed.n_patches
self.cls_token = nn.Parameter(torch.zeros(1, 1, dim)) # the "summary" token
self.pos_embed = nn.Parameter(torch.zeros(1, n + 1, dim)) # learned positions
self.blocks = nn.ModuleList([EncoderBlock(dim, n_heads) for _ in range(depth)])
self.norm = nn.LayerNorm(dim)
self.head = nn.Linear(dim, n_classes)
nn.init.trunc_normal_(self.pos_embed, std=0.02)
nn.init.trunc_normal_(self.cls_token, std=0.02)
def forward(self, x):
x = self.patch_embed(x)
cls = self.cls_token.expand(x.shape[0], -1, -1)
x = torch.cat([cls, x], dim=1) + self.pos_embed
for blk in self.blocks:
x = blk(x)
return self.head(self.norm(x)[:, 0]) # classify from the class token
Train it on CIFAR-10 for 15 epochs with AdamW and a cosine schedule:
import torch.nn.functional as F
from torchvision import datasets, transforms
device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"
norm = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
train = datasets.CIFAR10("data", train=True, download=True,
transform=transforms.Compose([transforms.ToTensor(), norm]))
loader = torch.utils.data.DataLoader(train, batch_size=128, shuffle=True)
model = VisionTransformer().to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.05)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=15)
for epoch in range(15):
model.train()
for x, y in loader:
x, y = x.to(device), y.to(device)
loss = F.cross_entropy(model(x), y)
opt.zero_grad(); loss.backward(); opt.step()
sched.step()
A 1.8M-parameter ViT reaches around 71% test accuracy. Respectable, but a three-block CNN with six times fewer parameters gets to 85% on the same data. That is not a bug in our ViT. It is the whole point, and the next sections explain why.
What Just Happened?
An image becomes a sequence of patch tokens
A Transformer consumes a sequence of vectors. A CNN consumes a grid. The first job of a ViT is to turn the grid into a sequence, and it does so by cutting the image into non-overlapping patches.
The elegant trick is that "cut into patches and linearly project each one" is exactly a convolution whose kernel size equals its stride. A nn.Conv2d(3, dim, kernel_size=4, stride=4) slides a 4x4 filter in non-overlapping 4-pixel jumps, so every patch is seen once and mapped to a dim-vector. One convolution does all the cutting and embedding at once.
patch_embed = PatchEmbedding(img_size=32, patch=4, dim=192)
x = torch.randn(1, 3, 32, 32) # one CIFAR image
tokens = patch_embed(x)
print(tokens.shape) # torch.Size([1, 64, 192])
The 32x32 image is now 64 tokens of width 192. From here the ViT does not know or care that the input was ever a picture.
The class token and positional embeddings
Two learned vectors get bolted on before the Transformer sees the sequence.
The class token is a single extra vector prepended to the 64 patch tokens, giving 65. It carries no image content at the start; its job is to act as a running summary. As it flows through the attention layers it gathers information from every patch, and the final classifier reads only from this one token. Think of it as a designated note-taker that attends the whole meeting and writes the verdict.
Positional embeddings solve a problem convolution never had. Self-attention is permutation-invariant: shuffle the patches and the maths is unchanged, because attention treats the input as a set, not a sequence. A CNN knows where every pixel is by construction; a ViT has to be told. So we add a learned vector to each position. The model discovers spatial structure from data, and as we will see later, it really does learn that patch 9 sits next to patch 10.
cls = model.cls_token.expand(x.shape[0], -1, -1)
x = torch.cat([cls, x], dim=1) + model.pos_embed # 65 tokens, each tagged with its position
Self-attention: every patch looks at every other
Convolution is local: a 3x3 kernel sees nine neighbours. Self-attention is global from the very first layer: every token computes a weighted average over all the others, where the weights say how relevant each token is.
Each token emits three vectors: a query (what am I looking for), a key (what do I contain), and a value (what I will pass on if attended to). The attention weight from token i to token j is the scaled dot product of i's query with j's key, softmaxed across all j:
Where:
- Q, K, V are the query, key, and value matrices (one row per token)
- QKᵀ is the table of all token-to-token relevance scores
- √(dₖ) rescales the scores so the softmax does not saturate as the head dimension grows
- multiplying by V mixes each token's representation toward the tokens it attends to
This is the same self-attention used in the KV-cache post and the LLM architecture-tricks post; a ViT is a language-model encoder pointed at image patches. Splitting Q, K, V into several heads lets the model attend to different relationships in parallel, then concatenate the results.
Because attention is global, a ViT can relate the top-left and bottom-right corners of an image in a single layer, something a shallow CNN cannot. The price is that it has no built-in notion that adjacent patches are special: it must learn every spatial relationship from scratch.
The encoder block
The rest is plumbing borrowed wholesale from the Transformer. Each block normalises, applies attention, and adds a residual, then normalises, applies a small MLP, and adds another residual. Stack six of them and you have the ViT body.
We can see what a trained ViT learns to look at by reading the attention weights from the class token back onto the patch grid:
With no convolutional prior telling it where to look, the trained model still concentrates its class-token attention on a few salient patches of the foreground rather than spreading it evenly across the image. It learned which regions carry the signal.
Going Deeper
Inductive bias is the whole story
An inductive bias is an assumption a model makes before it sees any data. A CNN's biases (locality and translation equivariance) are baked into the architecture: weight sharing across positions is the assumption that a feature means the same thing everywhere. Those assumptions are correct for natural images, so a CNN starts the race halfway to the finish line.
A ViT has almost none of this. The only spatial prior it ships with is the positional embedding, and even that is learned, not fixed. Self-attention can in principle represent convolution, but it can equally represent a thousand wrong functions, and with limited data it has no reason to prefer the right one. The ViT must learn locality and translation equivariance from examples, which costs data.
Train both from scratch on CIFAR-10's 50,000 images and the gap is stark:
The CNN reaches 85.2% test accuracy; the ViT, with six times the parameters, stalls at 70.8%. Augmentation helps, but it does not rescue the ViT. Random crops and horizontal flips synthesise exactly the translation invariance a CNN gets for free, so they lift the ViT more than the CNN (+1.5 points versus +0.9). Yet the effect is modest: the ViT climbs to 72.3% while the CNN reaches 86.2%, and the roughly 14-point gap barely narrows. At 50,000 images the ViT's flexibility is a liability, not an asset, because there is not enough data to pin down the spatial structure the CNN assumes from the start. This is the regime the original paper warned about, and the crossover only arrives with far more data.
Plot accuracy against parameter count and the efficiency gap is even clearer: the CNN reaches higher accuracy with a small fraction of the ViT's parameters, because every one of those convolutional parameters is spent on an assumption that happens to be true.
Where the positional embeddings end up
If the ViT truly learns spatial structure, its positional embeddings should encode 2D layout. They do. Plot the cosine similarity of each position's embedding against every other position and reshape it back to the 8x8 grid:
Each position is most similar to itself and to its immediate neighbours, with similarity fading by distance. The model was never told the patches form a grid. It recovered the geometry from the data, which is exactly the locality a CNN gets for free.
LayerNorm, not BatchNorm
Our CNN used batch normalisation; the ViT uses layer normalisation. The difference matters. BatchNorm normalises each channel across the batch, which couples examples together and behaves differently at train and test time. LayerNorm normalises each token across its features, independent of the batch, which suits variable-length sequences and the pre-norm residual structure Transformers rely on. Swapping one for the other in a ViT usually hurts.
Hyperparameters to tune
| Parameter | Our value | Effect |
|---|
| Patch size | 4 | Smaller patches = more tokens, finer detail, quadratically more attention cost |
| Embedding dim | 192 | Token width; more capacity per token |
| Depth | 6 | Number of encoder blocks; deeper sees more abstract relationships |
| Heads | 6 | Parallel attention subspaces |
| Weight decay | 0.05 | ViTs need heavy regularisation to fight overfitting on small data |
| Augmentation | crop + flip | Synthesises the translation invariance the CNN has built in |
Patch size is the lever unique to ViTs. Halving it from 4 to 2 quadruples the token count, and attention cost grows with the square of the token count, so it is the fastest way to blow up your compute budget.
When not to reach for a ViT
Use a CNN, not a ViT, when you are training from scratch on a small or medium dataset (anything short of hundreds of thousands of images), when compute is tight, or when you need a model that generalises from few examples. A ViT earns its keep when you have very large datasets, when you can start from pre-trained weights (the usual case in practice), or when you need to model long-range relationships a shallow CNN cannot reach. For tabular data, neither: gradient-boosted trees and tabular foundation models own that ground.
Where This Comes From
The Vision Transformer was introduced by Alexey Dosovitskiy and colleagues at Google in their 2020 paper "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale". The title is the thesis: a 16x16 image patch plays the role a word plays in a language model.
The paper's construction is the one we built. An image x ∈ ℝ^(H × W × C) is reshaped into N flattened patches, each of size P² · C, then linearly projected to dimension D:
Where E is the shared patch-embedding projection, x_class is the prepended class token, and E_pos are the learned positional embeddings. Each of the L encoder layers then applies multi-head self-attention (MSA) and an MLP, each with a residual connection and pre-layer-normalisation (LN):
The headline finding was about scale. The authors were blunt about the data dependence:
"When trained on mid-sized datasets such as ImageNet without strong regularization, these models yield modest accuracies of a few percentage points below ResNets of comparable size. This seemingly discouraging outcome may be expected: Transformers lack some of the inductive biases inherent to CNNs, such as translation equivariance and locality, and therefore do not generalize well when trained on insufficient amounts of data."
-- Dosovitskiy et al. (2020)
But when they pre-trained on JFT-300M, a Google-internal dataset of 300 million images, the picture flipped:
"However, the picture changes if the models are trained on larger datasets (14M-300M images). We find that large scale training trumps inductive bias."
-- Dosovitskiy et al. (2020)
That sentence, "large scale training trumps inductive bias", is the reason ViTs took over. Our CIFAR-10 experiment lives firmly in the first regime, where the CNN's biases win. The crossover only arrives with orders of magnitude more data than a laptop sees.
From ViT to today
The original ViT needed Google-scale data, which few teams have. The fix came fast. DeiT (Touvron et al., 2021) showed a ViT could be trained to competitive ImageNet accuracy on a single 8-GPU node in a few days, using aggressive augmentation and a distillation token, no JFT-300M required. The same augmentation idea you can watch helping in our training curves is what made ViTs practical.
Since then the patch-token recipe spread everywhere: hierarchical variants like Swin Transformers reintroduced a touch of locality, masked autoencoders (He et al., 2021) made ViTs self-supervised, and the patch-embedding front end became the standard way vision-language models such as the open VLMs of 2024-2026 turn pixels into tokens a language model can read. The model you built is the seed of all of it.
Further reading
- Attention Visualizer: The mechanism a ViT applies to image patches instead of words. Turn the causal mask off to get the bidirectional attention an encoder actually uses
- 2D Convolution Visualizer: See the local, weight-shared operation a ViT deliberately discards in favour of global attention
- CNN Layer Visualizer: The pooling hierarchy a ViT does away with. Stack conv and pool layers to watch a CNN trade resolution for receptive field, which patch embedding replaces in a single step
- Overfitting Explorer: ViTs overfit small datasets without heavy regularisation; explore the bias-variance tradeoff behind that
- Classification Metrics Calculator: Evaluate your ViT's CIFAR-10 predictions with precision, recall, and a confusion matrix
Related Posts
Frequently Asked Questions
A Vision Transformer splits an image into fixed-size patches, treats each patch as a token, and processes the sequence with a standard Transformer encoder, with no convolution. Use one when you have a very large training set or pre-trained weights to fine-tune; for training from scratch on small or medium datasets, a CNN is usually the better choice.
Why does a ViT need more data than a CNN?
A CNN has built-in inductive biases (locality and translation equivariance) that match how natural images work, so it generalises from fewer examples. A ViT has almost no such priors and must learn spatial structure from data, which only pays off once the dataset is large enough that flexibility beats built-in assumptions.
What is the class token for?
The class token is a single learnable vector prepended to the patch tokens. It accumulates information from every patch through the attention layers, and the final classifier reads only from it. It is a learned summary of the whole image.
Self-attention is permutation-invariant, so without positional information a ViT cannot tell where a patch came from. Adding a learned positional embedding to each token lets the model recover 2D spatial layout, which a CNN gets for free from its grid structure.
Smaller patches produce more tokens and capture finer detail, but attention cost grows with the square of the token count, so halving the patch size roughly quadruples compute and memory. Larger patches are cheaper but coarser. It is the main accuracy-versus-cost lever unique to ViTs.
Rarely from scratch. On small datasets the CNN's inductive biases win, as our CIFAR-10 experiment shows. A ViT can compete on small data when it starts from large-scale pre-trained weights and is fine-tuned, which is how ViTs are used in practice.