BiFormer: Vision Transformer with Bi-Level Routing Attention
摘要
作为视觉变换器的核心构建模块，注意力机制是一种强大的工具，用于捕获长距离依赖关系。然而，这种强大的功能是有代价的：它会导致巨大的计算负担和沉重的内存占用，因为需要在所有空间位置之间进行成对token交互计算。一系列工作试图通过引入手工制作和 content-agnostic 的稀疏性来缓解这个问题，例如将注意力操作限制在局部窗口、轴向条纹或扩张窗口内。与这些方法不同，本文提出了一种通过双层路由的新型动态稀疏注意力机制，以实现更灵活的内容感知计算分配。具体来说，对于查询，首先在粗粒度区域级别过滤掉无关的键值对，然后在剩余候选区域（即 routed regions）的并集中应用细粒度的 token-to-token 注意力。本文提供了一种简单而有效的双层路由注意力实现方法，该方法利用稀疏性来节省计算和内存，同时仅涉及 GPU 友好的密集矩阵乘法。基于所提出的双层路由注意力机制，一种新的通用视觉transformer，命名为BiFormer，随即提出。由于BiFormer以 query adaptive 的方式关注相关token的小子集，而不会受到其他无关token的干扰，因此它既具有良好的性能，又具有很高的计算效率，特别是在密集预测任务中。在图像分类、目标检测和语义分割等多个计算机视觉任务上的实验结果验证了本文设计的有效性。代码可在https://github.com/rayleizhu/BiFormer获取。
1 引言
Transformer具有许多适合构建强大的数据驱动模型的特性。首先，它能够捕捉数据中的长距离依赖关系[29_SwinTransformer, 42_Transformer]。其次，它几乎无归纳偏差，因此使模型更加灵活，能够拟合大量数据[15_Vit]。最后但同样重要的是，它具有高度并行性，这有利于大型模型的训练和推理[42_Transformer, 13_BERT, 33_GPT1, 36_DALL·E]。因此，Transformer不仅革新了自然语言处理，还在计算机视觉领域展现了非常promising的进展。
过去两年，计算机视觉领域见证了视觉transformer的爆炸式增长[29_SwinTransformer, 14_CSWinTransformer, 44_PVT, 15_Vit, 1_DETR, 46_CrossFormer]。在这些工作中，一个热门话题是改进核心构建模块，即attention机制。与本质上是一种局部操作的卷积不同，注意力的一个关键特性是全局感受野，这使视觉 Transformer 能够捕获长距离依赖关系[42_Transformer]。然而，这种特性是有代价的：由于注意力机制计算所有空间位置之间的 pairwise token affinity，它具有很高的计算复杂度，并导致巨大的内存占用。
为缓解这一问题，一个有前景的方向是将稀疏注意力[6_Sparse_Transformer]引入视觉 Transformer 中，使得每个查询仅关注一小部分键值对，而不是全部。