[Paper Review] Pre-training of Lightweight Vision Transformers on Small Datasets with Minimally Scaled Images

Jen Hong Tan|arXiv (Cornell University)|Feb 6, 2024

Advanced Neural Network Applications4 citations

TL;DR

This paper demonstrates that lightweight Vision Transformers (ViTs) can outperform ResNets on small datasets like CIFAR-10 and CIFAR-100 using Masked Auto-Encoder (MAE) pre-training with minimal image scaling. By applying MAE with learnable positional embeddings and a 75% masking ratio, the authors achieve state-of-the-art accuracy—96.41% on CIFAR-10 and 96.01% on CIFAR-100—using models with under 3.7 million parameters and minimal computational cost.

ABSTRACT

Can a lightweight Vision Transformer (ViT) match or exceed the performance of Convolutional Neural Networks (CNNs) like ResNet on small datasets with small image resolutions? This report demonstrates that a pure ViT can indeed achieve superior performance through pre-training, using a masked auto-encoder technique with minimal image scaling. Our experiments on the CIFAR-10 and CIFAR-100 datasets involved ViT models with fewer than 3.65 million parameters and a multiply-accumulate (MAC) count below 0.27G, qualifying them as 'lightweight' models. Unlike previous approaches, our method attains state-of-the-art performance among similar lightweight transformer-based architectures without significantly scaling up images from CIFAR-10 and CIFAR-100. This achievement underscores the efficiency of our model, not only in handling small datasets but also in effectively processing images close to their original scale.

Motivation & Objective

• Investigate whether lightweight Vision Transformers (ViTs) can match or exceed CNNs like ResNet on small datasets with minimal image scaling.
• Address the performance gap between ViTs and CNNs in low-data regimes by leveraging self-supervised pre-training.
• Demonstrate that ViTs can achieve state-of-the-art performance without architectural modifications like convolutional inductive biases or image upscaling.
• Evaluate the effectiveness of Masked Auto-Encoder (MAE) pre-training on small-scale images (32×32) with minimal augmentation.
• Establish a benchmark for efficient, small-parameter ViT models that are competitive with larger or more complex architectures on small datasets.

Proposed method

• Adopt a Masked Auto-Encoder (MAE) framework with a 75% patch masking ratio, where only 25% of patches are fed into the encoder.
• Use learnable positional embeddings in both the encoder and decoder instead of sinusoidal embeddings, improving spatial awareness and adaptability.
• Design a separate decoder with learnable positional embeddings and a flexible architecture independent from the encoder to enhance reconstruction quality.
• Compute a combined loss function that includes both masked patch reconstruction (using MSE) and a discounted loss on unmasked patches to stabilize training.
• Fine-tune the pre-trained ViT models on downstream classification tasks using standard cross-entropy loss and standard optimization protocols.
• Train models with minimal image scaling—only increasing resolution from 32×32 to 36×36—to preserve original data fidelity and reduce computational overhead.

Figure 1: Modified MAE: Implementation using separate learnable positional embeddings for Encoder and Decoder. The Encoder and the Decoder in this figure only consist of transformer layers, a layer of layerNorm and a linear projection layer.

Experimental results

Research questions

• RQ1 Can a lightweight Vision Transformer achieve superior performance to ResNet on small datasets like CIFAR-10 and CIFAR-100 without image upscaling?
• RQ2 Does MAE pre-training with minimal data augmentation and small image resolutions enable ViTs to overcome their lack of inductive bias in low-data regimes?
• RQ3 How does the use of learnable positional embeddings in both encoder and decoder affect reconstruction and downstream classification performance?
• RQ4 Can ViTs pre-trained on one dataset generalize effectively to another when fine-tuned, or is in-domain pre-training essential?
• RQ5 What is the impact of a combined loss function (masked + unmasked patch loss) on training stability and model generalization?

Key findings

• The proposed MAE-pretrained ViT models, Mae-ViT-C10 and Mae-ViT-C100, achieved 96.41% and 96.01% top-1 accuracy on CIFAR-10 and CIFAR-100 respectively, surpassing state-of-the-art lightweight ViT models.
• The ViT-Lite-7/4 model, which has similar parameter count (3.72M) and MACs (0.26G) to the MAE models, achieved only 93.57% and 73.94% accuracy, demonstrating the superiority of the MAE pre-training strategy.
• The models outperform other ViT variants such as CVT-7/4 (94.01% on CIFAR-100) and CCT-7/3×2 (95.04% on CIFAR-10), even without additional convolutional inductive biases.
• Pre-training on one dataset and fine-tuning on another led to faster convergence but slightly lower final performance, indicating that in-domain pre-training is more effective.
• Reconstruction quality improved progressively over training, with visible refinement in image details by epoch 3020, confirming effective feature learning.
• The use of a discounted loss on unmasked patches improved training stability and contributed to better generalization, especially in early training phases.

Figure 2: The architecture of our MAE. ’B’ stands for batch size, ’N’ for number of embeddings