BiFormer: Vision Transformer with Bi-Level Routing Attention

Abstract

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

introduction

Transformer 具有许多适合构建强大的数据驱动模型的属性。首先，它能够捕获数据中的远程依赖性 [29,42]。第二，它几乎没有归纳偏差，因此使模型更灵活地适应大量数据 [15]。最后但同样重要的是，它具有高并行性，这有利于大型模型的训练和推理 [13、33、36、42]。因此，transformer 不仅彻底改变了自然语言处理，而且在计算机视觉方面也显示出非常有前途的进展。

在过去的两年里，计算机视觉社区见证了vision transformer的爆炸式增长 [1、14、15、29、44、46]。在这些工作中，一个热门话题是提高核心构建块，即注意力。与本质上是局部运算符的卷积相反，注意力的一个关键属性是全局感受野，它使视觉转换器能够捕获远程依赖性 [42]。然而，这样的属性是有代价的：由于注意力计算跨所有空间位置的成对token affinity，因此它具有很高的计算复杂度并会产生大量内存占用。

为了缓解这个问题，一个有前途的方向是将稀疏注意力[6]引入视觉转换器，这样每个查询都会关注一小部分键值对而不是全部。以这种方式，已经探索了几种手工制作的稀疏模式，例如在局部窗口 [29]、扩张窗口 [41、46] 或轴向条纹 [46] 中限制注意力。另一方面，也有一些工作试图使稀疏性适应数据 [5, 48]。然而，虽然它们使用不同的策略来合并或选择键/值标记，但这些标记是查询不可知的，即它们由所有查询共享。尽管如此，根据预训练的 ViT 1 [15] 和 DETR 2 [1] 的可视化，不同语义区域中的查询实际上涉及完全不同的键值对。因此，强制所有查询都处理同一组标记可能不是最优的。

在本文中，我们寻求一种具有动态、查询感知稀疏性的注意力机制。基本上，我们的目标是让每个查询都关注语义最相关的键值对中的一小部分。第一个遇到的问题是如何分配需要关注的键值对。例如，如果我们像[17]中那样以每个查询的方式选择键值对，它仍然需要评估所有查询和键之间的成对亲和力，因此具有与普通注意力相同的复杂性。另一种可能性是根据每个查询的局部上下文预测注意力偏移 [10, 48]，因此避免了成对的亲和力计算。然而，以这种方式对长程依赖性进行建模是有问题的[48]。

为了高效地定位有价值的键值对以在全局范围内参与，我们提出了一种区域到区域的路由方法。我们的核心思想是在粗粒度的区域级别过滤掉最不相关的键值对，而不是直接在细粒度的令牌级别。这是通过首先构建区域级亲和图然后对其进行修剪以仅保留每个节点的前 k 个连接来完成的。因此，每个区域只需要关注前 k 个路由区域。确定了参与区域后，下一步是应用token到token的注意力，这很重要，因为现在假设键值对在空间上是分散的。对于这种情况，虽然稀疏矩阵乘法适用，但在现代 GPU 中效率低下。它依赖于合并的内存操作，即一次访问数十个连续字节的块 [31]。相反，我们通过收集键/值令牌提出了一个简单的解决方案，其中仅涉及硬件友好的密集矩阵乘法。我们将这种方法称为双层路由注意 (BRA)，因为它包含区域级路由步骤和标记级注意步骤。

通过使用 BRA 作为核心构建块，我们提出了 BiFormer，这是一种通用的vision transformer主干，可用于许多应用程序，例如分类、对象检测和语义分割。由于 BRA 使 BiFormer 能够以内容感知的方式处理每个查询最相关的键/值标记的一小部分，我们的模型实现了更好的计算性能权衡。例如，通过 4.6G FLOPs 计算，BiFormer-T 在 ImageNet-1K 分类上达到了 83.8% 的 top-1 准确率，这是我们所知在类似计算预算下最好的，无需使用外部数据进行训练或蒸馏 [23,40] . 这些改进也一致地显示在下游任务中，例如实例分割和语义分割。

总而言之，我们的贡献如下。我们为 vanilla attention 引入了一种新颖的双层路由机制，它以查询自适应的方式启用内容感知稀疏模式。使用双层路由注意作为基本构建块，我们提出了一个名为 BiFormer 的通用视觉转换器。包括图像分类、目标检测和语义分割在内的各种计算机视觉任务的实验结果表明，所提出的 BiFormer 在类似模型大小下比基线实现了显着更好的性能。

vision transformers。 Transformers 是一个神经网络家族，采用通道方式 MLP 块进行位置嵌入（通道混合）和注意 [42] 块进行交叉位置关系建模（空间混合）。Transformers 最初被提议用于自然语言处理 [13, 42]，然后通过 DETR [1] 和 ViT [15] 等开创性工作引入计算机视觉。与 CNN 相比，最大的区别在于 Transformer 使用注意力作为卷积的替代方法来启用全局上下文建模。然而，由于 vanilla attention 计算所有空间位置的成对特征亲和力，因此会产生高计算负担和大量内存占用，尤其是对于高分辨率输入。因此，一个重要的研究方向是寻求更有效的注意机制.

Efficient attention mechanisms。已经提出了大量的工作来通过利用稀疏连接模式 [6]、低秩近似 [43] 或循环操作 [11] 来减少普通注意力的计算和内存复杂性瓶颈。可以在 [39] 中找到对这些注意力变体的全面调查。在vision transformer的范围内，稀疏注意力最近由于 Swin 转换器 [29] 的巨大成功而受到欢迎。在Swin transformer中，注意力仅限于不重叠的局部窗口，引入shift window操作来实现相邻窗口之间的窗口间通信。为了在合理的计算预算下实现更大甚至准全局的感受野，一些后续工作引入了不同的手工稀疏模式，例如扩张窗口 [41,46] 或十字形窗口 [14]。也有尝试使稀疏模式自适应数据的工作，如 DAT [48]、TCFormer [53] 和 DPT，虽然这些工作通过不同的合并或选择策略减少了键/值标记的数量，但这些键/值标记由图像上的所有查询共享。相反，我们探索查询感知键/值标记选择。激发我们工作的关键观察是，根据预训练的 ViT [15] 和 DETR [1] 的可视化，不同查询的注意力区域可能会有很大差异。当我们以粗到细的方式实现查询自适应稀疏性的目标时，它与四叉树注意力机制有一些相似之处 [38]。与四叉树注意力不同，我们的双层路由注意力的目标是定位几个最相关的键值对，而四叉树注意力构建了一个令牌金字塔，并从不同粒度的各个级别组装消息。此外，四叉树需要深度递归来覆盖整个特征图，这会损害并行性，而我们的双层路由注意力可以通过键/值标记收集更有效地实现，然后进行密集矩阵乘法。因此，四叉树转换器比我们的 BiFormer 慢得多。

3. Our approach: Biformer

本节详细阐述了所提出的方法。我们首先简要总结第 3.1 节中的注意力机制。然后，我们在第 3.2 节中介绍了我们新颖的双层路由注意 (BRA) 机制，该机制支持动态和查询自适应稀疏性。在 3.3 节中，我们进一步表明 BRA 可以通过适当的区域分区大小实现
$$O((HW)^{4/3})$$
复杂度。最后，使用 BRA 作为核心构建块，我们在第 3.4 节中提出了一个新的分层视觉转换器，名为 BiFormer

3.1 预备知识：Attention

取

和
作为输入，注意力函数将每个查询转换为值的加权和，其中权重计算为查询和相应键之间的归一化点积。它可以形式化地定义为一个紧凑的矩阵形式，如:

这里，引入缩放因子以避免集中权重和梯度消失 [42]。

在 Transformer 中，实际使用的构建块是多头自注意力 (MHSA)。 “自注意力”意味着查询 Q、键 K 和值 V 是作为相同输入
的线性投影得出的。（对于vision transformer，X 是空间扁平化的特征图，即
，其中 H 和 W 分别是特征图的高度和宽度。）至于“多头”，它意味着沿着通道维度将输出分成 h 个块（即头），每个块使用一个独立的组投影权重。正式地，

其中是第i个attention head的输出。是对应的输入映射权重。一个用权重矩阵的额外的线性变换被使用来融合所有的头。

MHSA 的复杂度为 O(N2)，因为有 N 个查询，每个查询将处理 N 个键值对。
如此高的复杂性会导致严重的可扩展性问题 w.r.t. 输入的空间分辨率。

3.2 Bi-Level Routing Attention(BRA)

为了缓解 MHSA 的可扩展性问题，一些工作 [14、29、41、46、48] 提出了不同的稀疏注意力机制，其中每个查询只关注少量的键值对，而不是全部。然而，这些现有工作要么使用手工制作的静态模式，要么在所有查询中共享键值对的采样子集，如图 1 所示。在这项工作中，我们探索了一种动态的、查询感知的稀疏注意力机制。我们的关键思想是在粗区域级别过滤掉大部分不相关的键值对，以便只保留一小部分路由区域。然后，我们在这些路由区域的联合中应用细粒度的 token-to-token 注意力。为了简化符号，我们讨论了单头自注意力和单个输入的情况，尽管我们在实践中使用多头自注意力 [42] 和批量输入。整个算法在算法 1 中用 Pytorch-like [32] 伪代码进行了总结。我们给出详细解释如下。

Region partition and input projection
给定一个二维输入特征图，我们开始把它划分成
$$S\times S$$
个非重叠区域来使得每个区域包含
$$HW/S²$$
个特征向量。这一步通过把X调整成来完成。然后我们用线性映射推导query、key和value向量

其中是对于query、key和value对应的线性映射。

Region-to-region routing with directed graph。然后我们通过构建有向图找到参与关系（即，每个给定区域应该参与的区域）。具体来说，我们首先通过分别在 Q 和 K 上应用每个区域的平均值来导出区域级查询和密钥。然后我们推导出邻接矩阵通过 Qr 和转置 Kr 之间的矩阵乘法的区域到区域亲和图.

邻接矩阵 Ar 中的条目测量两个区域在语义上的相关程度。我们接下来执行的核心步骤是通过只保留每个区域的前 k 个连接来修剪亲和图。具体来说，我们使用逐行 topk 运算符推导路由索引矩阵

因此，Ir 的第 i 行包含第 i 个区域最相关区域的 k 个索引.

Token-to-token attention。使用区域到区域路由索引矩阵 Ir，我们可以应用细粒度的 token-to-token 注意力。对于区域 i 中的每个查询令牌，它将关注驻留在以为索引的 k 个路由区域的联合中的所有键值对。然而，有效地实施这一步骤并非易事，因为这些路由区域预计会分散在整个特征图上，而现代 GPU 依赖于一次加载数十个连续字节块的合并内存操作。因此，我们首先收集键张量和值张量，即:

其中是收集的键和值张量。然后我们可以将注意力集中在收集到的键值对上：

在这里，我们引入了一个局部上下文增强项 LCE((V)，如 [37] 中所述。函数 LCE(·) 使用深度卷积进行参数化，我们将内核大小设置为 5。

3.3. BRA 的复杂性分析

所提出的双层路由注意可以实现类似于vanilla注意的直接远程依赖建模。
然而，我们在这里表明，与具有 O((HW)2) 复杂度的普通注意力相比，BRA 具有低得多的复杂度 O((HW)4/3) 和适当的区域划分因子 S，并且quasi-global axial attention [14, 22]，其复杂度为 O((HW)3/2)

BRA 的计算由三部分组成：线性投影、区域到区域路由和令牌到令牌的注意力。
因此总的计算量为：

其中 C 是标记嵌入维度（即特征图的通道数），k 是要参与的区域数（“top-k”中的“k”）。在这里，应用了算术和几何平均值的不等式。等式中的平等。当且仅当成立。所以：

换句话说，如果我们缩放区域划分因子 S w.r.t，BRA 将实现复杂度。根据等式9的输入分辨率。

3.4. BiFormer 的架构设计

使用 BRA 作为基本构建块，我们提出了一种新的通用vision transformer BiFormer。如图 3 所示，我们遵循最近最先进的vision transformer[14,29,41] 使用四级金字塔结构。具体来说，在第 i 阶段，我们在第一阶段使用重叠块嵌入，在第二到第四阶段使用块合并模块 [25, 37] 来降低输入空间分辨率，同时增加通道数，然后是 Ni 连续 BiFormer 块来转换特征。
在每个 BiFormer 块中，我们遵循最近的工作 [7, 25, 41] 在开始时使用3×3深度卷积来隐式编码相对位置信息。然后，我们依次应用 BRA 模块和扩展率为 e 的 2 层 MLP 模块，分别用于交叉位置关系建模和每个位置嵌入。

我们通过缩放网络宽度（即基本通道 C 的数量）和深度（即每个阶段使用的 BiFormer 块的数量，Ni，i = 1，2，3，4）来实例化具有 3 种不同模型大小的 BiFormer ，如表 1 所列。它们共享其他配置。我们将每个注意力头设置为 32 个通道，MLP 扩展比率 e=3。对于 BRA，我们使用 topk = 1，4，16，S23为4个阶段，由于输入分辨率不同，分类/语义分割/目标检测任务的区域划分因子S = 7/8/16

4. 实验

我们通过实验评估我们提出的 BiFormer 在一系列主流计算机视觉任务上的有效性，包括图像分类（第 4.1 节）、对象检测和实例分割（第 4.2 节）和语义分割（第 4.3 节）。具体来说，我们在 ImageNet1K [12] 上从头开始训练图像分类。然后我们在 COCO [28] 上微调预训练骨干用于对象检测和实例分割，并在 ADE20K [55] 上微调预训练骨干用于语义分割。此外，我们还进行了消融研究，以验证所提出的双层路由注意和 BiFormer 的其他架构设计选择的有效性。 4.4. 最后，为了验证查询自适应、稀疏模式是通过双层路由注意实现的，我们在第4.5节中可视化了注意图。

4.1 Image Classification on ImageNet-1k

设置。我们在 ImageNet-1K [12] 数据集上进行图像分类实验，遵循 DeiT [40] 的实验设置进行公平比较。具体来说，每个模型都训练了 300 个 epoch，输入大小为 224×224。我们将 AdamW 作为权重衰减为 0.05 的优化器，并应用初始学习率为 0.001 的余弦衰减学习率计划，而前 5 个时期用于线性预热 [16]。批量大小设置为 1024。为了避免过度拟合，我们应用正则化技术，包括 RandAugment [9] (rand-m9-mstd0.5inc1)、MixUp [54] (prob = 0.8)、CutMix [52] (prob = 1.0）、随机擦除（概率 = 0.25）和增加随机深度 [21]（对于 BiFormer-T/S/B，概率分别为 0.1/0.15/0.4）。为了公平地比较使用标记标签 [23] 训练的模型，包括 Uniformer [25] 和 WaveViT [51]，我们还提供了一个使用 WaveViT 提供的相同配方训练的版本。

结果。我们将我们的方法与几种密切相关的方法和/或最近的最新技术进行比较。定量结果列于表 2，其中模型按计算量 (FLOPs) 分组。在所有 3 个组中，我们的模型始终优于其他比较模型。例如，对于最小组（2G FLOPs）中的模型，我们的 BiFormer-T 达到了 81.4% 的 top-1 精度，比最有竞争力的QuadTree-b1提高了1.4%。对于第二组（4G FLOPs）中的模型，BiFormer-S 达到了 83.8% 的 top-1 准确率。据我们所知，这是没有额外训练数据或训练技巧的最佳结果。此外，使用名为标记标记 [23] 的蒸馏技术，BiFormer-S 的准确率可以进一步提高到 84.3%，这意味着所提出的架构具有巨大的潜力。对于最大组（10G FLOPs）中的模型，BiFormer-B 的性能甚至优于现有模型，计算量高达 16G FLOPs，如 Swin-B [29]，CSWin-B [ 14] 和 CrossFormer-L [46]

4.2. 对象检测和实例分割设置

我们在 COCO 2017 [28] 上评估对象检测和实例分割模型。为了公平比较，所有实验都是使用 MMDetection [3] 工具箱进行的。 RetinaNet [27] 和 Mask R-CNN [18] 框架分别用于对象检测和实例分割。在对 COCO 进行训练之前，我们使用在 ImageNet1K 上预训练的权重初始化主干，同时让所有其他层随机初始化。这些模型使用 MMDetection 提供的标准 1 计划（12 个时期）进行训练，除了我们使用 AdamW 优化器 [30]，而不是 SGD。我们使用 1e } 4 的初始学习率和 16 的批量大小，而 RetinaNet 和 Mask R-CNN 的权重衰减分别设置为 1e-4 和 5e-2。在训练期间，我们通过将较短边固定为 800 像素同时保持较长边不超过 1,333 像素来调整输入图像的大小。

结果。我们在表 3 中列出了结果。对于使用 RetinaNet 进行的对象检测，我们报告了不同 IoU 阈值（50%、75%）和三种对象大小（即小型、中型和大号 (S/M/L)）。从结果中我们可以看出，虽然 BiFormer 的整体性能仅与一些最具竞争力的现有方法（如 WaveViT 和 QuadTree-B）相当，但小对象（APM）的性能明显优于这些方法。这可能是因为 BRA 通过稀疏采样而不是下采样来节省计算。因此，它保留了细粒度的细节，这对于小物体至关重要。用Mask R-CNN来做实力分割，我们在不同的 IoU 阈值（50%、75%）下报告边界框和掩码平均精度（APb 和 APm）。如表 3 所示，我们的方法在此任务中的所有指标上都显示出明显的优势

4.3. ADE20K 设置上的语义分割

在现有工作之后，我们基于 MMSegmentation [8] 在 ADE20K [55] 数据集上进行语义分割实验。我们在语义 FPN [24] 和 UperNet [49] 框架下进行比较。在这两种情况下，主干都使用 ImageNet-1K 预训练权重进行初始化，而其他层使用随机初始化。模型使用 AdamW 优化器进行优化，批量大小设置为 32。为了公平比较，我们的语义 FPN 实验使用与 PVT [44] 相同的设置来训练模型 80k步。上层网络实验使用和swin transformer相同设置训练160k次迭代。

结果。表 4 显示了两种不同框架的结果。它表明，使用语义 FPN 框架，我们的 BiFormer-S/B 分别达到了 48.9/49.9 mIoU，将 CSWin-T/S 提高了 0.7 mIoU。还观察到 UperNet 框架的类似性能提升

4.4. 消融研究 BRA 的有效性。

我们将 BRA 与几种现有的稀疏注意力机制进行了比较。在 [14] 之后，我们将宏架构设计与 Swin-T [29] 对齐以进行公平比较。具体来说，我们对四个阶段使用2、2、6、2块，非重叠的patch embedding，设置初始patch embedding维数C = 96和MLP扩展率e = 4。结果如表5所示。我们的bilevel 在图像分类和语义分割方面，路由注意力比现有的稀疏注意力机制具有明显更好的性能

其他网络结构设计选择 使用Swin-T结构作为baseline，我们总结了我们已经应用的其他修改，这些修改进一步将我们的 BiFormer-S 模型提升到 ImageNet-1K 数据集上的最先进性能。这些修改包括：(1) 将非重叠补丁嵌入 [29] 替换为重叠补丁嵌入 [14,37,45]，(2) 使用更深的布局（即在每个阶段堆叠更多块，同时将基本通道从 96 减少到 64 和 MLP 扩展比率从 4 到 3 以保持类似的 FLOPs。），（3）在 BiFormer 块的开头添加卷积位置编码 [7, 25]，以及（4）应用令牌标记 [23, 25, 51] 培训技术。如表 6 所示，简单地使用更深的布局可以显着提高性能。然而，这个因素通常不会在现有工作中讨论。

4.5. 注意图的可视化

为了进一步了解双层路由注意的工作原理，我们可视化路由区域和注意响应 w.r.t. 查询职位。对于此可视化，我们使用从第 3 阶段的最终 BiFormer 块中提取的路由索引和注意力分数，这是消耗最多计算的主要阶段。我们在图 4 中演示了两个场景。在这两种情况下，我们都可以清楚地观察到语义相关区域已成功定位。例如，在第一个场景中，街景，如果查询位置在建筑物或树上，则对应的路由区域覆盖相同或相似的实体。在第二个室内场景中，当我们将查询位置放在鼠标上时，路由区域包含主机、键盘和显示器的一部分，即使这些区域彼此不相邻。这意味着我们的双层路由注意可以捕获远程对象间关系。

5. 局限性和未来工作
与简单静态模式的稀疏注意力相比，我们引入了一个额外的步骤来定位要参与的区域，我们构建和修剪区域级图并从路由区域收集键值对。虽然此步骤不会产生太多计算，因为它在粗区域级别运行，但它不可避免地会导致额外的 GPU 内核启动和内存事务。因此，由于内核启动和内存瓶颈的开销，BiFormer 的吞吐量低于一些在 GPU 上具有类似 FLOPs 的现有模型。尽管如此，这个问题可以通过工程努力得到缓解，例如 GPU 内核融合。我们将在未来的工作中探索具有硬件意识的高效稀疏注意力和vision transformer。

6. 结论

我们提出了双层路由注意，以动态、查询感知的方式实现计算的有效分配。
BRA 的核心思想是在粗区域级别过滤掉最不相关的键值对。它是通过首先构建和修剪区域级有向图，然后在路由区域的联合中应用细粒度的token到token注意力来实现的。我们已经分析了 BRA 的计算复杂性，并证明它可以通过适当的区域分区大小实现。使用 BRA 作为核心构建块，我们提出了 BiFormer，这是一种新的视觉转换器，在四种流行的视觉任务、图像分类、对象检测、实例分割和语义分割方面表现出卓越的性能