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: This paper uses transformers to improve visual representation learning. The two view augmentation and the response generation solution from existing approaches are boosted by transformers for visual attention enhancement. The main contribution lies in a structure with two transformer branches to improve CNN attention and supervise the original CNN branches. I can see that the attention mechanism obtains gains due to the parallel design. The proposed method shows superiority, compared to existing work.