Self-supervised Models are Good Teaching Assistants for Vision Transformers

1 INTRODUCTION . Recently , Vision Transformers ( ViTs ) have been successfully used for computer vision tasks , including image recognition , object detection , semantic segmentation and so on . Remarkably , ViTs are capable to reach superior performance on image classification task when trained with large-scale datasets , e.g . JFT-300M ( Dosovitskiy et al . ( 2020 ) ) . However , ViTs achieve lower accuracies than Convolutional Neural Network ( CNN ) networks on medium-scale or small-scale datasets ( Dosovitskiy et al . ( 2020 ) ) . To alleviate the demand for data , DeiT ( Touvron et al . ( 2021 ) ) distills the inductive bias from a large CNN teacher by introducing an extra distillation token and shows satisfactory results . Self-supervised learning ( SSL ) and supervised learning ( SL ) are two different paradigms w.r.t . the way they construct training objectives . With the development of transformer , self-supervised learning for transformers has also attracted widespread attention from the community , and many approaches have been proposed ( Chen et al . ( 2021 ) ; Caron et al . ( 2021 ) ) . Caron et al . ( 2021 ) reported an interesting discovery that self-attention visualizations of self-supervised vision transformers and supervised vision transformers represent different tendentiousness . As shown in Figure 1 , vision transformers trained with supervised signal pay more attention to texture , while self-supervised counterparts focus on shape . In addition , when the size of the annotated training dataset is small , the supervised transformer is more prone to overfitting . For example , when the training dataset is ImageNet-1K ( Russakovsky et al . ( 2015 ) ) , self-supervised ViT can transfer better to downstream tasks than its counterpart ( Caron et al . ( 2021 ) ) . These observations motivate us to explore and exploit the differences between these two learning paradigms ( SSL v.s . SL ) applied on ViTs . We measured the similarity index between the feature layers of two randomly initialized supervised transformers and of a supervised transformer and a self-supervised transformer through Centered Kernel Alignment ( CKA ) indicator . The results are shown in Figure 2 . It can be seen that the similarity between the layers of two randomly ini- tialized supervised transformers significantly exceeds that between a supervised transformer and a self-supervised transformer , and the feature similarity of the last few layers is relatively low . Since the difference is quantitatively prominent and qualitatively compensating , we propose to introduce an self-supervised teaching assistant ( termed as SSTA ) besides the commonly used supervised teacher to further improve the performance of transformers . Specifically , we propose a head-level knowledge distillation method that selects the most important head of the supervised teacher and the self-supervised teaching assistant , and let the student to mimic the attention distribution of these two heads , so as to make the student focus on the relationship between tokens deemed by the teacher and the teacher assistant . Extensive experiments demonstrate that the proposed SSTA is a good compensation to the supervised teacher . Meanwhile , compared with supervised teaching assistant , SSTA with greater difference can bring more improvements . The success of SSTA prompted us to further reveal the otherness between the self-supervised ViTs and supervised ViTs . We explore the differences between two different teachers and the students distilled from different teachers on prediction , shape bias , robustness , and transferability to downstream tasks , some of which are counter-intuitive and are studied for the first time . Our contributions are summarized as follows : • By observing that self-supervised learning and supervised learning provide information from different perspectives , we exploit adding an self-supervised transformer as a teaching assistant to complement to commonly used supervised teacher , and firstly propose a headlevel knowledge distillation approach for data efficient vision transformer learning . • To effectively transfer the knowledge via heads , a heuristic head selection strategy is designed to choose most informative heads from teacher . Meanwhile , an early stop learning strategy is further derived to facilitate distillation . • Extensive experiments are conducted to demonstrate the advantage of the self-supervised teaching assistant . Besides , by comprehensive analyzing the variant combination of two teachers , some interesting findings , regarding the prediction , shape bias , robustness , and transferability , are studied for the first time . 2 RELATED WORK . 2.1 VISION TRANSFORMER . Recently , ViTs have made tremendous development , and various Transformer architectures for computer vision tasks have been proposed ( Dosovitskiy et al . ( 2020 ) ; Touvron et al . ( 2021 ) ; El-Nouby et al . ( 2021 ) ) . The Self-Attention mechanism allows transformers to capture long-distance relationships and become content-aware . Compared to CNN , ViTs are more robust to severe occlusions , perturbations , and domain shifts and significantly less biased towards textures ( Naseer et al . ( 2021 ) ) . However , ViTs are very hungry for data , when training on medium-scale or small-scale datasets , ViTs can ’ t exceed the results of CNN ( Dosovitskiy et al . ( 2020 ) ) . Therefore , a lot of works ( Touvron et al . ( 2021 ) ; Graham et al . ( 2021 ) ) introduce the inductive bias of a supervised pre-trained large CNN teacher through knowledge distillation , thereby alleviating the demand for annotated data . 2.2 KNOWLEDGE DISTILLATION . Knowledge Distillation ( KD ) was first proposed by Hinton et al . ( 2015 ) , which aims to transfer the knowledge of a larger teacher model to a smaller student model . Many approaches have achieved great success on CNN , e.g . Romero et al . ( 2014 ) ; Zagoruyko & Komodakis ( 2016 ) ; Park et al . ( 2019 ) , however due to the differences of transformers , few of them can be directly applied to transformers . DeiT ( Touvron et al . ( 2021 ) ) is the first work applying knowledge distillation to transformer , which adds an extra distillation token to transfer the inductive bias of a larger CNN to a relatively small transformer in the form of hard or soft output label . Ren et al . ( 2021 ) use different architectural inductive biases to co-advise the student transformer . These methods all rely on the inductive bias of other network structure , the knowledge needed by the student transformer and the effective transmission method are still to be explored . Furthermore , the teachers used in the existing distillation methods are all obtained by supervised training , and as far as we know , we are the first to try to use self-supervised representations to assist supervised training . 2.3 SELF-SUPERVISED LEARNING . Self-supervised Learning ( SSL ) is a generic framework that gets supervision from the data itself without any tags from human labor . Earlier methods heavily rely on constructing negative samples , e.g . SimCLR ( Chen et al . ( 2020a ; b ) ) , MoCo ( He et al . ( 2020 ) ; Chen et al . ( 2020c ) ) , while recent works eliminate the need for negative samples , e.g . BYOL ( Grill et al . ( 2020 ) ) , SimSiam ( Chen & He ( 2021 ) ) . With the development of vision transformer , some works ( Caron et al . ( 2021 ) , Chen et al . ( 2021 ) ) apply contrastive learning to vision transformers . Compared to supervised counterparts , self-supervised vision transformers exhibit some properties . As described in Caron et al . ( 2021 ) , self-supervised ViT features explicitly contain the scene layout and object boundaries . In this work , we show that the difference between self-supervised ViT representations and supervised ViT representations is far from that . 2.4 KD MEETS SSL . Recently , some works have combined KD and SSL . SSKD ( Xu et al . ( 2020 ) ) adds an SSL branch next to the supervisory branch and regards the information contained in the SSL task as additional dark knowledge . CRD ( Tian et al . ( 2019 ) ) proposes a contrastive-based objective for knowledge distillation , which allows student to capture more information in the teacher ’ s representations of data . SEED ( Fang et al . ( 2021 ) ) employs knowledge distillation as a means to improve the representation capability of small models in self-supervised learning . These methods are for CNN , and there is only one teacher with the same training paradigm as the student . While in our method , the teacher and the student are under different training paradigms and the two teachers are trained by different paradigms with obvious different tendentiousness . 3 METHODOLOGY . In this section , we introduce the proposed Self-Supervised Teacher Assistant ( SSTA ) . We first present the overall architecture in Section 3.1 , and then introduce the specific head-level distillation in detail in Section 3.2 . Finally , the entire training process is described in Section 3.3 . 3.1 OVERALL ARCHITECTURE . The framework of the proposed method is shown in Figure 3 , consisting of three transformer encoders . The Student in the middle is the encoder that we want to improve , the SL Teacher on the left is the pre-trained teacher obtained via supervised learning , and the SSTA on the right is the pretrained teaching assistant obtained through self-supervised learning . For each inputX ∈ RH×W×C , where H , W and C represents the height , width and channel of the image respectively , it is input to three encoders respectively . After patch embedding , the input image is projected to XPE ∈ RN×D where N is the number of tokens and D is the dimension of each token , and XPE is then fed into stacked layers . As shown in Figure 3 , each layer consists of LayerNorm ( Ba et al . ( 2016 ) ) , Multi-head Self Attention ( MSA ) , Multi-Layer Perceptron ( MLP ) and residual connections . For MSA , we first compute Q = XPE · WQ ∈ RN×h×d , K = XPE · WK ∈ RN×h×d and V = XPE ·WV ∈ RN×h×d via linear transformations WQ , WK , WV , where h is the number of heads , and d is the dimension of each head ( d = D/h ) . Figure 3 ( right ) shows the details of MSA , Q and K produce an attention matrix via inner product and then the matrix is rescaled by √ d and normalized with a softmax function . Finally the normalized attention matrix is multiplied by V to get the output of the MSA layer . The entire procedure can be formulated as : AttnMat = Softmax ( Q×KT / √ d ) , ( 1 ) Output = AttnMat× V , ( 2 ) note the dimension ofAttnMat is h×N ×N . For more details , please kindly refer to ( Dosovitskiy et al . ( 2020 ) ) . AttnMat describes the attention distribution , which is computed based on the similarity between tokens . The higher the value , the more the relevance . The attention distribution reflects the relationship between tokens , and the relationship between [ cls ] token and other patch tokens can further reflect where the model is focusing on , as shown in Figure 1 . Existing work ( Zagoruyko & Komodakis ( 2016 ) ) has demonstrated that the attention maps of a powerful teacher network are effective knowledge in CNN . As Transformers are based on attention mechanism , we consider to adopt attention distribution as knowledge to transfer . The specific knowledge transfer method is the newly proposed head-level distillation , which will be introduced in section 3.2 .

Summary: This paper adopts a self-supervised pretrained vision transformer as teaching assistant for improving the performance of student via head-level knowledge distillation. To further improve the effectiveness of knowledge distillation, a simple yet effective head selection strategy is proposed to transfer teachers’knowledge to the select critical heads in the student. Extensive experiments show the effectiveness of the proposed method, including on classification, robustness and transferability.