Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 METHODS
text: 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).

heading: 2.1 NETWORK ARCHITECTURES: FROM MONOLITHIC TO MULTI-STAGE VIT BACKBONE
text: 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.

heading: 2.2 PRE-TRAINING TASKS: DELVING INTO VIEWS WITH REGIONS
text: 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.

heading: 3 RELATED WORKS
text: Relation to mask prediction tasks We can consider the proposed LR as a proxy to mimick masked language modeling in BERT, where the “ground-truth” local token is a soft label provided by the
teacher network, while the student network makes predictions to match that target, based on the context of regions in a different augmented view. Importantly, our LR considers softmax with cross-entropy in the objective, rather than MSE as in MPP. A very sharp teacher distribution is used by choosing small temperatures. This encourages the model to focus on the salient dimensions, rather than waste modeling capability on training short-range dependencies and high-frequency details (Ramesh et al., 2021).
Relation to DenseCL The proposed LR mostly related to DenseCL (Wang et al., 2020b) in that the region correspondences in both methods are determined as the two most similar grid features. One critical difference is that DenseCL is a contrastive region-matching task, while our LR is a non-contrastive region-matching task, where no negative samples/queue is needed. This technical difference has a significant impact on the downstream task performance. We find that LR is particularly effective in serving our goal to improve image classification performance and build efficient & affordable SoTA SSL system; In contrast, DenseCL degrades the classification performance.
Relation to other region-level tasks The ideas of leveraging local region-level pre-training tasks for visual representation learning have been explored for ConvNets (Misra & Maaten, 2020; Xiong et al., 2020; Wang et al., 2020b; Xie et al., 2021a; Yang et al., 2021; Xie et al., 2021c). We summarize the differences in three aspects: (i) Motivation. Our region-matching task LR aims to recover the lost property of automatic correspondence learning in self-supervised monolithic Transformers, while most existing region-level tasks aim to improve dense visual prediction tasks. (ii) Technical difference. Our LR is a non-contrastive region-matching task, while others are contrastive learning. (iii) Empirical performance. Most region-level tasks improve dense visual prediction tasks but sacrifice their image classification performance, while LR consistently improves classification performance. Among them, EsViT training method achieves the best ImageNet linear probe performance with minimum computational overhead. For detailed comparisons, please refer to Table 7 in Appendix.
Self-supervised vision Transformers. The research on Transformer-based self-supervised representation learning just scratches the tip of the iceberg, and only a few attempts are made on this topic. ImageGPT (Chen et al., 2020b) and MoCo-v3 (Chen et al., 2021) dedicate huge compute resource with large models to exploring the frontier. DINO (Caron et al., 2021) achieves comparable performance of large self-supervised ConvNets using small/medium-size Transformers. The proposed EsViT further pursues efficient and affordable solutions to self-supervised vision Transformers. For more general related works on Transformers for vision tasks and self-supervised ConvNets, please refer to Section B in Appendix.

heading: 4 EXPERIMENTAL RESULTS
text: We describe the experimental settings in Appendix Section C.3, and evaluate the proposed EsViT to answer three questions: Q1: How does EsViT perform on standard ImageNet benchmark compared to SoTA methods? Q2: How effective EsViT is when transferring to downstream tasks? Q3: What are the design choices and empirical contributions of LR? Q4: When does the intriguing property of self-supervised Transformers exist, including learned correspondence and attentions?

heading: 4.1 COMPARISONS WITH PRIOR ART ON IMAGENET
text: We report top-1 linear probe and k-NN accuracy on the ImageNet validation set. Table 1 presents comparisons with SoTA SSL systems across various architectures. Please refer to Figure 1 for comparisons over scaling parameter counts and throughput. Our findings are summarized below.
Comparisons with self-supervised Transformers. The DINO- and MoCo-based ViT has higher accuracy and smaller models than iGPT, under the same linear probing protocol and training data. At the similar level of model size and compute complexity, the proposed EsViT improve SoTA methods DINO/MoCo-v3 by a large margin: EsViT (Swin-B) outperforms DINO (ViT-B/16) by 2.2% linear probe accuracy and 2.8% k-NN accuracy in absolute values. EsViT (Swin-B) even performs slightly better than DINO (ViT-B/8) (0.3% higher linear probe accuracy and 1.5% higher k-NN accuracy), with 4× higher throughput. MoBY (Xie et al., 2021b) is a con-current work that investigates multi-stage ViT in SSL. With the same architecture Swin-T, our EsViT pre-training tasks significantly outperform MoBY, showing 3% higher accuracy. In EsViT, longer sequences in self-attention is implemented by increasing the window size. We experiment this by considering a window size of W =14. Overall, the proposed EsViT (Swin-B/W=14) shows the best performance
(top-1 accuracy 81.3%, top-5 accuracy 95.5%, k-NN accuracy 79.3%), compared with all systems, and is 3.5× parameter-efficient and has at least 10× higher throughput than previous SoTA MoCo-v3.
Comparisons with big ConvNets. We compare with the SoTA big ResNets reported by SimCLRv2 (Chen et al., 2020d), BYOL (Grill et al., 2020) and SwAV (Caron et al., 2020). Among them, the best accuracy 79.8% under the linear probing protocol is reported by SimCLR-v2 with SK-ResNet, where Selective Kernel (SK) (Li et al., 2019c) is a form of attention to enhance CNNs. It is clear in Figure 1 (b) that all ConvNets-based SSL methods show an envelope in the regime of scaling up model sizes after passing 500M. EsViT achieves better accuracy than their highest envelope, with 16× less model parameters and 8× higher throughput.

heading: 4.2 TRANSFER LEARNING
text: We also conduct transfer learning in downstream tasks to evaluate the quality of learned representations. Two sets of tasks are considered:
• Classification on a suite of 18 small datasets. As exemplified in (Radford et al., 2021), it is a common and clean approach to evaluate a learned representation by fitting a linear classifier on the representation and measuring its performance across multiple datasets. We study 18 datasets used in (Radford et al., 2021). Automatic hyper-parameter tuning is considered to ensure fairness of comparison. Besides averaged scores, we report # wins as the number of datasets on which the model outperforms its supervised counterpart. Detailed dataset description and settings are in Appendix.
• Detection and segmentation on COCO. Different from previous monolithic self-supervised ViT, the multi-stage architecture in EsViT can be readily used for dense visual tasks that require hierarchical feature representations.
Comparison with supervised counterparts. We compare with the supervised-learning Swin, whose checkpoints are downloaded from the official codebase2. Figure 3 shows the classification results of Swin-S, EsViT consistently outperforms its supervised variant, often by a large margin.
2https://github.com/microsoft/Swin-Transformer

heading: Hat
text: APbb APbb50 AP bb 75
Sup. 46.0 68.1 50.3 EsViT 46.2 (46.2) 68.0 (67.9) 50.6 (50.5)
APmb APmb50 AP mb 75
Sup. 41.6 65.1 44.9 EsViT 41.6 (41.7) 64.9 (64.8) 44.8 (45.1)
Table 2: COCO Detection & Segmentation.
Similar conclusions are drawn for other model sizes. On COCO detection and segmentation task, however, EsViT shows comparable results with the variant with LV only (shown in parentheses) and the supervised counterpart (Swin-T trained with 3× schedule), as shown in Table 2. We hypothsize this is related to the non-constrastive nature of EsViT, as explained later.
Effects of larger, less-curated pre-train datasets. The performance of Transformer-based SSL research has thus far been limited to highly curated pre-train data such as ImageNet-1K. To push the frontier in leveraging large amounts of unlabeled data, we explore the effects of pre-training from larger, less-curated image datasets: WebVision-v1 (Li et al., 2017), OpenImages-v4 (Kuznetsova et al., 2020) and ImageNet-22K (Deng et al., 2009), described in Appendix. The pre-train epochs on different datasets are adjusted so that all models see a similar number of augmented views. We summarize the results in Table 3 and would like to emphasize the following findings. First, LR improves LV (shown in parentheses) on all datasets. Second, all EsViT pre-trained checkpoints outperform supervised checkpoint in downstream classification tasks, but performance varies a lot, with ImageNet-22K checkpoint showing the best transfer ability. Third, ImageNet-1K pre-trained model shows the best ImageNet-1K linear probe performance. We hypothesize that it is not only the size of pre-train dataset matters, but also the distribution of image classes matters: more diverse and well-balanced distribution results in a stronger generalization ability.

heading: 4.3 DISCUSSION ON THE NON-CONTRASTIVE REGION-MATCHING TASK
text: Compatibility with various network architectures. We investigate ResNet-50 and different efficient sparse Transformers in Table 4. DeiT is shown as a baseline reference. Batch size = 1024 in this experiment. To ensure fair comparison, we modify all into a 4-stage architecture with the number of Transformer layers in each stage as 2-2-6-2. We see that LR improves all network architectures, including ResNet-50, Swin (Liu et al., 2021), ViL (Zhang et al., 2021), CvT (Wu et al., 2021) and PvT (Wang et al., 2021). Though directly adding LR to monolithic ViT is computationally infeasible, we uniformly sampled top-layer grid features of DeiT and then add LR, but did not observe performance improvement. This is partly because the monolithic ViT itself already has a good corresponding ability, an extra region-matching task does not provide new learning signals. As compared in Appendix Table 12 with the ResNet-50 backbone, EsViT learning method shows the highest accuracy, compared with existing SSL methods.
Model scaling with LR. We compare the pre-training objective with and without LR in Table 1. Across different model scales and window sizes, the proposed region level LR can consistently improve the performance. The gains can be clearly seen by k-NN accuracy (around 1-2%), where no additional tuning is needed as in linear probe. Figure 4 demonstrates that LR helps model convergence, and can be used as a drop-in to improve models trained with the view level task.
Contrastive vs Non-contrastive region-matching tasks. The proposed LR adds a non-contrastive region-matching task to the non-contrastive view-level task LV ; On the contrary, DenseCL adds a contrastive region-matching task to the contrastive view-level task MoCo-v2. In Table 5, we compare four methods in the same setting with ResNet-50. DenseCL improves dense visual prediction performance, but hurts classification performance. LR improves both tasks, especially the classification performance. One limitation is that the non-contrastive methods show lower performance in dense prediction tasks, this is consistent with the observations for BYOL in (Wang et al., 2020b). The simple LR shows the best ImageNet accuracy compared with all sophisticated region-level tasks in this 200-epoch setting in Appendix Table 7, and the best overall accuracy in Table 12. It indicates that LR well serves our goal in building efficient SoTA SSL systems.
Design choices of LR. We ablate a couple of choices in constructing LR in Eq. (2). (i) Softmax vs MSE. One alternative way to measure the distance between two projected vectors is MSE, as employed in the popular non-contrastive SSL algorithm BYOL (Grill et al., 2020). When adding region-matching tasks to BYOL and pre-training 50 epochs, Softmax and MSE yield k-NN accuracy of 37.2% and 34.9%, while the baseline BYOL yields 33.1%. We also replace the region-matching metric in EsViT as MSE, yielding k-NN accuracy 72.6%, which lower than the view-level task only (74.2%). These results show that Softmax is essential in LR. (ii) Optimal Transport (OT) vs Simple Argmax. To avoid heavy computational overhead, a simple feature-level argmax solution is considered in Eq. (2) to pair two local regions. To study the impact of high region-matching quality, we consider OT. Empirically, we observe OT yields slightly higher k-NN accuracy at the early stage, but the gain is diminished in the end. Considering the extra computational cost of solving OT with an inner loop in sinkhorn algorithm (Cuturi, 2013), we opt for simple argmax in our experiments.

heading: 4.4 QUALITATIVE STUDIES
text: Visualization of correspondences. Given two views of the same image, we use the pre-trained backbone to extract the top-layer features z1 and z2. For each feature vector in z1, we find the feature vector in z2 that best matches it in terms of highest cosine similarity, as defined in Equation (2). In Figure 5, we show the top-10 correspondences between two views for three methods. In Figure 5 (b), EsViT with LV tends to identify pairs in the background as the most matched ones (and in a wrong way in this example). This could be a valid solution to LV , as the invariance in the level of aggregated global features does not necessarily induce invariances in the local region level. This is significantly alleviated with LR (shown in Figure 5 (c)), a task that implicitly requires local matching.
Surprisingly, DINO is able to learn good correspondences even without the region-level matching task. To the best of our knowledge, this is a previously unreported intriguing property of self-supervised Transformers with monolithic architectures: good semantic correspondences are automatically learned. We hypothesize that features at lower layers (image patch itself in the extreme case) can directly pass to higher layers, and the former regularizes the latter to remain discriminative. Nevertheless, the proposed LR can dramatically reduce the issue, and is good remedy to rescue the loss of semantic correspondence for the multi-stage architecture. In Appendix, we quantitatively measures the correspondence learning ability of these SSL methods on ImageNet validation dataset, the observations are consistent: LR improves the matching accuracy from 66% to 91%.
Visualization of attention maps. We look at the self-attention in the different heads of the last layer in Figure 6. A local region on the edge of the main object is employed as query, and the attended regions are highlighted in red for those the query’s top 60% mass are assigned. In Appendix, we visualize more examples with different query positions. DINO tends to automatically learn class-specific attention maps leading to foreground object segmentation, regardless of its query located in foreground or background. This is probably because main objects remain as the major invariance factor in different augmented views. This property is lost when a multi-stage architecture is employed, as shown in EsViT with LV . These patterns are consistent for different heads. After introducing LR for EsViT, we note that the attention maps become more diverse in different heads, i.e., entropy values of attentions get more skewed, and attended regions are more different. This is perhaps because LR requires each region to consider many matching tasks to regions in different augmented views, each head automatically learns to distribute the tasks and complete a few of them.

heading: 5 CONCLUSIONS
text: In this paper, we first discover the automatic correspondence learning property of self-supervised monolithic Transformers. Inspired by this, we present efficient self-supervised vision Transformers (EsViT) to with two major insights: a multi-stage Transformer architecture with sparse self-attentions, and a non-contrastive region-matching pre-training task. The synergy of both helps EsViT reach the SoTA performance of SSL vision systems with significantly less compute and smaller model size. Our study also reveals that exploration of effective solutions to learn from larger and less curated pre-training data in the wild is a key but less studied factor in paving the way toward the scaling success of SSL vision systems.

heading: ETHICS STATEMENT
text: Though self-supervised learning (SSL) has great potentials to learn powerful representation without human annotation, the existing techniques to build SoTA SSL vision systems tend to be Red AI (Schwartz et al., 2020): it could be environmentally unfriendly and the computational cost is extensively high. The required training resource is typically not accessible for a lab environment (thus raising barriers to participation in AI research). For example, the prior art MoCo-v3 has greatly pushes the performance limit of SSL system (Chen et al., 2021). The authors kindly reported that “it (MoCo-v3, ViT-H) takes 9.8 hours per 100 epochs using 512 TPUs. This is a gigantic scale of training: for the 300-epoch ViT-H, this amounts to ∼625 TPU days, or ∼1.7 TPU years of training.” The SoTA model MoCo-v3 with ViT-BN-L/7 should have a higher cost than this. Even for a smaller model ViT-B, “it takes 24 hours in 128 GPUs (vs. 2.1 hours in 256 TPUs)”. Hence, improving the efficiency of building SoTA SSL systems is of high value for the community and society to achieve Green AI (Schwartz et al., 2020).
To this end, we propose EsViT to provide more affordable and efficient solutions for the community to experiment and explore the directions of SoTA SSL in computer vision. Our EsViT model shows the best ImageNet linear probe performance compared with all existing SSL vision systems, and is 3.5× parameter-efficient and has 10× higher throughput than previous SoTA. This efficiency gain can significantly decrease its carbon footprint and increase its inclusivity, encouraging more researchers to participate the study of the SSL topic.

heading: REPRODUCIBILITY STATEMENT
text: Our paper provides comprehensive empirical studies on the EsViT algorithm. We provide PyTorchstyle pseudo-code in Appendix. We also include an example code with instruction as supplementary material to ensure the reproducibility. For empirical results on both various network architecture and large-scale datasets, we provide detailed hyper-parameter specifications. We will release the pre-trained checkpoints and codebase for the research community for reproducible research.

heading: A METHODS

heading: A.1 ALGORITHMS
text: We summarize the training algorithm procedure of EsViT with LV +LR in Algorithm 1. To clearly outline the main idea of the algorithm, we show the algorithm for two augmented views. For the full algorithm to deal with multi-crop, please refer to our codebase. In Algorithm 1, for a mini-batch of size n, the teacher/student network consists of three output variables: (1) p ∈ Rn×K is the probability vector for the view-level representation, output by an MLP head. (2) z ∈ Rn×T×P is the feature map, containing T region-level features of dimension P . (3) pz ∈ Rn×T×K are probability vectors of z, output by a different MLP head.
Algorithm 1: EsViT with LV +LR, pseudocode with 2-crop. # gs, gt: student and teacher networks # Cv, Cr: view and region center (K) # tmp s, tmp t: student and teacher temperatures # a, b: network and center momentum rates. # n: batch size, K: MLP-head-projected probability vector length, T: last layer feature map length, P: last layer feature vector length
1 gt.params = gs.params # The main training loop 2 for x in loader: 3 x1, x2 = augment(x), augment(x) # two random views 4 # student output, p:n×K, pz:n×T×K, z:n×T×P 5 p s1, pz s1, z s1 = gs(x1) 6 p s2, pz s2, z s2 = gs(x2) # teacher output, p:n×K, pz:n×T×K, z:n×T×P 7 p t1, pz t1, z t1 = gt(x1) 8 p t2, pz t2, z t2 = gt(x2) 9
# view-level loss 10 loss v = Hv(p s1, p t2)/2 + Hv(p s2, p t1)/2 # region-level loss 11 loss r = Hr(pz s1, pz t2, z s1, z t2)/2 + Hr(pz s2, pz t1, z s2, z t1)/2 12 loss = loss v/2 + loss r/2 13 loss.backward() # back-propagate 14 # update student, teacher and centers 15 update(gs) # AdamW for student 16 gt.params = a ∗ gt.params+ (1− a) ∗ gs.params # EMA for teacher 17 Cv = b ∗ Cv+ (1− b) ∗ cat([p t1,p t2].mean(0)) # EMA for view center 18 Cr = b ∗ Cr+ (1− b) ∗ cat([pz t1,pz t2].mean(0)) # EMA for region center 19
# The view-level loss function 20 def Hv(s, t): 21 t = t.detach() # stop gradient 22 s = softmax(s / tmp s, dim=-1) 23 t = softmax((t - Cv) / tmp t, dim=-1) 24 return - (t * log(s)).sum(dim=-1).mean()
# The region-level loss function 25 def Hr(ps, pt, zs, zt): 26 pt = pt.detach() # stop gradient 27 ps = softmax(ps / tmp s, dim=-1) # n×T×K 28 pt = softmax((pt - Cr) / tmp t, dim=-1) # n×T×K 29 sim matrix = torch.matmul(zs , zt.permute(0, 2, 1)) # n×T×T 30 sim idx = sim matrix.max(dim=-1)[1].unsqueeze(2) # n×T×1 31 pt idxed = torch.gather(pt, 1, sim idx.expand(-1, -1, pt.size(2))) 32 return - (pt idxed * log(ps)).sum(dim=-1).mean()

heading: A.2 NETWORK ARCHITECTURE CONFIGURATIONS AND IMPLEMENTATION DETAILS
text: Inspired by great successes of the multi-stage ConvNet architecture such as VGG (Simonyan & Zisserman, 2014)/ResNets (He et al., 2016) for computer vision, the multi-stage Transformer-based networks have been explored very recently in the supervised learning setting (Vaswani et al., 2021; Wang et al., 2021; Liu et al., 2021; Zhang et al., 2021; Wu et al., 2021). In multi-stage vision Transformers, since a larger number of patches is often produced at the early stages, an efficient Transformer with sparse self-attentions is considered to reduce the computational complexity. The basic idea is to split the feature maps into non-overlapping local windows (with size W×W ), and self-attention is performed within each local window. This however has one drawback that features in different local windows cannot interact. Various methods have been proposed to best approximate full-attention, with different trade-off between accuracy and efficiency.
We briefly describe three schemes as follows, and benchmark them in the experiments. (i) Swin Transformer (Liu et al., 2021): A shifted window partitioning approach is proposed, which alternates between two partitioning configurations in consecutive Transformer blocks, so that each local feature is grouped into different windows in self-attentions. (ii) Vision Longformer (ViL) (Zhang et al., 2021): Features in each local window are further allowed to attend all features in the 8-neighboring windows. (iii) Convolution vision Transformer (CvT) (Wu et al., 2021): Features in neighboring windows are considered in the convolutional projection in self-attentions.
The window size is set to W = 7 by default. The query dimension of each head in self-attentions is d = 32, and the hidden layer width of each MLP is 4× of its input’s width, for all experiments. The architecture configurations of model variants employed in the experiments are summarized in Table 6. Some notable implementation detailed are described as follows:
• The three configurations Swin-T, Swin-S and Swin-B indicate Tiny, Small, and Base models, respectively, which are almost identical to the original implementation (Liu et al., 2021), except that we add special treatments to deal with input augmented views of different resolutions, when the resolution (feature map size more specifically) is not divisible by the window size (i.e., resolution 96 and window size=7 or 14).
• Swin-T and Swin-S with window size W =14 are customized by us to allow full self-attention in stage 3 (where the majority of model capacity is allocated to) and stage 4, to study the impact of longer sequences in EsViT.
• In the original ViL (Zhang et al., 2021) and CvT (Wu et al., 2021) papers, different positional embedding strategies and multi-stage network configurations were employed. We modify them by only utilizing relative position bias and their proposed sparse self-attention mechanisms, and create a similar 4-stage architecture with Swin-T for fair comparison.
Relative Position Bias. To facilitate SSL, we consider relative position bias (Liu et al., 2021) to characterize the spatial information between features for the three efficient Transformers aforementioned, and do not use absolute position embeddings. This is because augmented views of varied resolutions can be cropped from anywhere in an image in SSL, maintaining the relative positions is easy in implementation, and is largely sufficient for invariance learning among these views.

heading: B RELATED WORK
text: Self-supervised ConvNets. ConvNets-based SSL has been extensively studied in the literature. Based on the pre-training tasks, they can be broadly categorized into three classes: Handcrafted pretext tasks (Doersch et al., 2015; Noroozi & Favaro, 2016; Pathak et al., 2016; Gidaris et al., 2018; Zhang et al., 2016; Larsson et al., 2016; Zhang et al., 2017; Pu et al., 2016; Donahue & Simonyan, 2019), contrastive learning (Dosovitskiy et al., 2015; Zhuang et al., 2019; Oord et al., 2018; Hjelm et al., 2018; Bachman et al., 2019; He et al., 2020; Chen et al., 2020c; Grill et al., 2020) and prototype learning (Caron et al., 2018; 2020; Li et al., 2020b; Xie et al., 2016; Yang et al., 2016; Ji et al., 2019; Zhan et al., 2020). It is also known that data augmentations play a crucial role in SSL pipeline (Chen et al., 2020e; Caron et al., 2020; Tian et al., 2020; Li et al., 2020a). The impact of pre-training dataset size/quality is explored for ConvNets in SSL (Goyal et al., 2021; Yonglong et al., 2021). To date, the search of best pre-taining tasks/datasets and augmentations are based on CNNs. Among them,
SimCLR-v2 (Chen et al., 2020d), BYOL (Grill et al., 2020) and SwAV (Caron et al., 2020) achieve the highest ImageNet linear probe performance with large ConvNet architectures. The performance tends to saturate with an increasingly growing model size, raising a question if ConvNets reach a limit in SSL.
Transformers for vision. Vision Transformers (ViT) (Dosovitskiy et al., 2021) shows the great potentials of generalizing Transformers for computer vision, by achieving compelling accuracy in supervised learning, especially with large-scale data and high capacity models. DeiT (Touvron et al., 2020) further provides an effective ViT training strategy to ease the adaption of Transformers for practitioners. Transformers have also been applied to other vision tasks, ranging from low-level tasks such as image generation (Parmar et al., 2018; Chen et al., 2020b) and enhancement (Chen et al., 2020a; Yang et al., 2020), to high-level tasks such as object detection (Carion et al., 2020; Zhu et al., 2020; Zheng et al., 2020; Dai et al., 2020) and segmentation (Wang et al., 2020a;c), and to vision-language tasks (Lu et al., 2019; Tan & Bansal, 2019; Chen et al., 2019; Su et al., 2019;
Li et al., 2019b;a; Zhou et al., 2020; Li et al., 2020c). Marrying Transformers with multi-stage architectures (Vaswani et al., 2021; Wang et al., 2021; Liu et al., 2021; Zhang et al., 2021; Wu et al., 2021) show higher classification accuracy in supervised learning, and enables applicability of Transformers for a broader range of vision tasks. Given these properties, we believe multi-stage ViT is a must-study baseline for SSL in computer vision.
Discussion with other region-level tasks. In Table 7, we compare LR against the existing regionlevel tasks, including PIRL (Misra & Maaten, 2020), DenseCL (Wang et al., 2020b), DetCo (Xie et al., 2021a), InstLoc (Yang et al., 2021), PixPro (Xie et al., 2021c). Most of these region-level tasks improve object detection tasks, but hurt the ImageNet classification accuracy. DetCo achieves the best trade-off: improving the performance of both tasks, but with a sophisticated multi-level, global-local interaction algorithm. With the same number of pre-training epochs and augmented views, EsViT achieve the best ImageNet linear probe accuracy among all region-level tasks, with as minimum computational overhead as possible. This well serves our goal of building efficient SSL SoTA image classification system.

heading: C EXPERIMENTS

heading: C.1 THROUGHPUT ESTIMATE AND CONVERSION
text: Since different papers report throughput on different hardwares, it is not ready to compare the numbers directly. Noting that all papers report the throughput for ViT-B/DeiT-B, we use this number to align and convert the throughput. In Table 8, we describe our process and results of standardizing the throughput.

heading: C.2 THE COMPUTATION AND MEMORY OVERHEAD OF THE PROPOSED LR
text: We emphasize that adding LR to the multi-stage transformer architectures yields acceptable extra computational cost, while adding LR directly to the monolithic transformer architectures has a huge computational overhead. To demonstrate this, we report the cost comparisons in Table 9. For each setting, we report [Memory Usage (MB) / Running time per iteration (second/iteration)]. In Table Table 9 (a), when the batch size is gradually increased (e.g., , batch-size=12), the memory cost increases nearly 4 times for monolithic architectures, while increases 1.6 times for multi-stage architectures. Similar trends are shown for training cost per iteration increase ratio (1.47 vs 1.15). This indicates LR can more naturally fit multi-stage architectures. Similarly, we compare computational cost comparisons [Memory Usage (MB) / Running time per iteration (second/iteration)] in Table 9 (b), for other network architecture configurations . From the increased cost ratio, we see that LR adds acceptable cost in terms of both memory and training time, compared with the baseline.

heading: C.3 EXPERIMENTAL SETTINGS OF PRE-TRAINING AND EVALUATION ON IMAGENET
text: We study unsupervised pre-training performed in ImageNet-1K dataset (Deng et al., 2009) without labels. The default training details are described as follows, mostly following (Caron et al., 2021). We train with the Adamw optimizer (Loshchilov & Hutter, 2018), a batch size of 512, and total epochs 300. Linear warmup of the learning rate is used during the first 10 epochs, with its base value
determined with the linear scaling rule (Goyal et al., 2017): lr = 0.0005 ∗ batchsize/256. After this warmup, the learning rate is decayed with a cosine schedule. We build our systems based on Swin Transformers (Liu et al., 2021) in our experiments. Swin-B has a model size and computation complexity similar to ViT-B/DeiT-B (patch size 16). We also considered Swin-T and Swin-S, which have the complexity that are similar to those of ResNet-50 (DeiT-S) and ResNet-101, respectively. The default window size is W =7.
One major common protocol to evaluate SSL is linear probe on ImageNet-1K, where features are extracted from a frozen backbone, and a supervised linear classifier is trained. For all Transformer models, we use the concatenation of view-level features z̄ in the last 4 layers (the results are similar to the use of last 3 or 5 layers in our initial experiments).

heading: C.4 COMPARISON WITH A RESNET-50 BACKBONE
text: To compare our EsViT learning method with other SSL algorithms, we conduct experiments with a ResNet-50 backbone, and show the results in Table 12.

heading: C.5 LINEAR PROBE ON A SUITE OF SMALL DATASETS
text: Datasets. Table 10 shows details and source of all datasets used for linear probe, including the number of classes, the size of training set and testing set, metrics used in evaluation, as well as a public source of the dataset. Note that original UCF101 dataset is a video dataset. Here the middle frame of each video is extracted to form a classification dataset. There are 3 train/val splits in Tensorflow, we use the first one.
Automatic hyper-parameter tuning. We rigorously follow (Radford et al., 2021) to conduct training and evaluation for linear probe on the downstream datasets. We train a logistic regression classifier using scikit-learn’s L-BFGS implementation, with maximum 1, 000 iterations, and report the corresponding metric for each dataset. We determine the L2 regularization strength λ using a hyperparameter sweep on the validation sets over the range between 10−6 and 106 , with 96 logarithmically spaced steps. To save compute required for the sweeps, we perform a parametric binary search that starts with λ = [10−6, 10−4, 10−2, 1, 102, 104, 106] and iteratively halves the interval around the peak until it reaches a resolution of 8 steps per decade. The hyperparameter sweeps are performed on a validation split of each dataset. For the datasets that contain a validation split in addition to a test split, we use the provided validation set to perform the hyperparameter search, and for the datasets that do not provide a validation split or have not published labels for the test data, we split the training dataset to perform the hyperparameter search. For the final result, we combine the validation split back with the training split and report the performance on the unused split.
Detailed results. Only the last layer feature is considered for all models for simplicity, though adding features from more layers may potentially improve the results. Table 11 shows the results for architectures at a similar scale of ResNet-50 or Swin-T. The first two columns are numbers from (Radford et al., 2021). CLIP with ResNet-50 is pre-trained on 400 million image-text pairs. Supervised ResNet-50 and Swin-T are pre-trained on ImageNet-1K, on which EsViT with Swin-T is pre-trained as well (Batch Size=512). EsViT outperforms its supervised counterpart, and is on par with the performance of CLIP in a similar image encoder architecture scale.

heading: C.6 PRE-TRAINING DATASETS
text: We describe the statistics and training schedule on larger and less curated datasets in Table 13. The pre-training epochs are chosen so that the model is trained with a similar number of augmented views.

heading: C.7 RESULTS ON CORRESPONDENCE LEARNING
text: We first quantitatively evaluate the correspondence learning results with 50K images in the ImageNet validation dataset. We create a simple evaluation dataset 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. Please see the 1st row of Figure 9 for one such example. The top-10 correspondences are used for evaluation. Two metrics are considered: (1) Accuracy measures the percentage of correctly matched region pairs, (2) distance error indicates the averaged `2 distance between the predicted matched region and ground-truth region (the value is 0 for perfect matching). The results are reported in Figure 8. DINO with monolithic Transformers shows surprisingly good performance on correspondence learning. The use of multi-stage Transformer architecture reduces this ability, shows a lack of good region correspondence. With LR, the region matching ability is significantly recovered. In Figure 9, we visualize the correspondences for more images. Overall, DINO with monolithic Transformers is able to discover most salient correspondences of semantic meaning in the mild augmentation conditions, even without an implicit region matching loss in training. We believe this previously underestimated property is whole-noting, and has potentials to enable more applications. However, this desired property gets dilated when changing from monolithic to multi-stage Transformer architecture (from column 1 to column 2), then the proposed region level task can alleviate this issue (from column 2 to column 3).
To more specifically analyze the correspondences, we note the following results. The first row shows a simple case, where only images of left-to-right flipped views are presented. The ground-truth
correspondences should be horizontal lines that link the two flipped regions. It reveals that the view-level pre-train task alone is insufficient to learn good correspondences for the multi-stage Transformer architecture, while region matching task can alleviate this issue significantly. Similar observations are shown in row 3 and row 4.
We further study more cases that requires real-world correspondences in row 2, row 5 and row 6. These views are not generated with data augmentation (as in model pre-training), but are often presented in more practical scenarios: one-to-many mappings, cartoon-to-toy, seasonal changing of the scene, respectively. The proposed region matching task can work particularly well in those cases.

heading: C.8 MORE VISUALIZATION RESULTS OF ATTENTION MAPS
text: We visualize attention maps at the top layer in Figure 10, 11, 12. With a monolithic Transformer architecture, DINO can automatically identify the main foreground objects. Unfortunately, changing from monolithic to the multi-stage Transformer architecture (From left column to middle column), this property gets lost. There are more heads in the multi-stage architecture than monolithic architecture (24 heads vs 6 heads in this case) in the last year. A fair number of heads in EsViT shows redundant patterns, this issue can be reduced when the region-level matching task is added (From middle column to right column).
We observed that DINO with monolithic Transformer architecture only learns to attend the foreground objects, even when the query is a background region (see Figure 12). This is perhaps because DINO models are trained to learn view-level invariance, the main objects in the pre-train dataset ImageNet tend to be the principle factor that remains invariant across different augmented views. Hence, all backgrounds are ignored, regardless of the query positions. This is improved in EsViT with the region-level pre-train task, as the model is trained to match individual regions.
DINO shows high entropy values in all of 6 heads (perhaps a required condition to cover all regions of the main object). In EsViT, LR plays an interesting role in modulating the entropy distributions among heads: it increases those with larger entropy values, while decreasing those with lower entropy values. In another word, it makes the attention patterns in different heads more diverse.

Review summaries:
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.

This paper develops an efficient self-supervised vision transformer for learning visual representations. It introduces a multi-stage architecture with sparse attentions to reduce computation complexity and proposes a new pretraining task of region matching to capture fine-grained region dependencies. The results on the ImageNet and 18 small datasets or downstream tasks are good and compared with other state-of-the-art approaches.