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: The authors propose a new ViT-based architecture for image recognition. The architecture is based on 3 main ideas: 1) using channel-wise conv to reduce the spatial dimension of feature maps to save computations. 2) using a stack of convs as "patch embedding", which is used at the beginning of each stage to shrink spatial size while increasing channel dimensions. 3) A depth-wise conv to replace the "position encoding" in the original ViT.   Overall the model achieves strong performance, outperforming Swin-Transformer at comparable computational cost.