Module 15 — Vision Transformer (ViT)

Introduction

The Vision Transformer (ViT) applies the Transformer architecture to computer vision by treating an image as a sequence of patches, similar to how words are treated as tokens in Natural Language Processing (NLP).

Instead of using convolutional filters like CNNs, ViT divides an image into fixed-size patches, converts each patch into an embedding, adds positional information, and processes the sequence using a standard Transformer encoder.

Since its introduction in the paper "An Image is Worth 16×16 Words" (2020), Vision Transformers have become one of the most influential architectures in computer vision.

Applications include:

• Image Classification
• Image Retrieval
• Object Detection
• Image Segmentation
• Medical Imaging
• Remote Sensing
• Autonomous Driving
• Vision-Language Models (VLMs)

In this module, you'll learn:

• Image Patch
• Patch Embedding
• CLS Token
• Position Embedding
• Vision Encoder
• Image Classification
• DeiT
• Swin Transformer
• ConvNeXt
• CIFAR-10 Classification

Vision Transformer Architecture

Unlike CNNs, ViT converts an image into a sequence of patches before applying Transformer encoder layers.

textCopy
1               Input Image
2                    │
3          Split into Patches
4                    │
5           Flatten Each Patch
6                    │
7            Linear Projection
8                    │
9           Patch Embeddings
10                    │
11          Add CLS Token + Position
12                    │
13          Transformer Encoder
14                    │
15         CLS Representation
16                    │
17       Classification Head
18                    │
19            Predicted Class

1. Image Patch

What is an Image Patch?

Instead of processing the entire image at once, ViT divides it into small non-overlapping patches.

Example

Suppose an image has size

textCopy
1224 × 224 × 3

Using a patch size of

textCopy
116 × 16

The number of patches becomes

textCopy
1224 / 16 = 14
2
314 × 14 = 196 patches

Each patch is treated like a word token.

Visualization

textCopy
1+----+----+----+----+
2| P1 | P2 | P3 | P4 |
3+----+----+----+----+
4| P5 | P6 | P7 | P8 |
5+----+----+----+----+
6| P9 |... |... |... |
7+----+----+----+----+

Extract Patches in PyTorch

pythonCopy
1import torch
2
3image = torch.randn(1, 3, 224, 224)
4
5patch_size = 16
6
7patches = image.unfold(2, patch_size, patch_size)
8
9patches = patches.unfold(3, patch_size, patch_size)
10
11print(patches.shape)

Output

textCopy
1torch.Size([1, 3, 14, 14, 16, 16])

2. Patch Embedding

Each image patch is flattened and projected into an embedding vector using a linear layer.

Formula

textCopy
1Patch
2
3↓
4
5Flatten
6
7↓
8
9Linear Layer
10
11↓
12
13Embedding

Suppose

textCopy
1Patch Size = 16×16
2
3Channels = 3
4
5Flatten Size
6
716 × 16 × 3 = 768

Each patch becomes a vector of dimension

textCopy
1768

PyTorch Example

pythonCopy
1import torch
2import torch.nn as nn
3
4embedding = nn.Linear(
5    16 * 16 * 3,
6    768
7)
8
9patch = torch.randn(1, 16 * 16 * 3)
10
11output = embedding(patch)
12
13print(output.shape)

Output

textCopy
1torch.Size([1, 768])

3. CLS Token

Like BERT, ViT prepends a Classification (CLS) Token.

Architecture

textCopy
1CLS
2
3↓
4
5Patch1
6
7↓
8
9Patch2
10
11↓
12
13Patch3
14
15↓
16
17...

After passing through the Transformer encoder, the CLS token contains information about the entire image.

The classification head uses only this token.

PyTorch Example

pythonCopy
1import torch
2
3cls = torch.randn(1, 1, 768)
4
5patch_embeddings = torch.randn(
6    1,
7    196,
8    768
9)
10
11sequence = torch.cat(
12    [cls, patch_embeddings],
13    dim=1
14)
15
16print(sequence.shape)

Output

textCopy
1torch.Size([1, 197, 768])

4. Position Embedding

Transformers do not understand spatial positions automatically.

Therefore, ViT adds learnable position embeddings.

Formula

textCopy
1Input =
2
3Patch Embedding
4
5+
6
7Position Embedding

Example

pythonCopy
1position = torch.randn(
2    1,
3    197,
4    768
5)
6
7output = sequence + position
8
9print(output.shape)

5. Vision Encoder

The sequence is passed through multiple Transformer Encoder layers.

Each encoder contains:

• Multi-Head Self-Attention
• Feed Forward Network
• Layer Normalization
• Residual Connections

Architecture

textCopy
1Patch Embeddings
2
3↓
4
5Encoder Layer 1
6
7↓
8
9Encoder Layer 2
10
11↓
12
13Encoder Layer 3
14
15↓
16
17...
18
19↓
20
21CLS Output

Using PyTorch Transformer Encoder

pythonCopy
1import torch.nn as nn
2
3encoder_layer = nn.TransformerEncoderLayer(
4    d_model=768,
5    nhead=12,
6    batch_first=True
7)
8
9encoder = nn.TransformerEncoder(
10    encoder_layer,
11    num_layers=12
12)
13
14output = encoder(sequence)
15
16print(output.shape)

6. Image Classification

The CLS token is fed into a linear classifier.

Architecture

textCopy
1CLS Token
2
3↓
4
5Linear Layer
6
7↓
8
9Softmax
10
11↓
12
13Class Prediction

PyTorch Example

pythonCopy
1classifier = nn.Linear(
2    768,
3    10
4)
5
6prediction = classifier(
7    output[:, 0]
8)
9
10print(prediction.shape)

Output

textCopy
1torch.Size([1, 10])

7. DeiT (Data-efficient Image Transformer)

What is DeiT?

DeiT improves Vision Transformer training using knowledge distillation.

Instead of requiring extremely large datasets, DeiT learns from a CNN teacher model.

Features

• Distillation Token
• Faster convergence
• Better performance on smaller datasets
• Efficient training

Load DeiT

pythonCopy
1from transformers import AutoImageProcessor
2from transformers import ViTForImageClassification
3
4processor = AutoImageProcessor.from_pretrained(
5    "facebook/deit-base-patch16-224"
6)
7
8model = ViTForImageClassification.from_pretrained(
9    "facebook/deit-base-patch16-224"
10)

8. Swin Transformer

What is Swin Transformer?

Swin Transformer introduces Shifted Window Attention to improve scalability.

Instead of computing attention over the entire image, it computes attention within local windows.

Advantages

• Lower computational cost
• Hierarchical feature maps
• Better for detection and segmentation
• Supports high-resolution images

Architecture

textCopy
1Image
2
3↓
4
5Window Partition
6
7↓
8
9Window Attention
10
11↓
12
13Shift Windows
14
15↓
16
17Merge Features

Load Swin

pythonCopy
1from transformers import SwinForImageClassification
2
3model = SwinForImageClassification.from_pretrained(
4    "microsoft/swin-tiny-patch4-window7-224"
5)

9. ConvNeXt

What is ConvNeXt?

ConvNeXt is not a Transformer. It is a convolutional neural network redesigned using ideas inspired by Vision Transformers.

Key ideas

• Large kernels
• Layer normalization
• Modern training recipes
• Hierarchical architecture

Advantages

• Strong CNN baseline
• Competitive with ViT
• Efficient inference

Load ConvNeXt

pythonCopy
1from transformers import ConvNextForImageClassification
2
3model = ConvNextForImageClassification.from_pretrained(
4    "facebook/convnext-tiny-224"
5)

Comparison of Vision Models

Practice — CIFAR-10 Classification

The following example fine-tunes a Vision Transformer on the CIFAR-10 dataset.

pythonCopy
1from transformers import (
2    AutoImageProcessor,
3    ViTForImageClassification,
4    Trainer,
5    TrainingArguments
6)
7from datasets import load_dataset
8
9dataset = load_dataset("cifar10")
10
11processor = AutoImageProcessor.from_pretrained(
12    "google/vit-base-patch16-224"
13)
14
15model = ViTForImageClassification.from_pretrained(
16    "google/vit-base-patch16-224",
17    num_labels=10
18)
19
20def preprocess(example):
21    image = example["img"].convert("RGB")
22    inputs = processor(image, return_tensors="pt")
23    example["pixel_values"] = inputs["pixel_values"][0]
24    return example
25
26dataset = dataset.map(preprocess)
27
28training_args = TrainingArguments(
29    output_dir="./vit-cifar10",
30    per_device_train_batch_size=16,
31    num_train_epochs=3,
32    evaluation_strategy="epoch",
33    save_strategy="epoch"
34)
35
36trainer = Trainer(
37    model=model,
38    args=training_args,
39    train_dataset=dataset["train"],
40    eval_dataset=dataset["test"]
41)
42
43trainer.train()

What You'll Learn

• Load the CIFAR-10 dataset
• Preprocess images using an image processor
• Fine-tune a pretrained Vision Transformer
• Evaluate image classification performance

Mini Project — Image Prediction

pythonCopy
1from PIL import Image
2from transformers import (
3    AutoImageProcessor,
4    ViTForImageClassification
5)
6import torch
7
8image = Image.open("cat.jpg").convert("RGB")
9
10processor = AutoImageProcessor.from_pretrained(
11    "google/vit-base-patch16-224"
12)
13
14model = ViTForImageClassification.from_pretrained(
15    "google/vit-base-patch16-224"
16)
17
18inputs = processor(
19    images=image,
20    return_tensors="pt"
21)
22
23with torch.no_grad():
24    outputs = model(**inputs)
25
26prediction = outputs.logits.argmax(dim=-1)
27
28print("Predicted Class ID:", prediction.item())

Module Summary

After completing this module, you will be able to:

• Explain how Vision Transformers process images as sequences of patches.
• Understand image patch extraction and patch embeddings.
• Describe the role of the CLS token and positional embeddings.
• Build and use a Vision Transformer encoder for image classification.
• Compare Vision Transformer, DeiT, Swin Transformer, and ConvNeXt.
• Fine-tune pretrained Vision Transformer models on image datasets.
• Build image classification systems using Hugging Face Transformers and PyTorch.
• Choose the appropriate vision architecture based on accuracy, efficiency, and downstream tasks.

Next Module: Module 16 – Multimodal Transformers, where you'll learn CLIP, BLIP, Flamingo, LLaVA, Qwen2.5-VL, Florence-2, vision-language modeling, image captioning, visual question answering (VQA), document understanding, and multimodal AI applications.