BEiT: BERT Pre-Training of Image Transformers

We introduce a self-supervised vision representation model BEIT , which stands for Bidirectional Encoder representation from Image Transformers . Following BERT ( Devlin et al. , 2019 ) 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 16× 16 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 . 1 INTRODUCTION . Transformer ( Vaswani et al. , 2017 ) has achieved promising performance in computer vision ( Dosovitskiy et al. , 2020 ; Touvron et al. , 2020 ) . However , empirical studies show that vision Transformers require more training data than convolutional neural networks . In order to solve the data-hungry issue ( Liu et al. , 2021a ) , 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 ( Chen et al. , 2021 ; Xie et al. , 2021 ) , and self-distillation ( Caron et al. , 2021 ) . Concurrently , BERT ( Devlin et al. , 2019 ) 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 BERTstyle 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 can not simply employ a softmax classifier to predict over all possible candidates for masked patches . In contrast , the language vocabulary , such as words and BPE ( Sennrich et al. , 2016 ) , 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 ( Ramesh et al. , 2021 ) . 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 ( Ramesh et al. , 2021 ) . 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 . visual tokens . We randomly mask some proportion of image patches ( gray patches in the figure ) and replace them with a special mask embedding [ M ] . Then the patches are fed to a backbone vision Transformer . The pre-training task aims at predicting the visual tokens of the original image based on the encoding vectors of the corrupted image . 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 . 2 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 . 2.1 IMAGE REPRESENTATIONS . 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 . 2.1.1 IMAGE PATCH . The 2D image is split into a sequence of patches ( Dosovitskiy et al. , 2020 ) , so that a standard Transformer can directly accept image data . Formally , we reshape the image x ∈ RH×W×C into N = HW/P 2 patches xp ∈ RN× ( P 2C ) , where C is the number of channels , ( H , W ) is the input image resolution , and ( P , P ) is the resolution of each patch . The image patches { xpi } Ni=1 are flattened into vectors and are linearly projected , which is similar to word embeddings in BERT ( Devlin et al. , 2019 ) . Image patches preserve raw pixels and are used as input features in BEIT . In our experiments , we split each 224× 224 image into a 14× 14 grid of image patches , where each patch is 16× 16 . 2.1.2 VISUAL TOKEN . Similar to natural language , we represent the image as a sequence of discrete tokens obtained by an “ image tokenizer ” , instead of raw pixels . Specifically , we tokenize the image x ∈ RH×W×C into z = [ z1 , . . . , zN ] ∈ Vh×w , where the vocabulary V = { 1 , . . . , |V| } contains discrete token indices . Following ( Ramesh et al. , 2021 ) , we use the image tokenizer learned by discrete variational autoencoder ( dVAE ) . There are two modules during visual token learning , namely , tokenizer and decoder . The tokenizer qφ ( z|x ) maps image pixels x into discrete tokens z according to a visual codebook ( i.e. , vocabulary ) . The decoder pψ ( x|z ) learns to reconstruct the input image x based on the visual tokens z . The reconstruction objective can be written as Ez∼qφ ( z|x ) [ log pψ ( x|z ) ] . Because the latent visual tokens are discrete , the model training is non-differentiable . Gumbel-softmax relaxation ( Jang et al. , 2017 ; Maddison et al. , 2017 ) is employed to train the model parameters . Moreover , a uniform prior is put on qφ during dVAE training . Refer to ( Ramesh et al. , 2021 ) for more training details of the image tokenizer . We tokenize each image to a 14× 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 |V| = 8192 . In our work , we directly use the publicly available1 image tokenizer described in ( Ramesh et al. , 2021 ) . We also compare it with a re-implemented tokenizer in Appendix C . 2.2 BACKBONE NETWORK : IMAGE TRANSFORMER . Following ViT ( Dosovitskiy et al. , 2020 ) , we use the standard Transformer ( Vaswani et al. , 2017 ) as the backbone network . So the results can be directly compared with previous work in terms of the network architecture . The input of Transformer is a sequence of image patches { xpi } Ni=1 . The patches are then linearly projected to obtain patch embeddings Expi , where E ∈ R ( P 2C ) ×D . Moreover , we prepend a special token [ S ] to the input sequence . We also add standard learnable 1D position embeddings Epos ∈ RN×D to patch embeddings . The input vectors H0 = [ e [ S ] , Expi , . . . , Ex p N ] +Epos is fed into Transformer . The encoder contains L layers of Transformer blocks H l = Transformer ( H l−1 ) , where l = 1 , . . . , L. The output vectors of the last layer HL = [ hL [ S ] , h L 1 , . . . , h L N ] are used as the encoded representations for the image patches , where hLi is the vector of the i-th image patch . 2.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 . Figure 1 shows the overview of our method . As presented in Section 2.1 , given an input image x , we split it into N image patches ( { xpi } Ni=1 ) , and tokenize it to N visual tokens ( { zi } Ni=1 ) . We 1https : //github.com/openai/DALL-E randomly mask approximately 40 % image patches , where the masked positions are denoted as M∈ { 1 , . . . , N } 0.4N . Next we replace the masked patches with a learnable embedding e [ M ] ∈ RD . The corrupted image patches xM = { xpi : i /∈ M } Ni=1 ⋃ { e [ M ] : i ∈ M } Ni=1 are then fed into the L-layer Transformer as described in Section 2.2 . The final hidden vectors { hLi } Ni=1 are regarded as encoded representations of the input patches . For each masked position { hLi : i ∈M } Ni=1 , we use a softmax classifier to predict the corresponding visual tokens pMIM ( z′|xM ) = softmaxz′ ( WchLi +bc ) , where xM is the corrupted image , Wc ∈ R|V|×D , and bc ∈ R|V| . The pre-training objective is to maximize the log-likelihood of the correct visual tokens zi given the corrupted image : max ∑ x∈D EM [ ∑ i∈M log pMIM ( zi|xM ) ] ( 1 ) where D is the training corpus , M represents randomly masked positions , and xM is the corrupted image that is masked according toM . Algorithm 1 Blockwise Masking Input : N ( = h× w ) image patches Output : Masked positionsM M← { } repeat s← Rand ( 16 , 0.4N − |M| ) . Block size r ← Rand ( 0.3 , 1 0.3 ) . Aspect ratio of block a← √ s · r ; b← √ s/r t← Rand ( 0 , h− a ) ; l← Rand ( 0 , w − b ) M←M ⋃ { ( i , j ) : i ∈ [ t , t+ a ) , j ∈ [ l , l + b ) } until |M| > 0.4N . Masking ratio is 40 % returnM Rather than randomly choosing patches for the masked positionsM , 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 ( Devlin et al. , 2019 ) , 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 ( Joshi et al. , 2020 ; Bao et al. , 2020 ; Raffel et al. , 2020 ) . 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 ( Ramesh et al. , 2021 ) . 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 .

Summary: This paper is one of the first to present state-of-the-art results for masked image modeling (BERT-style) self-supervised learning on images (contrastive approaches held the SOTA before). The whole system is based on a ViT encoder (e.g. with 16x16 pixels patches as input) to produce visual tokens. The pretraining objective consists in matching the visual tokens (one per image patch) from a provided visual tokenizer, with the ViT encoder. The tokenizer is a discrete VAE (with 8192 token types) from prior work called DALL-E (Ramesh et al., 2021). The results on ImageNet-1K are SOTA for comparable model sizes and pretraining data and settings. Idem for segmentation on ADE20K.