Quadtree Attention for Vision Transformers

1 INTROCUTION . Transformers can capture long-range dependencies by the attention module and have demonstrated tremendous success in natural language processing tasks . In recent years , transformers have also been adapted to computer vision tasks for image classification ( Dosovitskiy et al. , 2020 ) , object detection ( Wang et al. , 2021c ) , semantic segmentation ( Liu et al. , 2021 ) , feature matching ( Sarlin et al. , 2020 ) , and stereo ( Li et al. , 2021 ) , etc . Typically , images are divided into patches and these patches are flattened and fed to a transformer as word tokens to evaluate attention scores . However , transformers have quadratic computational complexity in terms of the number of tokens , i.e . number of image patches . Thus , applying transformers to computer vision applications requires careful simplification of the involved computation . To utilize the standard transformer in vision tasks , many works opt to apply it on low resolution or sparse tokens . ViT ( Dosovitskiy et al. , 2020 ) uses coarse image patches of 16 × 16 pixels to limit the number of tokens . DPT ( Ranftl et al. , 2021 ) up-samples low-resolution results from ViT to high resolution maps to achieve dense predictions . SuperGlue ( Sarlin et al. , 2020 ) applies transformer on sparse image keypoints . Focusing on correspondence and stereo matching applications , Germain et al . ( 2021 ) and Li et al . ( 2021 ) also apply transformers at a low resolution feature map . However , as demonstrated in several works ( Wang et al . ( 2021c ) ; Liu et al . ( 2021 ) ; Sun et al . ( 2021 ) ; Li et al . ( 2021 ) ) , applying transformers on high resolution is beneficial for a variety of tasks . Thus , many efforts have been made to design efficient transformers to reduce computational complexity . Linear approximate transformers ( Katharopoulos et al. , 2020 ; Wang et al. , 2020 ) approximate standard attention computation with linear methods . However , empirical studies ( Germain et al. , 2021 ; Chen et al. , 2021 ) show those linear transformers are inferior in vision tasks . To reduce the computational cost , the PVT ( Wang et al. , 2021c ) uses downsampled keys and values , which is harmful to capture pixel-level details . In comparison , the Swin Transformer ( Liu et al. , 2021 ) restricts the attention in local windows in a single attention block , which might hurt long-range dependencies , the most important merit of transformers . Unlike all these previous works , we design an efficient vision transformer that captures both fine image details and long-range dependencies . Inspired by the observation that most image regions are irrelevant , we build token pyramids and compute attention in a coarse to fine manner . In this way , we can quickly skip irrelevant regions in the fine level if their corresponding coarse level regions are not promising . For example , as in Figure 1 , at the 1st level , we compute the attention of the blue image patch in image A with all the patches in image B and choose the top K ( here , K = 2 ) patches which are also highlighted in blue . In the 2nd level , for the four framed sub-patches in image A ( which are children patches of the blue patch at the 1st level ) , we only compute their attentions with the sub-patches corresponding to the top K patches in image B at the 1st level . All the other shaded sub-patches are skipped to reduce computation . We highlight two sub-patches in image A in yellow and green . Their corresponding top K patches in image B are also highlighted in the same color . This process is iterated in the 3rd level , where we only show the sub-sub-patches corresponding to the green sub-patch at the 2nd level . In this manner , our method can both obtain fine scale attention and retain long-range connections . Most importantly , only sparse attention is evaluated in the whole process . Thus , our method has low memory and computational costs . Since a quadtree structure is formed in this process , we refer to our method as QuadTree Attention , or QuadTree Transformer . In experiments , we demonstrate the effectiveness of our quadtree transformer in both tasks requiring cross attention , e.g . feature matching and stereo , and tasks only utilizing self-attention , e.g . image classification and object detection . Our method achieves state-of-the-art performance with significantly reduced computation , comparing to relevant efficient transformers ( Katharopoulos et al . ( 2020 ) ; Wang et al . ( 2021c ) ; Liu et al . ( 2021 ) ) . In feature matching , we achieve 60.3 AUC @ 20◦ in ScanNet ( Dai et al. , 2017 ) , 2.7 higher than the linear transformer ( Katharopoulos et al. , 2020 ) but with similar flops . In stereo matching , we achieve a similar end-point-error as standard transformer , ( Li et al. , 2021 ) but with about 50 % flops reduction and 40 % memory reduction . In image classification , we achieve 82.6 % top-1 accuracy in ImageNet ( Deng et al. , 2009 ) , 6.2 % higher than ResNet ( He et al. , 2016 ) and 1.3 % higher than the Swin Transformer-T ( Liu et al. , 2021 ) , with fewer parameters and flops . In object detection , our QuadTree Attention + RetinaNet achieves 46.2 AP in COCO ( Lin et al. , 2014 ) , 1.6 higher than the backbone PVTv2 ( Wang et al. , 2021b ) but around 35 % reduction in flops . 2 RELATED WORK . Efficient Transformers . Transformers have shown great success in both natural language processing and computer vision . Due to the quadratic computational complexity , the computation of full attention is unaffordable when dealing with long sequence tokens . Therefore , many works design efficient transformers , aiming to reduce computational complexity ( Katharopoulos et al. , 2020 ; Choromanski et al. , 2020 ; Wang et al. , 2020 ; Lee et al. , 2019 ; Ying et al. , 2018 ) . Current efficient transformers can be categorized into three classes . 1 ) Linear approximate attention ( Katharopoulos et al. , 2020 ; Choromanski et al. , 2020 ; Wang et al. , 2020 ; Beltagy et al. , 2020 ; Zaheer et al. , 2020 ) approximates the full attention matrix by linearizing the softmax attention and thus can accelerate the computation by first computing the product of keys and values . 2 ) Inducing point-based linear transformers ( Lee et al. , 2019 ; Ying et al. , 2018 ) use learned inducing points with fixed size to compute attention with input tokens , thus can reduce the computation to linear complexity . However , these linear transformers are shown to have inferior results than standard transformers in different works ( Germain et al. , 2021 ; Chen et al. , 2021 ) . 3 ) Sparse attention , including Longformer ( Beltagy et al. , 2020 ) , Big Bird ( Zaheer et al. , 2020 ) , etc , attends each query token to part of key and value tokens instead of the entire sequence . Unlike these works , our quadtree attention can quickly skip the irrelevant tokens according to the attention scores at coarse levels . Thus , it achieves less information loss while keeps high efficiency . Vision Transformers . Transformers have shown extraordinary performance in many vision tasks . ViT ( Dosovitskiy et al. , 2020 ) applies transformers to image recognition , demonstrating the superiority of transformers for image classification at a large scale . However , due to the computational complexity of full attention , it is hard to apply transformers in dense prediction tasks , e.g . object detection , semantic segmentation , etc . To address this problem , Swin Transformer ( Liu et al. , 2021 ) restricts attention computation in a local window . Focal transformer ( Yang et al. , 2021 ) uses twolevel windows to increase the ability to capture long-range connection for local attention methods . Pyramid vision transformer ( PVT ) ( Wang et al. , 2021c ) reduce the computation of global attention methods by downsampling key and value tokens . Although these methods have shown improvements in various tasks , they have drawbacks either in capturing long-range dependencies ( Liu et al. , 2021 ) or fine level attention ( Wang et al. , 2021c ) . Different from these methods , our method simultaneously capture both local and global attention by computing attention from full image levels to the finest token levels with token pyramids in one single block . Besides , the K-NN transformers ( Wang et al. , 2021a ; Zhao et al. , 2019 ) aggregate messages from top K most similar tokens as ours , but they compute the attention scores among all pairs of query and key tokens , and thus still has quadratic complexity . Beyond self-attention , many tasks can largely benefit from cross attention . Superglue ( Sarlin et al. , 2020 ) processes detected local descriptors with self- and cross attention and shows significant improvement in feature matching . Standard transformers can be applied in SuperGlue because only sparse keypoints are considered . SGMNet ( Chen et al. , 2021 ) further reduces the computation by attending to seeded matches . LoFTR ( Sun et al. , 2021 ) utilizes linear transformer ( Katharopoulos et al. , 2020 ) on low-resolution feature maps to generate dense matches . For stereo matching , STTR ( Li et al. , 2021 ) applies self- and cross attention along epipolar lines and reduces the memory by gradient checkpointing engineering techniques . However , due to the requirement of processing a large number of points , these works either use linear transformers , which compromise performance , or a standard transformer , which compromises efficiency . In contrast , our transformer with quadtree attention achieves a significant performance boost compared with linear transformer or efficiency improvement compared with standard transformer . Besides , it can be applied to both self-attention and cross attention . 3 METHOD . We first briefly review the attention mechanism in transformers in Section 3.1 and then formulate our quadtree attention in Section 3.2 . 3.1 ATTENTION IN TRANSFORMER . Vision transformers have shown great success in many tasks . At the heart of a transformer is the attention module , which can capture long-range information between feature embeddings . Given two image embeddings X1 and X2 , the attention module passes information between them . Selfattention is the case when X1 and X2 are the same , while cross attention covers a more general situation when X1 and X2 are different . It first generates the query Q , key K , and value V by the Under review as a conference paper at ICLR 2022 level 3 qi Under review as a conference paper at ICLR 2022 QUADTREE ATTENTION FOR VISION TRANSFORMER Anonymous authors Paper under double-blind review Q K ⌦li ⌦l+1i ⌦l+2i mli ml+1i ml+2i li l+1i l+2i REFERENCES 1 Under review as a conference paper at ICLR 2022 QUADTREE ATTENTION FOR VISION TRANSFORMER Anonymous authors Paper under double-blind review Q K\V ⌦li ⌦l+1i ⌦l+2i mli ml+1i ml+2i li l+1i l+2i ⇤li ⇤l+1i ⇤l+2i REFERENCES 1 level 2level 1 QuadTree-A Under review as a conference paper at ICLR 2022 QUADTREE ATTENTION FOR VISION TRANSFORMER Anonymous authors Paper under double-blind review Q K\V ⌦1i ⌦2i ⌦3i m1i m2i m3i 1 2 3 ⇤li ⇤l+1i ⇤l+2i REFERENCES 1 QuadTree-B Under review as a conference paper at ICLR 2022 QUADTREE ATTENTION FOR VISION TRANSFORMER Anonymous authors Paper under double-blind review Q K\V ⌦1i ⌦2i ⌦3i m1i m2i m3i 1i 2i 3i ⇤li ⇤l+1i ⇤l+2i REFERENCES 1 ( b ) ( c ) ( a ) 0.1 0.2 0.3 0.4 0.25 0.30 0.10 0.05 0.10 0.05 0.050.10 Under review as a conference paper at ICLR 2022 QUADTREE ATTENTION FOR VISION TRANSFORMER Anonymous authors Paper under double-blind review Q K\V ⌦1i ⌦2i ⌦3i m1i m2i m3i 1i 2i 3i ⇤li ⇤l+1i ⇤l+2i level 1 level 2 level 3 m1i m2i m3i REFERENCES 1 Under review as a conference paper at ICLR 2022 QUADTREE ATTENTION FOR VISION TRANSFORMER Anonymous authors Paper under double-blind review Q K\V ⌦1i ⌦2i ⌦3i m1i m2i m3i 1i 2i 3i ⇤li ⇤l+1i ⇤l+2i level 1 level 2 level 3 m1i m2i m3i REFERENCES 1 Under review as a conference paper at ICLR 2022 QUADTREE ATTENTION FOR VISION TRANSFORMER Anonymous authors Paper under double-blind review Q K\V ⌦1i ⌦2i ⌦3i m1i m2i m3i 1i 2i 3i ⇤li ⇤l+1i ⇤l+2i level 1 level 2 level 3 m1i m2i m3i REFERENCES 1 Under review as a conference paper at ICLR 2022 QUADTREE ATTENTION FOR VISION TRANSFORMER Anonymous authors Paper under double-blind review Q K\V ⌦1i ⌦2i ⌦3i m1i m2i m3i 1i 2i 3i ⇤li ⇤l+1i ⇤l+2i level 1 level 2 level 3 m1i m2i m3i REFERENCES 1 Figure 2 : Illustration of quadtree message aggregation for a query token qi . ( a ) shows the token pyramids and involved key/value tokens in each level . Attention scores are marked in the first two levels for clarification , and the top K scores are highlighted in red . ( b ) shows message aggregation for QuadTree-A architecture . The message is assembled from different levels along a quadtree . ( c ) shows message aggregation for QuadTree-B architecture . The message is collected from overlapping regions from different levels . following equation , Q = WqX1 , K = WkX2 , V = WvX2 , where Wq , Wk and Wv are learnable parameters . Then , it performs message aggregation by computing the attention scores between query and key as following , Y = softmax ( QKT√ C ) V , ( 1 ) where C is the embedding channel dimension . The above process has O ( N2 ) computational complexity , where N is the number of image patches in a vision transformer . This quadratic complexity hinders transformers from being applied to tasks requiring high resolution output . To address this problem , PVT ( Wang et al . ( 2021c ) ) downsamples K and V , while Swin Transformer ( Liu et al . ( 2021 ) ) limits the attention computation within local windows .

Summary: This paper proposes a new attention mechanism in vision transformers, which is called quadtree attention. This quadtree attention mechanism builds token pyramids and computes the attention in a coarse-to-fine manner. At each level, only top K regions with the highest attention scores from query and key are selected for further attention computation on finer tokens, while other regions simply use (current) coarse attention as a part of output. In the empirical study, the proposed quadtree attention achieves good performance on a wide range of vision tasks (e.g. feature matching, stereo matching, classification and object detection) with less computation.