Transformed CNNs: recasting pre-trained convolutional layers with self-attention

INTRODUCTION Since the success of AlexNet in 2012 ( Krizhevsky et al. , 2017 ) , the field of Computer Vision has been dominated by Convolutional Neural Networks ( CNNs ) ( LeCun et al. , 1998 ; 1989 ) . Their local receptive fields give them a strong inductive bias to exploit the spatial structure of natural images ( Scherer et al. , 2010 ; Schmidhuber , 2015 ; Goodfellow et al. , 2016 ) , while allowing them to scale to large resolutions seamlessly . Yet , this inductive bias limits their ability to capture long-range interactions . In this regard , self-attention ( SA ) layers , originally introduced in language models ( Bahdanau et al. , 2014 ; Vaswani et al. , 2017 ; Devlin et al. , 2018 ) , have gained interest as a building block for vision Ramachandran et al . ( 2019 ) ; Zhao et al . ( 2020 ) . Recently , they gave rise to a plethora of Vision Transformer ( ViT ) models , able to compete with state-of-the-art CNNs in various tasks Dosovitskiy et al . ( 2020 ) ; Touvron et al . ( 2020 ) ; Wu et al . ( 2020 ) ; Touvron et al . ( 2021 ) ; Liu et al . ( 2021 ) ; Heo et al . ( 2021 ) while demonstrating better robustness ( Bhojanapalli et al. , 2021 ; Mao et al. , 2021 ) . However , capturing long-range dependencies necessarily comes at the cost of quadratic complexity in input size , a computational burden which many recent directions have tried to alleviate ( Bello , 2021 ; Wang et al. , 2020 ; Choromanski et al. , 2020 ; Katharopoulos et al. , 2020 ) . Additionally , ViTs are generally harder to train ( Zhang et al. , 2019 ; Liu et al. , 2020 ) , and require vast amounts of pre-training ( Dosovitskiy et al. , 2020 ) or distillation from a convolutional teacher ( Hinton et al. , 2015 ; Jiang et al. , 2021 ; Graham et al. , 2021 ) to match the performance of CNNs . Faced with the dilemma between efficient CNNs and powerful ViTs , several approaches have aimed to bridge the gap between these architectures . On one side , hybrid models append SA layers onto convolutional backbones ( Chen et al. , 2018 ; Bello et al. , 2019 ; Graham et al. , 2021 ; Chen et al. , 2021 ; Srinivas et al. , 2021 ) , and have already fueled successful results in a variety of tasks ( Carion et al. , 2020 ; Hu et al. , 2018 ; Chen et al. , 2020 ; Locatello et al. , 2020 ; Sun et al. , 2019 ) . Conversely , a line of research has studied the benefit of introducing convolutional biases in Transformer architectures to ease learning ( d ’ Ascoli et al. , 2021 ; Wu et al. , 2021 ; Yuan et al. , 2021 ) . Despite these interesting compromises , modelling long-range dependencies at low computational cost remains a challenge for practitioners . Contributions At a time when pre-training on vast datasets has become common practice , we ask the following question : does one need to train the SA layers during the whole learning process ? Could one instead learn cheap components such as convolutions first , leaving the SA layers to be learnt at the end ? In this paper , we take a step in this direction by presenting a method to fully reparameterize any pre-trained convolutional layer as a Gated Positional Self-Attention ( GPSA ) layer ( d ’ Ascoli et al. , 2021 ) . The latter is initialized to reproduce the mapping of the convolutional layer , but is then encouraged to learn more general mappings which are not accessible to the CNN by adjusting positional gating parameters . We leverage this method to reparametrize pre-trained CNNs as functionally equivalent hybrid models . After only 50 epochs of fine-tuning , the resulting Transformed CNNs ( T-CNNs ) boast significant performance and robustness improvements as shown in Fig . 1 , demonstrating the practical relevance of our method . We analyze the inner workings of the T-CNNs , and show that they learn more robust representations by combining convolutional heads and attentional heads in a complementary way . Related work Our work mainly builds on two pillars . First , the idea that SA layers can express any convolution , introduced by Cordonnier et al . ( 2019 ) . This idea was recently leveraged by d ’ Ascoli et al . ( 2021 ) , which initialize the SA layers of end-to-end Transformers as random convolutions to imbue them with a local inductive bias and improve their sample efficiency . Our approach leverages the opposite idea : giving an end-to-end CNN the freedom to escape locality by learning self-attention at late times . Second , we exploit the following learning paradigm : train a simple and fast model , then reparameterize it as a more complex model for the final stages of learning . This approach was studied from a scientific point of view in d ’ Ascoli et al . ( 2019 ) , which shows that reparameterizing a CNN as a fully-connected network ( FCN ) halfway through training can lead the FCN to outperform the CNN . Yet , the practical relevance of this method is limited by the vast increase in number of parameters required by the FCN to functionally represent the CNN . In contrast , our reparameterization hardly increases the parameter count of the CNN , making it easily applicable to any state-of-the-art CNN . Note that these reparameterization methods can be viewed an informed version of dynamic architecture growing algorithms such as AutoGrow ( Wen et al. , 2020 ) . In the context of hybrid models , various works have studied the performance gains obtained by introducing MHSA layers in ResNets with minimal architectural changes ( Srinivas et al. , 2021 ; Graham et al. , 2021 ; Chen et al. , 2021 ) . However , the MHSA layers used in these works are initialized randomly and need to be trained from scratch . Our approach is different , as it makes use of GPSA layers , which can be initialized to represent the same function as the convolutional layer it replaces . We emphasize that the novelty in our work is not in the architectures used , but in the unusual way they are blended together . 1 BACKGROUND . Multi-head self-attention The SA mechanism is based on a trainable associative memory with ( key , query ) vector pairs . To extract the semantic interpendencies between the L elements of a sequence X ∈ RL×Din , a sequence of “ query ” embeddings Q = WqryX ∈ RL×Dh is matched against another sequence of “ key ” embeddingsK = WkeyX ∈ RL×Dh using inner products . The result is an attention matrix whose entry ( ij ) quantifies how semantically relevantQi is toKj : A = softmax ( QK > √ Dh ) ∈ RL×L . ( 1 ) Multi-head SA layers use several SA heads in parallel to allow the learning of different kinds of dependencies : MSA ( X ) : = Nh∑ h=1 [ SAh ( X ) ] W hout , SAh ( X ) : = A hXW hval , ( 2 ) whereW hval ∈ RDin×Dv andW hout ∈ RDv×Dout are two learnable projections . To incorporate positional information , ViTs usually add absolute position information to the input at embedding time , before propagating it through the SA layers . Another possibility is to replace the vanilla SA with positional SA ( PSA ) , by including a position-dependent term in the softmax ( Ramachandran et al. , 2019 ; Shaw et al. , 2018 ) . Although there are several ways to parametrize the positional attention , we use encodings rij of the relative position of pixels i and j as in ( Cordonnier et al. , 2019 ) : Ahij : = softmax ( QhiK h > j + v h > posrij ) . ( 3 ) Each attention head learns an embedding vhpos ∈ RDpos , and the relative positional encodings rij ∈ RDpos only depend on the distance between pixels i and j , denoted denoted as a twodimensional vector δij . Self-attention as a generalized convolution Cordonnier et al . ( 2019 ) shows that a multi-head PSA layer ( Eq . 3 ) with Nh heads and dimension Dpos ≥ 3 can express any convolutional layer of filter size √ Nh , with Din input channels and min ( Dv , Dout ) output channels , by setting the following : vhpos : = −αh ( 1 , −2∆h1 , −2∆h2 , 0 , . . . 0 ) rδ : = ( ‖δ‖2 , δ1 , δ2 , 0 , . . . 0 ) Wqry = Wkey : = 0 ( 4 ) In the above , the center of attention ∆h ∈ R2 is the position to which head h pays most attention to , relative to the query pixel , whereas the locality strength αh > 0 determines how focused the attention is around its center ∆h . When αh is large , the attention is focused only on the pixel located at ∆h ; when αh is small , the attention is spread out into a larger area . Thus , the PSA layer can achieve a convolutional attention map by setting the centers of attention ∆h to each of the possible positional offsets of a √ Nh × √ Nh convolutional kernel , and sending the locality strengths αh to some large value . 2 APPROACH . In this section , we introduce our method for mapping a convolutional layer to a functionally equivalent PSA layer with minimal increase in parameter count . To do this , we leverage the GPSA layers introduced in d ’ Ascoli et al . ( 2021 ) . Loading the filters We want each head h of the PSA layer to functionally mimic the pixel h of a convolutional filterWfilter ∈ RNh×Din×Dout , where we typically have Dout ≥ Din . Rewriting the action of the MHSA operator in a more explicit form , we have MHSA ( X ) = Nh∑ h=1 AhX W hvalW h out︸ ︷︷ ︸ W h∈RDin×Dout ( 5 ) In the convolutional configuration of Eq . 4 , AhX selects pixel h of X . Hence , we need to set W h = W hfilter . However , as a product of matrices , the rank ofWh is bottlenecked by Dv . To avoid this being a limitation , we need Dv ≥ Din ( since Dout ≥ Din ) . To achieve this with a minimal number of parameters , we choose Dv = Din , and simply set the following initialization : W hval = I , W h out = W h filter . ( 6 ) Note that this differs from the usual choice made in SA layers , where Dv = bDin/Nhc . However , to keep the parameter count the same , we share the sameW hval across different heads h , since it plays a symmetric role at initialization . Note that this reparameterization introduces three additional matrices compared to the convolutional filter : Wqry , Wkey , Wval , each containing Din × Din parameters . However , since the convolutional filter contains Nh ×Din ×Dout parameters , where we typically have Nh = 9 and Dout ∈ { Din , 2Din } , these additional matrices are much smaller than the filters and hardly increase the parameter count . This can be seen from the model sizes in Tab . 3 . Gated Positional self-attention Recent work ( d ’ Ascoli et al. , 2021 ) has highlighted an issue with standard PSA : the fact that the content and positional terms in Eq . 3 are potentially of very different magnitudes , in which case the softmax ignores the smallest of the two . This can typically lead the PSA to adopt a greedy attitude : choosing the form of attention ( content or positional ) which is easiest at a given time then sticking to it . To avoid this , the authors suggest to sum the content and positional terms after the softmax , with their relative importances governed by a learnable gating parameter λh ( one for each attention head ) . The resulting Gated Positional Self-Attention ( GPSA ) layers are parametrized as follows : Ahij : = ( 1− σ ( λh ) ) softmax ( QhiK h > j ) + σ ( λh ) softmax ( vh > posrij ) , ( 7 ) where σ : x 7→ 1/ ( 1+e−x ) is the sigmoid function . In the positional part , the encodings rij are fixed rather than learnt ( see Eq . 4 ) , which makes changing input resolution straightforward ( see SM . C ) and leaves only 3 learnable parameters per head : ∆1 , ∆2 and α1 . How convolutional should the initialization be ? The convolutional initialization of GPSA layers involves two parameters , determining how strictly convolutional the behavior is : the initial value of the locality strength α , which determines how focused each attention head is on its dedicated pixel , and the initial value of the gating parameters λ , which determines the importance of the positional information versus content . If λh 0 and α 1 , the T-CNN will perfectly reproduce the input-output function of the CNN , but may want to greedily stay in the convolutional configuration . Conversely , if λh 0 and α 1 , the T-CNN will forget about the input-output function of the CNN . Hence , we choose α = 1 and λ = 1 to lie in between these two extremes , encouraging the T-CNNs to escape locality throughout training . Architectural details To make our setup as canonical as possible , we focus on ResNet architectures ( He et al. , 2016 ) , which contain 5 stages , with spatial resolution halfed and number of channels doubled at each stage . Our method involves reparameterizing 3×3 convolutions as GPSA layers with 9 attention heads . However , global SA is too costly in the first layers , where the spatial resolution is large . We therefore only reparameterize the last stage of the architecture , while replacing the first stride-2 convolution by a stride-1 convolution , exactly as in ( Srinivas et al. , 2021 ) . We also add explicit padding layers to account for the padding of the original convolutions . 1Since α represents the temperature of the softmax , its value must stay positive at all times . To ensure this , we instead learn a rectified parameter α̃ using the softplus function : α = 1 β log ( 1 + e−βα̃ ) , with β = 5 .

Summary: This work proposes an approach to bridge CNNs and vision transformers for image recognition. The idea is to replace the last convolutional stage of a ResNet by a self-attention layer which is initialized from the weights of the convolutional stage. The approach is shown to improve the performance of the CNN, especially in terms of robustness to adversarial examples, corruptions in the input images, and domain shift.