BiFormer:基于双层路由注意力的视觉Transformer

文章目录

  • 摘要
  • 1、简介
  • 2、相关工作
  • 3、我们的方法:BiFormer
    • 3.1、预备知识:注意力
    • 3.2、双层路由注意(BRA)
    • 3.3、BRA的复杂性分析
  • 4、实验
    • 4.1、ImageNet-1K图像分类
    • 4.2. 目标检测与实例分割
    • 4.3. 基于ADE20K的语义分割
    • 4.4、消融研究
    • 4.5、注意图可视化
  • 5、局限性和未来工作
  • 6、结论

摘要

论文链接:https://arxiv.org/abs/2303.08810
代码链接:https://github.com/rayleizhu/BiFormer

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

1、简介

Transformer具有许多适合构建强大的数据驱动模型的属性。首先,它能够捕捉数据中的长程依赖关系[29,42]。其次,它几乎没有归纳偏差,从而使模型更灵活地拟合成吨的[15]数据。最后但并非最不重要的是,它具有高度的并行性,有利于大型模型的训练和推理[13,33,36,42]。因此,transformer不仅彻底改变了自然语言处理,而且在计算机视觉方面取得了非常有希望的进展。

在过去的两年中,计算机视觉社区见证了视觉transformer的爆炸式增长[1,14,15,29,44,46]。在这些工作中,一个流行的话题是改进核心构建块,即注意力。卷积本质上是一种局部算子,与之相比,注意力的一个关键特性是全局感受野,它使视觉transformer能够捕捉长程依赖性[42]。然而,这样的属性是有代价的:由于注意力计算所有空间位置成对的token亲和性,它具有很高的计算复杂度,并会产生大量的内存占用。

为缓解该问题,一个有希望的方向是将稀疏注意力[6]引入视觉transformer,使每个查询处理一小部分而不是全部键值对。以这种方式,人们探索了几种手工制作的稀疏模式,例如将注意力限制在本地窗口[29]、膨胀窗口[41,46]或轴向条纹[46]中。另一方面,也有工作试图使稀疏性自适应数据[5,48]。然而,虽然它们使用不同的策略来合并或选择键值对标记,但这些标记是查询不可知的,即它们被所有查询共享。尽管如此,根据预训练的ViT[15]和DETR[1]的可视化,不同语义区域的查询实际上涉及非常不同的键值对。因此,强制所有查询处理同一组令牌可能不是最佳选择。

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

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

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

综上所述,我们的贡献如下。我们为普通注意力引入了一种新的双层路由机制,它以一种查询自适应的方式支持内容感知的稀疏模式。利用双层路由注意作为基本构建块,我们提出了一种通用的视觉转换器,称为BiFormer。在图像分类、目标检测和语义分割等各种计算机视觉任务上的实验结果表明,在相似模型尺寸下,所提出的BiFormer在基线上获得了明显更好的性能。

2、相关工作

视觉transformers。transformers是一个神经网络家族,它采用通道级MLP块进行位置嵌入(通道混合),注意[42]块进行跨位置关系建模(空间混合)。transformers最初被提出用于自然语言处理[13,42],然后通过诸如DETR[1]和ViT[15]等开创性工作引入计算机视觉。与cnn相比,最大的不同在于transformer使用注意力作为卷积的替代方案来实现全局上下文建模。然而,由于普通注意在所有空间位置上计算成对的特征亲和力,它会导致很高的计算负担和沉重的内存占用,特别是对于高分辨率的输入。因此,寻求更有效的注意机制是一个重要的研究方向。

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

3、我们的方法:BiFormer

本节详细介绍所提议的方法。我们从3.1节中简要总结注意力机制开始。然后,我们在第3.2节介绍了我们的新型双层路由注意(BRA)机制,它支持动态和查询自适应稀疏性。在3.3节中,我们进一步证明了BRA在适当的区域分区大小下可以达到 O ( ( H W ) 4 3 ) O\left((H W)^{\frac{4}{3}}\right) O((HW)34)复杂度。最后,使用BRA作为核心构建块,我们在3.4节中提出了一个新的分层视觉转换器,名为BiFormer。

3.1、预备知识:注意力

将查询 Q ∈ R N q × C \mathbf{Q} \in \mathbb{R}^{N_{q} \times C} QRNq×C、键 K ∈ R N k v × C \mathbf{K} \in \mathbb{R}^{N_{k v} \times C} KRNkv×C和值 V ∈ R N k v × C \mathbf{V} \in \mathbb{R}^{N_{k v} \times C} VRNkv×C作为输入,注意力函数将每个查询转换为值的加权和,其中权重以查询和相应键之间的标准化点积计算。它可以用紧凑矩阵形式正式定义,如:
Attention ⁡ ( Q , K , V ) = softmax ⁡ ( Q K T C ) V (1) \operatorname{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V})=\operatorname{softmax}\left(\frac{\mathbf{Q K}^{T}}{\sqrt{C}}\right) \mathbf{V} \tag{1} Attention(Q,K,V)=softmax(C QKT)V(1)
这里引入标量因子 C \sqrt{C} C 来避免权重集中和梯度消失[42]。

在transformers中,实际使用的构件是多头自注意(MHSA)。通过“自注意力”,它意味着查询 Q \mathbf{Q} Q、键 K \mathbf{K} K和值 V \mathbf{V} V被导出为相同输入 X ∈ R N × C \mathbf{X} \in \mathbb{R}^{N \times C} XRN×C的线性投影。(对于视觉transformers,X为空间平面化的特征图,即N = H × W,其中H和W分别为特征图的高度和宽度。)至于“多头”,它意味着沿着信道维度将输出分成h个块(即头部),每个块使用一组独立的投影权重。在形式上,
MHSA ⁡ ( X ) = Concat ⁡ ( head ⁡ 0 , h e a d 1 , … , h e a d h ) W o , \operatorname{MHSA}(\mathbf{X})=\operatorname{Concat}\left(\operatorname{head}_{0}\right. , head _{1}, \ldots , head \left._{h}\right) \mathbf{W}^{\mathbf{o}} , MHSA(X)=Concat(head0,head1,,headh)Wo,
其中 h e a d i ∈ R N × C h head_{i} \in \mathbb{R}^{N \times \frac{C}{h}} headiRN×hC是第i个注意头的输出。 W i q , W i k , W i v ∈ R C × C h \mathbf{W}_{i}^{q}, \mathbf{W}_{i}^{k}, \mathbf{W}_{i}^{v} \in \mathbb{R}^{C \times \frac{C}{h}} Wiq,Wik,WivRC×hC为对应的输入投影权值。一个额外的线性变换与权矩阵 W o ∈ R C × C \mathbf{W}^{o} \in \mathbb{R}^{C \times C} WoRC×C用于组合所有heads。

MHSA的复杂度为 O ( N 2 ) O\left(N^{2}\right) O(N2),因为有N个查询,每个查询将处理N个键值对。如此高的复杂性导致了严重的可伸缩性问题,即输入的空间分辨率。

3.2、双层路由注意(BRA)

为了缓解MHSA的可伸缩性问题,一些著作[14,29,41,46,48]提出了不同的稀疏注意机制,其中每个查询只关注少量的键-值对,而不是所有键-值对。但是,这些现有的工作要么使用手工制作的静态模式,要么在所有查询中共享抽样的键-值对子集,如图1所示。在这项工作中,我们探索了一种动态的、查询感知的稀疏注意机制。我们的关键思想是在粗区域级别过滤掉大多数不相关的键-值对,以便只保留一小部分路由区域。然后,我们在这些路由区域的联合中应用细粒度的令牌对令牌注意。为了简化符号,我们讨论了单输入单头自注意的情况,尽管我们在实践中使用批处理输入的多头自注意[42]。算法1中采用仿pytorch的[32]伪代码对整个算法进行了总结。具体解释如下。
BiFormer:基于双层路由注意力的视觉Transformer_第1张图片

区域划分和输入投影。给定二维输入特征图 X ∈ R H × W × C \mathbf{X} \in \mathbb{R}^{H \times W \times C} XRH×W×C,我们首先将其划分为S×S个不重叠的区域,这样每个区域都包含 H W S 2 \frac{H W}{S^{2}} S2HW特征向量。这一步通过将X重塑为 X r ∈ R S 2 × H W S 2 × C \mathbf{X}^{r} \in \mathbb{R}^{S^{2} \times \frac{H W}{S^{2}} \times C} XrRS2×S2HW×C来完成。然后我们用线性投影导出查询,键,值的张量, Q , K , V ∈ R S 2 × H W S 2 × C \mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{S^{2} \times \frac{H W}{S^{2}} \times C} Q,K,VRS2×S2HW×C

Q = X r W q , K = X r W k , V = X r W v , (3) \mathbf{Q}=\mathbf{X}^{r} \mathbf{W}^{q}, \mathbf{K}=\mathbf{X}^{r} \mathbf{W}^{k}, \quad \mathbf{V}=\mathbf{X}^{r} \mathbf{W}^{v}, \tag{3} Q=XrWq,K=XrWk,V=XrWv,(3)

W q , W k , W v ∈ R C × C \mathbf{W}^{q}, \mathbf{W}^{k}, \mathbf{W}^{v} \in \mathbb{R}^{C \times C} Wq,Wk,WvRC×C分别是查询、键、值的投影权重。
BiFormer:基于双层路由注意力的视觉Transformer_第2张图片

基于有向图的区域到区域路由。然后,通过构建有向图来找到关注关系(即,每个给定区域应该被关注的区域)。具体来说,我们首先推导出区域级查询和键 Q r , K r ∈ R S 2 × C \mathbf{Q}^{r}, \mathbf{K}^{r} \in \mathbb{R}^{S^{2} \times C} Qr,KrRS2×C,通过对Q和K分别应用每个区域的平均值。然后,我们通过Qr和转置的Kr之间的矩阵乘法推导出区域到区域的亲和图的邻接矩阵 A r ∈ R S 2 × S 2 \mathbf{A}^{r} \in \mathbb{R}^{S^{2} \times S^{2}} ArRS2×S2
A r = Q r ( K r ) T . (4) \mathbf{A}^{r}=\mathbf{Q}^{r}\left(\mathbf{K}^{r}\right)^{T} . \tag{4} Ar=Qr(Kr)T.(4)

邻接矩阵 A r \mathbf{A}_{r} Ar中的条目衡量两个区域在语义上有多少关联。接下来我们执行的核心步骤是通过仅为每个区域保留top-k个连接来修剪亲和性图。具体来说,我们用逐行topk运算符推导出一个路由索引矩阵 I r ∈ N S 2 × k \mathbf{I}_{r} \in \mathbb{N}^{S^{2} \times k} IrNS2×k
I r = topkIndex ⁡ ( A r ) . (5) \mathbf{I}^{r}=\operatorname{topkIndex}\left(\mathbf{A}^{r}\right) . \tag{5} Ir=topkIndex(Ar).(5)

因此, I r \mathbf{I}^{r} Ir的第i行包含了与第i个区域最相关的k个区域的索引。

Token-to-token注意力。通过区域到区域的路由索引矩阵 I r \mathbf{I}^{r} Ir,我们可以应用细粒度的令牌对令牌注意力。对于region i中的每个查询令牌,它将关注k个路由region的并集中的所有键值对,这些路由region都被 I ( i , 1 ) r , I ( i , 2 ) r , … , I ( i , k ) r \mathbf{I}_{(i, 1)}^{r}, \mathbf{I}_{(i, 2)}^{r}, \ldots, \mathbf{I}_{(i, k)}^{r} I(i,1)r,I(i,2)r,,I(i,k)r。然而,高效地实现这一步并非易事,因为这些路由区域预计分散在整个特征映射上,而现代gpu依赖于合并内存操作,一次性加载数十个连续字节的块。因此,我们首先收集键值张量,即:
K g = gather ⁡ ( K , I r ) , V g = gather ⁡ ( V , I r ) , (6) \mathbf{K}^{g}=\operatorname{gather}\left(\mathbf{K}, \mathbf{I}^{r}\right), \quad \mathbf{V}^{g}=\operatorname{gather}\left(\mathbf{V}, \mathbf{I}^{r}\right), \tag{6} Kg=gather(K,Ir),Vg=gather(V,Ir),(6)

K g , V g ∈ R S 2 × k H W S 2 × C \mathbf{K}^{g}, \mathbf{V}^{g} \in \mathbb{R}^{S^{2} \times \frac{k H W}{S^{2}} \times C} Kg,VgRS2×S2kHW×C是键值张量集合。然后,我们可以对收集到的keyvalue对应用注意力机制:
O = Attention ⁡ ( Q , K g , V g ) + LCE ⁡ ( V ) (7) \mathbf{O}=\operatorname{Attention}\left(\mathbf{Q}, \mathbf{K}^{g}, \mathbf{V}^{g}\right)+\operatorname{LCE}(\mathbf{V}) \tag{7} O=Attention(Q,Kg,Vg)+LCE(V)(7)

在这里,我们引入了一个局部上下文增强术语 LCE ⁡ ( ( V ) \operatorname{LCE}((\mathbf{V}) LCE((V),如[37]。函数 LCE ⁡ ( ⋅ ) \operatorname{LCE}(\cdot) LCE()通过深度卷积进行参数化,我们将内核大小设置为5。

3.3、BRA的复杂性分析

所提出的双层路由注意力实现了类似于普通注意力的直接长距离依赖建模。然而,本文表明,与复杂度为 O ( ( H W ) 4 3 ) O\left((H W)^{\frac{4}{3}}\right) O((HW)34)的普通注意力和复杂度为 O ( ( H W ) 2 ) O\left((H W)^{2}\right) O((HW)2)的准全局轴向注意力[14,22]相比,在适当的区域划分因子S下,BRA的复杂度为 O ( ( H W ) 3 2 ) O\left((H W)^{\frac{3}{2}}\right) O((HW)23)

BRA的计算由3部分组成:线性投影、区域到区域路由和token到token注意力。因此总的计算量如下:
F L O P s = F L O P s p r o j + F L O P s routing  + F L O P s attn  = 3 H W C 2 + 2 ( S 2 ) 2 C + 2 H W k H W S 2 C = 3 H W C 2 + C ( 2 S 4 + k ( H W ) 2 S 2 + k ( H W ) 2 S 2 ) ≥ 3 H W C 2 + 3 C ( 2 S 4 ⋅ k ( H W ) 2 S 2 ⋅ k ( H W ) 2 S 2 ) 1 3 = 3 H W C 2 + 3 C k 2 3 ( 2 H W ) 4 3 , (8) \begin{aligned} \mathrm{FLOPs} & =\mathrm{FLOPs}_{p r o j}+\mathrm{FLOPs}_{\text {routing }}+\mathrm{FLOPs}_{\text {attn }} \\ & =3 H W C^{2}+2\left(S^{2}\right)^{2} C+2 H W k \frac{H W}{S^{2}} C \\ & =3 H W C^{2}+C\left(2 S^{4}+\frac{k(H W)^{2}}{S^{2}}+\frac{k(H W)^{2}}{S^{2}}\right) \\ & \geq 3 H W C^{2}+3 C\left(2 S^{4} \cdot \frac{k(H W)^{2}}{S^{2}} \cdot \frac{k(H W)^{2}}{S^{2}}\right)^{\frac{1}{3}} \\ & =3 H W C^{2}+3 C k^{\frac{2}{3}}(2 H W)^{\frac{4}{3}}, \end{aligned} \tag{8} FLOPs=FLOPsproj+FLOPsrouting +FLOPsattn =3HWC2+2(S2)2C+2HWkS2HWC=3HWC2+C(2S4+S2k(HW)2+S2k(HW)2)3HWC2+3C(2S4S2k(HW)2S2k(HW)2)31=3HWC2+3Ck32(2HW)34,(8)

其中C是token嵌入维度(即特征图的通道数量),k是要参加的区域数量(“top- k”中的“k”)。这里应用了算术平均数和几何平均数的不等式。等式8中的等式当且仅当 2 S 4 = k ( H W ) 2 S 2 2 S^{4}=\frac{k(H W)^{2}}{S^{2}} 2S4=S2k(HW)2成立。因此:
S = ( k 2 ( H W ) 2 ) 1 6 (9) S=\left(\frac{k}{2}(H W)^{2}\right)^{\frac{1}{6}} \tag{9} S=(2k(HW)2)61(9)

换句话说,如果我们根据等式9的输入分辨率将区域划分因子S按比例缩放,BRA的复杂度为 O ( ( H W ) 4 3 ) O\left((H W)^{\frac{4}{3}}\right) O((HW)34)

BiFormer:基于双层路由注意力的视觉Transformer_第3张图片

4、实验

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

4.1、ImageNet-1K图像分类

设置。我们在ImageNet-1K [12] 数据集上进行图像分类实验,遵循DeiT[40]的实验设置进行公平比较。每个模型训练300 epoch,输入大小为224 × 224。我们以权重衰减为0.05的AdamW作为优化器,采用初始学习率为0.001的余弦衰减学习率调度,前5个epoch用于线性预热[16]。批处理大小设置为1024。为了避免过拟合,我们应用正则化技术,包括RandAugment [9] (rand-m9-mstd0.5- inc1)、MixUp [54] (prob = 0:8)、CutMix [52] (prob = 1:0)、随机擦除(prob = 0:25)和增加随机深度[21](分别为BiFormer-T/S/B的prob =0:1 =0:15=0:4)。为了公平地比较使用令牌标记[23]训练的模型,包括Uniformer[25]和WaveViT[51],我们还提供了一个使用WaveViT提供的相同配方训练的版本。

BiFormer:基于双层路由注意力的视觉Transformer_第5张图片

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

4.2. 目标检测与实例分割

设置。我们在COCO 2017[28]上评估了目标检测和实例分割模型。为了公平的比较,所有的实验都是用MMDetection[3]工具箱进行的。分别使用RetinaNet[27]和Mask R-CNN[18]框架进行对象检测和实例分割。在COCO上训练之前,我们用ImageNet- 1K上预训练的权重初始化主干,同时让所有其他层随机初始化。模型使用MMDetection提供的标准1×时间表(12 epoch)进行训练,除了我们使用AdamW优化器[30],而不是SGD。我们使用初始学习率为1e−4,批处理大小为16,而RetinaNet和Mask R-CNN的权重衰减分别设置为1e−4和5e−2。在训练过程中,我们调整输入图像的大小,将较短的一侧固定为800像素,同时保持较长的一侧不超过1333像素。
BiFormer:基于双层路由注意力的视觉Transformer_第6张图片

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

4.3. 基于ADE20K的语义分割

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

结果。表4显示了两种不同框架的结果。结果表明,在语义FPN框架下,我们的BiFormer-S/B分别达到48.9/49.9 mIoU, CSWin-T/S提高了0.7 mIoU。在UperNet框架上也观察到类似的性能提升。
BiFormer:基于双层路由注意力的视觉Transformer_第7张图片

4.4、消融研究

BRA的有效性。我们将BRA与现有的几种稀疏注意机制进行了比较。在[14]之后,我们将宏观架构设计与Swin-T[29]进行比较。具体来说,我们使用2、2、6、2个块进行四个阶段,不重叠的补丁嵌入,设置初始补丁嵌入维数C = 96, MLP扩展比e = 4。结果如表5所示。我们的双层路由注意在图像分类和语义分割方面都明显优于现有的稀疏注意机制。
BiFormer:基于双层路由注意力的视觉Transformer_第8张图片

其他架构设计选择。使用Swin-T布局作为基线,我们总结了我们应用的其他修改,这进一步提高了我们的BiFormer-S模型在ImageNet-1K数据集上的最先进性能。这些修改包括:(1)将非重叠的贴片嵌入[29]替换为重叠的贴片嵌入[29][14,37,45],(2)采用更深层次的布局(即在每个阶段堆叠更多的块,同时将基通道从96减少到64,MLP扩展比从4减少到3,以保持相似的FLOPs.),(3)在BiFormer块的开始添加卷积位置编码[7,25],以及(4)应用令牌标记[23,25,51]训练技术。如表6所示,只需使用较深的布局就可以显著提高性能。然而,这一因素在现有的著作中通常没有被讨论。

BiFormer:基于双层路由注意力的视觉Transformer_第9张图片

4.5、注意图可视化

为了进一步理解双层路由注意是如何工作的,我们通过查询位置可视化路由区域和注意响应。对于这个可视化,我们使用从第三阶段的最后BiFormer块中提取的路由指数和注意力分数,这是消耗最多计算的主要阶段。我们在图4中演示了两个场景。在这两种情况下,我们可以清楚地观察到语义相关的区域被成功定位。例如,在第一个场景(街景)中,如果查询位置在建筑物或树上,则对应的路由区域覆盖相同或相似的实体。在第二个室内场景中,当我们将查询位置放在鼠标上时,路由区域包含主机、键盘和显示器的一部分,尽管这些区域彼此并不相邻。这意味着我们的双层路由注意可以捕获长期的对象间关系。
BiFormer:基于双层路由注意力的视觉Transformer_第10张图片

5、局限性和未来工作

与使用简单静态模式的稀疏注意相比,我们引入了一个额外的步骤来定位要关注的区域,在那里我们构建和修剪区域级图,并从路由区域中收集键值对。虽然这一步在粗略的区域级别上操作,不会引起太多的计算,但它不可避免地会引起额外的GPU内核启动和内存事务。因此,由于内核启动和内存瓶颈的开销,BiFormer在GPU上的吞吐量比一些现有的具有类似FLOPs的模型要低。尽管如此,这个问题可以通过工程上的努力来缓解,比如GPU内核融合。我们将在未来的工作中探索具有硬件感知的高效稀疏注意和视觉转换器。

6、结论

我们提出了两级路由注意,以一种动态的、查询感知的方式实现计算的有效分配。BRA的核心思想是在粗区域级别过滤掉最不相关的键值对。首先构建和修剪区域级有向图,然后在路由区域的并集中应用细粒度的令牌对令牌注意。我们分析了BRA的计算复杂度,并证明在适当的区域分区大小下,它可以达到O((HW)4 3)。以BRA为核心构建块,我们提出了一种新的视觉转换器BiFormer,它在图像分类、目标检测、实例分割和语义分割这四种流行的视觉任务上表现出了优越的性能。

你可能感兴趣的:(transformer,深度学习,计算机视觉)