ResT: An Efficient Transformer for Visual Recognition

1 Introduction . Deep learning backbone architectures have been evolved for years and boost the performance of computer vision tasks such as classification [ 5 , 26 , 33 , 11 ] , object detection [ 2 , 41 , 18 , 25 ] , and instance segmentation [ 10 , 24 , 31 ] , etc . There are mainly two types of backbone architectures most commonly applied in computer vision : convolutional network ( CNN ) architectures [ 11 , 38 ] and Transformer ones [ 6 , 5 , 33 , 39 ] . Both of them capture feature information by stacking multiple blocks . The CNN block is generally a bottleneck structure [ 11 ] , which can be defined as a stack of 1 × 1 , 3 × 3 , and 1 × 1 convolution layers with residual learning ( shown in Figure 1a ) . The 1× 1 layers are responsible for reducing and then increasing channel dimensions , leaving the 3× 3 layer a bottleneck with smaller input/output channel dimensions . The CNN backbones are generally faster and require less inference time thanks to parameter sharing , local information aggregation , and dimension reduction . However , due to the limited and fixed receptive field , CNN blocks may be less effective in scenarios that require modeling long-range dependencies . For example , in instance segmentation , being able to collect and associate scene information from a large neighborhood can be useful in learning relationships across objects [ 23 ] . To overcome these limitations , Transformer backbones are recently explored for their ability to capture long-distance information [ 5 , 33 , 26 , 19 ] . Unlike CNN backbones , the Transformer ones first split an image into a sequence of patches ( i.e. , tokens ) , then sum these tokens with positional encoding to represent coarse spatial information , and finally adopt a stack of Transformer blocks to capture feature information . A standard Transformer block [ 28 ] comprises a multi-head self-attention ( MSA ) that employs a query-key-value decomposition to model global relationships between sequence 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . tokens , and a feed-forward network ( FFN ) to learn wider representations ( shown in Figure 1b ) . As a result , Transformer blocks can dynamically adapt the receptive field according to the image content . Despite showing great potential than CNNs , the Transformer backbones still have four major shortcomings : ( 1 ) It is difficult to extract the low-level features which form some fundamental structures in images ( e.g. , corners and edges ) since existing Transformer backbones direct perform tokenization of patches from raw input images . ( 2 ) The memory and computation for MSA in Transformer blocks scale quadratically with spatial or embedding dimensions ( i.e. , the number of channels ) , causing vast overheads for training and inference . ( 3 ) Each head in MSA is responsible for only a subset of embedding dimensions , which may impair the performance of the network , particularly when the tokens embedding dimension ( for each head ) is short , making the dot product of query and key unable to constitute an informative function . ( 4 ) The input tokens and positional encoding in existing Transformer backbones are all of a fixed scale , which are unsuitable for vision tasks that require dense prediction . In this paper , we proposed an efficient general-purpose backbone ResT ( named after ResNet [ 11 ] ) for computer vision , which can remedy the above issues . As illustrated in Figure 2 , ResT shares exactly the same pipeline of ResNet , i.e. , a stem module applied for extracting low-level information and strengthening locality , followed by four stages to construct hierarchical feature maps , and finally a head module for classification . Each stage consists of a patch embedding , a positional encoding module , and multiple Transformer blocks with specific spatial resolution and channel dimension . The patch embedding module creates a multi-scale pyramid of features by hierarchically expanding the channel capacity while reducing the spatial resolution with overlapping convolution operations . Unlike the conventional methods which can only tackle images with a fixed scale , our positional encoding module is constructed as spatial attention which is conditioned on the local neighborhood of the input token . By doing this , the proposed method is more flexible and can process input images of arbitrary size without interpolation or fine-tune . Besides , to improve the efficiency of the MSA , we build an efficient multi-head self-attention ( EMSA ) , which significantly reduce the computation cost by a simple overlapping Depth-wise Conv2d . In addition , we compensate short-length limitations of the input token for each head by projecting the interaction across the attention-heads dimension while keeping the diversity ability of multi-heads . We comprehensively validate the effectiveness of the proposed ResT on the commonly used benchmarks , including image classification on ImageNet-1k and downstream tasks , such as object detection , and instance segmentation on MS COCO2017 . Experimental results demonstrate the effectiveness and generalization ability of the proposed ResT compared with the recently state-of-the-art Vision Transformers and CNNs . For example , with a similar model size as ResNet-18 ( 69.7 % ) and PVT-Tiny ( 75.1 % ) , our ResT-Small obtains a Top-1 accuracy of 79.6 % on ImageNet-1k . 2 ResT . As illustrated in Figure 2 , ResT shares exactly the same pipeline as ResNet [ 11 ] , i.e. , a stem module applied to extract low-level information , followed by four stages to capture multi-scale feature maps . Each stage consists of three components , one patch embedding module ( or stem module ) , one positional encoding module , and a set of L efficient Transformer blocks . Specifically , at the beginning of each stage , the patch embedding module is adopted to reduce the resolution of the input token and expanding the channel dimension . The positional encoding module is fused to restrain position information and strengthen the feature extracting ability of patch embedding . After that , the input token is fed to the efficient Transformer blocks ( illustrated in Figure 1c ) . In the following sections , we will introduce the intuition behind ResT . 2.1 Rethinking of Transformer Block . The standard Transformer block consists of two sub-layers of MSA and FFN . A residual connection is employed around each sub-layer . Before MSA and FFN , layer normalization ( LN [ 1 ] ) is applied . For a token input x ∈ Rn×dm , where n , dm indicates the spatial dimension , channel dimension , respectively . The output for each Transformer block is : y= x′ + FFN ( LN ( x′ ) ) , and x′= x +MSA ( LN ( x ) ) ( 1 ) MSA . MSA first obtains query Q , key K , and value V by applying three sets of projections to the input , each consisting of k linear layers ( i.e. , heads ) that map the dm dimensional input into a dk dimensional space , where dk = dm/k is the head dimension . For the convenience of description , we assume k = 1 , then MSA can be simplified to single-head self-attention ( SA ) . The global relationship between the token sequence can be defined as SA ( Q , K , V ) = Softmax ( QKT√ dk ) V ( 2 ) The output values of each head are then concatenated and linearly projected to form the final output . The computation costs of MSA are O ( 2dmn2 + 4d2mn ) , which scale quadratically with spatial dimension or embedding dimensions according to the input token . FFN . The FFN is applied for feature transformation and non-linearity . It consists of two linear layers with a non-linearity activation . The first layer expands the embedding dimensions of the input from dm to df and the second layer reduce the dimensions from df to dm . FFN ( x ) = σ ( xW1 + b1 ) W2 + b2 ( 3 ) where W1 ∈ Rdm×df and W2 ∈ Rdf×dm are weights of the two Linear layers respectively , b1 ∈ Rdf and b2 ∈ Rdm are the bias terms , and σ ( · ) is the activation function GELU [ 12 ] . In standard Transformer block , the channel dimensions are expanded by a factor of 4 , i.e. , df = 4dm . The computation costs of FFN are 8nd2m . 2.2 Efficient Transformer Block . As analyzed above , MSA has two shortcomings : ( 1 ) The computation scales quadratically with dm or n according to the input token , causing vast overheads for training and inference ; ( 2 ) Each head in MSA only responsible for a subset of embedding dimensions , which may impair the performance of the network , particularly when the tokens embedding dimension ( for each head ) is short . To remedy these issues , we propose an efficient multi-head self-attention module ( illustrated in Figure 3 ) . Here , we make some explanations . ( 1 ) Similar to MSA , EMSA first adopt a set of projections to obtain query Q . ( 2 ) To compress memory , the 2D input token x ∈ Rn×dm is reshaped to 3D one along the spatial dimension ( i.e. , x̂ ∈ Rdm×h×w ) and then feed to a depth-wise convolution operation to reduce the height and width dimension by a factor s. To make simple , s is adaptive set by the feature map size or the stage number . The kernel size , stride and padding are s+ 1 , s , and s/2 respectively . ( 3 ) The new token map after spatial reduction x̂ ∈ Rdm×h/s×w/s is then reshaped to 2D one , i.e. , x̂ ∈ Rn′×dm , n′ = h/s× w/s . Then x̂ is feed to two sets of projection to get key K and value V. ( 4 ) After that , we adopt Eq . 4 to compute the attention function on query Q , K and value V. EMSA ( Q , K , V ) = IN ( Softmax ( Conv ( QKT√ dk ) ) ) V ( 4 ) Here , Conv ( · ) is a standard 1 × 1 convolutional operation , which model the interactions among different heads . As a result , attention function of each head can depend on all of the keys and queries . However , this will impair the ability of MSA to jointly attend to information from different representation subsets at different positions . To restore this diversity ability , we add an Instance Normalization [ 27 ] ( i.e , IN ( · ) ) for the dot product matrix ( after Softmax ) . ( 5 ) Finally , the output values of each head are then concatenated and linearly projected to form the final output . The computation costs of EMSA are O ( 2dmn 2 s2 + 2d 2 mn ( 1 + 1 s2 ) + dmn ( s+1 ) 2 s2 + k2n2 s2 ) , much lower than the original MSA ( assume s > 1 ) , particularly in lower stages , where s is tend to higher . Also , we add FFN after EMSA for feature transformation and non-linearity . The output for each efficient Transformer block is : y= x′ + FFN ( LN ( x′ ) ) , and x′= x + EMSA ( LN ( x ) ) ( 5 ) 2.3 Patch Embedding . The standard Transformer receives a sequence of token embeddings as input . Take ViT [ 5 ] as an example , the input image x ∈ R3×h×w is split with a patch size of p× p. These patches are flattened into 2D ones and then mapped to latent embeddings with a size of c , i.e , x ∈ Rn×c , where n = hw/p2 . However , this straightforward tokenization is failed to capture low-level feature information ( such as edges and corners ) [ 33 ] . In addition , the length of tokens in ViT are all of a fixed size in different blocks , making it unsuitable for downstream vision tasks such as object detection and instance segmentation that require multi-scale feature map representations . Here , we build an efficient multi-scale backbone , calling ResT , for dense prediction . As introduced above , the efficient Transformer block in each stage operates on the same scale with identical resolution across the channel and spatial dimensions . Therefore , the patch embedding modules are required to progressively expand the channel dimension , while simultaneously reducing the spatial resolution throughout the network . Similar to ResNet , the stem module ( can be seen as the first patch embedding module ) are adopted to shrunk both the height and width dimension with a reduction factor of 4 . To effectively capture the low-feature information with few parameters , here we introduce a simple but effective way , i.e , stacking three 3× 3 standard convolution layers ( all with padding 1 ) with stride 2 , stride 1 , and stride 2 , respectively . Batch Normalization [ 14 ] and ReLU activation [ 7 ] are applied for the first two layers . In stage 2 , stage 3 , and stage 4 , the patch embedding module is adopted to down-sample the spatial dimension by 4× and increase the channel dimension by 2× . This can be done by a standard 3× 3 convolution with stride 2 and padding 1 . For example , patch embedding module in stage 2 changes resolution from h/4× w/4× c to h/8× w/8× 2c ( shown in Figure 2 ) .

Summary: This work presents a computationally efficient vision transformer network for image recognition. In particular, this work propose memory-efficient self-attention module to replace the original one in the original transformer. Second, it modifies the inflexible position encoding (that only allows fixed length of tokens) to accept variable size of input images. Lastly, the patch embedding is achieved by stacking overlapping convolution operations with strides. Extensive experiments are conducted to show the superiority of the proposed ResT.