CrossFormer: A Versatile Vision Transformer Hinging on Cross-scale Attention

1 INTRODUCTION . It turns out that transformer ( Vaswani et al. , 2017 ; Devlin et al. , 2019 ; Brown et al. , 2020 ) has achieved great success in the field of natural language processing ( NLP ) . Benefitting from its selfattention module , transformer is born with the key ability to build long-distance dependencies . Since long-distance dependencies are also needed by a number of vision tasks ( Zhang & Yang , 2021 ; Chu et al. , 2021 ) , a surge of research work ( Dosovitskiy et al. , 2021 ; Touvron et al. , 2021 ; Wang et al. , 2021 ) has been conducted to explore various transformer-based vision architectures . A transformer requires a sequence of embeddings2 ( e.g. , word embeddings ) as input . To adapt this requirement to typical vision tasks , most existing vision transformers ( Dosovitskiy et al. , 2021 ; Touvron et al. , 2021 ; Wang et al. , 2021 ; Liu et al. , 2021b ) produce embeddings by splitting an input image into equal-sized patches . For example , a 224 × 224 image can be split into 56 × 56 patches of size 4 × 4 , and these patches are projected through a linear layer to yield an embedding sequence . Inside a certain transformer , self-attention is engaged to build the interactions between any two embeddings . Thus , the computational or memory cost of the self-attention module is O ( N2 ) , where N is the length of an embedding sequence . Such a cost is too big for a visual input because its embedding sequence is much longer than that of NLP . Therefore , the recently proposed vision transformers ( Wang et al. , 2021 ; Liu et al. , 2021b ; Lin et al. , 2021 ) develop multiple substitutes to approximate the vanilla self-attention module with a lower cost . Though the aforementioned vision transformers have made some progress , they do not explicitly utilize features of different scales , whereas multi-scale features are very vital for a lot of vision tasks . For example , an image often contains many objects of different scales , and to fully understand 1The code has been submitted along with our supplementary material . 2In this paper , we also use “ embeddings ” to represent the input of each layer . the image , building the interactions among those objects is helpful . Besides , some particular tasks such as instance segmentation need the interactions between large-scale ( coarse-grained ) features and small-scale ( fine-grained ) ones . Existing vision transformers fail to deal with the above cases due to two reasons : ( 1 ) The embeddings are generated from equal-sized patches . Though these patches theoretically have a chance to extract any scale features if only the receptive field is large enough , it is difficult to promise that they can learn appropriate multi-scale features automatically in practice . ( 2 ) Inside the self-attention module , adjacent embeddings are often grouped together and merged ( Wang et al. , 2021 ) . Since the number of groups is smaller than that of embeddings , such behavior can reduce the computational budget of the self-attention . In this case , however , even if embeddings have both small-scale and large-scale features , merging operations will lose the small-scale features of each individual embedding , thereby disabling the cross-scale attention . To enable the building of cross-scale interactions , we co-design a novel embedding layer and selfattention module as follows . 1 ) Cross-scale Embedding Layer ( CEL ) – Following Wang et al . ( 2021 ) , we also employ a pyramid structure for our transformer , which naturally splits the vision transformer model into multiple stages . CEL appears at the start of each stage , which receives last stage ’ s output ( or an input image ) as input and samples patches with multiple kernels of different scales ( e.g. , 4×4 or 8 × 8 ) . Then , each embedding is constructed by projecting and concatenating these patches . Through this way , we enforce some dimensions ( e.g. , dimensions from 4 × 4 patches ) to focus on small-scale features only , while others ( e.g. , those from 8× 8 patches ) have a chance to learn largescale features , leading to an embedding with explicitly cross-scale features . 2 ) Long Short Distance Attention ( LSDA ) – We propose a substitute of the vanilla self-attention , but to preserve small-scale features , the embeddings will not be merged . In contrast , we split the self-attention module into Short Distance Attention ( SDA ) and Long Distance Attention ( LDA ) . SDA builds the dependencies among neighboring embeddings , while LDA takes charge of the dependencies among embeddings far away from each other . The proposed LSDA can also reduce the cost of the self-attention module like previous studies ( Wang et al. , 2021 ) , but different from them , LSDA does not undermine either small-scale or large-scale features . As a result , attention with cross-scale interactions is enabled . Besides , following prior work ( Shaw et al. , 2018 ; Liu et al. , 2021b ) , we employ a relative position bias for embeddings ’ position representations . The Relative Position Bias ( RPB ) only supports fixed image/group size3 . However , image size for many vision tasks such as object detection is variable , so does group size for many architectures , including ours . To make the RPB more flexible , we further introduce a trainable module called Dynamic Position Bias ( DPB ) , which receives two embeddings ’ relative distance as input and outputs their position bias . The DPB module is optimized end-to-end in the training phase , inducing an ignorable cost but making RPB apply to variable image/group size . All our proposed modules can be implemented with about ten lines of code . Based on them , we construct four versatile vision transformers of different sizes , dubbed CrossFormers . Other than image classification , the proposed CrossFormer can handle a variety of tasks with variable-sized inputs such as object detection . Experiments on four representative vision tasks ( i.e. , image classification , object detection , instance segmentation , and semantic segmentation ) demonstrate that CrossFormer outperforms the other state-of-the-art vision transformers on all the tasks . Remarkably , the performance gains brought by CrossFormer are substantially significant on dense prediction tasks , e.g. , object detection and instance/semantic segmentation . It is worth highlighting our contributions as follows : • We propose cross-scale embedding layer ( CEL ) and long short distance attention ( LSDA ) , which together compensate for existing transformers ’ incapability of building cross-scale attention . • The dynamic position bias module ( DPB ) is further proposed to make the relative position bias more flexible , i.e. , accommodating variable image size or group size . • Multiple CrossFormers with different sizes are constructed , and we corroborate their effective- ness through sufficient experiments on four representative vision tasks . 2 BACKGROUND . Vision Transformers . Motivated by the transformers developed for NLP , researchers design specific visual transformers for vision tasks to take full advantage of their powerful attention mechanism . In 3Some vision transformers split input embeddings into several groups . Group size means the number of embeddings in a group . particular , ViT and DeiT transfer the original transformer Vaswani et al . ( 2017 ) to vision tasks ( Touvron et al. , 2021 ; Dosovitskiy et al. , 2021 ) , achieving impressive performance . Later , PVT ( Wang et al. , 2021 ) , HVT ( Pan et al. , 2021 ) , Swin ( Liu et al. , 2021b ) , and ViTAE ( Xu et al. , 2021b ) introduce a pyramid structure into the visual transformers , greatly decreasing the number of patches in the later layers of a respective model . They also extend the visual transformers to other vision tasks like object detection and segmentation ( Wang et al. , 2021 ; Liu et al. , 2021b ) . Cross-scale Feature Extraction . Szegedy et al . ( 2015 ) ; Tan & Le ( 2019 ) use multi-scale convolutional kernels at every layer of models for cross-scale features , while the largest kernel size can only be 7 × 7 , and larger kernel size will induce unaccpetable computational budget . For vision transformers , CoaT ( Xu et al. , 2021a ) uses multi-scale features at later layers of models by mixing features from different layers ; CViT ( Chen et al. , 2021a ) keeps embeddings of different scales in different branch , and introduces cross-scale interaction through self-attention between branches . More discussions can be seen in the appendix C.1 . Sparse Self-attention . Many substitutes have been proposed ( Liu et al. , 2021b ; Wang et al. , 2021 ; Chu et al. , 2021 ; Child et al. , 2019 ; Beltagy et al. , 2020 ; Zaheer et al. , 2020 ) to alleviate the cost of the vanilla self-attention module . Instead of going self-attention among all embeddings , each embedding only interact with part of other embeddings , which change dense self-attention to a sparse one . More discussions can be seen in the appendix C.2 . Position Representations . To make the respective model aware of position information of embeddings , many different position representations ( Vaswani et al. , 2017 ) are proposed . For example , Dosovitskiy et al . ( 2021 ) directly add the embeddings with the vectors that contain absolute position information . In contrast , Relative Position Bias ( RPB ) ( Shaw et al. , 2018 ) resorts to position information indicating the relative distance of two embeddings . 3 CROSSFORMER . The overall architecture of CrossFormer is plotted in Figure 1 . Following ( Wang et al. , 2021 ; Liu et al. , 2021b ; Lin et al. , 2021 ) , CrossFormer also employs a pyramid structure , which naturally splits the transformer model into four stages . Each stage consists of a cross-scale embedding layer ( CEL , Section 3.1 ) and several CrossFormer blocks ( Section 3.2 ) . A CEL receives last stage ’ s output ( or an input image ) as input and generates cross-scale embeddings . In this process , CEL ( except that in Stage-1 ) reduces the number of embeddings to a quarter while doubles their dimensions for a pyramid structure . Then , several CrossFormer blocks , each of which involves long short distance attention ( LSDA ) and dynamic position bias ( DPB ) , are set up after CEL . A specialized head ( e.g. , the classification head in Figure 1 ) follows after the final stage accounting for a specific task . 3.1 CROSS-SCALE EMBEDDING LAYER ( CEL ) . Cross-scale embedding layer ( CEL ) is leveraged to generate input embeddings for each stage . Figure 2 takes the first CEL , which is ahead of Stage-1 , as an example . It receives an image as input , then sampling patches using four kernels of different sizes . The stride of four kernels is kept the same so that they generate the same number of embeddings4 . As we can observe in Figure 2 , every four corresponding patches have the same center but different scales , and all these four patches will be projected and concatenated as one embedding . In practice , the process of sampling and projecting can be fulfilled through four convolutional layers . For a cross-scale embedding , one problem is how to set the projected dimension of each scale . The computational budget of a convolutional layer is proportional to K2D2 , where K and D represent kernel size and input/output dimension , respectively ( assuming the input dimension equals to the output dimension ) . Therefore , given the same dimension , a large kernel consumes more budget than a smaller one . To control the total budget of the CEL , we use a lower dimension for large kernels while a higher dimension for small kernels . Figure 2 provides the specific allocation rule in its subtable , and a 128 dimensional example is given . Compared with allocating the dimension equally , our scheme saves much computational cost but does not explicitly affect the model ’ s performance . The cross-scale embedding layers in other stages work in a similar way . As shown in Figure 1 , CELs for Stage-2/3/4 use two different kernels ( 2 × 2 and 4 × 4 ) . Further , to form a pyramid structure , the strides of CELs for Stage-2/3/4 are set as 2× 2 to reduce the number of embeddings to a quarter .

Summary: The paper describes a pyramidal vision transformer that introduces cross-scale patch embeddings, dynamic relative position biases and "Long-Short Attention". The model shows superior performance on imagenet classification and especially on dense prediction tasks. Ablations on classification are provided hinting that the multi-scale embeddings are the most important factor contributing to the overall superior results to similar, previous architectures.