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 presents the observation that supervised and SSL-trained Vision Transformers (VIT) do not encode the exact same information. Based on this observations, the authors propose to include a distillation loss to guide the training of supervised classifiers. By distilling information from the SSL-trained network, the final classifier obtains better validation performance.