BE I T: BERT P RE -T RAINING OF I MAGE T RANSFORMERS
Anonymous authors
Paper under double-blind review
A BSTRACT
We introduce a self-supervised vision representation model BE I T , which stands for B idirectional E ncoder representation from I mage T ransformers.
Following
BERT (Devlin et al., 2019) developed in the natural language processing area, we propose a masked image modeling task to pretrain vision Transformers.
Speciﬁcally, 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 ﬁrst “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 BE I T, we directly ﬁne-tune the model parameters on downstream tasks by appending task layers upon the pretrained encoder.
Experimental results on image classiﬁcation and semantic segmentation show that our model achieves competitive results with previous pre-training methods.
I NTRODUCTION
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 ﬁrst 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 cannot simply employ a softmax classiﬁer 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-deﬁned 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 BE I T , which stands for
B idirectional E ncoder representation from I mage T ransformers.
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.
334211 322 433+ + + + + + + + + + + + + + + + +
BE I T Encoder
BlockwiseMasking        9 10 11 12 13 14 15
Flatten
Tokenizer
Decoder
Position Embedding
Patch Embedding
Original Image
Image Patches
Visual Tokens 𝐡 2L 𝐡 3L 𝐡 6L 𝐡 7L 𝐡 14L
Masked Image Modeling Head
Reconstructed Image Unused During Pre-Training 234876 [S] [M] [M] [M] [M] [M]
We perform self-supervised learning and then ﬁne-tune the pretrained BE I T on two downstream tasks, i.e., image classiﬁcation, and semantic segmentation.
Experimental results indicate that BE I T outperforms both from-scratch training and previous strong self-supervised models.
Moreover, BE I T is complementary to supervised pre-training.
Performance of BE I T can be further improved by intermediate ﬁne-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 ﬁne-tuning reduce training costs on end tasks.
In addition, we demonstrate that self-supervised BE I T 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 BE I T and conduct extensive ﬁne-tuning experiments on downstream tasks, such as image classiﬁcation, and semantic segmentation.
• We present that the self-attention mechanism of self-supervised BE I T learns to distinguish semantic regions and object boundaries, although without using any human annotation.
M ETHODS
Given an input image x , BE I T encodes it to contextualized vector representations.
As shown in Figure 1, BE I T 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 classiﬁcation, and semantic segmentation), we append task layers upon pretrained BE I T and ﬁne-tune the parameters on the speciﬁc datasets.
2. I MAGE R EPRESENTATIONS
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. I MAGE P ATCH
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 ∈ R H × W × C into
N = HW/P 2 patches x p ∈ R N ×
( P 2 C ) , 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 { x pi } Ni =1 are ﬂattened 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 BE I T.
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. V ISUAL T OKEN
Similar to natural language, we represent the image as a sequence of discrete tokens obtained by an “image tokenizer”, instead of raw pixels.
Speciﬁcally, we tokenize the image x ∈ R H × W × C into z =
[ z 1 , . .
, z N ] ∈ V h × 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 E z ∼ 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 available 1 image tokenizer described in (Ramesh et al., 2021).
We also compare it with a re-implemented tokenizer in Appendix C.
2. B ACKBONE N ETWORK : I MAGE T RANSFORMER
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 { x pi } Ni =1 .
The patches are then linearly projected to obtain patch embeddings Ex pi , where E ∈ R ( P 2 C ) × D .
Moreover, we prepend a special token [S] to the input sequence.
We also add standard learnable 1D position embeddings
E pos ∈ R N × D to patch embeddings.
The input vectors H 0 =
[ e [S] , Ex pi , . . .
, Ex pN ] + E pos 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 H L =
[ h L [S] , h L 1 , . .
, h LN ] are used as the encoded representations for the image patches, where h Li is the vector of the i -th image patch.
P RE -T RAINING BE I T:
M ASKED I MAGE M ODELING
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 ( { x pi } Ni =1 ), and tokenize it to N visual tokens ( { z i } Ni =1 ).
We randomly mask approximately 40% image patches, where the masked positions are denoted as
M ∈ { 1 , . . .
, N } 0 .
4 N .
Next we replace the masked patches with a learnable embedding e [M] ∈ R D .
The corrupted image patches x M = { x pi : i / ∈ M} Ni =1 (cid:83) { e [M] : i ∈ M} Ni =1 are then fed into the
L -layer Transformer as described in Section 2.2.
The ﬁnal hidden vectors { h Li } Ni =1 are regarded as encoded representations of the input patches.
For each masked position { h Li : i ∈ M} Ni =1 , we use a softmax classiﬁer to predict the corresponding visual tokens p MIM ( z (cid:48) | x M ) =
softmax z (cid:48)
( W c h Li + b c ) , where x M is the corrupted image, W c ∈ R |V|× D , and b c ∈ R |V| .
The pre-training objective is to maximize the log-likelihood of the correct visual tokens
z i given the corrupted image:
[Equation]
where D is the training corpus, M represents randomly masked positions, and x M is the corrupted image that is masked according to M .
[Table]
} until
[Table]
is
[Table]
% return M
Rather than randomly choosing patches for the masked positions 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 .
4 N , 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).
BE I T 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 signiﬁcantly outperforms pixel-level auto-encoding.
2. F ROM THE P ERSPECTIVE OF V ARIATIONAL A UTOENCODER
The BE I T pre-training can be viewed as variational autoencoder (Kingma & Welling, 2014) training.
Let x denote the original image, ˜ x the masked image, and z the visual tokens.
Considering the evidence lower bound (ELBO) of the log-likelihood p ( x | ˜ x ) , i.e., recovering the original image from its corrupted version:
[Equation]
where (1) q φ ( z | x ) denotes the image tokenizer that obtains visual tokens; (2) p ψ ( x | z ) decodes the original image given input visual tokens; (3) p θ ( z | ˜ x ) recovers the visual tokens based on the masked image, which is our MIM pre-training task.
We learn the model following a two-stage procedure similar to (van den Oord et al., 2017; Razavi et al., 2019).
In the ﬁrst stage, we obtain the image tokenizer as a discrete variational autoencoder (Ramesh et al., 2021).
Speciﬁcally, the ﬁrst stage minimizes the reconstruction loss − E z i ∼ q φ ( z | x i ) [log p ψ ( x i
| z i )] with an uniform prior as described in Equation (2).
In the second stage, we learn the prior p θ while keeping q φ and p ψ ﬁxed.
We simplify q φ ( z | x i ) to a one-point distribution with the most likely visual tokens ˆ z i = arg max
z q φ ( z | x i ) .
Then Equation (2) can be rewritten as:
[Equation]
where the second term is our BE I T pre-training objective.
2. P RE -T RAINING S ETUP
The network architecture of BE I T follows that of ViT-Base (Dosovitskiy et al., 2020) 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 × 16 input patch size.
We directly borrow the image tokenizer trained by Ramesh et al.
The vocabulary size of visual tokens is 8192 .
We pretrain BE I T on the training set of ImageNet-1K (Russakovsky et al., 2015), which contains about 1 . 2 M images.
Our augmentation policy includes random resized cropping, horizontal ﬂipping, color jittering (Wu et al., 2018).
Notice that we do not use the labels for self-supervised learning.
We use the 224 × 224 resolution in our experiments.
So the input is split to 14 × 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 (Loshchilov &
Hutter, 2019) with β 1 = 0 .
9 , β 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 .
We employ stochastic depth (Huang et al., 2016) with a 0 . 1 rate, and disable dropout.
The 500 k training steps take about ﬁve days using 16 Nvidia Telsa V100 32GB GPU cards.
2. F INE -T UNING BE I T ON D OWNSTREAM V ISION T ASKS
After pre-training BE I T, we append a task layer upon the Transformer, and ﬁne-tune the parameters on downstream tasks, like BERT.
We take image classiﬁcation and semantic segmentation as examples in our work.
It is straightforward to leverage the pre-training-then-ﬁne-tuning paradigm on other vision tasks with BE I T.
Image classiﬁcation.
For image classiﬁcation tasks, we directly employ a simple linear classiﬁer as the task layer.
Speciﬁcally, we use average pooling to aggregate the representations, and feed the global to a softmax classiﬁer.
The category probabilities are computed as softmax(avg( { h Li } Ni =1 W c )) , where h Li is the ﬁnal encoding vector of the i -th image patch,
W c ∈ R D
× C is a parameter matrix, and C is the number of labels.
We maximize the likelihood of labeled data by updating the parameters of BE I T and the softmax classiﬁer.
Semantic segmentation.
For semantic segmentation, we follow the task layer used in SETRPUP (Zheng et al., 2020).
To be speciﬁc, we use pretrained BE I T as a backbone encoder, and incorporate several deconvolution layers as decoder to produce segmentation.
The model is also end-to-end ﬁne-tuned similar to image classiﬁcation.
Intermediate ﬁne-tuning.
After self-supervised pre-training, we can further train BE I T on a datarich intermediate dataset (i.e., ImageNet-1K in our work), and then ﬁnetune the model on the target downstream tasks.
Such intermediate ﬁne-tuning is the common practice of BERT ﬁne-tuning in
NLP (Pruksachatkun et al., 2020).
We directly follow the method for BE I T.
E XPERIMENTS
We conduct full ﬁne-tuning experiments on image classiﬁcation and semantic segmentation.
Moreover, we present various ablation studies for pre-training and analyze the representations learned by BE I T.
We also report linear probes on ImageNet in Appendix D.
3. I MAGE C LASSIFICATION
The image classiﬁcation task classiﬁes input images to various categories.
We evaluate BE I T on the ILSVRC-2012 ImageNet dataset (Russakovsky et al., 2015) with 1k classes and 1 .
3 M images.
We directly follow the most of hyperparameters of DeiT (Touvron et al., 2020) in our ﬁne-tuning experiments for a fair comparison.
We reduce ﬁne-tuning epochs compared with training from scratch, as BE I T has been pre-trained.
Accordingly, we use a larger learning rate with layer-wise decay.
The detailed hyperparameters are summarized in Appendix H.
[Table] [Table]
× 384 .
† : iGPT-1.36B contains 1 .
36 billion parameters, while others are base-size models.
‡ : ViT 384 -B-JFT300M is pretrained with the “masked patch prediction” task on Google’s in-house 300 M images, while others use ImageNet.
Table 1 reports top-1 accuracy on image classiﬁcation.
We compare BE I T 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 ﬁnd that pre-trained BE I T signiﬁcantly improves performance on both datasets.
BE I T improves the performance on ImageNet, which shows the effectiveness under the rich-resource setting.
Moreover, we compare BE I T with previous state-of-the-art self-supervised methods for Transformer, such as DINO (Caron et al., 2021), and MoCo v3 (Chen et al., 2021).
Our proposed method outperforms previous models on ImageNet ﬁne-tuning.
Among them, iGPT-1.36B (Chen et al., 2020a) uses much more parameters (i.e., 1 .
36 B vs 86 M), and ViT-JFT300M (Dosovitskiy et al., 2020) 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.
Speciﬁcally, 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 BE I T and DINO in a multi-task learning manner, which is presented in
Appendix E.
In addition, we evaluate our proposed method with intermediate ﬁne-tuning.
In other words, we ﬁrst pretrain BE I T in a self-supervised manner, and then ﬁne-tune the pretrained model on ImageNet with labeled data.
The results show that BE I T is complementary to supervised pre-training, achieving additional gain after intermediate ﬁne-tuning on ImageNet.
Fine-tuning to 384 × 384 resolution.
After ﬁne-tuning with resolution 224 × 224 , we additionally ﬁne-tune the model on 384 × 384 images by 10 more epochs.
We follow the standard higher-resolution setting of DeiT (Touvron et al., 2020), except using fewer epochs.
Notice that we keep patch size the same for both 224 × 224 and 384 × 384 images.
So the input sequence length of Transformers
[Figure]
DeiT from scratch and ﬁne-tuning BE I T on
ImageNet-1K.
[Table]
becomes longer for higher resolutions.
Table 1 shows that higher resolution improves the BE I T results by 1+ points on ImageNet.
More importantly, BE I T 384 pretrained on ImageNet-1K even outperforms supervised pre-training ViT 384 that uses ImageNet-22K, when they use the same input resolution.
Scaling up to larger size.
We further scale up BE I T to the large size (same as ViT-L).
As shown in
Table 1, ViT 384 -L is worse than ViT 384 on ImageNet, when training from scratch.
The results veriﬁes the data-hungry issue of vision Transformers.
Supervised pre-training on ImageNet-22K partially relieves the issue, where ViT 384 -L ﬁnally outperforms ViT 384 by 1 . 2 .
In comparison, BE I T-L is better than BE I T by 2 . 0 , and BE I T 384 -L outperforms BE I T 384 by 1 . 7 .
In other words, the beneﬁts of scaling up BE I T from base to large are greater than supervised pre-training with ImageNet-22K.
More importantly, comparing between BE I T 384 with ViT 384 that conducts supervised pre-training on ImageNet-22K, the improvements of BE I T become greater along with scaling the size from base (i.e., 0 . 6 ) to large (i.e., 1 . 1 ).
The results suggest that BE I T tends to help more for extremely larger models (such as 1B, or 10B), especially when labeled data are insufﬁcient 2 to conduct supervised pre-training 3 for such large models.
Convergence curves.
Figure 2 compares the convergence curves of the training-from-scratch and pre-training-then-ﬁne-tuning paradigms.
We ﬁnd that ﬁne-tuning BE I T not only achieves better performance, but also converging much faster than training DeiT from scratch.
Moreover, ﬁne-tuning
BE I T can reach reasonable numbers within very few epochs.
3. S EMANTIC S EGMENTATION
Semantic segmentation aims to predict a corresponding class for each pixel of the input image.
We evaluate BE I T on the ADE20K benchmark (Zhou et al., 2019) with 25 K images andsemantic 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 (Zheng et al., 2020).
On ADE20K, we use
Adam (Loshchilov & Hutter, 2019) as the optimizer.
The learning rate is set to 1e-3 with layer-wise decay similar to image classiﬁcation.
We conduct ﬁne-tuning for 160K steps.
The batch size is 16 .
The detailed hyperparameters are described in Appendix I.
As shown in Table 3, we compare BE I T with supervised pre-training that relies on labeled data of ImageNet.
We ﬁnd that our proposed method achieves better performance than supervised pretraining, although BE I T does not require manual annotations for pre-training.
Moreover, we employ intermediate ﬁne-tuning for BE I T on ImageNet, i.e., we ﬁrst ﬁne-tune pretrained BE I T on ImageNet, and then ﬁne-tune the model on ADE20K.
The results indicate that intermediate ﬁne-tuning further improves BE I T on semantic segmentation.
[Table]
3. A BLATION S TUDIES
We conduct ablation studies to analyze the contributions of each component in BE I
T.
The models are evaluated on image classiﬁcation (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 ﬁnd that blockwise masking is beneﬁcial 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 signiﬁcantly 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 BE I
T. Third, we ablate the usage of visual tokens and blockwise masking together.
We ﬁnd 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 BE I T with different training steps.
Pre-training the model longer can further improve performance on downstream tasks.
3. A NALYSIS OF S ELF -A TTENTION M AP
We show that the self-attention mechanism in BE I T can separate objects, even though our pre-training does not rely on any manual annotation at all.
Similar properties are also observed by Caron et al.
The probing images are taken from the MS COCO (Lin et al., 2014) 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, BE I T learns to distinguish semantic regions using self-attention heads, without any task-speciﬁc supervision.
The property partially indicates the reason why BE I T is able to help downstream tasks.
Such knowledge acquired by BE I T potentially improves the generalization ability of ﬁne-tuned models, especially on small-scale datasets.
R ELATED W ORK
Self-supervised visual representation learning.
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 (Noroozi & Favaro, 2016), colorization (Zhang et al., 2016), and predicting rotation angles (Komodakis & Gidaris, 2018).
In addition, Trinh et al.
(2019) 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 pretraining (Noroozi & Favaro, 2016).
The recent strand of research follows contrastive paradigm (Wu et al., 2018;
Oord et al., 2018; Hjelm et al., 2019; Bachman et al., 2019;
He et al., 2020; Chen et al.,
2020b;c).
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 (Wu et al., 2018;
He et al., 2020) or large batch size (Chen et al., 2020b).
BYOL (Grill et al., 2020) and SimSiam (Chen & He, 2020) further eliminate the requirement of negative samples, et et al
Self-supervised vision Transformers.
Pre-training vision Transformers has received signiﬁcant attention recently due to the data-hungry issue.
iGPT (Chen et al., 2020a)
ﬁrst 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 (Dosovitskiy et al., 2020) conducts a preliminary exploration with the masked patch prediction task, which predicts the 3-bit mean color of the masked patches.
Dosovitskiy et al.
(2020) 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 (Chen et al., 2021; Xie et al., 2021), and self-distillation (Caron et al., 2021).
In comparison, BE I T can achieve several times of improvement in terms of pre-training throughput (Appendix E), and memory consumption.
The advantages make BE I T appealing to scale up vision Transformers.
C ONCLUSION
We introduce a self-supervised pre-training framework for vision Transformers, achieving strong ﬁne-tuning results on downstream tasks, such as image classiﬁcation, 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 BE I T pre-training in terms of data size and model size.
Moreover, we will conduct multi-modal pre-training in a more uniﬁed way, using the similar objectives and the shared architecture for texts and images.
A A RCHITECTURE V ARIANTS OF V ISION T RANSFORMER
We use the standard vision Transformer (ViT; Dosovitskiy et al.
in the experiments for fair comparisons.
In addition, we ﬁnd that LayerScale (Touvron et al., 2021) and relative position bias (Bao et al., 2020; Raffel et al., 2020) 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 classiﬁcation and ADE20K semantic segmentation.
We denote the improved architecture as BE I T + 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.
Architecture
[Table]
For ADE20K, we use UperNet (Xiao et al., 2018) as the task layer, and report mIoU scores of single-scale inference.
B C OMPARISON WITH L ARGE -S CALE S UPERVISED P RE -T RAINING
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 BE I T 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 BE I T + 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 ﬁne-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]
[Table] [Table]
Table 6 compares BE I T with previous state-of-the-art supervised pre-training (Dosovitskiy et al., 2020; Zhai et al., 2021) on ImageNet ﬁne-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 BE
I T pretraining can catch up with only ImageNet-22k (14M).
Speciﬁcally, BE I T-L ﬁne-tuned on ImageNet22K achieves comparable performance with ViT-L trained on Google JFT-3B.
Moreover, BE I T-L obtains 89 .
5% top-1 accuracy on ImageNet after intermediate ﬁne-tuning on an in-house 70M dataset.
The results indicate that BE I T 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 ﬁne-tuning results on the ADE20K semantic segmentation benchmark.
Following Swin (Liu et al., 2021b), we use the same task layer (i.e., UperNet; Xiao et al.
and evaluate the models at the resolution 640 × 640 .
The BE I T-L model obtains state-of-the-art performance on ADE20K.
[Table]
semantic segmentation.
We follow Swin-L (Liu et al., 2021b) to use UperNet (Xiao et al., 2018) as the task layer and evaluate at resolution 640 × 640 .
C A BLATION S TUDIES OF I MAGE T OKENIZER
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 (Ramesh et al., 2021).
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 ﬁne-tuning performance compared with the off-the-shelf DALL-E tokenizer.
[Table]
D L INEAR P ROBES ON I MAGE N ET
We evaluate linear probes on ImageNet for various pretrained vision Transformers.
We compare
BE I T with two main strands of work, namely discriminative and generative self-supervised learning.
The ﬁrst one applies discriminative learning for pre-training, such as contrastive learning (Chen et al., 2021), and self distillation (Caron et al., 2021).
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 (Chen et al., 2020a) and ours, usually do not pretrain such global feature aggregation, which tends to make linear probes difﬁcult.
Following iGPT (Chen et al., 2020a), 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 ﬁnal layer.
Similarly, we ﬁnd that the best layer lies in 9 -th layer for BE I T-B, and 14 -th layer for
BE I T-L.
To be speciﬁc, we use AdamW (Loshchilov & Hutter, 2019) 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 (Caron et al., 2021), which uses random resize crops and horizontal ﬂips augmentation during training and evaluates on central crops.
[Table]
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 ﬁxed and only update the linear layer.
So the pre-training of global aggregation of image-level features is beneﬁcial to linear probing in DINO and MoCo v3, although full ﬁne-tuning eliminates the gap.
Moreover, the results indicate that increasing the model size from base (86M) to large (304M) signiﬁcantly improves accuracy for our proposed method.
In contrast, the gap between base- and large-size MoCo v3 is smaller.
We also ﬁnd that BE I T outperforms iGPT by a large margin even using much fewer parameters.
E M ULTI -T ASK P RE -T RAINING WITH DINO
We train the pre-training tasks of BE I T and DINO (Caron et al., 2021) 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 classiﬁcation.
Moreover,
BE I T is more efﬁcient in terms of pre-training speed, as DINO has two copies of Transformer parameters for self-distillation and multi-crop augmentation (Caron et al., 2020).
For the throughput comparisons between BE I T and BE I T+DINO, we set batch size to the same.
Because BE I T is also more memory-efﬁcient, we can use larger batch size to fully utilize GPU cards, which obtains greater speedup in practice than the reported numbers.
[Table]
measures the speed, where larger numbers indicate faster pre-training.
F I MAGE C LASSIFICATION ON CIFAR-
In addition to ImageNet classiﬁcation
[Table]
we conduct ﬁne-tuning experiments on the CIFAR100 (Krizhevsky & Hinton, 2009) 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 (Chen et al., 2021).
In comparison, BE I T achieves 90 .
1% with the help of pre-training.
The results indicate that BE I T can greatly reduce the requirement of annotation efforts.
BE I T also outperforms MoCo v3.
Moreover, intermediate ﬁne-tuning on ImageNet-1K further improves the results on CIFAR-100.
Models
[Table]
G H YPERPARAMETERS FOR P RE -T RAINING
[Table]
H H YPERPARAMETERS FOR I MAGE C LASSIFICATION F INE -T UNING Hyperparameters
CIFAR
[Table]
Peak learning rate {
[Table]
4e
[Table]
Fine-tuning
[Table]
Batch
[Table]
Warmup
[Table]
Layer
[Table]
rate
[Table]
Label smoothing ε
[Table]
Stoch
[Table]
Dropout
[Table]
clipping
[Table]
Erasing
[Table] [Table]
resolution
[Table]
Augment
[Table]
Mixup
[Table] [Table]
Cutmix
[Table]
I H YPERPARAMETERS FOR ADE20K S EMANTIC S EGMENTATION F INE -T UNING
[Table]