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: The paper proposed an attention approach to handle the global or long-range attention by leveraging the idea of quadtree structure, meanwhile reducing the quadratic complexity of original attention operation to linear. Experiments are performed on several tasks, e.g. feature matching, image classification, and objection detection. Superior results are achieved on these tasks.