摘要

作为视觉Transformer的核心构建模块，注意力机制能够有效捕捉长距离依赖关系。然而，这种强大能力也伴随着高昂的代价：由于需要计算所有空间位置之间成对的标记（token）交互，导致其带来巨大的计算开销和沉重的内存占用。为缓解这一问题，已有大量研究通过引入人工设计且与内容无关的稀疏性来优化注意力机制，例如将注意力操作限制在局部窗口、轴向条带或扩张窗口内。与这些方法不同，本文提出一种新颖的动态稀疏注意力机制，通过双层路由（bi-level routing）实现更具灵活性且具备内容感知能力的计算资源分配。具体而言，对于每个查询（query），我们首先在粗粒度区域层面过滤掉无关的键值对，随后在剩余候选区域的并集（即路由区域）内执行细粒度的标记到标记注意力计算。我们设计了一种简单而高效的双层路由注意力实现方案，充分利用稀疏性在显著降低计算量和内存消耗的同时，仅依赖GPU友好的密集矩阵乘法操作，从而保持了良好的硬件可扩展性。基于所提出的双层路由注意力机制，我们进一步构建了一种新型通用视觉Transformer——BiFormer。由于BiFormer能够以查询自适应的方式仅关注一小部分相关标记，避免了无关信息的干扰，因此在保持优异性能的同时，展现出极高的计算效率，尤其在密集预测任务（如语义分割、目标检测）中表现突出。在图像分类、目标检测和语义分割等多个计算机视觉任务上的实验结果充分验证了该设计的有效性。相关代码已开源，地址为：https://github.com/rayleizhu/BiFormer。