Vision Transformer (ViT) in Computer Vision - Deep Dive

Overview - Vision Transformer (ViT)

What is it?

Vision Transformer (ViT) is a type of machine learning model designed to understand images by breaking them into small patches and processing these patches like words in a sentence. Instead of using traditional methods that look at pixels in grids, ViT treats image patches as a sequence and uses a transformer architecture originally made for language. This approach allows the model to learn complex patterns and relationships in images. It has shown strong performance in image recognition tasks.

Why it matters

ViT exists because traditional image models like convolutional neural networks (CNNs) have limits in capturing long-range relationships in images. Without ViT, models might miss important connections between distant parts of an image, reducing accuracy. ViT enables better understanding of global image context, improving tasks like object recognition and classification. This helps technologies like self-driving cars, medical imaging, and photo search become more accurate and reliable.

Where it fits

Before learning ViT, you should understand basic image processing and convolutional neural networks (CNNs). Knowing how transformers work in language models helps too. After ViT, learners can explore advanced vision transformers, hybrid models combining CNNs and transformers, and applications in video and 3D data.

Mental Model

Core Idea

Vision Transformer breaks an image into patches and treats them like words in a sentence, using transformer attention to learn relationships across the whole image.

Think of it like...

Imagine reading a picture like a book made of small tiles, where each tile is a word. Instead of reading line by line, you look at all tiles at once and understand how they connect to tell the story.

Image
┌───────────────┐
│               │
│  ┌───┐ ┌───┐  │
│  │P1 │ │P2 │  │  P1, P2, ... are patches
│  └───┘ └───┘  │
│  ┌───┐ ┌───┐  │
│  │P3 │ │P4 │  │
│  └───┘ └───┘  │
│               │
└───────────────┘

Patches → Flatten → Linear Projection → Add Position Embeddings → Transformer Encoder → Classification Head

Build-Up - 7 Steps

1

Foundation

Understanding Image Patches

Concept: Images can be split into smaller square pieces called patches to simplify processing.

An image is a grid of pixels. Instead of looking at the whole image at once, we cut it into small patches, like cutting a photo into puzzle pieces. Each patch contains a small part of the image, for example, a 16x16 pixel square. These patches are easier to handle and can be processed one by one or as a sequence.

Result

The image is now represented as a list of patches, each containing pixel data from a small area.

Understanding patches helps us convert images into a format that transformers, which work on sequences, can process.

2

Foundation

Basics of Transformer Architecture

Concept: Transformers use attention to focus on important parts of a sequence and learn relationships between elements.

Transformers were first made for language, where words in a sentence relate to each other. They use a mechanism called self-attention to weigh how much each word matters to others. This helps the model understand context and meaning. The transformer has layers that process sequences and learn complex patterns.

Result

A powerful way to analyze sequences by focusing on relevant parts and ignoring less important ones.

Knowing how transformers work with sequences is key to applying them to image patches.

3

Intermediate

Converting Patches to Tokens

🤔Before reading on: do you think image patches are used as raw pixels or transformed before input to the transformer? Commit to your answer.

Concept: Each image patch is flattened and projected into a vector called a token, similar to word embeddings in language models.

After cutting the image into patches, each patch's pixels are flattened into a single long vector. Then, a linear layer (a simple math operation) converts this vector into a fixed-size token. These tokens represent patches in a way the transformer can understand. Position embeddings are added to keep track of where each patch was in the original image.

Result

A sequence of tokens representing image patches with position information, ready for the transformer.

Transforming patches into tokens bridges the gap between images and sequence models, enabling the use of transformers.

4

Intermediate

Self-Attention Across Image Patches

🤔Before reading on: does self-attention in ViT only look at nearby patches or all patches globally? Commit to your answer.

Concept: Self-attention allows the model to consider relationships between all patches, not just neighbors.

In ViT, self-attention computes how much each patch should pay attention to every other patch. This means the model can learn connections between distant parts of the image, like how a dog's head relates to its tail. This global view helps capture the full context of the image.

Result

The model understands complex patterns by relating all parts of the image simultaneously.

Global attention is what gives ViT an edge over traditional models that focus only on local areas.

5

Intermediate

Training Vision Transformer Models

Concept: ViT models require large datasets and careful training to perform well.

Because ViT has many parameters and no built-in image-specific biases like CNNs, it needs a lot of training data to learn effectively. Training involves feeding many labeled images, adjusting model weights to reduce errors. Techniques like data augmentation and regularization help prevent overfitting.

Result

A trained ViT model that can classify images accurately on new data.

Knowing the training needs helps set realistic expectations and guides dataset preparation.

6

Advanced

Comparing ViT to Convolutional Networks

🤔Before reading on: do you think ViT always outperforms CNNs on small datasets? Commit to your answer.

Concept: ViT and CNNs have different strengths; ViT excels with large data and global context, CNNs with local features and smaller data.

CNNs use filters that scan small areas and build up features hierarchically, which works well with limited data. ViT treats images as sequences and learns global relationships but needs more data to avoid overfitting. Hybrid models combine both approaches. Understanding these differences helps choose the right model for a task.

Result

Clear understanding of when to use ViT or CNNs based on data and task.

Knowing model strengths prevents misuse and guides better architecture choices.

7

Expert

Scaling and Efficiency in Vision Transformers

🤔Before reading on: do you think increasing patch size always improves ViT performance? Commit to your answer.

Concept: Scaling ViT involves trade-offs between patch size, model size, and computational cost, with innovations to improve efficiency.

Larger patches reduce sequence length, making computation cheaper but lose fine details. Smaller patches capture more detail but increase computation. Techniques like hierarchical transformers, sparse attention, and distillation help scale ViT efficiently. Understanding these trade-offs is crucial for deploying ViT in real-world systems with limited resources.

Result

Ability to design and optimize ViT models balancing accuracy and efficiency.

Recognizing scaling trade-offs is key to practical ViT applications and innovation.

Under the Hood

ViT works by first splitting an image into fixed-size patches, flattening each patch into a vector, and projecting it into a token embedding. Position embeddings are added to retain spatial information. These tokens form a sequence input to a standard transformer encoder, which uses multi-head self-attention layers to compute relationships between all patches simultaneously. The output tokens are pooled and passed to a classification head. Unlike CNNs, ViT does not use convolutional filters but relies entirely on attention mechanisms to learn image features.

Why designed this way?

ViT was designed to leverage the success of transformers in language, applying their powerful sequence modeling to images. Traditional CNNs have strong inductive biases like locality and translation invariance, which help with small data but limit global context. ViT removes these biases to allow learning more flexible representations, especially when large datasets are available. This design choice trades off data efficiency for model expressiveness and scalability.

Image → Patch Split → Flatten → Linear Projection → + Position Embeddings → Transformer Encoder (Multi-head Self-Attention + Feedforward Layers) → Classification Token → MLP Head → Output

┌─────────────┐
│   Image     │
└─────┬───────┘
      │
┌─────▼───────┐
│  Patching   │
└─────┬───────┘
      │
┌─────▼───────┐
│ Flattening  │
└─────┬───────┘
      │
┌─────▼───────┐
│ Linear Proj │
└─────┬───────┘
      │
┌─────▼───────────────┐
│ Add Position Embed   │
└─────┬───────────────┘
      │
┌─────▼───────────────┐
│ Transformer Encoder  │
│ (Self-Attention + FF)│
└─────┬───────────────┘
      │
┌─────▼───────────────┐
│ Classification Head  │
└───────────────┬──────┘
                │
           Output Label

Myth Busters - 4 Common Misconceptions

Quick: Does ViT use convolutional filters like CNNs? Commit to yes or no before reading on.

Common Belief: ViT is just a CNN with a different name and uses convolutional filters internally.

Reality: ViT does not use any convolutional filters; it relies entirely on transformer self-attention mechanisms to process image patches.

Why it matters: Confusing ViT with CNNs leads to misunderstanding its strengths and weaknesses, causing poor model design and training choices.

Quick: Can ViT perform well with very small datasets without special techniques? Commit to yes or no before reading on.

Common Belief: ViT works well on small datasets just like CNNs without extra tricks.

Reality: ViT generally requires large datasets or pretraining because it lacks CNNs' built-in image biases, making it prone to overfitting on small data.

Why it matters: Ignoring data requirements causes poor model performance and wasted resources.

Quick: Does ViT only consider local patch neighbors during attention? Commit to yes or no before reading on.

Common Belief: ViT's attention is limited to nearby patches to reduce computation.

Reality: ViT's self-attention is global, meaning every patch attends to every other patch in the image.

Why it matters: Misunderstanding attention scope leads to wrong assumptions about ViT's ability to capture global context.

Quick: Is increasing patch size always better for ViT accuracy? Commit to yes or no before reading on.

Common Belief: Larger patches always improve ViT performance by simplifying the input.

Reality: Larger patches reduce detail and can hurt accuracy; smaller patches capture more detail but increase computation.

Why it matters: Wrong patch size choices degrade model accuracy or efficiency.

Expert Zone

1

ViT's lack of convolutional inductive biases means it learns spatial relationships purely from data, which can be both a strength and a weakness depending on dataset size.

2

Position embeddings in ViT are crucial; without them, the model loses spatial order information, making it unable to understand image structure.

3

The classification token (CLS token) in ViT acts as a summary of the entire image, and its learned representation is what the final classifier uses.

When NOT to use

ViT is not ideal for small datasets or real-time applications with limited compute due to its data hunger and computational cost. In such cases, CNNs or hybrid CNN-transformer models are better alternatives. For tasks requiring fine-grained local feature extraction, CNNs may outperform ViT.

Production Patterns

In production, ViT models are often pretrained on large datasets like ImageNet or JFT and then fine-tuned on specific tasks. Hybrid models combining CNN feature extractors with transformer layers are common to balance efficiency and accuracy. Techniques like knowledge distillation and pruning are used to reduce model size and latency.

Connections

Natural Language Processing Transformers

ViT builds directly on the transformer architecture developed for language, applying the same sequence modeling to image patches.

Understanding language transformers helps grasp how ViT processes image data as sequences, showing the power of attention beyond text.

Convolutional Neural Networks (CNNs)

ViT and CNNs are alternative approaches to image understanding, with ViT focusing on global attention and CNNs on local filters.

Knowing CNNs clarifies what ViT changes and why, highlighting trade-offs in model design.

Human Visual Attention

ViT's self-attention mechanism loosely mimics how humans focus on different parts of a scene to understand it holistically.

Connecting ViT to human attention helps appreciate why global context matters in vision tasks.

Common Pitfalls

#1 Using ViT on small datasets without pretraining or augmentation.

Wrong approach:

model = VisionTransformer()
model.train(small_dataset) # No pretraining or data augmentation

Correct approach:

model = VisionTransformer()
model.load_pretrained_weights()
model.train(small_dataset_with_augmentation)

Root cause:

Misunderstanding ViT's need for large data or pretraining leads to poor generalization and overfitting.

#2 Ignoring position embeddings in ViT input tokens.

Wrong approach:

tokens = patch_embeddings # No position embeddings added
output = transformer(tokens)

Correct approach:

tokens = patch_embeddings + position_embeddings
output = transformer(tokens)

Root cause:

Forgetting position embeddings causes the model to lose spatial order, harming performance.

#3 Choosing too large patch size for detailed images.

Wrong approach:

patch_size = 64 # Very large patches for small objects
model = VisionTransformer(patch_size=patch_size)

Correct approach:

patch_size = 16 # Smaller patches to capture details
model = VisionTransformer(patch_size=patch_size)

Root cause:

Misjudging patch size reduces image detail representation, lowering accuracy.

Key Takeaways

Vision Transformer (ViT) processes images by splitting them into patches and treating these patches as a sequence for transformer models.

Self-attention in ViT allows the model to learn global relationships between all parts of an image, unlike CNNs which focus locally.

ViT requires large datasets or pretraining because it lacks the built-in image biases of CNNs, making training on small data challenging.

Position embeddings are essential in ViT to maintain spatial information about where each patch belongs in the image.

Choosing the right patch size and understanding ViT's computational trade-offs are key to building effective and efficient models.

Practice

1. What is the main purpose of splitting an image into patches in a Vision Transformer (ViT)?

easy

A. To reduce the image size by cropping

B. To convert the image into smaller parts that the transformer can process as tokens

C. To apply convolution filters on each patch separately

D. To increase the image resolution for better detail

Solution

1. Step 1: Understand ViT input processing

   ViT splits images into fixed-size patches to treat each patch like a word token in language models.

2. Step 2: Purpose of patch splitting

   This allows the transformer to process image patches as a sequence, enabling attention mechanisms to learn relationships.

3. Final Answer:

   To convert the image into smaller parts that the transformer can process as tokens -> Option B

4. Quick Check:

   Image patches = tokens for transformer [OK]

Hint: Think of patches as words in a sentence for the transformer [OK]

Common Mistakes:

• Confusing patch splitting with image resizing
• Thinking patches are processed by convolution
• Assuming patches increase image resolution

2. Which of the following is the correct way to add a class token to the patch embeddings in ViT using Python-like pseudocode?

easy

A. patches = torch.cat([class_token, patches], dim=1)

B. patches = torch.cat([patches, class_token], dim=1)

C. patches = torch.cat([patches, class_token], dim=0)

D. patches = torch.cat([class_token, patches], dim=0)

Solution

1. Step 1: Understand tensor concatenation dimension

   Patch embeddings are sequences along dimension 1 (batch, seq, embed); class token must be prepended along this dimension.

2. Step 2: Correct concatenation syntax

   Using torch.cat with dim=1 adds class_token at the start of the sequence correctly.

3. Final Answer:

   patches = torch.cat([class_token, patches], dim=1) -> Option A

4. Quick Check:

   Class token prepended along sequence dim = patches = torch.cat([class_token, patches], dim=1) [OK]

Hint: Class token goes first, concat along sequence dimension (dim=1) [OK]

Common Mistakes:

• Concatenating along wrong dimension (dim=0)
• Appending class token at the end instead of start
• Mixing order of tensors in concat

3. Given the following simplified ViT patch embedding code, what is the shape of patch_embeddings after processing a batch of 8 images of size 32x32 with patch size 8 and embedding dimension 64?

patch_size = 8
embedding_dim = 64
batch_size = 8
image_size = 32
num_patches = (image_size // patch_size) ** 2
patch_embeddings = torch.randn(batch_size, num_patches, embedding_dim)

medium

A. (16, 8, 64)

B. (8, 64, 16)

C. (8, 8, 64)

D. (8, 16, 64)

Solution

1. Step 1: Calculate number of patches

   Number of patches = (32 / 8)^2 = 4^2 = 16 patches per image.

2. Step 2: Determine patch_embeddings shape

   Shape is (batch_size, num_patches, embedding_dim) = (8, 16, 64).

3. Final Answer:

   (8, 16, 64) -> Option D

4. Quick Check:

   Batch=8, patches=16, embed=64 [OK]

Hint: Calculate patches as (image/patch)^2, then batch x patches x embed [OK]

Common Mistakes:

• Mixing embedding dimension and patch count order
• Calculating patches incorrectly
• Confusing batch size with patch count

4. You have this ViT code snippet that throws an error:

class_token = torch.randn(1, 1, 64)
patches = torch.randn(8, 16, 64)
input_seq = torch.cat([class_token, patches], dim=1)

What is the cause of the error?

medium

A. Embedding dimensions do not match

B. Wrong concatenation dimension; should be dim=0

C. class_token shape should be (8, 1, 64) to match batch size

D. Dimension mismatch because class_token sequence size is 1 but patches sequence size is 16

Solution

1. Step 1: Check batch size compatibility

   class_token has batch size 1, patches have batch size 8; they must match for concatenation.

2. Step 2: Fix class_token shape

   class_token should be repeated or created with shape (8, 1, 64) to match patches batch size.

3. Final Answer:

   class_token shape should be (8, 1, 64) to match batch size -> Option C

4. Quick Check:

   Batch sizes must match for concat [OK]

Hint: Match batch sizes before concatenating tensors [OK]

Common Mistakes:

• Ignoring batch size mismatch
• Changing wrong concat dimension
• Assuming embedding dims cause error

5. In a Vision Transformer model, why is the class token important for image classification tasks?

hard

A. It aggregates information from all patches via attention to produce a final image representation

B. It stores the positional information of patches

C. It applies convolution to patches before transformer layers

D. It normalizes the patch embeddings before feeding to the transformer

Solution

1. Step 1: Understand class token role

   The class token is a special token that attends to all patch tokens and gathers their information.

2. Step 2: Use in classification

   After transformer layers, the class token embedding is used as the image's summary representation for classification.

3. Final Answer:

   It aggregates information from all patches via attention to produce a final image representation -> Option A

4. Quick Check:

   Class token = image summary for classification [OK]

Hint: Class token collects info from patches for final decision [OK]

Common Mistakes:

• Confusing class token with positional encoding
• Thinking class token applies convolution
• Assuming class token normalizes embeddings