Vision Transformer › Pre-training Strategies

To overcome the data inefficiency inherent in the absence of inductive biases, ViT relies on large-scale pre-training. The original paper explored two main approaches, and subsequent research has expanded these significantly.

The most direct approach involves pre-training ViT on large supervised image classification datasets. The model is trained with standard cross-entropy loss to predict image labels from the CLS token representation. After pre-training on a large dataset, the classification head is replaced or fine-tuned for specific downstream tasks.

The two primary supervised pre-training datasets used in the original ViT work were:

• ImageNet-21k: 14 million images across 21,000 classes. This provides broad visual diversity. Pre-training on this dataset was enough for ViT to match or exceed well-tuned ResNets on ImageNet-1k classification.
• JFT-300M: 300 million images across 18,000 classes, an internal Google dataset. Training at this scale demonstrated ViT's strongest performance and clearest advantages over CNNs in terms of accuracy and scaling behavior.

The key practical concern with supervised pre-training is the requirement for labeled data at scale. While ImageNet-21k is publicly available, collecting and labeling 300 million images requires substantial resources. This limitation motivated the development of self-supervised pre-training approaches that can learn from unlabeled data.

An important pre-training detail for supervised ViT is that the classification head used during pre-training differs from the one used during fine-tuning. During pre-training on JFT-300M, ViT uses an MLP head with one hidden layer and tanh activation. For fine-tuning on downstream datasets, a single linear layer is used instead, initialized to zero. The zero initialization of the fine-tuning head encourages the model to preserve the pre-trained representations at the start of fine-tuning, avoiding the disruption that a randomly initialized large head might cause.

More recent approaches eliminate the need for manual labels by learning representations directly from unlabeled image data. The Masked Autoencoder (MAE) approach, introduced by He et al. in 2021, draws direct inspiration from masked language modeling in BERT (which we covered in Masked Language Modeling).

MAE works by randomly masking a high proportion of input patches (typically 75%) and training the model to reconstruct the pixel values of the masked patches from the visible ones. The high masking ratio is intentional: it creates a challenging task that prevents the model from solving the reconstruction problem through simple linear interpolation or nearest-neighbor copying. To recover a heavily masked image, the model must develop an internal representation that captures the semantic and structural regularities of natural images.

The architecture uses an asymmetric encoder-decoder design that provides an efficiency advantage:

The encoder is a standard ViT that processes only the visible (unmasked) patches. Because 75% of patches are masked, the encoder handles only 25% of the sequence. This reduces the computational cost of the encoder by roughly $4\times$ compared to processing all patches, enabling larger batch sizes and faster training iterations.

The decoder is a much smaller transformer that takes both the encoded visible patches and learned mask tokens (placeholders for the missing patches) as input. The mask tokens are shared across all masked positions (the same vector is used for every masked patch). The decoder reconstructs all patches in their original spatial positions, with the loss computed only on the masked patches using mean squared error between predicted and actual pixel values.

The asymmetry is strategically important: making the decoder lightweight means the encoder bears the full burden of learning good visual representations. The decoder, being simple and narrow, cannot compensate for a weak encoder through its own capacity. This forces the encoder to develop rich semantic representations rather than relying on decoder complexity to smooth over impoverished features.

Other self-supervised approaches applied to ViT include:

DINO (Self-Distillation with No Labels): Uses knowledge distillation between a student network and a teacher network that is a momentum-averaged version of the student. The student and teacher both receive different augmented views of the same image, and the student is trained to predict the teacher's output distribution. DINO-trained ViTs develop particularly strong feature representations that exhibit striking emergent properties: the self-attention maps of DINO-trained models often segment objects from background without any segmentation supervision, revealing that the model has learned object boundaries as a side effect of self-supervised visual understanding.

BEiT (BERT Pre-training of Image Transformers): Instead of reconstructing raw pixel values, BEiT converts each patch into a discrete token from a visual vocabulary learned by a separate vector quantization model (a visual tokenizer). The model then performs masked patch prediction in this discrete token space, much like BERT predicts masked word tokens. This approach provides a cleaner target signal than raw pixel reconstruction, which may focus the model on learning semantically meaningful features rather than low-level texture details.

Some architectures combine CNNs and transformers to capture the benefits of both paradigms. A CNN stem, typically using the early layers of a ResNet, processes the image first to extract low-level features while preserving spatial inductive biases. The resulting feature maps are then patched and fed into the transformer for global modeling. This hybrid approach offers a useful engineering compromise: the CNN stem provides strong inductive bias that enables efficient learning from limited data, while the transformer backbone enables global reasoning and flexible scaling with more data. Hybrid ViT models trained on ImageNet-1k alone can outperform pure ViT models trained on the same data, since the CNN stem compensates for ViT's lack of visual priors in the data-scarce regime.