MST: Masked Self-Supervised Transformer for Visual Representation

1 Introduction . As Yann LeCun said , “ if intelligence is a cake , the bulk of the cake is unsupervised learning ” . This sentence reflects that Un-/Self-supervised Learning played a central role in the resurgence of deep learning . Common approaches focus on designing different pretext tasks [ 10 , 29 , 14 , 3 , 4 , 6 , 5 , 13 , 1 , 36 ] and aim to learn useful representations of the input data without relying on human annotations . It then uses those representations in downstream tasks , such as image classification , objection detection , and semantic segmentation . ∗Work done as an intern at SenseTime Research . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In computer vision , previous methods focus on designing different pretext tasks . One of the most promising directions among them is contrastive learning/instance discrimination [ 17 , 23 ] , which regards each instance in the training dataset as a single category . Based on instance discrimination [ 14 , 4 , 6 , 5 , 13 , 1 ] , some methods show the effectiveness in the image classification task . They successfully bridge the performance gap between self-supervised and full-supervised methods . However , almost all of self-supervised learning methods , which formulate the learning as image-level prediction using global features , are suboptimal in the pixel-level predictions [ 14 , 1 , 13 ] , such as object detection and semantic segmentation . Also , InfoMin [ 35 ] finds that high-level features do not truly matter in transferring to dense prediction tasks . Here , current self-supervised learning may overfit to image classification while not being well tamed for downstream tasks requiring dense prediction . Meanwhile , large-scale pre-trained models have become the prevailing formula for a wide variety of Natural Language Processing ( NLP ) tasks due to its impressive empirical performance . These models typically abstract semantic information from massive unlabeled corpora in a self-supervised manner . The Masked Language Modeling ( MLM ) [ 10 ] has been widely utilized as the objective for pre-training language models . In the MLM setup , a certain percentage of tokens within the input sentence are randomly masked , and the objective is to predict the original information of the masked tokens based only on its context . In NLP tasks , we found that the different mask strategies used in the MLM framework had a great impact on the performance of the model . However , in the field of vision , images have higher-dimensional , noisy , and redundant format compared to text . The main information of input images is randomly distributed in tokens . If tokens are randomly masked , it will lead to poor performance . Some of previous methods use random tokens , such as iGPT [ 3 ] and ViT [ 11 ] . iGPT trains self-supervised Transformers using an amount of 6801M parameters and achieves 72.0 % Top-1 accuracy on ImageNet by masking and reconstructing pixels , while ViT trains ViT-B model on the JFT-300M dataset , and the result is significantly lower than the supervised model . The random MLM is prone to mask the tokens of crucial region for images , resulting in misunderstanding , and is not suitable for directly applying to self-supervised vision Transformers . In order to avoid masking the tokens of crucial region , we propose a masked token strategy based on the multi-head self-attention map , which dynamically masks some tokens of patches without damaging the crucial structure for self-supervised learning . Notably , the strategy would not increase the training time . Also , predicting original tokens alone may cause the model to over-emphasize local region , and therefore suppress the ability to recognize objects . Hence , in this paper , we present a novel Masked Self-supervised Transformer approach named MST , which can explicitly capture the local context of an image while preserving the global semantic information . In addition , a global image decoder is further exploited to recover the spatial information of the image and is thus more friendly to the downstream dense prediction tasks . We validate our method on multiple visual tasks . In particular , on the ImageNet linear evaluation protocol , we reach 76.9 % top-1 accuracy with DeiT-S and achieve the state-of-the-art performance . Overall , we make the following contributions : • We propose a new masked self-supervised transformer approach called MST . It makes full use of self-attention map to guide the masking of local patches , thus enhancing the understanding of local context semantics in pre-training without damaging the crucial structure . • Our method can effectively recover the spatial information of the image by a global image decoder , which is vital for the downstream dense prediction task and greatly improves the versatility and scalability of the pre-training model . • Extensive experiments demonstrate the effectiveness and transfer ability of our method . Specifically , the results on ImageNet [ 9 ] , MS COCO [ 18 ] and Cityscapes [ 8 ] show that our method outperforms previous state-of-the-art methods . 2 Related Works . 2.1 Self-supervised visual representation learning . Following MLM paradigm in NLP [ 10 , 25 ] , iGPT [ 3 ] trains self-supervised Transformers by masking and reconstructing pixels , while ViT [ 11 ] masks and reconstructs patches . Recently , the most competitive pretext task for self-supervised visual representation learning is instance discrimination [ 14 , 4 , 6 , 5 , 13 , 1 ] . The learning objective is simply to learn representations by distinguishing each image from others , and this approach is quite intractable for large-scale datasets . MoCo [ 14 ] improves the training of instance discrimination methods by storing representations from a momentum encoder instead of the trained network . SimCLR [ 4 ] shows that the memory bank can be entirely replaced with the elements from the same batch if the batch is large enough . In order to avoid comparing every pair of images and incur overfitting , BYOL [ 13 ] directly bootstraps the representations by attracting the different features from the same instance . SwAV [ 1 ] maps the image features to a set of trainable prototype vectors and proposes multi-crop data augmentation for self-supervised learning to increase the number of views of an image . MoCov3 [ 7 ] and DINO [ 2 ] apply the self-supervised learning methods of computer vision to Transformers and achieve superior performance in image classification task . These works achieve comparable results compared to supervised ImageNet [ 9 ] pre-training . The success of these methods suggest that it is of central importance to learn invariant features by matching positive samples . However , almost all of these self-supervised learning methods formulate the learning process as image-level prediction using global features , so they lack the ability to pay attention to local features . 2.2 Self-supervised dense prediction learning . Based on the existing instance discrimination , some researchers propose self-supervised dense prediction methods . Self-EMD [ 19 ] adopts Earth Mover ’ s Distance ( EMD ) to compute the similarity between two embedding . Insloc [ 33 ] pastes image instances at various locations and scales onto background images . The pretext task is to predict the instance category given the composited images as well as the foreground bounding boxes . PixPro [ 31 ] directly applies contrastive learning at the pixel level . DenseCL [ 28 ] presents dense contrastive learning by optimizing a pairwise contrastive loss at the pixel level between two views of input images . VADeR [ 24 ] and FlowE [ 32 ] also learn dense image representations for downstream detection and segmentation tasks . Meanwhile , HED [ 34 ] focuses on downstream segmentation task . These methods also show the effectiveness in detection and segmentation tasks but get poor performance on image classification tasks . In a word , these methods overfit a single task and can not train a general pre-training model . 3 Methods . The pipeline of our proposed MST is shown in Figure 1 . We propose a Masked Self-supervised Transformer ( MST ) approach , which creatively introduces attention-guided mask strategy and uses it to complete image restoration task . Our method is combined with some classical components of instance discrimination , such as the momentum design , asymmetric data augmentations , and multi-crop strategies . Here , we first review the basic instance discrimination method in 3.1 . Then , the mechanism and effect of our attention-guided mask strategy are explained in 3.2 . Finally , we describe the reconstruction branch and the training target of our method in 3.3 . 3.1 The basic instance discrimination method . As noted in prior works [ 4 , 14 , 13 , 29 , 1 ] , many existing augmentation policies adopt random resized cropping , horizontal flipping , color jittering and so on . We generate multiple views for each image x under random data augmentation according to multi-crop [ 1 ] . This operation can acquire two standard resolution crops x1 and x2 representing the global view and sample N low-resolution crops indicating partial view . They are encoded by two encoders , teacher network ft and student network fs , parameterized by θt and θs respectively , and outputting vectors Ot and Os . Both encoder fs and ft consist of a Transformer backbone and a projection head [ 5 ] , which share the same architecture with different parameters . The parameters θt of fixed encoder ft is updated by the moving-average of θs according to Eq ( 1 ) . θt = m ∗ θt + ( 1−m ) ∗ θs ( 1 ) Given a fixed teacher network ft , the student network fs learns the parameters θs by minimizing cross entropy loos as Eq ( 2 ) . LCE ( θs ) = ∑ i∈ { 1,2 } N+2∑ j=1 j 6=i −ft ( θt ; xi ) log ( fs ( θs ; xj ) ) ( 2 ) 3.2 Masked token strategy . Random mask strategy . Inspired of the MLM strategy for natural language pre-training , we apply the random mask strategy to self-supervised learning . Given a dataset Q without manual annotations , and e = ( e1 , ... , en ) denote a image of n tokens , where i = 1 , ... , n. Let m = ( m1 , ... , mn ) denote a binary vector of length n , where mi ∈ { 0 , 1 } , representing the mask over image . According to BERT [ 10 ] , the m can be obtained with probability p by Eq ( 3 ) , and the p is 0.15 by default . mi = { 1 , probi < p 0 , otherwise ( 3 ) Algorithm 1 Pseudo code of attention-guided mask strategy in a PyTorch-like style . # l ( ) : linear projection # f_s : backbone + projection head # f_t : backbone + projection head # mask_embedding : learnable token # p : mask probability e_t = f_t.l ( x ) # linear projection patch_attention , _ = f_t.Transformer ( e_t ) importance = measure_importance ( patch_attention ) # acquire threshold e_s = f_s.l ( x ) # linear projection mask = M ( e_s , patch_attention , importance ) # generate mask e_s_masked = ( 1-mask ) * e_s + mask * mask_embedding _ , _ = f_s.Transformer ( e_s_masked ) def M ( embedding , patch_attention , importance ) : B , L , _ = embedding.shape mask_tokeep = zeros ( ( B , L ) ) mask_remove = bernoulli ( ones ( ( B , L ) ) * p ) mask = where ( importance > patch_attention , mask_tokeep , mask_remove ) return mask According to Eq ( 3 ) , the tokens of crucial and nonessential regions have the same probability of being masked . As shown in Figure 2 ( c ) , we observe that the random mask strategy may eliminate tokens of crucial regions that are responsible for recognizing objects , resulting in indistinguishable semantic features for input images . The random mask strategy is prone to mask crucial regions for images , and suppress the ability of network to recognize objects . It is not suitable to directly apply this strategy to self-supervised vision Transformers and the overall performance would deteriorate if the mask strategy is not properly modulated . Attention-guided mask strategy . In this section , we propose our attention-guided mask strategy for dynamically controlling the fidelity of masked tokens and thereby decreasing the probability of masking crucial regions in self-supervised Transformer . Meanwhile , our strategy does not increase additional time consumption . Our algorithm is shown as Alg . 1 . Our framework consists of two networks , teacher network ft and student network fs , with the same transformer architecture . Let x denotes the input image . It is firstly projected to a sequence of n 1-d tokens e = e1 , ... , en , and then processed by several self-attention layers . Each self-attention layer [ 27 ] owns three groups of embeddings for one token , denoted as Qi ( query ) , Ki ( key ) , Vi ( value ) . The attention map is calculated as the correlation between the query embedding of class token Qcls and key embeddings of all other patches K. It is averaged for all heads as Eq ( 4 ) . We output the attention map from the last layer in the teacher network to guide our strategy . Attn = 1 H H∑ h=1 Softmax ( Qclsh · KTh√ d ) ( 4 ) We sort the attention of different patches for each image in ascending order , and take the sorted attention value of 1/num of total tokens as the threshold τ , where num is the hyperparameter for selection . This means that the lowest 1/num of total tokens are selected as the masked candidates . The student model receives the importance of different patches and generates the mask m with probability p , according to the Bernoulli distribution as Eq ( 5 ) . probi refers to the probability of randomly generated . mi = { 1 , probi < p and Attni < τ 0 , otherwise ( 5 ) We use m e to denote the final masked tokens as Eq ( 6 ) . Follow the BERT [ 10 ] , the masked regions are filled with a learnable mask embedding [ MASK ] . Our strategy can ensure the patches with the highest scores are always presented ( in Figure 2 ) . ( m e ) = { [ MASK ] , mi = 1 ei , mi = 0 ( 6 ) The attention-guided mask strategy can benefit pre-training models in two ways : 1 . The models utilize contextual information to understand the relationship of different patches , thus preserving the global semantic information of the image while paying more attention to the local details of the image . 2 . Our strategy can avoid masking crucial regions while replacing nonessential regions with the learnable mask embedding , making the models focus on the crucial regions .

Summary: This paper presents a new self-supervised learning method based on a combination of mean-teacher distillation (in this case DINO) and masked patch modelling (like masked language modelling in NLP). They use this to train transformer based vision models and show downstream results for ImageNet classification using linear probing and knn eval, MS COCO object detection and instance segmentation using Mask-RCNN framework and Cityscapes semantic segmentation using SERT framework.  Their key technical contributions are (a) the overall training framework combining DINO with masked patch modelling. The latter requires a CNN decoder and reconstructs the original global image and not just the masked patches and (b) a smart mask generation scheme based on the attention map of the CLS token which turns out to be crucial for good results. They use the attention maps corresponding to the CLS token to select patches that are not important. The idea is that if important regions of the image are masked out then the model will end up modelling only local statistics.  At a deeper level, their method is trying to retain spatial structure in the latent representation while capturing global image properties. If I understood correctly, the spatial structure comes from masked patch modelling and global image properties come from DINO.