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: The paper presents a new objective called masked image modeling (MIM) to pre-train vision transformers, making the model predict a visual token from the masked patches.   The visual tokens are obtained by discrete VAE, which is trained on 250M images. (Though, the authors showed that training discrete VAE with only 1M images (imagenet-1K) is enough to demonstrate the power of the proposed objective.)   The authors pre-train the ViT with MIM and fine-tune the pre-trained weights (BEiT) on two visual downstream tasks: image classification and semantic segmentation, and showed superior performances compared to previous methods, including DINO and MoCo v3.   The authors also propose an additional trick called blockwise masking to improve BEiT further.