QuadTree Attention for Vision Transformers

TL;DR

QuadTree Attention introduces a linear-complexity attention mechanism for vision transformers, enabling efficient dense predictions and achieving state-of-the-art results across multiple vision tasks.

Contribution

It proposes a novel quadtree-based attention method that reduces complexity from quadratic to linear, improving efficiency and performance in vision transformers.

Findings

Achieves 4.0% improvement in feature matching on ScanNet

Reduces FLOPs by about 50% in stereo matching

Improves top-1 accuracy on ImageNet by 0.4-1.5%

Abstract

Transformers have been successful in many vision tasks, thanks to their capability of capturing long-range dependency. However, their quadratic computational complexity poses a major obstacle for applying them to vision tasks requiring dense predictions, such as object detection, feature matching, stereo, etc. We introduce QuadTree Attention, which reduces the computational complexity from quadratic to linear. Our quadtree transformer builds token pyramids and computes attention in a coarse-to-fine manner. At each level, the top K patches with the highest attention scores are selected, such that at the next level, attention is only evaluated within the relevant regions corresponding to these top K patches. We demonstrate that quadtree attention achieves state-of-the-art performance in various vision tasks, e.g. with 4.0% improvement in feature matching on ScanNet, about 50% flops reduction in stereo matching, 0.4-1.5% improvement in top-1 accuracy on ImageNet classification, 1.2-1.8% improvement on COCO object detection, and 0.7-2.4% improvement on semantic segmentation over previous state-of-the-art transformers. The codes are available at https://github.com/Tangshitao/QuadtreeAttention.