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: The paper explores a hybrid type of model architectures that combines the recently popular transformer based model architectures with already well established convolutional neural networks. Specifically, the proposed hybrid model follows a two-stage training strategy where first a CNN model is trained and the pre-trained weights are used to further improve the training by recasting the weights into a new transformed CNN model termed int he paper as T-CNN. The proposed T-CNN model architecture is able to outperform its CNN counterparts and other models compared in the paper. The representations learnt by the T-CNN are also analyzed in the experimental analysis section.