Revitalizing CNN Attention via Transformers in Self-Supervised Visual Representation Learning

1 Introduction . Learning visual features effectively has a profound influence on the recognition performance [ 5 , 53 ] . Upon handling large-scale natural images , self-supervised visual representation learning benefits downstream recognition tasks via pretext feature training . Existing SSL methods typically leverage two branches to measure the similarity between different view representations derived from the same input image . By maximizing the similarity between the correlated views within one image ( e.g. , BYOL [ 24 ] , SimSiam [ 14 ] and Barlow Twins [ 73 ] ) , or minimizing the similarity between views from different images ( e.g. , MoCo [ 27 ] and SimCLR [ 12 ] ) , these methods are shown to be effective towards learning self-supervised visual representations . SSL evolves concurrently with the transformer . Debuting at natural language processing [ 62 , 17 ] , transformers have shown their advantages to process large-scale visual data since ViT [ 19 ] . The encoder-decoder architecture in this vision transformer consistently explores global attention without convolution . This architecture is shown effective for visual recognition with [ 7 , 76 , 54 ] or without CNN integration [ 38 , 21 ] . Inspired by these achievements via supervised learning , studies [ 15 , 10 , 4 , 71 ] arise recently to train transformers in a self-supervised manner . These methods maintain most of ∗Y . Song is the corresponding author . The code is available at https : //github.com/ChongjianGE/CARE 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . the SSL pipeline ( i.e. , encoder , projector , and predictor ) utilized for training CNN encoders . Without significant framework alteration , original SSL methods for CNN encoders can be adapted to train transformer encoders and achieve favorable performance . The success of using transformer encoders indicates that visual attention benefits encoder backbones in SSL . On the other hand , in supervised learning , CNN attention is usually developed via network supervision [ 47 ] . However , we observe that existing SSL methods do not incorporate visual attention within CNN encoders . This motivates us to explore CNN attention in SSL . We expect CNN encoders to maintain similar visual attention to transformers for recognition performance improvement with lower computational complexity and less memory consumption . In this paper , we propose a CNN Attention REvitalization framework ( CARE ) to make CNN encoder attentive via transformer guidance . Fig . 1 ( d ) shows an intuitive illustration of CARE and compares it with other state-of-the-art SSL frameworks . There are two streams ( i.e. , C-stream and T-stream ) in CARE where each stream contains two branches . C-stream is similar to existing SSL frameworks with two CNN encoders , two projectors , and one predictor . T-stream consists of two transformers , two projectors , and one predictor . T-stream takes CNN encoder features as input and improves feature attention via transformers . During the training process , we perform SSL in both streams simultaneously and use the T-stream output to supervise C-stream . The self-supervised learning in T-stream ensures attentive features produced by transformers are suitable for this SSL scenario . Meanwhile , we use the attention supervision on C-stream . This supervision enables both C-stream and T-stream to produce similar features . The feature representation of CNN encoders is improved by visual attention from transformers . As a result , the pre-trained CNN encoder produces attentive features , which benefits downstream recognition scenarios . Experiments on standard image classification , object detection , and semantic segmentation benchmarks show that the proposed CARE framework improves prevalent CNN encoder backbones to the state-of-the-art performance . 2 Related works . In the proposed CARE framework , we introduce transformers into self-supervised visual representation learning . In this section , we perform a literature survey on related works from the perspectives of visual representation learning as well as vision transformers . 2.1 Visual representation learning . There is an increasing need to learn good feature representations with unlabeled images . The general feature representation benefits downstream visual recognition scenarios . Existing visual representation learning methods can be mainly categorized as generative and discriminative methods . The generative methods typically use an auto-encoder for image reconstruction [ 63 , 49 ] , or model data and representation in a joint embedding space [ 18 , 6 ] . The generative methods focus on image pixel-level details and are computationally intensive . Besides , further adaption is still required for downstream visual recognition scenarios . The discriminative methods formulate visual representation learning as sample comparisons . Recently , contrastive learning is heavily investigated since its efficiency and superior performance . By creating different views from images , SSL obtains positive and negative sample pairs to constitute the learning process . Examples include memory bank [ 70 ] , multi-view coding [ 58 , 61 ] , predictive coding [ 30 , 60 ] , pretext invariance [ 42 ] , knowledge distillation [ 22 , 10 ] and information maximization [ 32 ] . While negative pairs are introduced in MoCo [ 27 ] and SimCLR [ 12 ] , studies ( e.g. , BYOL [ 24 ] and SimSiam [ 14 ] ) show that using only positive pairs are effective . Also , clustering methods [ 8 , 9 ] construct clusters for representation learning . The negative pairs are not introduced in these methods . Besides these discriminative methods focusing on image data , there are similar methods learning representations from either video data [ 66 , 26 , 33 , 67 , 45 , 65 , 64 ] or multi-modality data [ 1 , 2 , 3 ] . Different from these SSL methods , we use the transformer architectures to improve CNN encoders attention . 2.2 Vision transformers . Transformer is proposed in [ 62 ] where self-attention is shown effective for natural language processing . BERT [ 17 ] further boosts its performance via self-supervised training . The sequential modeling of transformer has activated a wide range of researches in natural language processing [ 57 ] , speech processing [ 56 ] , and computer vision [ 25 ] . In this section , we only survey transformer-related works from the computer vision perspective . There are heavy researches on transformers in both visual recognition and generation . ViT [ 19 ] has shown that CNN is not a must in image classification . DETR [ 7 ] , Deformable DETR [ 78 ] and RelationNet++ [ 16 ] indicate that transformers are able to detect objects with high precisions . SETR [ 76 ] brings transformers into semantic segmentation while VisTR [ 69 ] has shown transformers are able to perform video object segmentation . TrackFormer [ 40 ] introduces transformers into multiple object tracking . A general form of transformer is formulated in NLM [ 68 ] for video classification . Furthermore , transformers have been show effective in image generation [ 46 ] and image processing [ 11 ] scenarios . Examples include image super-resolution [ 72 ] , video inpainting [ 74 ] , and video captioning [ 77 ] . There are several emerging studies [ 15 , 10 , 4 , 71 ] on how to use selfsupervised learning to improve a transformer backbone . The learning paradigm for CNN encoders is adapted to the transformer without significant alteration . Different from existing methods that focus on learning a transformer encoder backbone with supervised or self-supervised learning . We explore how to use the transformers as guidance to enhance CNN visual attention . The pretrained CNN encoder benefits downstream recognition scenarios . 3 Proposed method . Our CARE framework consists of C-stream and T-stream . Fig . 2 shows an overview of the pipeline . We first illustrate the detailed structure of these streams . Then , we illustrate the network training process . The CNN encoder features are visualized as well for attention display . 3.1 CNN-stream ( C-stream ) . Our C-stream is similar to the existing SSL framework [ 24 ] where there are two CNN encoders , two projectors , and one predictor . The structures of the two encoders are the same , and the structures of the two projectors are the same . Given a training image x , we process it with a set of random augmentations to create two different augmented views . We feed these two views to C-stream and obtain corresponding outputs f1 ( x ) and f2 ( x ) , respectively . Then , we compute a loss Lc to penalize the dissimilarity of the outputs . This loss term is the mean squared error of the normalized feature vectors and can be written as what follows : Lc = 2− 2 · ⟨f1 ( x ) , f2 ( x ) ⟩ ∥f1 ( x ) ∥2 · ∥f2 ( x ) ∥2 ( 1 ) where ∥ · ∥2 is the ℓ2 normalization , and the ⟨ , ⟩ is the dot product operation . As the inputs of C-stream are from one image , the outputs of C-stream are supposed to become similar during the training process . 3.2 Transformer-stream ( T-stream ) . The T-stream takes the output feature maps of the CNN encoders as its inputs , which are set in parallel to the C-stream . It consists of two transformers , two projectors , and one predictor . The structures of the projectors and the predictor are the same as those in the C-stream . The structures of two transformers share the same architecture , which consists of n consecutive attention blocks where each block contains two Multilayer Perception ( MLP ) layers with one multi-head self-attention ( MHSA ) layer in between . We mainly follow the design of [ 54 ] to construct the attention block in transformer as shown in Fig . 3 . The input feature map ( denoted as s ) of an attention block is first processed by the a MLP layer for dimension mapping , and then passes the self-attention layer and finally the another MLP layer . MHSA consists of multiple attention heads that process the input features in parallel . In one attention head , as illustrated on the right of Fig . 3 , the input feature map is mapped to the query feature ( q ) , the key feature ( k ) , and the value feature ( k ) via 3 different MLP layers wq , wk , and wv , respectively . As detailed in Eq . ( 2 ) , the query q and key k are multiplied to form the content-based attention , and q and the position encoding p are multiplied to form the position-based attention . Attention ( q , k , v ) = softmax ( qpT + qkT√ dk ) v ( 2 ) where dk is the dimension of the query and the key . There are learnable parameters in the positional encoding module [ 46 ] to output p that follows the shape of s. In Eq . ( 2 ) , we perform matrix multiplication between q and pT , q and kT , and the softmax output and v by treating each pixel as a token [ 62 ] ( i.e. , for a feature map with c channels and spatial dimension of h × w , it forms h · w c-dimensional tokens and thus obtains a matrix of size h · w × c ) . Besides , we perform matrix addition between qpT and qkT . The output of the second MLP layer is added to the original input feature map s via a residual connection [ 29 ] , and finally passes a ReLU activation layer . In T-stream , the outputs of the two transformers are feature maps with 2D dimensions , which are then average-pooled and sent to the projectors and the predictor . We denote the outputs of T-stream as f3 ( x ) and f4 ( x ) . Following the dissimilarity penalty in Sec . 3.1 , we compute Lt as follows : Lt = 2− 2 · ⟨f3 ( x ) , f4 ( x ) ⟩ ∥f3 ( x ) ∥2 · ∥f4 ( x ) ∥2 . ( 3 ) Besides introducing SSL loss terms in both streams , we use the T-stream output to supervise the C-stream . This attention supervision loss term can be written as : Latt = ∥f1 ( x ) − f3 ( x ) ∥2 + ∥f2 ( x ) − f4 ( x ) ∥2 ( 4 ) where the C-stream outputs are required to resemble the T-stream outputs during the training process . Note that the network supervision in Eq . 4 can not be simply considered as the knowledge distillation ( KD ) process . There are several differences from 3 perspectives : ( 1 ) The architecture design between ours and KD is different . In KD [ 31 ] , a large teacher network is trained to supervise a small student network . In contrast , the CNN backbones are shared by two similar networks in our method . The different modules are only transformers , lightweight projectors , and predictor heads . ( 2 ) The training paradigm is different . In KD , the teacher network is typically trained in advance before supervising the student network . In contrast , two branches of our method are trained together from scratch for mutual learning . ( 3 ) The loss function in KD is normally the cross-entropy loss while we adopt mean squared error . During KD , supervision losses are also computed between feature map levels . While our method only computes losses based on the network outputs .

Summary: A self-supervised visual representation learning method is proposed in this paper. Besides the prevalent two-branch CNN architecture design, the two-branch transformer structure is introduced to improve CNN backbone features via attention enhancement. The proposed method is evaluated on the several visual recognition benchmarks and shown effective through extensive experiments.