分享

如何将二维空间先验注入到ViT中? UMA&港理工&阿里提出SP-ViT,为视觉Transformer学习2D空间先验知识!

 我爱计算机视觉 2022-08-02 发布于江苏

详细信息如下:

  • 论文地址:https:///abs/2206.07662

  • 代码地址:未开源



      01      

摘要

最近,Transformer 在图像分类方面显示出巨大的潜力,并在 ImageNet 基准上建立了最先进的结果。然而,与 CNN 相比,transformer 收敛速度较慢,并且由于缺乏空间归纳偏置,在低数据状态下容易过度拟合。这种空间归纳偏差可能特别有益,因为输入图像的 2D 结构在Transformer 中没有得到很好的保留。

在这项工作中,作者提出了空间先验增强自注意 (SP-SA),这是一种为视觉Transformer量身定制的普通自注意 (SA) 的新变体。空间先验(SPs)是本文提出的一系列归纳偏置,它突出了某些空间关系组。与被迫只关注硬编码局部区域的卷积归纳偏置不同,本文提出的 SP 是由模型本身学习的,并考虑了各种空间关系。

具体来说,注意力分数的计算重点是每个头的某些类型的空间关系,并且这些学习到的空间焦点可以相互补充。基于 SP-SA,作者提出了 SP-ViT 系列,它始终优于其他具有类似 GFlops 或参数的 ViT 模型。本文最大的模型 SP-ViT-L 实现了 86.3% 的 Top-1 准确度,与之前的最先进模型相比,参数数量减少了近 50%

      02      

Motivation

在主导自然语言处理 (NLP) 任务之后,Transformers最近在图像分类方面取得了令人兴奋的成果。所有Transformer的核心是所谓的自注意力机制,它全局捕获所有输入token对之间的内容关系,并有选择地关注相关对。与卷积相比,自注意力更灵活,卷积是通过硬编码专门捕获局部依赖关系的。


这可能为Transformer模型配备更大的容量和更大的计算机视觉任务潜力。正如最近的工作所报道的那样,当在大型数据集上进行预训练时,Transformer的性能优于卷积神经网络 (CNN),并通过预训练 CNN 的知识蒸馏或伪标签来促进。

尽管如此,CNN 的泛化能力和收敛速度都比 Vision Transformers (ViT) 更好。这表明卷积中使用的某些类型的归纳偏置仍然可以有益于解决视觉任务。因此,最近的许多研究提出以不同的方式将卷积归纳偏置纳入 ViT,并且都证明了性能的提升。卷积的有效性依赖于自然图像的相邻像素高度相关的事实,但在卷积滤波器的局部感受野之外可能存在其他高度相关的内容被忽略。

因此,作者提出同时利用各种归纳偏置,就像人类在日常生活中所做的那样,例如,如果我们看到一个水平物体的一部分,我们自然会顺着它的方向看,而不是将我们的视线限制在局部范围内。

在这项工作中,作者通过称为空间先验增强自注意 (SP-SA) 的普通自注意 (SA) 的扩展,将名为空间先验 (SPs) 的新型归纳偏置系列引入 ViT。SP-SA 根据 key 和 query patch 的相对位置在每个注意力头突出显示一组 2D 空间关系。它有助于在此类空间关系的上下文中计算注意力分数。由于适当空间先验的构建和验证非常费力,因此作者引入了可学习空间先验的概念。

具体地说,作者只是将弱的先验知识强加给模型,即不同的相对距离应该被区别对待。然而,作者并不强制模型预先支持任何类型的空间关系,例如,既不是局部的也不是非局部的。有效空间先验(SP)应该由模型本身在训练阶段发现。为此,SP 由一系列数学函数表示,这些函数将相对坐标映射到抽象分数,称为空间关系函数。

为了寻找理想的空间关系函数,作者通过神经网络对这些函数进行参数化,并与 ViT 一起优化它们。因此,该模型可以学习类似于卷积中诱导的空间先验,但它也可以学习更大距离上的空间关系。

如上图b所示,不同的注意力头呈现不同的互补模式,从而分别处理不同类型的空间关系。同时,一个全局感受野是通过考虑所有头来近似组成的。

事实上,卷积归纳偏置可以看作是一种特殊的空间先验:它们首先将坐标空间关系分为两类,即关注局部邻域的和关注非局部区域的。然后他们学习局部邻域的先验并忽略非局部关系。为了比较,将这种卷积偏置与 ViT 相结合的一些现有方法如上图a所示。基于局部窗口,分别通过改变纵横比或移动中心,为 CSWin-Transformer和 Swin-Transformer提出了新的变体。但是,这种窗口的设计非常直观,关注局部关系的主要思想没有改变。

总的来说,本文做出以下贡献

  1. 本文提出了一系列针对 ViT 的归纳偏置,这些偏置集中在不同类型的空间关系上,称为空间先验 (SP)。SP 将局部受限的卷积归纳偏置推广到局部和非局部相关。通过神经网络参数化,SP 在训练期间自动学习,无需预先对任何硬编码区域施加偏好。

  2. 作者提出了 SP-SA,这是一种新颖的自注意变体,可以自动学习有益的 2D 空间归纳偏置。基于本文提出的 SP-SA,作者进一步构建了一个称为 SP-ViT 的 ViT 变体。SP-ViT 在没有额外数据的情况下在 ImageNet Benchmark 上实现了最先进的结果。

  3. 本文提出的 SP 与各种输入大小兼容,因为它们来自每对patch之间的相对坐标,而不是它们的绝对位置。当在更高分辨率上进行微调时,SP-ViT 还表现出比图像分类baseline模型更好的性能。


      03      

方法

3.1 Revisiting Multi-Head Self-Attention

Self-attention 接收输入序列并输出相同长度的新序列,其中计算每个元素作为线性变换的输入元素的加权和:


每个权重系数或注意力分数是根据两个元素之间的语义依赖关系确定的,通过将 softmax 函数应用于线性变换元素的缩放点积,在本文中将其称为内容得分:


多头自注意力并行使用几个这样的操作来学习不同类型的相互依赖关系。最终输出是通过对来自每个头的concat输出应用线性变换来获得的。

3.2 Spatial Prior-enhanced Self-Attention


由于观察到某些类型的空间关系归纳偏置可能对Transformer有益,作者提出了通过结合学习的 2D 空间先验 (SP) 来增强自注意的扩展,称为空间先验增强自注意 (SP-SA)。将SP-SA嵌入到ViT,可以形成本文提出的SP-ViT,其结构如上图所示。

每个 SP形成一种用于计算注意力分数的空间上下文,它源自输入元素对之间的坐标空间关系,即 ViT 的键和查询patch之间的相对位置。因此,SP 具有完全相同形状的注意力分数,可以简单地通过乘法将其整合到Attention计算中:

3.2.1 2D Spatial Prior

以查询块i为参考点,可以得到图像块 j 的相对坐标。然后为所有查询和键patch对使用共享映射,命名为空间关系函数:


输出一起形成所谓的 2D SP 矩阵 Ω。

3.2.2 Learnable 2D Spatial Priors

为了使模型能够自动学习所需的归纳偏置,作者使用多层感知器(MLP)来参数化从 2D 相对坐标到 Ω 的映射:


因此,Ω可以学习查询 和键的注意力分数的权重,该权重仅取决于它们的相对坐标并以非线性方式应用,即在 softmax 之前。作者通过向每个头添加一个独特的网络将 SP-SA 扩展到其多头版本。这种设计遵循与多头自注意力相同的动机,并假设不同 SP 的组合应该提高性能。

3.3 Relation to Other Methods

3.3.1 Relation to Local Windows

之前工作中使用的方形和十字形窗口在实践中可以看作是本文提出的空间关系函数的一种特殊形式:


其中分别控制移位、窗口宽度和高度。如果,则生成方形窗口,否则生成十字形窗口。两项工作都只对整个网络采用了一些硬编码模式,而本文的方法提出从各种有益的 2D 结构中受益。

3.3.2 Relation to PSA

Positional Self-Attention (PSA)也可以看作是手动设计的空间关系函数族:


其中是可学习的参数,它们按照明确的规则初始化以近似卷积效果。
它们的主要贡献是所谓的局部/卷积初始化,它将头的数量限制为整数的平方,并且的初始值都需要额外的超参数调整。为了与他们的方法进行比较,作者采用具有 9 个头的 ViT baseline进行消融分析。

3.3.3 Relation to Relative Positional Embeddings

之前也有一些工作考虑了相对位置编码:


其中是一个可学习的嵌入表,从中获取相对位置嵌入。然后它与查询进行乘法交互。

与 1D 相对位置嵌入不同,作者提出的空间关系函数在建模图像的 2D 结构方面具有更大的能力。如果扩展到 2D,他们的方法相当于将线性变换应用于相对距离的 one-hot 表示。对于 one-hot 表示,距离的大小被忽略,而相对坐标则不是这样。

更重要的是,相对位置编码只需将相对位置嵌入添加到注意力方程中的键和值。自注意力的修改是相当直观的,但缺乏明确的物理解释。与位置嵌入相反,本文的方法不仅提供了中性位置信息,而且还学习了有益的归纳偏置并将其注入模型中。


      04      

实验

上表展示了本文方法在ImageNet上和SOTA方法的对比,可以看出,本文方法在性能上具有一定的优势。


在上图中,作者使用最近的 Transformer Explainability 对几个图像展示目标类激活图的可视化,以展示 SPViT 的行为。虽然 DeiT 模型只显示目标类区域的一小部分的类激活,例如“澳洲鹦鹉”的头部、“埃及猫”的皮毛或“美洲短吻鳄”的下巴,但提出的 SP- ViT 模型显示了更广泛的目标类区域上的类激活。


如上图 a 所示,与 DeiT baseline相比,用 SP-SA 替换多个 SA 层可以提高 Top-1 精度。一般来说,随着更多的层被替换,性能会提高。对于总共有 12 层的模型,当替换 10 层时性能最佳。当用本文的 SP-SA 替换除最后一个 SA 层之外的所有层时,性能略有下降。

作者认为这很可能是因为分类token只涉及最后一层,因此在这种情况下没有充分提取类特定的特征。作者进一步研究了上图b 中具有更多层的模型(总共 16 层),并发现了类似的趋势。当第一层到倒数第二层被替换时,性能最佳。


上表展示了消除在 ImageNet-100 上后些层插入分类token的影响。可以看出,分类token的后期插入确实对分类结果有积极的影响。


作者发现一个简单的 2 层 MLP 可以很好地作为空间关系函数。在上表中,作者研究其隐藏维度对性能的影响。在 Top-1 准确度中,性能从 16 到 32 的隐藏维度提高了 0.4%,但隐藏维度的进一步增加并没有进一步改善结果。隐藏维度为 32 的结果是最好的。


上表展示了本文方法和其他基于相对位置方法的对比,可以看出,本文的方法具有明显的优势。


为了进一步验证本文提出的空间关系函数的重要性,在上表中,作者将其与相对空间信息的乘法交互的不同方式进行了比较。结果表明强大的非线性函数对于学习有效的空间先验是必要的。


为了验证组合各种 SP 的好处,作者将本文的 SP-SA 与它的两个变体进行了比较:一个只为每一层采用一个 SP,另一个为整个网络学习相同的 SP。结果在上表中报告,与为每一层学习单个 SP 相比,整个网络的共享 SP 提供了更好的结果。然而,每层和头具有唯一可学习 SP 的设置表现最佳,这证明了组合不同 SP 的好处。


      05      

总结
在本文中,作者提出了一种名为 Spatial Prior-enhanced Self-Attention (SP-SA) 的 Vanilla self-attention (SA) 变体,以促进具有自动学习空间先验的视觉Transformer。基于SP-SA,作者进一步提出了SP-ViT,并通过实验证明了本文方法的有效性。

作者提出的不同大小的 SP-ViT 都为仅在 ImageNet-1K 上训练的模型建立了最先进的结果。例如,与之前最先进的 LV-ViT-M 相比,SP-ViT-M 的准确度提高了 0.8%。本文的 SP-SA 能够激发更多关于为视觉Transformer设计和利用适当归纳偏置的研究。

最后,学习到的空间先验可以被二值化并用于设计计算效率更高的Transformer。尽管具有出色的性能,但视觉Transformer依赖于大量的数据增强技术,并且与卷积相比,注意力机制在计算上效率低下,仍然需要进一步研究改善 ViT。

参考资料

[1]https:///abs/2206.07662

▊ 作者简介
研究领域:FightingCV公众号运营者,研究方向为多模态内容理解,专注于解决视觉模态和语言模态相结合的任务,促进Vision-Language模型的实地应用。
知乎/公众号:FightingCV

END



加入「Transformer交流群👇备注:TFM



    转藏 分享 献花(0

    0条评论

    发表

    请遵守用户 评论公约

    类似文章