BEiT: BERT Pre-Training of Image Transformers

Hangbo Bao, Li Dong, Songhao Piao, Furu Wei

Introduction

Transformer has achieved promising performance in computer vision . However, empirical studies show that vision Transformers require more training data than convolutional neural networks. In order to solve the data-hungry issue , self-supervised pre-training is a promising solution to leverage large-scale image data. Several strands of methods have been explored for vision Transformers, such as contrastive learning , and self-distillation .

Concurrently, BERT has achieved great success in natural language processing. Its masked language modeling task first randomly masks some proportion of tokens within a text, and then recovers the masked tokens based on the Transformer encoding results of the corrupted text. Motivated by BERT, we turn to the denoising auto-encoding idea to pretrain vision Transformers, which has not been well studied by the vision community. It is challenging to directly apply BERT-style pre-training for image data. First of all, there is no pre-exist vocabulary for vision Transformer’s input unit, i.e., image patches. So we cannot simply employ a softmax classifier to predict over all possible candidates for masked patches. In contrast, the language vocabulary, such as words and BPE , is well-defined and eases auto-encoding prediction. A straightforward alternative is regarding the task as a regression problem, which predicts the raw pixels of masked patches. However, such pixel-level recovery task tends to waste modeling capability on pre-training short-range dependencies and high-frequency details . Our goal is to overcome the above issues for pre-training of vision Transformers.

In this work, we introduce a self-supervised vision representation model BEiT, which stands for Bidirectional Encoder representation from Image Transformers. Inspired by BERT, we propose a pre-training task, namely, masked image modeling (MIM). As shown in Figure 1, MIM uses two views for each images, i.e., image patches, and visual tokens. We split the image into a grid of patches that are the input representation of backbone Transformer. Moreover, we “tokenize” the image to discrete visual tokens, which is obtained by the latent codes of discrete VAE . During pre-training, we randomly mask some proportion of image patches, and feed the corrupted input to Transformer. The model learns to recover the visual tokens of the original image, instead of the raw pixels of masked patches.

We perform self-supervised learning and then fine-tune the pretrained BEiT on two downstream tasks, i.e., image classification, and semantic segmentation. Experimental results indicate that BEiT outperforms both from-scratch training and previous strong self-supervised models. Moreover, BEiT is complementary to supervised pre-training. Performance of BEiT can be further improved by intermediate fine-tuning with ImageNet labels. Ablation studies show that our proposed techniques are critical to the effectiveness of BERT-style pre-training for image data. Apart from performance, the improvements of convergence speed and stability of fine-tuning reduce training costs on end tasks. In addition, we demonstrate that self-supervised BEiT can learn reasonable semantic regions via pre-training, unleashing the rich supervision signals contained in images.

Our contributions are summarized as follows:

We propose a masked image modeling task to pretrain vision Transformers in a self-supervised manner. We also provide a theoretical explanation from the perspective of variational autoencoder.

We pretrain BEiT and conduct extensive fine-tuning experiments on downstream tasks, such as image classification, and semantic segmentation.

We present that the self-attention mechanism of self-supervised BEiT learns to distinguish semantic regions and object boundaries, although without using any human annotation.

Methods

Given an input image $x$ , BEiT encodes it to contextualized vector representations. As shown in Figure 1, BEiT is pretrained by the masked image modeling (MIM) task in a self-supervised learning manner. MIM aims at recovering the masked image patches based on encoding vectors. For downstream tasks (such as image classification, and semantic segmentation), we append task layers upon pretrained BEiT and fine-tune the parameters on the specific datasets.

The images have two views of representations in our method, namely, image patch, and visual tokens. The two types serve as input and output representations during pre-training, respectively.

In our experiments, we split each $224\times 224$ image into a $14\times 14$ grid of image patches, where each patch is $16\times 16$ .

1.2 Visual Token

We tokenize each image to a $14\times 14$ grid of visual tokens. Notice the number of visual tokens and the number of image patches for one image are the same. The vocabulary size is set to $|\mathcal{V}|=8192$ . In our work, we directly use the publicly availablehttps://github.com/openai/DALL-E image tokenizer described in . We also compare it with a re-implemented tokenizer in Appendix C.

2 Backbone Network: Image Transformer

Following ViT , we use the standard Transformer as the backbone network. So the results can be directly compared with previous work in terms of the network architecture.

3 Pre-Training BEiT: Masked Image Modeling

We propose a masked image modeling (MIM) task. We randomly mask some percentage of image patches, and then predict the visual tokens that are corresponding to the masked patches.

where $\mathcal{D}$ is the training corpus, $\mathcal{M}$ represents randomly masked positions, and $x^{\mathcal{M}}$ is the corrupted image that is masked according to $\mathcal{M}$ .

Rather than randomly choosing patches for the masked positions $\mathcal{M}$ , we employ blockwise masking in our work. As summarized in Algorithm 1, a block of image patches is masked each time. For each block, we set the minimum number of patches to $16$ . Then we randomly choose an aspect ratio for the masking block. We repeat the above two steps until obtaining enough masked patches, i.e., $0.4N$ , where $N$ is the total number of image patches, and $0.4$ is masking ratio.

The MIM task is greatly inspired by masked language modeling , which is one of the most successful pre-training objective in natural language processing. Moreover, blockwise (or n-gram) masking is also widely applied in BERT-like models . However, directly using pixel-level auto-encoding (i.e., recovering the pixels of masked patches) for vision pre-training pushes the model to focus on short-range dependencies and high-frequency details . BEiT overcomes the above issue by predicting discrete visual tokens, which summarizes the details to high-level abstractions. Ablation studies in Section 3.3 show that our proposed method significantly outperforms pixel-level auto-encoding.

4 From the Perspective of Variational Autoencoder

where the second term is our BEiT pre-training objective.

5 Pre-Training Setup

The network architecture of BEiT follows that of ViT-Base for a fair comparison. We use a $12$ -layer Transformer with $768$ hidden size, and $12$ attention heads. The intermediate size of feed-forward networks is $3072$ . We employ the default $16\times 16$ input patch size. We directly borrow the image tokenizer trained by . The vocabulary size of visual tokens is $8192$ .

We pretrain BEiT on the training set of ImageNet-1K , which contains about $1.2$ M images. Our augmentation policy includes random resized cropping, horizontal flipping, color jittering . Notice that we do not use the labels for self-supervised learning. We use the $224\times 224$ resolution in our experiments. So the input is split to $14\times 14$ image patches, and the same amount of visual tokens. We randomly mask at most $75$ patches (i.e., roughly $40\%$ of total image patches).

The pre-training runs for about $500$ k steps (i.e., $800$ epochs) with $2$ k batch size. Adam with $\beta_{1}=0.9,\beta_{2}=0.999$ is employed for optimization. The learning rate is set to 1.5e-3, with a warmup of $10$ epochs, and cosine learning rate decay. The weight decay is $0.05$ . We employ stochastic depth with a $0.1$ rate, and disable dropout. The $500$ k training steps take about five days using $16$ Nvidia Telsa V100 32GB GPU cards.

We find that proper initialization is important to stabilize Transformer, especially for large-scale pre-training. We first randomly initialize all the parameters within a small range, such as $[-0.02,0.02]$ . Then, for the $l$ -th Transformer layer, we rescale the output matrices (i.e., the last linear projection within each sub-layer) of the self-attention module and the feed-forward network by $\frac{1}{\sqrt{2l}}$ .

6 Fine-Tuning BEiT on Downstream Vision Tasks

After pre-training BEiT, we append a task layer upon the Transformer, and fine-tune the parameters on downstream tasks, like BERT. We take image classification and semantic segmentation as examples in our work. It is straightforward to leverage the pre-training-then-fine-tuning paradigm on other vision tasks with BEiT.

For semantic segmentation, we follow the task layer used in SETR-PUP . To be specific, we use pretrained BEiT as a backbone encoder, and incorporate several deconvolution layers as decoder to produce segmentation. The model is also end-to-end fine-tuned similar to image classification.

After self-supervised pre-training, we can further train BEiT on a data-rich intermediate dataset (i.e., ImageNet-1K in our work), and then finetune the model on the target downstream tasks. Such intermediate fine-tuning is the common practice of BERT fine-tuning in NLP . We directly follow the method for BEiT.

Experiments

We conduct full fine-tuning experiments on image classification and semantic segmentation. Moreover, we present various ablation studies for pre-training and analyze the representations learned by BEiT. We also report linear probes on ImageNet in Appendix D.

The image classification task classifies input images to various categories. We evaluate BEiT on the ILSVRC-2012 ImageNet dataset with 1k classes and $1.3$ M images. We directly follow the most of hyperparameters of DeiT in our fine-tuning experiments for a fair comparison. We reduce fine-tuning epochs compared with training from scratch, as BEiT has been pre-trained. Accordingly, we use a larger learning rate with layer-wise decay. The detailed hyperparameters are summarized in Appendix H.

Table 1 reports top-1 accuracy on image classification. We compare BEiT with vision Transformers trained by random initialization, supervised pre-training, and previous self-supervised learning methods. All the compared models are base-size, except iGPT has 1.36B parameters. Pre-training is conducted on ImageNet for the comparison purpose, except ViT-JFT300M is pretrained on Google’s in-house $300$ M images.

Compared with the models trained by random initialization, we find that pre-trained BEiT significantly improves performance on both datasets. BEiT improves the performance on ImageNet, which shows the effectiveness under the rich-resource setting.

Moreover, we compare BEiT with previous state-of-the-art self-supervised methods for Transformer, such as DINO , and MoCo v3 . Our proposed method outperforms previous models on ImageNet fine-tuning. Among them, iGPT-1.36B uses much more parameters (i.e., $1.36$ B vs $86$ M), and ViT-JFT300M is pretrained on larger corpus (i.e., $300$ M vs $1.3$ M), while others pretrain ViT-Base on ImageNet-1K. iGPT-1.36B and ViT-JFT300M are the most comparable methods, which also follows auto-encoding pre-training for vision Transformer. Specifically, iGPT uses clustered image tokens as both input and output for image GPT or image BERT. In contrast, we use image patches as input to preserve raw pixels, and employ discrete visual tokens as a prediction bottleneck. ViT-JFT300 predicts the mean, 3-bit color of each masked patch, rather than visual tokens learned by discrete VAE. We also pretrain the self-supervised tasks of BEiT and DINO in a multi-task learning manner, which is presented in Appendix E.

In addition, we evaluate our proposed method with intermediate fine-tuning. In other words, we first pretrain BEiT in a self-supervised manner, and then fine-tune the pretrained model on ImageNet with labeled data. The results show that BEiT is complementary to supervised pre-training, achieving additional gain after intermediate fine-tuning on ImageNet.

After fine-tuning with resolution $224\times 224$ , we additionally fine-tune the model on $384\times 384$ images by $10$ more epochs. We follow the standard higher-resolution setting of DeiT , except using fewer epochs. Notice that we keep patch size the same for both $224\times 224$ and $384\times 384$ images. So the input sequence length of Transformers becomes longer for higher resolutions. Table 1 shows that higher resolution improves the BEiT results by $1+$ points on ImageNet. More importantly, BEiT384 pretrained on ImageNet-1K even outperforms supervised pre-training ViT384 that uses ImageNet-22K, when they use the same input resolution.

We further scale up BEiT to the large size (same as ViT-L). As shown in Table 1, ViT384-L is worse than ViT384 on ImageNet, when training from scratch. The results verifies the data-hungry issue of vision Transformers. Supervised pre-training on ImageNet-22K partially relieves the issue, where ViT384-L finally outperforms ViT384 by $1.2$ . In comparison, BEiT-L is better than BEiT by $2.0$ , and BEiT384-L outperforms BEiT384 by $1.7$ . In other words, the benefits of scaling up BEiT from base to large are greater than supervised pre-training with ImageNet-22K. More importantly, comparing between BEiT384 with ViT384 that conducts supervised pre-training on ImageNet-22K, the improvements of BEiT become greater along with scaling the size from base (i.e., $0.6$ ) to large (i.e., $1.1$ ). The results suggest that BEiT tends to help more for extremely larger models (such as 1B, or 10B), especially when labeled data are insufficient report that supervised pre-training of a 1.8B-size vision Transformer requires billions of labeled images. to conduct supervised pre-trainingAppendix B shows that BEiT fine-tuned on ImageNet-22K (14M) can match the performance of supervised pre-training on Google’s in-house JFT-3B , while using 214x less labels. We also demonstrate that large-size BEiT fine-tuned on 70M labeled images can achieve $89.5\%$ top-1 accuracy on ImageNet and $58.4\%$ mIoU on ADE20K, creating new state-of-the-art results for large-size vision Transformers. for such large models.

Figure 3 compares the convergence curves of the training-from-scratch and pre-training-then-fine-tuning paradigms. We find that fine-tuning BEiT not only achieves better performance, but also converging much faster than training DeiT from scratch. Moreover, fine-tuning BEiT can reach reasonable numbers within very few epochs.

2 Semantic Segmentation

Semantic segmentation aims to predict a corresponding class for each pixel of the input image. We evaluate BEiT on the ADE20K benchmark with $25$ K images and $150$ semantic categories. We report the metric of mean Intersection of Union (mIoU) averaged over all semantic categories. As presented in Section 2.6, we directly follow the task layer and the most of hyperparameters described in SETR-PUP . On ADE20K, we use Adam as the optimizer. The learning rate is set to 1e-3 with layer-wise decay similar to image classification. We conduct fine-tuning for 160K steps. The batch size is $16$ . The detailed hyperparameters are described in Appendix I.

As shown in Table 3, we compare BEiT with supervised pre-training that relies on labeled data of ImageNet. We find that our proposed method achieves better performance than supervised pre-training, although BEiT does not require manual annotations for pre-training. Moreover, we employ intermediate fine-tuning for BEiT on ImageNet, i.e., we first fine-tune pretrained BEiT on ImageNet, and then fine-tune the model on ADE20K. The results indicate that intermediate fine-tuning further improves BEiT on semantic segmentation.

3 Ablation Studies

We conduct ablation studies to analyze the contributions of each component in BEiT. The models are evaluated on image classification (i.e., ImageNet) and semantic segmentation (i.e., ADE20K). We set the default pre-training steps to $300$ epochs for the ablation studies, which is $37.5$ % of the total steps used in the previous experiments.

Table 4 reports the results of various model variants. First, we ablate blockwise masking by randomly sample masked positions. We find that blockwise masking is beneficial on both tasks, especially on semantic segmentation. Second, we ablate the usage of visual tokens by predicting the raw pixels of masked patches, i.e., the pre-training task becomes a pixel regression problem to recover masked patches. Our proposed masked image modeling task significantly outperforms naive pixel-level auto-encoding. Compared with the results in Table 1, the ablation result is worse than training vision Transformer from scratch on two tasks. The results indicate that the prediction of visual tokens is the key ingredient of BEiT. Third, we ablate the usage of visual tokens and blockwise masking together. We find that blockwise masking is even more helpful for pixel-level auto-encoding, which relieves the suffering of short-distance dependency. Forth, recovering all the visual tokens harms performance on downstream tasks. Fifth, we compare BEiT with different training steps. Pre-training the model longer can further improve performance on downstream tasks.

4 Analysis of Self-Attention Map

We show that the self-attention mechanism in BEiT can separate objects, even though our pre-training does not rely on any manual annotation at all. Similar properties are also observed by . The probing images are taken from the MS COCO corpus to avoid appearing in the pre-training data.

As shown in Figure 2, we plot the self-attention map for different reference points within an image. The visualizations are produced by attention scores computed via query-key product in the last layer. For each reference point, we use the corresponding patch as query, and show which patch it attends to. After pre-training, BEiT learns to distinguish semantic regions using self-attention heads, without any task-specific supervision. The property partially indicates the reason why BEiT is able to help downstream tasks. Such knowledge acquired by BEiT potentially improves the generalization ability of fine-tuned models, especially on small-scale datasets.

Related Work

Various methods have been introduced over the years to pretrain vision models in a self-supervised manner. Pioneering works design clever pretext tasks, such as predicting the patch orderings , colorization , and predicting rotation angles . In addition, propose to mask some patches within an image, and classify whether the masked patches are real or fake for each masked position. The method is similar to the masked version of Jigsaw pre-training . The recent strand of research follows contrastive paradigm . The models typically regard various data augmentations as different views of an image, and then make the representations of positive pairs similar while pushing negative pairs away. In order to obtain enough informative negative samples in contrastive learning, the methods usually rely on large memory banks or large batch size . BYOL and SimSiam further eliminate the requirement of negative samples, using various techniques to avoid representation collapse. Another strand of methods use clustering to organize image examples .

Pre-training vision Transformers has received significant attention recently due to the data-hungry issue. iGPT first creates a 9-bit color palette by k-means clustering RGB pixels, and then uses the clustered tokens to represent images. Next iGPT uses the tasks of BERT and GPT to pretrain Transformers. In comparison, our proposed method uses image patches as input without losing pixel-level information. Moreover, our visual tokens are obtained by discrete VAE instead of clustering. ViT conducts a preliminary exploration with the masked patch prediction task, which predicts the 3-bit mean color of the masked patches. also report that pixel-level auto-encoding performs worse, although it is the most straightforward translation of BERT from NLP to CV. Rather than using heuristically designed pre-training tasks, our proposed model leverages visual tokens learned by discrete VAE, which not only achieves better performance but also is better theoretically motivated. Apart from masked auto-encoding, other mainstream research works use contrastive learning , and self-distillation . In comparison, BEiT can achieve several times of improvement in terms of pre-training throughput (Appendix E), and memory consumption. The advantages make BEiT appealing to scale up vision Transformers.

Conclusion

We introduce a self-supervised pre-training framework for vision Transformers, achieving strong fine-tuning results on downstream tasks, such as image classification, and semantic segmentation. We show that the proposed method is critical to make BERT-like pre-training (i.e., auto-encoding with masked input) work well for image Transformers. We also present the intriguing property of automatically acquired knowledge about semantic regions, without using any human-annotated data. In the future, we would like to scale up BEiT pre-training in terms of data size and model size. Moreover, we will conduct multimodal pre-training in a more unified way, using the similar objectives and the shared architecture for texts and images.

We would like to acknowledge Yue Cao, Han Hu, Hang Hua, Jingdong Wang, Zheng Zhang for the helpful discussions, and Yaru Hao for some analysis experiments using .

References

Appendix A Architecture Variants of Vision Transformer

We use the standard vision Transformer (ViT) in the experiments for fair comparisons. In addition, we find that LayerScale and relative position bias improve ViTs on downstream tasks. We employ the same setting as in Section 3.3 for ablation studies, which pretrains base-size models for $300$ epochs on ImageNet-1K.

As shown in Table 5, both LayerScale and relative position bias improve performance on ImageNet classification and ADE20K semantic segmentation. We denote the improved architecture as BEiT+ and use it for the experiments in Appendix B. We empirically notice that vanilla Transformer is the most stable when scaling up the model to billions of parameters, so we do not use LayerScale for extra-large models.

Appendix B Comparison with Large-Scale Supervised Pre-Training

We compare with state-of-the-art supervised pre-training at scale. In addition to using ImageNet-1K for fair comparisons with previous work, we pretrain BEiT on ImageNet-22K to boost performance. We employ the architecture improvements (i.e., LayerScale, and relative position bias) as described in Appendix A, which is denoted as BEiT+ in Table 6 and Table 7. We follow the same pre-training setup as in Section 2.5, except we pretrain $150$ epochs on ImageNet-22K. After self-supervised pre-training, we conduct intermediate fine-tuning on ImageNet-22K for $90$ epochs. Moreover, we use an in-house dataset that has about 70M labeled images as a drop-in replacement of ImageNet-22K.

Table 6 compares BEiT with previous state-of-the-art supervised pre-training on ImageNet fine-tuning. Rather than heavily relying on extremely large-size labeled data (such as Google’s in-house JFT-300M and JFT-3B), we demonstrate that BEiT pre-training can catch up with only ImageNet-22k (14M). Specifically, BEiT-L fine-tuned on ImageNet-22K achieves comparable performance with ViT-L trained on Google JFT-3B. Moreover, BEiT-L obtains $89.5\%$ top-1 accuracy on ImageNet after intermediate fine-tuning on an in-house 70M dataset. The results indicate that BEiT pre-training greatly reduces the required labeling efforts and advances the new state of the art for large-size vision Transformers.

As shown in Table 7, we report the fine-tuning results on the ADE20K semantic segmentation benchmark. Following Swin , we use the same task layer (i.e., UperNet) and evaluate the models at the resolution $640\times 640$ . The BEiT-L model obtains state-of-the-art performance on ADE20K.

Appendix C Ablation Studies of Image Tokenizer

For comparison, we re-train the image tokenizer on ImageNet-1K. The reimplementation is based on https://github.com/lucidrains/DALLE-pytorch. We use the same codebook size 8K as in DALL-E . Then we plug the tokenizer into our pre-training process. We follow the same experimental setup of ablation studies as in Section 3.3. Table 8 shows that our reimplemented tokenizer obtains comparable reconstruction loss and ImageNet fine-tuning performance compared with the off-the-shelf DALL-E tokenizer.

Appendix D Linear Probes on ImageNet

We evaluate linear probes on ImageNet for various pretrained vision Transformers. We compare BEiT with two main strands of work, namely discriminative and generative self-supervised learning. The first one applies discriminative learning for pre-training, such as contrastive learning , and self distillation . The above methods typically learn to aggregate the image-level features into a global vector, which is relatively suitable for linear probing. In contrast, the second strand of methods, such as iGPT and ours, usually do not pretrain such global feature aggregation, which tends to make linear probes difficult.

Following iGPT , we use average pooling to aggregate the hidden states of each image patches, and add the probing layer at the middle layer of Transformer instead of always at the final layer. Similarly, we find that the best layer lies in $9$ -th layer for BEiT-B, and $14$ -th layer for BEiT-L. To be specific, we use AdamW to update the linear probe layer for $50$ epochs. The learning rate is 4e-3 with cosine decay. The batch size is 1024. The weight decay is set to 1e-4. We follow data augmentation used in DINO , which uses random resize crops and horizontal flips augmentation during training and evaluates on central crops.

As shown in Table 9, we evaluate linear probes on ImageNet-1K for self-supervised learning. Overall, discriminative methods perform better than generative pre-training on linear probing. Linear probes keep the Transformer parameters fixed and only update the linear layer. So the pre-training of global aggregation of image-level features is beneficial to linear probing in DINO and MoCo v3, although full fine-tuning eliminates the gap. Moreover, the results indicate that increasing the model size from base (86M) to large (304M) significantly improves accuracy for our proposed method. In contrast, the gap between base- and large-size MoCo v3 is smaller. We also find that BEiT outperforms iGPT by a large margin even using much fewer parameters.

Appendix E Multi-Task Pre-Training with DINO

We train the pre-training tasks of BEiT and DINO together in a multi-task manner. As shown in Table 10, augmenting masked image modeling with DINO improves semantic segmentation on ADE20K, and obtains comparable results on ImageNet classification. Moreover, BEiT is more efficient in terms of pre-training speed, as DINO has two copies of Transformer parameters for self-distillation and multi-crop augmentation . For the throughput comparisons between BEiT and BEiT+DINO, we set batch size to the same. Because BEiT is also more memory-efficient, we can use larger batch size to fully utilize GPU cards, which obtains greater speedup in practice than the reported numbers.

Appendix F Image Classification on CIFAR-100

In addition to ImageNet classification, we conduct fine-tuning experiments on the CIFAR-100 benchmark with 100 classes and 60k images. The experimental setup is the same as in Section 3.1.

Table 11 reports the top-1 accuracy on CIFAR-100. Notably, on the smaller CIFAR-100 dataset, ViT trained from scratch only reaches $48.5\%$ accuracy . In comparison, BEiT achieves $90.1\%$ with the help of pre-training. The results indicate that BEiT can greatly reduce the requirement of annotation efforts. BEiT also outperforms MoCo v3. Moreover, intermediate fine-tuning on ImageNet-1K further improves the results on CIFAR-100.