PaperlayerView exact revision on arXiv

Source-linked AI summary

BEiT: BERT Pre-Training of Image Transformers

Hangbo Bao, Li Dong, Songhao Piao, Furu Wei

arXiv:2106.08254v2cs.CVcs.LG

TL;DR

Vision Transformers need more training data, while BERT-style masked pretraining remains underexplored for images. BEiT addresses this with masked image modeling over image patches and discrete visual tokens, achieving strong downstream results in image classification and semantic segmentation.

  • Problem

    Vision Transformers require more training data, and BERT-style denoising pretraining is difficult because image patches lack a pre-existing prediction vocabulary.

  • Method

    BEiT masks image patches and trains a Transformer to recover discrete visual tokens produced by a learned image tokenizer.

  • Results

    BEiT achieves strong fine-tuning results, outperforming from-scratch training and previous self-supervised models on image classification and semantic segmentation.

  • Takeaways & Limitations

    BEiT makes BERT-like masked pretraining effective for vision Transformers and learns semantic regions without human-annotated data.

Abstract

from arXiv · show

We introduce a self-supervised vision representation model BEiT, which stands for Bidirectional Encoder representation from Image Transformers. Following BERT developed in the natural language processing area, we propose a masked image modeling task to pretrain vision Transformers. Specifically, each image has two views in our pre-training, i.e, image patches (such as 16x16 pixels), and visual tokens (i.e., discrete tokens). We first "tokenize" the original image into visual tokens. Then we randomly mask some image patches and fed them into the backbone Transformer. The pre-training objective is to recover the original visual tokens based on the corrupted image patches. After pre-training BEiT, we directly fine-tune the model parameters on downstream tasks by appending task layers upon the pretrained encoder. Experimental results on image classification and semantic segmentation show that our model achieves competitive results with previous pre-training methods. For example, base-size BEiT achieves 83.2% top-1 accuracy on ImageNet-1K, significantly outperforming from-scratch DeiT training (81.8%) with the same setup. Moreover, large-size BEiT obtains 86.3% only using ImageNet-1K, even outperforming ViT-L with supervised pre-training on ImageNet-22K (85.2%). The code and pretrained models are available at https://aka.ms/beit.

1 Introduction

BEIT addresses the data demands of vision Transformers with self-supervised masked image modeling inspired by BERT. It pretrains on corrupted image patches to recover discrete visual tokens, then fine-tunes for downstream vision tasks.

  • Motivation: Vision Transformers require more training data than convolutional neural networks, motivating self-supervised pre-training to leverage large-scale image data.
  • Motivation: Inspired by BERT’s masked language modeling, BEIT introduces masked image modeling to pretrain vision Transformers through denoising auto-encoding.
  • Method: BEIT represents each image as patches and discrete visual tokens, masks some patches, and trains a Transformer to recover the original visual tokens.The discrete visual tokens are obtained from a latent-code discrete VAE.
  • Results: BEIT is self-supervised pretrained and then fine-tuned for image classification and semantic segmentation, outperforming from-scratch training and previous strong self-supervised models.The model is also complementary to supervised pre-training, with further gains from intermediate fine-tuning using ImageNet labels.
  • Contributions: Self-supervised BEIT’s self-attention distinguishes semantic regions and object boundaries without using human annotations.

2 Methods

BEiT pretrains a standard Transformer with self-supervised masked image modeling, using raw image patches as inputs and discrete visual tokens as prediction targets. It then fine-tunes the pretrained encoder for downstream tasks.

  • Downstream adaptation: For downstream tasks, task-specific layers are appended to pretrained BEiT and the model parameters are fine-tuned on the target datasets.The paper names image classification and semantic segmentation as examples.
  • Image representations: BEiT uses two image views: raw patches provide input representations, while discrete visual tokens provide pretraining targets.Image patches preserve raw pixels; visual tokens are obtained by an image tokenizer.
  • Transformer encoder: The image is split into patches, flattened, and linearly projected before entering a standard Transformer backbone with positional embeddings.In experiments, each 224 × 224 image uses a 14 × 14 grid of 16 × 16 patches.
  • Two-stage training: BEiT training follows two stages: a discrete variational autoencoder learns the tokenizer and decoder, then the visual-token prior is learned while those components remain fixed.The first stage minimizes visual-token reconstruction loss; the second stage is the BEiT pretraining objective.
  • Masked image modeling: BEiT randomly masks approximately 40% of image patches, replaces them with a learnable embedding, and predicts their corresponding visual tokens.The model applies a softmax classifier to masked-position representations and maximizes the log-likelihood of the correct tokens.
  • Masked image modeling: The method uses blockwise masking, repeatedly masking blocks of at least 16 patches with random aspect ratios until reaching 0.4N masked patches.This design follows masked language modeling while avoiding pixel-level reconstruction that emphasizes short-range dependencies and high-frequency details.

3 Experiments

Experiments show that BEiT improves image classification and semantic segmentation, converges faster than training from scratch, and benefits from intermediate fine-tuning and larger input or model sizes. Ablations and attention visualizations further examine the method’s components and learned representations.

  • Image classification: BEiT outperforms randomly initialized and prior self-supervised Transformer models on ImageNet fine-tuning.Comparisons include DeiT-style training, DINO, and MoCo v3 under the reported experimental settings.
  • Image classification: Intermediate supervised fine-tuning provides additional gains for BEiT on ImageNet, demonstrating complementarity between self-supervised and supervised pre-training.BEiT is first pretrained self-supervised, then fine-tuned with labeled ImageNet data.
  • Image classification: Higher-resolution fine-tuning and scaling to larger models further improve BEiT’s ImageNet classification results.The experiments retain patch size when increasing resolution and compare base- and large-size variants.
  • Convergence: Fine-tuning BEiT both achieves better performance and converges faster than training DeiT from scratch.The model reaches reasonable results within very few epochs.
  • Semantic segmentation: BEiT achieves better semantic segmentation performance than supervised pre-training on ADE20K without manual annotations during pre-training.Intermediate fine-tuning on ImageNet further improves BEiT for semantic segmentation.
  • Ablation studies: Blockwise masking benefits both evaluated tasks, while predicting visual tokens outperforms naive pixel-level reconstruction.The ablations evaluate image classification on ImageNet and semantic segmentation on ADE20K.

4 Related Work

Prior self-supervised vision work uses handcrafted pretext tasks, masked-patch discrimination, or contrastive learning. For vision Transformers, iGPT uses clustered color tokens with BERT/GPT objectives, whereas BEiT uses image patches as input without losing pixel-level information.

  • Self-supervised visual representation learning: Earlier self-supervised methods predict patch orderings, colorization, rotation angles, or whether masked patches are real or fake.The masked-patch method classifies each masked position as real or fake and is similar to masked Jigsaw pre-training.
  • Self-supervised visual representation learning: Contrastive methods treat augmented images as different views, bringing positive-pair representations together while separating negative pairs.They typically rely on large memory banks or large batches to obtain informative negative samples.
  • Self-supervised vision Transformers: iGPT represents images with 9-bit color-palette tokens from k-means-clustered RGB pixels and pretrains Transformers using BERT and GPT tasks.BEiT instead uses image patches as input without losing pixel-level information.

5 Conclusion

BEiT provides a self-supervised pre-training framework for vision Transformers that enables strong fine-tuning on image classification and semantic segmentation. The authors identify masked-input auto-encoding as critical for BERT-like pre-training of image Transformers and propose scaling and multimodal extensions.

  • 5 Conclusion: BEiT is a self-supervised pre-training framework for vision Transformers with strong fine-tuning results on image classification and semantic segmentation.The conclusion explicitly names both downstream tasks.
  • 5 Conclusion: Masked-input auto-encoding is critical for making BERT-like pre-training work well for image Transformers.This conclusion frames the masked-input objective as central to the method’s effectiveness.
  • 5 Conclusion: Future work will scale BEIT pre-training in data and model size.The passage identifies both data size and model size as scaling directions.
  • 5 Conclusion: The authors plan unified multimodal pre-training for text and images using similar objectives and a shared architecture.The proposed direction applies related objectives and common architecture across modalities.

A Architecture Variants of Vision Transformer

The experiments use standard ViT for fair comparisons, while LayerScale and relative position bias improve downstream performance. The resulting BEIT+ architecture is used subsequently, except that extra-large models retain vanilla Transformers for scaling stability.

  • Architecture variants: Standard ViT is used in experiments to enable fair comparisons, with base-size models pretrained for 300 epochs on ImageNet-1K for ablations.The ablation setting matches Section 3.3.
  • Architecture variants: LayerScale and relative position bias both improve ImageNet classification and ADE20K semantic segmentation, defining the BEIT+ architecture.ADE20K results use UperNet and report single-scale mIoU.
  • Architecture variants: Vanilla Transformer is more stable when scaling to billions of parameters, so extra-large models omit LayerScale.BEIT+ is used for the experiments in Appendix B, but not for extra-large models.

B Comparison with Large-Scale Supervised Pre-Training

BEIT is compared with large-scale supervised pre-training using ImageNet-22K, showing competitive ImageNet fine-tuning and state-of-the-art ADE20K segmentation. Its self-supervised pre-training also outperforms supervised Swin baselines on segmentation.

  • Comparison setup: BEIT uses ImageNet-22K pre-training for large-scale comparisons, with BEIT+ adding LayerScale and relative position bias.The ImageNet-22K pre-training runs for 150 epochs under the Section 2.5 setup.
  • ImageNet classification: BEIT-L fine-tuned on ImageNet-22K achieves comparable performance with ViT-L trained on Google JFT-3B.The comparison emphasizes achieving this with ImageNet-22K rather than extremely large labeled datasets such as JFT-300M and JFT-3B.
  • ADE20K semantic segmentation: Swin-B and Swin-L reach 50.0 and 52.1 mIoU (%) with supervised ImageNet-22K pre-training, below the corresponding BEIT-B+ and BEIT-L+ results.Their multi-scale mIoU (%) values are 51.7 and 53.5, respectively.
  • ADE20K semantic segmentation: 53.6 mIoU (%) is achieved by BEIT-B+ and 56.7 mIoU (%) by BEIT-L+ on ADE20K after ImageNet-22K intermediate fine-tuning.Their multi-scale mIoU (%) values are 54.2 and 57.0, respectively.
  • ADE20K semantic segmentation: 57.9 mIoU (%) and 58.4 multi-scale mIoU (%) are achieved by BEIT-L+ after intermediate fine-tuning on In-House-70M.The model obtains state-of-the-art performance on ADE20K using UperNet at 640 × 640 resolution.

C Ablation Studies of Image Tokenizer

The study re-trains an image tokenizer on ImageNet-1K and evaluates it within the BEiT pre-training process. The reimplemented tokenizer uses an 8K codebook and achieves reconstruction loss and ImageNet fine-tuning performance comparable to the off-the-shelf DALL-E tokenizer.

  • Tokenizer setup: The image tokenizer is re-trained on ImageNet-1K and integrated into the BEiT pre-training process using the same ablation setup as Section 3.3.The reimplementation is based on lucidrains/DALLE-pytorch.
  • Tokenizer setup: 8K codebook size is used for the reimplemented tokenizer, matching DALL-E.The tokenizer is trained on ImageNet-1K without labels.
  • Ablation results: Comparable reconstruction loss and ImageNet fine-tuning performance are obtained relative to the off-the-shelf DALL-E tokenizer.Reconstruction is measured by mean absolute error of normalized RGB values, while downstream performance is reported as ImageNet-1K top-1 accuracy.

D Linear Probes on ImageNet

The section evaluates ImageNet-1K linear probes for pretrained vision Transformers, comparing discriminative and generative self-supervised methods. It finds that discriminative methods perform better, while global feature aggregation benefits linear probing but not full fine-tuning.

  • Evaluation setup: BEiT linear probes average-pool patch hidden states and attach the probing layer at an intermediate Transformer layer.The best layer is the 9-th layer for BEIT-B and the 14-th layer for BEIT-L.
  • Results: Discriminative self-supervised methods perform better than generative pre-training on ImageNet-1K linear probing.The comparison includes contrastive learning, self-distillation, iGPT, and BEiT.
  • Results: Global aggregation of image-level features benefits linear probing in DINO and MoCo v3, but full fine-tuning eliminates the gap.Linear probing keeps Transformer parameters fixed and updates only the linear layer.

E Multi-Task Pre-Training with DINO · F Image Classification on CIFAR-100

The paper evaluates joint BEiT–DINO pre-training and shows improved ADE20K segmentation with comparable ImageNet classification, while BEiT substantially improves CIFAR-100 classification over scratch training.

  • E Multi-Task Pre-Training with DINO: BEiT and DINO are trained jointly in a multi-task pre-training scheme.The combined objective augments masked image modeling with DINO.
  • E Multi-Task Pre-Training with DINO: Joint training improves semantic segmentation on ADE20K while obtaining comparable ImageNet classification results.Table 10 evaluates fine-tuned ImageNet-1K classification and ADE20K semantic segmentation.
  • E Multi-Task Pre-Training with DINO: BEiT pre-training is more efficient because DINO uses two Transformer parameter copies for self-distillation and multi-crop augmentation.Throughput comparisons measure pre-training speed, with larger numbers indicating faster training.
  • F Image Classification on CIFAR-100: CIFAR-100 experiments use a 100-class, 60k-image benchmark with the same experimental setup as Section 3.1.Table 11 reports top-1 accuracy at 224 × 224 resolution, except ViT384 at 384 × 384; results are generally from base-size models.
  • F Image Classification on CIFAR-100: 48.5% top-1 accuracy is achieved by ViT trained from scratch on CIFAR-100, whereas 90.1% top-1 accuracy is achieved by BEiT with pre-training.The comparison is made on the smaller CIFAR-100 dataset.
  • F Image Classification on CIFAR-100: BEiT outperforms MoCo v3 on CIFAR-100 and further improves results through intermediate fine-tuning on ImageNet-1K.These results indicate that pre-training can reduce annotation requirements.

G Hyperparameters for Pre-Training · H Hyperparameters for Image Classification Fine-Tuning · I Hyperparameters for ADE20K Semantic Segmentation Fine-Tuning

This block documents the hyperparameter settings for BEIT pre-training, image-classification fine-tuning, and ADE20K semantic-segmentation fine-tuning. The settings are organized in Tables 12–14 by training stage and dataset.

  • G Hyperparameters for Pre-Training: Table 12 specifies hyperparameters for pre-training BEIT on ImageNet-1K.It covers the pre-training stage.
  • G Hyperparameters for Pre-Training: The pre-training hyperparameters are associated with ImageNet-1K.The table is explicitly labeled for BEIT pre-training on this dataset.
  • H Hyperparameters for Image Classification Fine-Tuning: Table 13 specifies hyperparameters for fine-tuning BEIT on ImageNet-1K and CIFAR-100.It covers image-classification fine-tuning across both datasets.
  • H Hyperparameters for Image Classification Fine-Tuning: ImageNet-1K is one dataset covered by the image-classification fine-tuning hyperparameters.The table also includes CIFAR-100.
  • H Hyperparameters for Image Classification Fine-Tuning: CIFAR-100 is the second dataset covered by the image-classification fine-tuning hyperparameters.These settings are presented together with ImageNet-1K in Table 13.
  • I Hyperparameters for ADE20K Semantic Segmentation Fine-Tuning: Table 14 specifies hyperparameters for fine-tuning BEIT on ADE20K.It covers semantic-segmentation fine-tuning.
Loading 2106.08254v2…