Vision Transformer, Part 4 — Building the ViT on CIFAR-10

Abstract. Everything from the first three parts, now in about 200 lines of PyTorch. This part builds the Vision Transformer end to end — dataset pipeline, patch embedding, encoder, classification head, and training loop — and trains it from scratch on CIFAR-10 to ~81% accuracy on a single GPU. It closes with a full per-class and confusion-matrix analysis of where the model succeeds, where it fails, and why.

Part 4 of the series. Parts 1–3 explained the mechanics with no code (embeddings, attention, the full block); this part builds and trains the model. Part 5 — Do overlapping patches help? A controlled experiment then runs a controlled experiment on overlapping patches.

In this 6-part series of Vision Transformer:

1. Part 1 — Embeddings: from pixels to a token sequence
2. Part 2 — Attention: how tokens relate to each other (single head)
3. Part 3 — The full block: multi-head, LayerNorm, residuals, MLP, and the prediction
4. Part 4 — Building the ViT in PyTorch, trained on CIFAR-10 — you are here
5. Part 5 — Do overlapping patches help? A controlled experiment
6. Part 6 — How it learns: loss, backprop, the chain rule, the weight update

Why this series

I wanted two things out of this project:

1. Understand a Vision Transformer by building every piece myself — no timm, no pretrained weights — on a dataset small enough to train on one GPU.
2. Run a clean experiment (Part 5): what happens if the image patches overlap instead of tiling the image edge-to-edge?

This part is the build. By the end you’ll have a ~1.2M-parameter ViT that reaches 81% on CIFAR-10 from scratch, and — more importantly — you’ll understand every number.

The one-paragraph idea

A Transformer eats a sequence of tokens. An image isn’t a sequence — so a ViT makes it one. Chop the image into fixed-size patches, flatten each patch and linearly project it into an embedding (a “visual token”). Prepend a learnable [CLS] token, add positional embeddings so the model knows where each patch came from, and feed the whole sequence through a standard Transformer encoder. The final state of the [CLS] token goes to a small classifier head.

Here’s the whole model as a functional block diagram:

Figure 1: The full ViT Forward Pass

Notice the shape contract flowing left to right: an image becomes [B, 65, 128] — 65 tokens of 128 dimensions — which is exactly the [batch, seq_len, dim] shape a text Transformer expects. Reaching that shape is the whole trick.

The dataset: CIFAR-10

CIFAR-10 is 60,000 tiny 32×32 RGB images across 10 classes (airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck), split 50k train / 10k test. It’s small enough to train quickly but real enough to be interesting — and it downloads automatically via torchvision.

The pipeline is deliberately simple: augment the training set (random crop + horizontal flip — ViTs overfit small data, so this matters), and normalize both splits with CIFAR-10’s channel statistics.

CIFAR10_MEAN = (0.4914, 0.4822, 0.4465)
CIFAR10_STD  = (0.2470, 0.2435, 0.2616)

def build_transforms(train: bool):
    if train:
        return transforms.Compose([
            transforms.RandomCrop(32, padding=4),   # small translations
            transforms.RandomHorizontalFlip(),       # left/right invariance
            transforms.ToTensor(),
            transforms.Normalize(CIFAR10_MEAN, CIFAR10_STD),
        ])
    return transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(CIFAR10_MEAN, CIFAR10_STD),
    ])

Each batch is a tensor of shape [B, 3, 32, 32] with integer labels [B]. That [B, 3, 32, 32] is the input contract for everything that follows.

A caveat worth stating up front: 32×32 is tiny. To your own eyes many of these images are ambiguous — some CIFAR “cats” could pass for dogs. This resolution limit turns out to be the ceiling on everything we do — a point that returns with force in Part 5.

Step 1 — Patch embedding: turning an image into tokens

This is the heart of a ViT. We need [B, 3, 32, 32] → [B, 65, 128] (64 patch tokens + 1 CLS, each 128-dimensional).

Get Chandrashekhar Padole’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

The elegant trick: a single Conv2d with kernel_size = stride = patch_size does "cut into patches + flatten + linearly project" in one fused operation. Here is the data flow for a single patch:

Figre 2: 4x4 patch to one token transformation

A 4×4 patch is 16 pixels × 3 channels = 48 raw numbers; the conv’s weight tensor is [128, 3, 4, 4] = a learned [48 → 128] matrix. Each output dimension is a dot product of those 48 values with a learned weight vector — identical to flatten → Linear(48, 128), just fused and faster.

class PatchEmbed(nn.Module):
    def __init__(self, img_size=32, patch_size=4, in_chans=3, embed_dim=128):
        super().__init__()
        self.proj = nn.Conv2d(in_chans, embed_dim,
                              kernel_size=patch_size, stride=patch_size)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))

    def forward(self, x):
        x = self.proj(x)                     # [B, 128, 8, 8]
        x = x.flatten(2).transpose(1, 2)     # [B, 64, 128]
        cls = self.cls_token.expand(x.shape[0], -1, -1)
        x = torch.cat((cls, x), dim=1)       # [B, 65, 128]
        x = x + self.pos_embed               # add positional info
        return x

Two things worth internalizing:

• The 48 → 128 projection is not compression or new information — it re-represents the raw pixels in a wider, learned feature space (edges, textures, color blobs), trained end-to-end by backprop. It’s the visual analog of an LLM’s token embedding, except the vector is computed from pixels rather than looked up by token id.
• Positional embeddings are learned ([65, 128]). Self-attention is permutation-invariant — without positions the model couldn't tell a top-left patch from a bottom-right one.

Step 2 — The Transformer encoder

The token sequence flows through 6 identical blocks. Each block has two sub-layers — multi-head self-attention and an MLP — each wrapped in a pre-norm residual. Here’s one block:

Figure 3: Stack of 6 layers, each layer has one block shown in left above

Two things make this stack work:

• Every layer is shape-preserving — [B, 65, 128] in, [B, 65, 128] out. That's precisely why blocks can be stacked to any depth without adapters between them.
• The 6 layers are not weight-tied. They share a structure, not parameters: each holds its own 198,272 weights (6 × 198,272 = the encoder’s 1,189,632). So early layers can learn local edge/texture relations while later ones assemble object-level structure — impossible if the weights were shared. Contrast this with the MLP inside a layer, which genuinely is shared across all 65 tokens.

Self-attention is where the “understanding” happens: every token forms Query, Key, and Value vectors and attends to all other tokens. Unlike GPT, there is no causal mask — a ViT classifies a whole image that exists all at once, so a bottom-right patch should see the top-left one. The full N×N attention matrix is computed in one shot:

Figure 4: Single-head self-attention

def forward(self, x):
    B, N, D = x.shape
    qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
    qkv = qkv.permute(2, 0, 3, 1, 4)
    q, k, v = qkv[0], qkv[1], qkv[2]           # each [B, heads, N, head_dim]

    attn = (q @ k.transpose(-2, -1)) * self.scale   # [B, heads, N, N]
    attn = attn.softmax(dim=-1)                      # rows sum to 1
    out = attn @ v                                   # [B, heads, N, head_dim]

    out = out.transpose(1, 2).reshape(B, N, D)       # concat heads
    return self.proj(out)

And the block, in code, is exactly the diagram above:

def forward(self, x):
    x = x + self.attn(self.norm1(x))   # patches share information
    x = x + self.mlp(self.norm2(x))    # per-token feature refinement
    return x

A key property: the MLP is token-wise and weight-shared — the same MLP is applied to every token independently. And the shapes are invariant to sequence length: N (tokens) can change, but D = 128 per token never does. (That invariance is exactly what lets Part 5's overlap experiment "just work" downstream, no matter how many patches we create.)

Step 3 — Classification head and full assembly

The head reads only the CLS token (index 0) — its final state is the whole-image summary the encoder accumulated via attention — applies a LayerNorm, and projects to 10 class logits.

class ViT(nn.Module):
    def forward(self, x):
        x = self.patch_embed(x)   # [B, 65, 128]
        x = self.encoder(x)       # [B, 65, 128]
        cls = x[:, 0]             # [B, 128]  — CLS token only
        cls = self.norm(cls)
        return self.head(cls)     # [B, 10]

The full model is ~1.2M parameters — patch embedding 14.7K, encoder 1.19M, head 1.5K. Tiny by modern standards, trainable in minutes.

One loss per image. Only the CLS token is supervised — exactly one cross-entropy loss per image, between the CLS logits and the label. The 64 patch tokens never get their own targets; they’re trained indirectly, via gradients flowing back through attention because CLS attended to them. (Contrast GPT, which has a loss term at every token position.)

Step 4 — Training

The recipe is standard and unremarkable — which is the point.

Figure 5: Batch Training Loop

• Loss: cross-entropy with label smoothing 0.1
• Optimizer: AdamW (lr 3e-4, weight decay 0.05) — decoupled weight decay, the same optimizer family used for GPT-style models
• Schedule: linear warmup → cosine decay to zero
• Extras: mixed precision (AMP), gradient clipping at 1.0
• 50 epochs, batch size 128

for images, labels in train_loader:
    with torch.amp.autocast("cuda", enabled=use_amp):
        logits = model(images)
        loss = criterion(logits, labels)
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer); scaler.update(); scheduler.step()

Putting it together: batch → image → token

It’s easy to lose track of what happens at which level. This nested view shows all three at once — the batch loop on the outside (where loss and weight updates live), one image inside it, and the per-token operations innermost.

Attention is drawn here as a single head to keep the picture clean. The model we built actually uses 4 heads of 32 dims — same operation, just partitioned (see Step 2). Everything else in the diagram is unchanged by that.

Figure 6: Detailed view of loss driven learning

Three things this makes obvious:

• Weights update once per batch — not per image and certainly not per token. A batch of 128 images produces one scalar loss and one optimizer step.
• Attention is the only operation that mixes tokens. LayerNorm, the Q/K/V projection, and the MLP are all applied to each token independently with shared weights; attention is where the 65 tokens of an image finally talk to each other. And they only ever talk within their own image — tokens never attend across images in the batch.
• Each image contributes exactly one loss term, from its CLS token. The 64 patch tokens are never given targets — they learn only through gradients flowing back from that single prediction.

Step 5 — Measuring classification

A single accuracy number hides a lot. Here’s the full picture for the plain model on the 10,000-image test set.

Overall test accuracy: 81.25%. For a small from-scratch ViT on CIFAR-10 with no pretraining, that’s right in the expected range. (CNNs reach ~95% here; ViTs are notoriously data-hungry without pretraining — a known trade-off, not a bug.)

Per-class accuracy, worst to best:

The confusion matrix (rows = true class, columns = predicted) shows the model makes sensible mistakes — a sign it’s genuinely learning structure:

              airp  auto  bird   cat  deer   dog  frog  hors  ship  truc
  airplane     861    17    14    15     7     1     6     9    42    28
  automobile    19   900     2     5     0     0     1     3     7    63
  bird          52     5   746    41    44    33    33    32     3    11
  cat           24     5    47   633    36   164    42    25    11    13
  deer          12     3    63    31   749    20    32    79     8     3
  dog           12     5    23   121    24   750    14    44     4     3
  frog           5     6    37    41    16    19   860     8     2     6
  horse         24     1    16    18    20    41     1   870     4     5
  ship          65    20     9     6     2     1     3     1   875    18
  truck         32    51     3     5     1     2     2     8    15   881

The single most common confusion is cat → dog, 164 times. Rigid, distinctive shapes (automobile, truck, ship) are easy; semantically similar animal pairs (cat↔dog, deer↔horse) are hard — exactly what you’d expect, and exactly the kind of confusion a higher resolution would help resolve.

Remember that number, 164 — it’s the yardstick for Part 5, where we ask whether overlapping the patches can shrink it.

What’s next

We now have a complete, working ViT: image → patches → embeddings → attention → CLS → prediction, at 81.25% on CIFAR-10.

In Part 5, we change exactly one thing — let the patches overlap — and run a controlled sweep to find out whether it helps. The result is a clean lesson in diminishing returns… and a cautionary tale about a confound that nearly inverted the conclusion.

Code and full experiment notes were built and documented from scratch — the goal was never to beat a benchmark, but to understand every number.