A comparative study of Vision Transformers (ViT) and Convolutional Neural Networks (CNNs) for image classification. This project explores how model performance evolves across datasets of increasing complexity, highlighting the transition from CNN dominance on simple tasks to ViT advantages on more challenging visual data.
This repository also showcases the Vision Transformer (ViT), providing insight into its core architectural ideas, training dynamics, and its ability to capture global relationships through self-attention. It offers a practical view of the strengths and limitations of ViTs, and how they differ from traditional convolutional approaches.
Run a training experiment in one command:
python main.py --model vit --dataset tiny_imagenet| Dataset | CNN |
ViT |
Best Model | ||||
|---|---|---|---|---|---|---|---|
| Params | Accuracy (%) | Loss | Params | Accuracy (%) | Loss | ||
| MNIST | 843,850 | 99.35 | 0.5305 | 18,955,274 | 93.27 | 0.6838 | CNN 🏆 |
| CIFAR-10 | 1,204,874 | 72.28 | 1.1652 | 18,979,338 | 62.26 | 1.4040 | CNN 🏆 |
| Tiny ImageNet | 6,545,224 | 22.67 | 3.9591 | 19,175,112 | ... | ... | ViT 🏆 |
| Dataset | Model | Parameters | Accuracy (%) | Loss | Best Model |
|---|---|---|---|---|---|
| MNIST | CNN |
843,850 | 99.35 | 0.5305 | CNN |
ViT |
18,955,274 | 93.27 | 0.6838 | ||
| CIFAR-10 | CNN |
1,204,874 | 72.28 | 1.1652 | CNN |
ViT |
18,979,338 | 62.26 | 1.4040 | ||
| Tiny ImageNet | CNN |
6,545,224 | 22.67 | 3.9591 | ViT |
ViT |
19,175,112 | ... | ... |
Follow the steps below to set up and run the project.
pip install -r requirements.txtpython main.py --model <model> --dataset <dataset>Choose a model and dataset:
Models:
vit— Vision Transformercnn— Convolutional Neural Network
Datasets:
mnist— MNIST datasetcifar10— CIFAR-10 datasettiny_imagenet— Tiny ImageNet dataset
After training:
- Metrics are saved to
results/ - Training and validation plots are generated automatically
- GPU is used if available
- Datasets are downloaded automatically if not present
- Hyperparameters can be modified in
config.py
The Vision Transformer (ViT) is a deep learning architecture that adapts the Transformer, originally developed for natural language processing, to image recognition tasks. Introduced by Dosovitskiy et al. in “An Image is Worth 16x16 Words” (2020), ViT replaces traditional convolutional feature extractors with a sequence of image patches processed by self-attention. This approach demonstrated that, with sufficient data and compute, Transformers can outperform convolutional neural networks (CNNs) in computer vision benchmarks, paving the way for a broad family of vision transformer models.
In practice, ViT transforms an image into a sequence of smaller patches, which are then processed using the same self-attention mechanism that made Transformers successful in language tasks. By modeling relationships between patches directly, ViT captures both local details and long-range dependencies within an image, offering a flexible alternative to the strictly hierarchical representations of CNNs. Positional information is incorporated to maintain awareness of spatial structure, and a dedicated representation is used for classification. This design shifts the focus from handcrafted inductive biases toward a more data-driven approach, where the model learns to interpret visual structure primarily from large-scale training data.
A key step in the Vision Transformer (ViT) is the patch embedding stage, which transforms an image into a sequence suitable for a Transformer. Instead of processing pixels individually or relying on convolutional filters, the input image is divided into fixed-size patches (for example, 16×16 pixels). Each patch is then flattened into a vector and projected through a linear layer to a chosen embedding space. The result is a sequence of patch embeddings that can be treated similarly to word tokens in natural language processing, allowing the Transformer to apply self-attention mechanisms across the entire image.
The [CLS] token is a learnable embedding prepended to the sequence of patch embeddings in a Vision Transformer. Its primary purpose is to serve as a global representation of the entire image. During the forward pass, the Transformer encoder processes the sequence of patch embeddings along with the [CLS] token, allowing the self-attention mechanism to integrate information from all patches into this special token. After the final encoder block, the [CLS] token contains a summary of the image’s content and is typically fed into the classification head to produce the output logits. By using the [CLS] token in this way, ViT can perform classification based on a single learned representation rather than aggregating information from all patch embeddings.
Since Transformers process input sequences without any inherent notion of order, it is necessary to provide information about the position of each patch in the image. In the Vision Transformer, this is achieved through positional encoding, which adds a vector to each patch embedding to indicate its location within the image. Unlike the fixed sinusoidal encodings used in the original Transformer for NLP, ViT often uses learnable positional embeddings, which are initialized randomly and updated during training. These learnable embeddings allow the model to adaptively encode spatial relationships between patches, helping the self-attention mechanism capture both local and global structure in the image.
The Vision Transformer (ViT) is built on the Transformer encoder architecture, which processes images as a sequence of patch embeddings. Each encoder block combines multi-head self-attention and a feed-forward network, with normalization and residual connections to stabilize training.
The key component is self-attention, which allows each image patch to interact with every other patch. Unlike convolutional layers that focus on local neighborhoods, self-attention captures global relationships across the entire image. This enables the model to learn long-range dependencies and complex visual patterns more effectively.
By stacking multiple encoder blocks, ViT builds increasingly rich representations of the input, integrating both local details and global context for tasks such as image classification.
For a deeper explanation of the Transformer architecture, see: Simple Transformer
Convolutional Neural Networks (CNNs), introduced with LeNet-5 and popularized by AlexNet, have long been the dominant approach in computer vision. They use convolutional filters over local receptive fields, along with pooling and fully connected layers, to build hierarchical representations. This design encodes strong inductive biases such as locality and translation equivariance, making CNNs highly effective and data-efficient for many visual tasks.
To learn more about CNNs, see Simple CNN Guide.
The Vision Transformer (ViT), introduced by Dosovitskiy et al. (2020), replaces convolutions with a Transformer encoder. It represents an image as a sequence of fixed-size patches, which are embedded and processed using multi-head self-attention. This allows the model to capture global relationships between all patches directly, with a [CLS] token used to aggregate information for classification.
In terms of architecture, CNNs build hierarchical representations through stacked convolutions, gradually expanding their receptive fields and excelling at capturing local patterns such as edges and textures. In contrast, Vision Transformers operate in patch-embedding space, where self-attention provides a global receptive field from the first layer.
This difference leads to distinct trade-offs: CNNs benefit from strong inductive biases, making them data-efficient and effective on smaller datasets, while ViTs rely more on large-scale data to learn spatial relationships. As a result, CNNs tend to perform well in data-limited settings, whereas ViTs scale more effectively with increased data and model size, often achieving superior performance on more complex tasks.
The experiments are conducted on three standard image classification datasets of increasing complexity: MNIST, CIFAR-10, and Tiny ImageNet. MNIST consists of simple grayscale handwritten digits, CIFAR-10 introduces low-resolution RGB images of natural objects, and Tiny ImageNet provides higher-resolution images with a large number of classes and greater visual diversity. Together, these datasets enable a progressive evaluation of model performance from simple to more complex visual recognition tasks.
MNIST is a grayscale image dataset of handwritten digits (0–9), where each image has a resolution of 28×28 pixels. It contains 60,000 training samples and 10,000 test samples, with relatively low variability and simple patterns. Due to its simplicity and strong local features, it is well-suited for benchmarking basic models and is typically dominated by convolutional neural networks.
CIFAR-10 is a dataset of 32×32 RGB images across 10 object classes such as animals and vehicles. It consists of 50,000 training samples and 10,000 test samples, with moderate variability and more complex visual structure than MNIST. While still relatively small in resolution, it provides a more realistic benchmark for image classification and highlights the limitations of models that rely only on local features.
Tiny ImageNet is a scaled-down version of ImageNet, containing 64×64 RGB images across 200 classes, with 100,000 training samples and 10,000 validation samples. It introduces higher resolution, greater class diversity, and increased semantic complexity compared to CIFAR-10. This makes it a more challenging dataset where models that capture global context, such as Vision Transformers, begin to show advantages.The three datasets differ primarily in image complexity, resolution, and semantic diversity. MNIST contains simple, centered grayscale digits with minimal variation, making it an easy task dominated by local features. CIFAR-10 introduces color, background noise, and object variability at a small resolution, increasing the difficulty and requiring better feature extraction. Tiny ImageNet further increases resolution and the number of classes, adding significant intra-class variation and global structure, which makes the task more challenging and better suited for models that capture broader contextual relationships.
BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding
An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale