BiFormer Vision Transformer with Bi-Level Routing Attention

BiFormer:具有双电平路由注意的视觉变压器

  摘要作为视觉转换器的核心组成部分，注意力是捕捉长期依赖关系的有力工具。然而，这样的能力是有代价的:当计算跨所有空间位置的成对令牌交互时，它会产生巨大的计算负担和沉重的内存占用。一系列的作品试图通过引入手工制作和内容不可知的注意力稀疏性来缓解这个问题，例如将注意力操作限制在局部窗口，轴向条纹或扩展窗口内。与这些方法相比，我们提出了一种新的动态稀疏注意，通过双层路由实现更灵活的内容感知计算分配。具体来说，对于查询，首先在粗区域级别过滤掉不相关的键值对，然后在剩余候选区域(即路由区域)的联合中应用细粒度的令牌到令牌注意。我们提供了一个简单而有效的实现所提出的双层路由注意，它利用稀疏性来节省计算和内存，同时只涉及gpu友好的密集矩阵乘法。在此基础上，提出了一种新的通用视觉变压器——BiFormer。由于BiFormer以查询自适应的方式处理一小部分相关令牌，而不会分散其他不相关的令牌，因此它具有良好的性能和较高的计算效率，特别是在密集的预测任务中。多个计算机视觉任务(如图像分类、目标检测和语义分割)的经验结果验证了我们设计的有效性。代码可从https://github.com/rayleizhu/BiFormer获得。

引言

  Transformer有许多适合于构建强大的数据驱动模型的属性。首先，它能够捕获数据中的远程依赖关系[29,42]。第二,它几乎没有诱导偏差，从而使模型更灵活地拟合大量数据[15]。最后但并非最不重要的是，它具有高并行性，这有利于大型模型的训练和推理[13,33,36,42]。因此，transformer不仅彻底改变了自然语言处理，而且在计算机视觉方面也显示出非常有希望的进展。
  在过去的两年里，计算机视觉社区见证了视觉变压器的爆炸式增长[1,14,15,29,44,46]。在这些工作中，一个热门的话题是提高核心构件，即注意力。卷积本质上是一个局部算子，与之相反，注意力的一个关键属性是全局接受场，它使视觉转换器能够捕获远程依赖[42]。然而，这样的属性是有代价的:当注意力跨所有空间位置成对地计算令牌关联时，它具有很高的计算复杂性，并导致大量内存占用。
  为了缓解这个问题，一个有希望的方向是将稀疏关注[6]引入到视觉转换中，这样每个查询只关注一小部分键值对，而不是全部。以这种方式，研究人员探索了几种手工制作的稀疏图案，例如将注意力限制在局部窗口[29]、扩展窗口[41,46]或轴向条纹[46]。另一方面，也有人试图使稀疏度适应数据[5,48]。然而，虽然它们使用不同的策略来合并或选择键/值令牌，但这些令牌是查询无关的，也就是说，它们由所有查询共享。然而，根据预训练的ViT 1[15]和DETR 2[1]的可视化，不同语义区域的查询实际上关注的键值对是完全不同的。因此，强制所有查询处理同一组令牌可能不是最优的。
  在本文中，我们寻求一种具有动态、查询感知稀疏性的注意机制。基本上，我们的目标是让每个查询都关注语义上最相关的键值对的一小部分。第一个问题是如何解决查找要参加的这些键值对。例如，如果我们像[17]中那样以每个查询的方式选择键值对，它仍然需要评估所有查询和键之间的配对亲和力，因此具有相同的香草注意复杂性。另一种可能性是基于每个查询的本地上下文来预测注意力偏移量[10,48]，因此可以避免两两关联计算。然而，通过这种方式，对远程依赖关系进行建模是有问题的[48]。

  为了高效地定位有价值的键值对，我们提出了一种区域到区域路由方法。我们的核心思想是在粗粒度的区域级别过滤掉最不相关的键值对，而不是直接在细粒度的令牌级别。这是通过首先构造一个区域级关联图，然后对其进行修剪，使每个节点只保留top-k连接来实现的。因此，每个区域只需要关注top-k路由区域。确定了参与区域后，下一步是应用令牌到令牌的注意，这是非常重要的，因为现在假定键值对在空间上是分散的。对于这种情况，虽然稀疏矩阵乘法是适用的，但它在现代gpu中是低效的，它依赖于合并内存操作，即一次访问数十个连续字节的块[31]。相反，我们提出一个简单的通过收集键/值令牌来解决，其中只涉及硬件友好的密集矩阵乘法。我们将这种方法称为双级路由注意(BRA)，因为它包含一个区域级路由步骤和一个令牌级注意步骤。
  通过使用BRA作为核心构建块，我们提出了BiFormer，这是一个通用的视觉变压器骨干，可用于许多应用，如分类，目标检测和语义分割。由于BRA使BiFormer能够以内容感知的方式为每个查询处理最相关的键/值令牌的一小部分，因此我们的模型实现了更好的计算性能权衡。例如，在4.6G FLOPs的计算下，BiFormer-T在ImageNet-1