Vision Transformer with Deformable Attention

论文地址:https://arxiv.org/abs/2201.00520

代码地址:https://github.com/LeapLabTHU/DAT

Abstract

        Transformers 最近在各种视觉任务上表现出了卓越的性能。大的,有时甚至是全局的感受野赋予了 Transformer 模型比 CNN 模型更高的表示能力。然而,简单地扩大感受野也会引起一些问题。一方面,使用密集注意力,例如在 ViT 中,会导致过多的内存和计算成本,并且特征可能会受到超出兴趣区域的不相关部分的影响。另一方面,PVT 或 Swin Transformer 中采用的稀疏注意力与数据无关,可能会限制对远程关系建模的能力。为了缓解这些问题,我们提出了一种新颖的可变形自注意模块,其中以数据相关的方式选择自注意中 key 和 value 位置。这种灵活的方案使自注意模块能够专注于相关区域并捕获更多信息特征。在此基础上,我们提出了Deformable Attention Transformer, 具有可变形注意力的通用骨干模型,用于图像分类和密集预测任务。大量实验表明,我们的模型在综合基准测试中取得了更好的结果。 

1. Introduction 

        最初引入 Transformer [34] 是为了解决自然语言处理任务。它最近在计算机视觉领域显示出巨大的潜力 [12,26,36]。开创性的工作,Vision Transformer [12] (ViT),堆叠多个 Transformer 块来处理非重叠图像块(即视觉标记)序列,从而产生用于图像分类的无卷积模型。与它们的 CNN 对应物 [18,19] 相比,基于 Transformer 的模型具有更大的感受野并且擅长建模长程依赖性,这被证明在大量训练数据和模型参数。然而,视觉识别中多余的注意力是一把双刃剑,有很多缺点。具体来说,每个查询补丁需要过多的键会产生高计算成本和缓慢的收敛,并增加过度拟合的风险。

        为了避免过多的注意力计算,现有作品 [6、11、26、36、43、49] 已经利用精心设计的高效注意力模式来降低计算复杂度。作为其中两种具有代表性的方法,Swin Transformer [26] 采用基于窗口的局部attention 限制局部窗口的注意力,而 Pyramid Vision Transformer (PVT) [36] 对键和值特征映射进行下采样以节省计算。虽然有效,但手工制作的注意力模式与数据无关,可能不是最佳的。很可能相关的键/值被删除,而不太重要的键/值仍然保留。

        理想情况下,人们会期望给定查询的候选键/值集是灵活的并且能够适应每个单独的输入,这样可以缓解手工制作的稀疏注意力模式中的问题。事实上,在 CNN 的文献中,学习卷积滤波器的可变形感受野已被证明可以有效地在数据依赖的基础上选择性地关注更多信息区域 [9]。最著名的工作是可变形卷积网络 [9],它在许多具有挑战性的视觉任务上取得了令人印象深刻的结果。这促使我们探索 Vision Transformers 中的可变形注意力模式。然而,这个想法的天真实现会导致不合理的高内存/计算复杂度:可变形偏移引入的开销是补丁数量的二次方。因此,尽管最近的一些工作 [7, 46, 54] 研究了变形金刚中可变形机制的想法,但没有人将其视为基本的 由于高计算成本,构建强大的骨干网络(如 DCN)的基石。相反,它们的可变形机制要么在检测头中采用[54],要么用作预处理层为后续骨干网络[7]采样补丁。 

        在本文中,我们提出了一个简单高效的可变形自注意力模块,配备了一个强大的金字塔主干,称为可变形注意力转换器 (DAT),用于图像分类和各种密集预测任务。与为整个特征图中的不同像素学习不同偏移量的 DCN 不同,我们建议学习几组查询不可知偏移量以将键和值移动到重要区域(如图 1(d)所示),基于观察在 [3,52] 中,全局注意力通常会导致不同查询的注意力模式几乎相同。这种设计既保持了线性空间复杂度,又为 Transformer 骨干引入了可变形的注意力模式。具体来说,对于每个注意力模块,参考点首先生成为统一网格,这些网格在输入数据中是相同的。然后,偏移网络将查询特征作为输入并为所有参考点生成相应的偏移量。 通过这种方式,候选键/值被转移到重要区域,从而以更高的灵活性和效率增强原始的自注意力模块,以捕获更多信息特征。 

        总而言之,我们的贡献如下:我们提出了第一个用于视觉识别的可变形自我注意主干,其中数据依赖的注意模式en-

dows 更高的灵活性和效率。在 ImageNet [10]、ADE20K [51] 和 COCO [25] 上进行的大量实验表明,我们的模型始终优于包括 Swin Transformer 在内的竞争基线,在图像分类的 top-1 准确度上领先 0.7,在语义的 mIoU 上领先 1.2分割,1.1 关于 box AP 和 mask AP 的对象检测。在小型和大型物体上的优势更加明显,差距为 2.1。

2. Related Work 

        Transformer vision backbone. 自从引入 ViT [12] 以来,改进 [6、11、26、28、36、43、49] 一直专注于学习多尺度特征以用于密集预测任务和高效注意力机制。这些注意力机制包括窗口注意力 [11、26]、全局标记 [6、21、32]、焦点注意力 [43] 和动态标记大小 [37]。最近,基于卷积的方法被引入到 Vision Transformer 模型中。其中现有的研究侧重于用卷积运算来补充 transformer 模型以引入额外的归纳偏置。 CvT [39] 在标记化过程中采用卷积,并利用步长卷积来降低自注意力的计算复杂度。 ViT with convolutional stem [41] 提出在早期加入卷积以实现更稳定的训练。 CSwin Transformer [11] 采用基于卷积的位置编码技术,并对下游任务进行了改进。其中许多基于卷积 技术可以潜在地应用于 DAT 之上以进一步提高性能。

        Deformable CNN and attention. 可变形卷积 [9, 53] 是一种强大的机制,可以处理以输入数据为条件的灵活空间位置。最近它已被应用于 Vision Transformers [7,46,54]。可变形 DETR [54] 通过为 CNN 主干顶部的每个查询选择少量键来提高 DETR [4] 的收敛性。它的可变形注意力不适合用于特征提取的视觉主干,因为缺少键限制了表示能力。此外,可变形 DETR 中的注意力来自简单学习的线性投影,并且键不在查询令牌之间共享。 DPT [7] 和 PS-ViT [46] 构建可变形模块来改进视觉标记。具体来说,DPT 提出了一个可变形的补丁嵌入来改进跨阶段的补丁,PS-ViT 在 ViT 主干之前引入了一个空间采样模块来改进视觉标记。它们都没有包含可变形的 注意力进入视觉骨干。相比之下,我们的可变形注意力采用强大而简单的设计来学​​习一组在视觉标记之间共享的全局键,并且可以用作各种视觉任务的通用骨干。我们的方法也可以看作是一种空间自适应机制,已在各种工作中证明是有效的 [16,38]。  

图 2. 可变形注意机制的图示。 (a) 呈现可变形注意力的信息流。在左侧部分,一组参考点均匀地放置在特征图上,其偏移量是通过偏移网络从查询中学习到的。然后根据变形点从采样的特征中投影出变形的键和值,如右图所示。相对位置偏差也由变形点计算,增强了输出变形特征的多头注意力。我们只展示了 4 个参考点以进行清晰的展示,实际实施中还有更多参考点。 (b) 揭示了偏移生成网络的详细结构,标有特征图的大小。

3. Deformable Attention Transformer

3.1. Preliminaries

        我们首先重新审视最近的 Vision Transformers 中的注意力机制。以扁平化的特征映射 x \in R^{N\times C} 作为输入,具有 M 个头的多头自注意力 (MHSA) 块表示为:

其中σ(·)表示softmax函数,d = C/M是每个头的维度。 z^{(m)} 表示第 m 个注意力头的嵌入输出,q^{(m)},k^{(m)},v^{(m)}\in R^{(N\times C)} 分别表示query、key和value。W_q,W_k,W_v,W_o\in R^{C\times C} 是投影矩阵。为了构建 Transformer 块,通常采用具有两个线性变换和 GELU 激活的 MLP 块来提供非线性。

        使用归一化层和身份捷径,第 l 个 Transformer 块表示为:

where LN is Layer Normalization [1].

3.2. Deformable Attention

        现有的分层 Vision Transformers,特别是 PVT [36] 和 Swin Transformer [26] 试图解决过度关注的挑战。前者的下采样技术导致严重的信息丢失,而后者的 shiftwindow attention 导致感受野的增长缓慢得多,这限制了对大型对象建模的潜力。因此,需要依赖数据的稀疏注意力来灵活地对相关特征进行建模,从而产生 DCN [9] 中提出的可变形机制。然而,简单地在 Transformer 模型中实现 DCN 是一个非常重要的问题。在 DCN 中,特征图上的每个元素单独学习其偏移量,其中 H × W ×C 特征图上的 3 × 3 可变形卷积具有 9HWC 的空间复杂度。如果我们直接在注意力模块中应用相同的机制,空间复杂度将急剧上升到 N_qN_kC,其中 N_q,N_k 是query和key的数量,通常与特征图大小 HW 具有相同的规模, 带来近似于双二次复杂度。尽管 Deformable DETR [54] 已经设法通过在每个尺度上设置较少数量的 N_k = 4 的密钥来减少这种开销,并且作为检测头工作得很好,但由于不可接受的信息丢失(参见附录中的详细比较)。同时,[3,52] 中的观察表明,不同的query在视觉注意力模型中具有相似的注意力图。因此,我们选择一个更简单的解决方案,为每个query共享移位键和值(shifted keys and values),以实现有效的权衡。 

        具体来说,我们提出可变形注意力机制,以在特征图中重要区域的指导下有效地对标记之间的关系进行建模。这些聚焦区域由多组变形采样点确定,这些变形采样点是通过偏移网络从query中学习到的。我们采用双线性插值从特征映射中采样特征,然后将采样的特征馈送到key和value投影以获得可变形的key和value。最后,应用标准的多头注意力。此外,变形点的位置提供了更强大的相对位置偏差,以促进可变形注意的学习,这将在以下部分中讨论。

Deformable attention module.如图 2(a) 所示,给定输入特征图 x\in R^{H\times W\times C},生成点 p\in R^{H_G\times W_G \times2} 的均匀网格作为参考。 具体来说,网格大小从输入特征图大小下采样一个因子 r,H_G=H/r,W_G=W/r。参考点的值是线性间隔的二维坐标 { (0, 0), . . . ,

(H_G-1,W_G-1) } , 然后我们根据网格形状 H_G\times W_G 将它们归一化到范围 [−1, 1] ,其中 (−1, −1) 表示左上角( 1, 1) 表示右下角。为了获得每个参考点的偏移量,将特征图线性投影到query标记 q=xW_q ,然后馈送到轻量级子网络 \theta_{offset}(\cdot ) 以生成偏移量 \Delta p=\theta_{offset}(q) 为了稳定训练过程,我们通过一些预定义因子 s 来缩放 Δp 的幅度,以防止偏移量过大,即 \Delta p \leftarrow s\ tand(\Delta p)。然后在变形点的位置对特征进行采样作为key和value,然后是投影矩阵: 

 \tilde{k} 和 \tilde{v} 分别表示变形后的key, value embeddings。具体来说,我们将采样函数 φ (·;·) 设置为双线性插值以使其可微分:

其中 g(a, b) = max(0, 1 − | a − b | ) 和 (r_x , r_y ) 索引 z\in R^{H\times W\times C} 上的所有位置。由于 g 仅在最接近 (p_x , p_y) 的 4 个积分点上非零,因此它简化了公式 (8) 为 4 个位置的加权平均值。与现有方法类似,我们对 q、k、v 执行多头注意力,并采用相对位置偏移 R。注意力头的输出公式为:

其中 \o (\hat{B};R)\in R^{HW\times H_GW_G} 对应于先前工作 [26] 之后的位置嵌入,同时进行了一些改编。详细信息将在本节稍后说明。每个头部的特征连接在一起并通过 W_o 投影以获得最终输出 z,如公式(3)。 

Offset generation. 正如我们所说,使用query特征设计的子网络用语生成偏移量并分别输出参考点的偏移值。考虑到每个参考点覆盖一个局部s×s区域(s是偏移量的最大值),生成网络也应该具有局部特征的感知,以学习合理的偏移量。因此,我们将子网络实现为两个具有非线性激活的卷积模块,如图 2(b) 所示。输入特征首先通过 5×5 深度卷积来捕获局部特征。然后,采用 GELU 激活和 1×1 卷积来获得 2D 偏移量。同样值得注意的是,1×1 卷积中的偏差被降低以减轻所有位置的强制偏移。

Offset groups. 为了促进变形点的多样性,我们在 MHSA 中遵循类似的范例,并将特征通道分成 G 组。来自每个组的特征使用共享子网络分别生成相应的偏移量。在实践中,注意模块的头数 M 设置为偏移组 G 大小的倍数,以确保将多个注意头分配给一组变形的key和value。

Deformable relative position bias. 相对位置偏差对每对query和key之间的相对位置进行编码,从而通过空间信息增强普通注意力。考虑一个形状为 H×W 的特征图,其相对坐标位移在两个维度上分别位于 [−H, H] 和 [−W, W] 范围内。在 Swin Transformer [26] 中,构建了一个相对位置偏差表 \hat{B} \in R^{ (2H-1)\times (2W -1)},通过用两个方向的相对位移对表进行索引来获得相对位置偏差 B。由于我们的可变形注意力具有连续的键位置,我们计算归一化范围 [−1, 1] 中的相对位移,然后在参数化偏差表中内插 φ (B; R) 通过连续的相对位移来覆盖所有可能的偏移值。

Computational complexity. 可变形多头注意力 (DMHA) 与 PVT 或 Swin Transformer 中的对应物具有相似的计算成本。唯一的额外开销来自用于生成偏移量的子网络。整个模块的复杂度可以概括为:

其中 N_s = H_G W_G = HW/r^2 是采样点的数量。可以立即看出偏移网络的计算成本具有线性复杂度关于通道大小,这相对于注意力计算的成本来说是较小的。通常,考虑用于图像分类的 Swin-T [26] 模型的第三阶段,其中 H =W =14,Ns =49,C =384,单个块中注意模块的计算成本为 79.63M FLOPs。如果配备我们的可变形模块(k = 5),额外开销为5.08M Flops,仅占整个模块的6.0%。此外,通过选择额外开销较大的下采样因子 r,复杂性将进一步降低,这使其对具有更高分辨率输入的任务(例如对象检测和实例分割)更友好。

图 3. DAT 架构图解。 N 1 到N 4 是堆叠的连续局部注意和移位窗口/可变形注意块的数量。 k 和 s 表示补丁嵌入中卷积层的内核大小和步幅。 

3.3. Model Architectures

        我们在 Transformer(公式(4))中用我们的可变形注意力替换标准 MHSA,并将其与 MLP(公式(5))结合以构建可变形视觉变换器块。在网络架构方面,我们的模型 Deformable Attention Transformer 与 [7, 26, 31, 36] 具有相似的金字塔结构,广泛适用于需要多尺度特征图的各种视觉任务。如图 3 所示,形状为 H × W × 3 的输入图像首先嵌入步幅为 4 的 4×4 非重叠卷积,然后是归一化层以获得 \frac{H}{4}\times \frac{W}{4}\times 4 补丁嵌入。为了构建一个层次化的特征金字塔,主干包括 4 个步幅逐渐增加的阶段。在两个连续的阶段之间,有一个步长为 2 的非重叠 2×2 卷积对特征图进行下采样,从而将空间大小减半并将特征维度加倍。在分类任务中,我们首先对上一阶段输出的特征图进行归一化,然后采用具有池化特征的线性分类器来预测逻辑。在目标检测、实例分割和语义分割任务中,DAT 在集成视觉模型中扮演骨干角色,以提取多尺度特征。我们在每个阶段的特征中添加一个归一化层,然后将它们输入以下模块,例如对象检测中的 FPN [23] 或语义分割中的解码器。 

        我们在 DAT 的第三和第四阶段引入了连续的局部注意力和可变形注意力块。特征图首先由基于窗口的局部注意力处理以在局部聚合信息,然后通过可变形注意力块对局部增强标记之间的全局关系进行建模。这种具有局部和全局感受野的注意力块的交替设计有助于模型学习强大的表征,在 GLiT [5]、TNT [15] 和 Pointformer [29]。由于前两个阶段主要学习局部特征,因此这些早期阶段的可变形注意力不太受欢迎。此外,前两个阶段的键和值具有相当大的空间大小,这大大增加了可变形注意力中点积和双线性插值的计算开销。因此,为了在模型容量和计算负担之间取得权衡,我们只将可变形注意力放在第三和第四阶段,并采用 Swin Transformer [26] 中的 shift-window 注意力,以便在早期阶段有更好的表现。我们在不同的参数和 FLOP 中构建了三种 DAT 变体,以便与其他 Vision Transformer 模型进行公平比较。我们通过在第三阶段堆叠更多块并增加隐藏维度来改变模型大小。详细架构在表 1 中报告。请注意,DAT 的前两个阶段还有其他设计选择,例如 PVT 中的 SRA 模块。我们展示比较结果见表7。 

猜你喜欢

转载自blog.csdn.net/like_jmo/article/details/128114780