As the core building block of vision transformers, attention is a powerful tool to capture long-range dependency. However, such power comes at a cost: it incurs a huge computation burden and heavy memory footprint as pairwise token interaction across all spatial locations is computed. A series of works attempt to alleviate this problem by introducing handcrafted and content-agnostic sparsity into attention, such as restricting the attention operation to be inside local windows, axial stripes, or dilated windows. In contrast to these approaches, we propose a novel dynamic sparse attention via bi-level routing to enable a more flexible allocation of computations with content awareness. Specifically, for a query, irrelevant key-value pairs are first filtered out at a coarse region level, and then fine-grained token-to-token attention is applied in the union of remaining candidate regions (\ie, routed regions). We provide a simple yet effective implementation of the proposed bi-level routing attention, which utilizes the sparsity to save both computation and memory while involving only GPU-friendly dense matrix multiplications. Built with the proposed bi-level routing attention, a new general vision transformer, named BiFormer, is then presented. As BiFormer attends to a small subset of relevant tokens in a \textbf{query adaptive} manner without distraction from other irrelevant ones, it enjoys both good performance and high computational efficiency, especially in dense prediction tasks. Empirical results across several computer vision tasks such as image classification, object detection, and semantic segmentation verify the effectiveness of our design. Code is available at \url{https://github.com/rayleizhu/BiFormer}.
翻译:作为视觉Transformer的核心构建模块,注意力机制是捕获长程依赖关系的强大工具。然而,这种能力伴随着代价:由于需要计算所有空间位置上的成对令牌交互,导致巨大的计算负担和沉重的内存占用。一系列研究尝试通过在手工艺设计的、与内容无关的稀疏性注意力机制(例如将注意力操作限制在局部窗口、轴向条纹或膨胀窗口内)来缓解这一问题。与这些方法不同,我们提出一种基于双层路由的新型动态稀疏注意力机制,以实现更具灵活性的内容感知计算分配。具体而言,对于某个查询,首先在粗粒度区域级别过滤掉无关的键值对,然后在剩余候选区域(即路由区域)的并集中应用细粒度的令牌到令牌注意力。我们为所提出的双层路由注意力提供了简单有效的实现,它利用稀疏性节省计算与内存,同时仅涉及GPU友好的密集矩阵乘法。基于所提出的双层路由注意力,我们构建了名为BiFormer的新型通用视觉Transformer。由于BiFormer以查询自适应方式关注相关令牌的小子集,且不受无关令牌干扰,它在密集预测任务中兼具优异性能与高计算效率。在图像分类、目标检测和语义分割等多项计算机视觉任务上的实验结果验证了我们设计的有效性。代码见:\url{https://github.com/rayleizhu/BiFormer}。