DeepViT Towards Deeper Vision Transformer

发布 : 2021-06-25 分类 : 深度学习 浏览 :

摘要

最近,视觉变换器(ViTs)被成功应用于图像分类任务。在本文中,我们表明,与卷积神经网络(CNN)可以通过堆叠更多的卷积层而得到改善不同,ViTs的性能在扩展到更深的时候会迅速饱和。更具体地说,我们从经验上观察到,这种扩展的困难是由注意力塌陷问题引起的:随着变换器的深入,注意力图逐渐变得相似,甚至在某些层之后变得基本相同。换句话说,在深度ViT模型的顶层,特征图往往是相同的。这一事实表明,在ViT的较深层,自我注意机制不能学习有效的概念来进行表征学习,阻碍了模型获得预期的性能增益。基于上述观察,我们提出了一种简单而有效的方法,即Re-attention,来重新生成注意力图,以增加不同层的多样性,其计算和内存成本可以忽略不计。所提出的方法使训练更深的ViT模型成为可能,并通过对现有的ViT模型进行微小的修改,使其性能不断提高。值得注意的是,当用32个transformer block训练一个深度ViT模型时,在ImageNet上的Top-1分类精度可以提高1.6%。代码可在https://github.com/zhoudaquan/dvit_repo 公开。

1. 引言

最近的研究[7, 37]表明,transformer[38]可以成功地应用于视觉任务[18],与卷积神经网络(CNN)[9, 35]相比,性能具有竞争力。 视觉网络(CNN)[9, 35]。与通过堆叠多个卷积(如3×3)来聚集全局信息的CNN不同[9, 10],视觉变换器(ViTs)[7]利用自我注意(SA)机制的优势[38]来捕捉空间模式和非局部依赖。这使得ViTs能够聚集丰富的全局信息,而无需像CNN那样手工制作逐层的局部特征提取。CNNs,从而实现了更好的性能。例如,如[37]所示,在ImageNet分类中,具有22M可学习参数的12块ViT模型比具有30多个瓶颈卷积块的ResNet-101模型取得更好的结果。
深度CNN模型的最新进展主要是通过训练具有大量层的深度模型来推动的,这是由新颖的模型架构设计所促成的[9, 41, 31, 20, 47]。这是因为更深的CNN可以为输入图像学习更丰富、更复杂的表征,并在视觉任务上提供更好的性能[1, 45, 29].因此,如何有效地扩展CNN的深度是最近深度学习领域的一个重要主题,这刺激了像残差学习[9]的技术。
考虑到浅层ViTs的显著性能,一个自然的问题出现了:我们能不能像CNN一样,通过使ViTs变得更深来进一步提高性能?虽然乍一看似乎很简单,但答案可能并不简单,因为ViT与CNN的本质区别在于它对自我注意机制的严重依赖。为了解决这个问题,我们在这项工作中详细调查了ViTs沿着深度的可扩展性。
image.png
我们首先在ImageNet上进行了试验性研究,以调查ViT的性能如何随着模型深度的增加而变化。在图1中,我们展示了ViTs[7]在不同块数(绿线)下的表现。不同的块数(绿线),范围从12到32。如图所示,随着变压器块数的增加,模型的性能并没有相应提高。令我们惊讶的是,有32个变压器块的ViT模型的性能甚至比有24个变压器块的模型还要差。这意味着像CNN[9]那样直接堆叠更多的变换器块在增强ViT模型方面是低效的。然后,我们深入研究这一现象的原因。我们根据经验观察到,随着ViTs深度的增加,用于聚合每个transformer block特征的注意力图在某些层之后往往会过于相似,这使得表征在某些层之后停止发展。我们把这个具体问题命名为注意力崩溃。这表明,随着ViT的深入,自我注意机制对于产生不同的注意来捕捉丰富的表征变得不那么有效。
为了解决注意力崩溃的问题,并有效地扩大视觉转换器的深度,我们提出了一个简单而有效的自我注意力机制,命名为Re-attention。我们的Re-attention利用了多头自我注意(MHSA)结构的优势,通过以可学习的方式交换不同注意头的信息来重新生成注意图。实验表明,在没有任何额外的增强和正则化策略的情况下,只需用Re-attention替换ViTs中的MHSA模块,我们就可以用甚至32个转化器块来训练非常深的视觉转化器,并取得一致的改进,如图1所示。此外,我们还提供了消减分析,以帮助更好地理解Re-attention在扩展视觉变换器方面的作用。
总的来上,我们的贡献如下:

  • 我们深度研究了视觉Transformer的行为,观察他们不能像CNNs那样从持续堆叠更多层中受益。我们进一步确定了这种违反直觉的现象背后的根本原因,并首次将其总结为注意力崩溃。
  • 我们提出了Re-attention,一种简单而有效的注意力机制,它考虑到了不同注意力头之间的信息交流。
  • 据我们所知,我们是第一个成功地在ImageNet-1k上从头开始训练32块ViT,并且性能持续提高。我们表明,通过用我们的Re-attention替换自我注意模块,在ImageNet-1k数据集上可以取得最先进的结果,而不需要在更大的数据集上进行任何预训练。



3. 重温视觉Transformer

image.png
一个视觉变换器 (ViT) 模型 [37, 7],如图 2(a) 所示,由三个主要部分组成:用于补丁嵌入的线性层(即,将高分辨率输入图像映射到低分辨率特征图),一堆具有多头自注意和前馈层的transformer block,用于特征编码,以及一个用于分类分数预测的线性层。在这一节中,我们首先回顾了它独特的transformer block,特别是自注意力机制,然后我们提供了关于自注意力崩溃问题的研究。

3.1 多头自注意力

Transformers[38]在自然语言中被广泛使用,用于将输入词标记的序列编码为嵌入序列。在处理图像时,为了符合这种序列到序列的学习结构,ViTs首先将输入图像均匀地分成多个块,并将每个块编码为一个标记嵌入。然后,所有这些标记,连同一个类标记,被送入Transformers块的堆栈。
每个转换块由一个多头自注意(MHSA)层和一个前馈多层感知器(MLP)组成。MHSA通过对输入进行线性转换,将query(Q)和一对key(K)-value(V)生成一个可训练的关联存储器,并将其输出。在数学上,MHSA的输出是通过以下方式计算的。
image.png
其中√d是一个基于网络深度的缩放系数。然后,MHSA的输出被归一化,并被送入MLP以生成下一个区块的输入。在上述自注意中,Q和K相乘生成注意图,它代表了每层内所有tokrn之间的关联性。它被用来检索和组合值V中的嵌入。在下文中,我们特别分析注意力图在扩展ViT的作用。为了方便起见,我们用$\mathbf{A} \in \mathbb{R}^{H \times T \times T}$来表示注意力图,其中H是SA头的数量,T是token的数量。对于第h个SA头,注意力图被计算为$\mathbf{A}_{h,:,:}=\operatorname{Softmax}\left(Q_{h} K_{h}^{\top} / \sqrt{d}\right)$,其中$Q_h$和$K_h$来自于对应的头。在上下文清楚的情况下,我们省略了 下标h。

3.2 注意力崩溃

受深度CNN的成功启发[9, 30, 35, 36],我们对ViTs的性能随着深度的增加而变化进行了系统研究。在不丧失一般性的前提下,我们首先按照[37]中的常见做法,将隐藏维度和头数分别固定为384和121。然后,我们将不同数量的转化器块(从12到32不等)堆叠起来,以建立对应于不同深度的多个ViT模型。在ImageNet[18]上对图像分类的整体性能进行了评估,并在图1中进行了总结。从性能曲线上看,我们惊讶地发现,随着模型的深入,分类的准确性提高得很慢,而且很快就饱和了。更具体地说,我们可以观察到,在采用了24个转化器块之后,改进就停止了。这一现象表明,现有的ViTs很难从更深的架构中获得好处。
这样的问题是相当反直觉的,值得探索,因为类似的问题(即如何有效地训练一个更深的模型)在CNN的早期发展阶段也被观察到[9],但后来得到了妥善解决[9,10]。通过对transfromer架构的深入研究,我们想强调的是,自注意力机制在ViTs中起着关键作用,这使得它与CNN有很大的不同。与CNN的显著不同。因此,我们首先调查了自注意力,或者更具体地说,生成的注意图$\mathbf{A}$是如何随着模型的深入而变化的。
image.png
为了衡量各层注意图的演变,我们计算了不同层的注意图之间的跨层相似性。 不同层的注意图之间的跨层相似性:
image.png
其中$M^{p,q}$是p层和q层的注意力图之间的余弦相似性矩阵,每个元素$M_{h,t}^{p,q}$衡量head h和token t的注意力相似性。考虑到一个特定的自注意层和它的第h个头,$\mathbf{A}_{h,:, t}^{*}$是一个T维向量,代表输入token t对T个输出token中的每个标记的贡献程度。因此,$M_{h, t}^{p, q}$提供了一个适当的衡量标准,说明一个token的贡献在p层和q层之间如何变化。当$M_{h, t}^{p, q}$等于1时,意味着token t在p层和q层中对自注意力起着完全相同的作用。
给出公式(2),然后我们在ImageNet-1k上训练一个有32个转换块的ViT模型,并研究所有注意力图之间的上述相似性。如图3(a)所示,第17个块之后的M中相似注意力图的比例大于90%。这表明之后学习的注意力图谱是相似的,转化器块可能退化为MLP。因此,进一步堆叠这种退化的MHSA可能会引入模型等级退化的问题(即,将各层参数相乘所得的模型参数张量的等级会降低),并限制了模型的学习能力。这也被我们对所学特征的退化的分析所验证,如下图所示。这种观察到的注意力塌陷可能是观察到的ViTs的性能饱和的原因之一。为了进一步验证不同深度的ViTs是否存在这种现象,我们分别对具有12、16、24和32个转换块的ViTs进行了同样的实验,并计算具有类似注意力图的块数。图3(b)所示的结果清楚地表明,当增加更多的transformer块时,类似注意力图块的数量与总块数的比例增加。
为了了解注意力崩溃如何伤害ViT模型的性能,我们进一步研究它如何影响深层的特征学习。对于一个特定的32块ViT模型,我们通过调查它们的余弦相似度来比较最终输出特征和每个中间转化块的输出。图4中的结果表明,相似度相当高,学到的特征在第20个块之后就停止了演变。注意力相似性的增加与特征相似性之间存在着密切的关联。这一观察表明,注意力崩溃是造成ViTs的不可扩展问题的原因。
image.png

4. Re-attention for Deep ViT

正如上文所揭示的,将ViT扩展到更深层次的一个主要障碍是注意力崩溃问题。在本节中,我们提出了两种解决方法,一种是增加计算自我注意力的隐藏维度,另一种是一种新的Re-attention机制。

4.1 高维空间中的自注意力

征服注意力崩溃的一个直观的解决方案是增加每个符号的嵌入维度。这将增强每个token embedding的表示能力,以编码更多的信息。因此,由此产生的注意力图谱可以更加多样化,每个区块的注意力图谱之间的相似性可以得到教育。在不丧失一般性的前提下,我们通过进行一组基于ViT模型的实验,用12个区块进行快速实验来验证这种方法。按照以前基于变压器的工作[38, 7],选择了四个嵌入尺寸,从256到768不等。详细的配置和结果见表1。
image.png
从图5和表1可以看出,通过增加嵌入维度,具有相似注意力图的区块数量减少了,注意力塌陷得到了缓解。因此,模型的性能也相应提高。这验证了我们的核心假设–注意力塌陷是扩展ViT的主要瓶颈。尽管它很有效,但增加嵌入维度也大大增加了计算成本,带来的性能改善也趋于减少。此外,一个更大的模型(具有更高的嵌入维度)通常需要更多的数据进行训练,遭受过度拟合的风险并降低效率。

4.2 Re-attention

在第3节中已经证明,不同transformer block的注意图之间的相似性很高,特别是对于深层。然而,我们发现来自同一transformer block的不同头部的注意力图的相似性相当小,如图3(c)所示。很明显,来自同一自我注意层的不同头部关注的是输入标记的不同方面。基于这一观察,我们建议建立跨头通信来重新生成注意力图,并训练深度ViTs,使其表现更好。
具体来说,我们使用头部的注意力图作为基础,并通过动态地聚合它们来生成一组新的注意力图谱。为了实现这一目标,我们定义了一个可学习的变换矩阵$\Theta \in \mathbb{R}^{H \times H}$,然后用它将多头的注意力图混合成重新生成的新图,再与V相乘。具体来说,重新注意是通过以下方式实现:
image.png
其中转换矩阵Θ沿头部维度乘以自我注意图A。这里Norm是一个归一化函数,用于减少层间的差异。Θ是端到端可学习的。
优点:所提出的Re-attention的优点有两个方面。首先,与其他可能的注意力增强方法相比,如随机丢弃
的一些元素,或者调整SoftMax的温度,我们的Re-attention利用了不同注意力头之间的相互作用来收集它们的互补信息,更好地提高了注意力图的多样性。这一点也被我们下面的实验所验证。此外,我们的Re-attention是有效的,易于实现的。它只需要几行代码,计算量可以忽略不计。与原始的自我注意相比,它只需要几行代码和可以忽略不计的计算开销。因此,它比增加嵌入维度的方法要有效得多。

本文作者 : HeoLis
原文链接 : http://ishero.net/DeepViT%20Towards%20Deeper%20Vision%20Transformer.html
版权声明 : 本博客所有文章除特别声明外,均采用 CC BY-NC-SA 4.0 许可协议。转载请注明出处!

学习、记录、分享、获得

微信扫一扫, 向我投食

微信扫一扫, 向我投食