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 proposed a novel multi-stage vision transformer architecture, named CrossFormer, which aimed at building interaction among features from different scales. The proposed architecture follows the layout of Swin transformer but replaces the shifted windows transformer block with the proposed short-distance attention and long-distance attention. A cross-scale module is proposed to replace the patch embedding or patch merging layer in the Swin transformer. Comprehensive experiments are conducted on image classification, objection detection, and segmentation tasks. Results shows the effectiveness of the proposed architecture.