Efficient Self-supervised Vision Transformers for Representation Learning

This paper investigates two techniques for developing efficient self-supervised vision transformers ( EsViT ) for visual representation learning . First , we show through a comprehensive empirical study that multi-stage architectures with sparse self-attentions can significantly reduce modeling complexity but with a cost of losing the ability to capture fine-grained correspondences between image regions . Second , we propose a new pre-training task of region matching which allows the model to capture fine-grained region dependencies and as a result significantly improves the quality of the learned vision representations . Our results show that combining the two techniques , EsViT achieves 81.3 % top-1 accuracy on the ImageNet linear probe evaluation , outperforming prior arts with around an order magnitude of higher throughput . When transferring to downstream linear classification tasks , EsViT outperforms its supervised counterpart on 17 out of 18 datasets . The code and models will be publicly available . 1 INTRODUCTION . Self-supervised learning ( SSL ) with Transformers ( Vaswani et al. , 2017 ) has become a de facto standard of model choice in natural language processing ( NLP ) . The dominant approaches such as GPT ( Radford et al. , 2018 ) and BERT ( Devlin et al. , 2019 ) are pre-training on a large text corpus and then fine-tuning to various smaller task-specific datasets , showing superior performance . Larger Transformers pre-trained with larger-scale language datasets often lead to a stronger generalization ability , demonstrated by improved performance in downsteam tasks ( with no sign of performance saturation yet ) , as exemplified in GPT-3 ( Brown et al. , 2020 ) . In computer vision ( CV ) , however , self-supervised visual representation learning is still dominated by convolutional neural networks ( CNNs ) . Sharing a similar goal/spirit with NLP , SSL in CV aims to learn general-purpose image features from raw pixels without relying on manual supervisions , and the learned networks are expected to serve as the backbone of various downstream tasks such as classification , detection and segmentation . Recently , impressive performance have been achieved by CNN-based SSL , outperforming state-of-the-art ( SoTA ) fully-supervised pre-training methods ( He et al. , 2020 ; Caron et al. , 2020 ) on tasks with a limited number of labels . The key to success is view-level learning : maximizing agreement of learned representations between differently augmented views of the same example . Recent works , including SimCLR-v2 ( Chen et al. , 2020d ) , BYOL ( Grill et al. , 2020 ) and SwAV ( Caron et al. , 2020 ) , have scaled up the CNN-based models to hundreds of millions of parameters . However , SSL has not enjoyed the same scaling success in CV as that in NLP . Several attempts have been made to close the gap by combining SSL with Transformer and selfattention architectures . Early works include Selfie ( Trinh et al. , 2019 ) , which generalizes the concept of masked language modeling of BERT for images . The idea has been recently revisited in Vision Transformer ( ViT ) ( Dosovitskiy et al. , 2021 ) via pre-training on a much larger scale dataset , e.g. , JFT300M . ImageGPT ( iGPT ) ( Chen et al. , 2020b ) generalizes the concept of auto-regressive language modeling of GPT for images , showing encouraging ImageNet recognition accuracy with a large model size . Contrastive learning with ViT has also been studied very recently in DINO ( Caron et al. , 2021 ) and MoCo-v3 ( Chen et al. , 2021 ) , where new SoTA result by linear probe evaluation on ImageNet-1K is achieved , by exhaustively consuming computation resource on full self-attention operators with long sequences of split image patches . Aiming to improve the efficiency of Transformer-based SSL , this paper presents Efficient selfsuperivsed Vision Transformers ( EsViT ) , by using a multi-stage architecture and a region-based pre-training task for self-supervised representation learning . Our main findings and contributions can be summarized as follows : ( 1 ) An intriguing property of self-supervised monolithic Transformers is firstly reported in our paper : automatic discovery of semantic correspondence between local regions . ( 2 ) We present the first comprehensive empirical study to show the pros and cons of multi-stage vision Transformer architectures for SSL . Though greatly reducing compute complexity , we find that the multi-stage architecture causes the loss of the property in ( 1 ) . ( 3 ) A region matching pre-train task is proposed to alleviate the issue in ( 2 ) , and further improve the learned representations and attentions . ( 4 ) We validate the new EsViT , which combines the two techniques , on a range of tasks . It significantly reduces the cost in building SoTA SSL vision systems , as summarized in Figure 1 , and shows better scaling performance on accuracy vs. throughput and model size . Under the linear evaluation protocol , EsViT achieves 81.3 % top-1 accuracy , showing the best performance compared with all systems , and is 3.5× parameter-efficient and has at least 10× higher throughput than previous SoTA ( 81.0 % , MoCo-v3 with ViT-BN-L/7 ( Chen et al. , 2021 ) ) . Compared with its supervised counterpart Swin Transformers ( Liu et al. , 2021 ) , EsViT shows superior performance on 17 out 18 datasets , when transferring the learned representations to downstream linear classification tasks . 2 METHODS . Transformer-based SSL methods emerge very recently to lead the state-of-the-art performance on the ImageNet linear probe task ( Chen et al. , 2021 ; Caron et al. , 2021 ) . It inherits the successes from ( 1 ) monolithic Transformer architectures that dominate in NLP ( Devlin et al. , 2019 ; Radford et al. , 2018 ) , and ( 2 ) instance-level contrastive learning objectives that demonstrate arguably the best SSL performance in computer vision ( Chen et al. , 2020c ) . Though simple and effective , the existing Transformer-based SSL methods require a large amount of compute resources ( e.g. , > 1.7 TPU years of training ) to reach SoTA performance . We believe that the SSL system efficiency is highly related to two ingredients : the network architecture and the pre-train task . To strike for a better tradeoff between accuracy and efficiency , we present EsViT , showing better synergy of networks ( a multi-stage Transformer architecture ) and pre-train tasks ( a non-contrastive region-matching task ) . 2.1 NETWORK ARCHITECTURES : FROM MONOLITHIC TO MULTI-STAGE VIT BACKBONE . Multi-stage ViT . This paper presents the first empirical study of multi-stage Transformer architectures ( Vaswani et al. , 2021 ; Wang et al. , 2021 ; Liu et al. , 2021 ; Zhang et al. , 2021 ; Wu et al. , 2021 ) for SSL . Each stage consists of a patch merging/embedding module , and a Transformer with sparse self-attention module . ( i ) The patch merging module plays a slightly different roles in different stages . In the first stage , it splits an input RGB image into non-overlapping patches . Each patch is treated as a “ token ” , constructed as a concatenation of the raw pixel RGB values , which is further projected into a C-dimension feature . In the later stage , the patch merging module concatenates the features of each group of 2× 2 neighboring patches , and applies a linear layer on the 4C-dimensional concatenated features . This reduces the number of tokens by a multiple of 2× 2 = 4 , and the output dimension is set to 2C . ( ii ) A Transformer with sparse self-attention module are then employed to enable interactions among the merged features . The two modules above are repeated for multiple times , typically 4 times , resulting in a multi-stage ViT . As a result , a hierarchical representation is generated : the number of tokens is reduced and the feature dimension ( and the number of heads in self-attentions ) of each token is increased , as the network gets deeper . An overview comparison of the monolithic and multi-stage Transformer architectures for SSL is illustrated in Figure 7 in Appendix . An intriguing property of self-supervised monolithic ViT . Though straightforward in implementation , changing from monolithic to multi-stage architecture without careful treatments may lose some desirable properties of self-supervised Transformers In out study , we first empirically note an intriguing property of self-supervised monolithic ViT ( Caron et al. , 2021 ) : the pre-trained model exhibits a very strong ability to automatically discovers correspondences , even without a region-level matching objective specified in training . We quantitatively evaluate the correspondence learning to illustrate this property , as discussed in the following process . ( i ) Simulated benchmark . Based on 50K images in the ImageNet validation dataset , we create a simple evaluation benchmark with mild augmentations : For a center-crop image , we apply HorizontalFlip , then ColorJitter and RandomGrayscale to create a new augmented view . In this way , ground-truth correspondences are created . ( ii ) Evaluation process . Given two views of the same image , we use the pre-trained backbone to extract the top-layer features , and find the feature vector in one view that best matches the other in terms of highest cosine similarity . The accuracy is measured as the averaged percentage of correctly identifying the region-to-region correspondences . Please see details in Section C.7 in Appendix . ( iii ) Results . We quantitatively show that a self-supervised monolithic ViT yields 95 % accuracy . However , simply replacing the network with a multi-stage Transformer yields only 66 % accuracy . This significant degradation ( absolute 29 % accuracy drop ) reveals the loss of the correspondence learning property . We first raise this critical problem , and believe that it has a large impact on the pre-trained model ’ s performance in various downstream tasks . 2.2 PRE-TRAINING TASKS : DELVING INTO VIEWS WITH REGIONS . We employ a non-contrastive learning framework to build our SSL method . Specifically , Selfdistillation with no labels ( DINO ) ( Caron et al. , 2021 ) is considered . It leverages the knowledge distillation learning paradigm where a student network gθs is trained to match the output of a given teacher network gθt , parameterized by θs and θt respectively . The neural network g is composed of a backbone f ( e.g. , Transformers or ConvNets ) , and of a projection head h : g = h ◦ f . The features used in downstream tasks are the output of backbone f . In SSL , different augmented views x̃ of an image x are fed into backbone network to obtain feature maps z = f ( x̃ ) . Two MLP heads followed by softmax per network further convert the feature vectors z ∈ z into probability vectors p = h ( z ) ; one head for view-level and the other head for region-level , respectively . More precisely , from a given image , we generate a set V of different views1 following ( Caron et al. , 2021 ) . The resulting feature map at the top layer for each view is z = [ z1 , . . . , zT ] , where T is the sequence length , and zi is a region-level representation for the local patch at position i . Average pooling is applied to obtain the view-level representation z̄ = avg-pool ( z ) . View-level task Given the augmented view set for student V and teacher V∗ , a set of pairs P = { ( s , t ) |x̃s ∈ V , x̃t ∈ V∗ and s 6= t } is constructed to perform cross-view prediction tasks . We consider the pre-training task at the view level proposed by ( Caron et al. , 2021 ) : LV = 1 |P| ∑ ( s , t ) ∈P MV ( s , t ) , with MV ( s , t ) = −ps log pt , ( 1 ) where ps = h ( z̄s ) and pt = h ( z̄t ) are the probability output of an MLP head h over the view-level representations z̄s and z̄t , learned by student and teacher , respectively . In DINO , ViT/DeiT are considered , hence the view-level representation is the feature of the [ CLS ] token . 1This set often contains views of two different resolutions V = [ Vg , Vl ] , where Vg = { x̃gi |i = 1 , 2 } is a global-view set of higher resolution , and Vl = { x̃li |i = 1 , . . . , 8 } is a local-view set of lower resolution . All views V are passed through the student while only the global views Vg are passed through the teacher . Region-level task In ( Caron et al. , 2021 ) , the LV encourages “ local-to-global ” correspondences only at a coarse level : the large crop and the small crop are matched in the view level , leaving region-to-region correspondence unspecified . In monolithic Transformers , the drop paths and skip connections from low-level features to high-level features help the the latter to remain discriminative , thus maintain good region-matching performance . However , such a property gets diluted due to the merging operators in multi-stage Transformers . As shown in our experiments later , training a multi-stage network with LV only indeed results in sub-optimal representations . Further , it could be a waste of computation not to leverage region-level features z that are computed in the process of extracting view-level feature . Inspired by the success of masked language modeling task in BERT , we argue that it is important to have region-level pre-training task for computer vision , so that the model can ( 1 ) amortize the computation and fully leverage the extracted region-level features , and ( 2 ) take into account the co-occurrences/structures between local features . Unfortunately , directly performing masked patch prediction ( MPP ) for the multi-stage Transformer architecture is infeasible , as the one-to-one correspondences between the input visual tokens and output features get diluted due to the merging operation . Even for monolithic architectures , MPP has not been proved effective in computer vision , as empirically shown in ( Dosovitskiy et al. , 2021 ) . To address this problem , we propose a non-contrastive , region-matching method that directly works at the level of local features by taking into account their correspondences : LR = 1 |P| ∑ ( s , t ) ∈P MR ( s , t ) , with MR ( s , t ) = − 1 T T∑ i=1 pj∗ log pi , j ∗ = arg max j zTi zj ‖zi‖‖zj‖ , ( 2 ) where pi = h′ ( zi ) and pj = h′ ( zj ) are the probability outputs of a new MLP head h′ over the local features of student zi ∈ zs and teacher zj ∈ zt , respectively . j∗ is the index of the feature in zt that best matches the i-th feature in zs , in the sense of highest cosine similarity . Note that zi and zj∗ are contextualized features of two best matched regions from different augmentated views , minimizing LR encourages different contexts ( i.e. , surrounding regions ) to learn invariant features , and thus captures the region-dependency . The overall pre-training objective of EsViT is L = LR + LV , we learn to match the feature distributions at both the view and region levels by minimizing the cross-entropy loss w.r.t . the parameters of the student network gθs . A visual illustration is in Figure 2 , and the full algorithm is in Appendix . We updates teacher/student network alternatively : ( i ) Given a fixed teacher network , the student network is updated by minimizing the full cross-entropy loss : θs ← arg minθs L ( s , t ; θs ) . ( ii ) The teacher model is updated as an exponential mov- ing average ( EMA ) of the student weights θt ← λθt + ( 1−λ ) θs , with λ following a cosine schedule from 0.996 to 1 during training . By default , the full objective L is used from the beginning . One can also load a checkpoint trained by LV only , and add LR for continual pre-training , which is shown effective in boosting performance in our experiments . Computational overhead Note that applying LR on the traditional monolithic Transformer architecture can be prohibitively computationally expensive , as it requires O ( T 2 ) to compute LR . For a typical image of resolution 224×224 , the feature map length of ViT/DeiT ( with patch size 16 ) at the top layer is T = 196 , while the multi-stage architecture yields T = 49 , which requires 3 times less compute in computing LR . To empirically illustrate this , we show in Appendix Section C.2 that LR adds acceptable extra memory and computational cost ( around 1.2 and 1.05 × , respectively ) for multi-stage Transformers , while it will quickly go out-of-memory for monolithic Transformers when the batch size is increased .

Summary: The paper investigates how to use self-supervised learning for multi-stage visual transformer models. Previous works have shown that SSL can learn image correspondences and lead to performant pre-trained models, while the multi-stage models can reduce the computation cost dramatically. This work tries to merge these two trends together. The solution is a new region-based loss that can be applied to the local features. The comprehensive experiments show the advantages of the resulting models on multiple tasks.