搞懂 Vision Transformer 原理和代码，看这篇技术综述就够了（八）

本文目录

17 DeepViT: 解决注意力坍塌以构建深层ViT
(来自 新加坡国立大学, 字节跳动AI Lab(美国))
17.1 DeepViT原理分析

18 CaiT：Going deeper with Image Transformers
(来自 Facebook)
18.1 CaiT原理分析
18.2 CaiT代码解读

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型，现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制，不采用 RNN 的顺序结构，使得模型可以并行化训练，而且能够拥有全局信息。

本文介绍的两个工作工作的初衷是为了加深Transformer模型。在模型加深的过程中可以同时保证稳定训练和精度提升。深度学习的基本思想之一就是让模型更深一点，这在CNN中已经很好的体现。例如2012年图像分类冠军AlexNet就是一个更深的模型。著名的ResNet本质上也是一种如何训练更深的CNN模型的方法。

但是反观基于Transformer的模型ViT，DeiT，它们不论是base，small还是tiny尺寸的模型，大小都是只有12个layers。如果直接把depth加深的话，模型的性能会迅速地饱和。那么一个自然而然的问题就是：能不能像CNN一样，采取一些什么手段来加深Transformer模型呢？

17 DeepViT: 解决注意力坍塌以构建深层ViT

论文名称：17 DeepViT: Towards Deeper Vision Transformer

论文地址：

DeepViT: Towards Deeper Vision Transformer

https://arxiv.org/abs/2103.11886

• 17.1 Deep-ViT原理分析：

Motivation

这个工作的初衷是为了加深Transformer模型。为什么这么做的原因论文中没有明显的交代，但是深度学习的基本思想之一就是让模型更深一点，这在CNN中已经很好的体现。例如2012年图像分类冠军AlexNet就是一个更深的模型。著名的ResNet本质上也是一种如何训练更深的CNN模型的方法。

但是反观基于Transformer的模型ViT，DeiT，它们不论是base，small还是tiny尺寸的模型，大小都是只有12个layers。如果直接把depth加深的话，模型的性能会迅速地饱和，甚至32层ViT的性能还不如24层的模型，如下图1所示。那么一个自然而然的问题就是：能不能像CNN一样，采取一些什么手段来加深Transformer模型呢？

为了解答这个问题，首先要弄清楚的是：为啥简单地加深ViT性能会出现饱和呢？

原因是作者发现：模型的注意力图 (attention maps)，在深层时会非常相似，甚至是变成一样的。这表明，随着ViT的深入，self-attention 在生成不同注意力以捕获rich representation方面变得不那么有效了。作者称之为注意力坍塌 (attention collapse)。

具体来讲，作者衡量不同layer直接attention map的相似性的方法是计算余弦相似度：

式中  代表第  个layer，  代表第  个head，  代表token。

是个  维度的张量，  是个  维的向量，代表第  个head的第  个输入的token对  个output token的注意力大小。所以  指的是第  层的第  个head的第  个输入的这个值，而  指的是第  层的第  个head的第  个输入的这个值。二者做cosine值以后记得到的  为  层和  层的第  个head的第  个输入的这个值的相似度。 当  越大时，代表相似度越高；当  越小时，代表相似度越低。当  时，代表token  在  层和  层所起的作用完全一致。

根据这样的定义，作者使用 ImageNet-1k 训练了一个32层的ViT模型。下图2的意义论文里讲得很含糊，我推测是每个点代表一种ratio，这个ratio是指这个层的token与周围  个层的对应token的相似性。该图表明了随着深度的增加，attention map和  个附近的block的attention map越来越相似。

下图3代表的是这32个层里面，随着层数的加深，具有相似的attention map的block的数量(红色线)，以及跨层的相似度(黑色线)。比如最右边的红色点表示在到了第32层的时候，相似的block就达到了16个，而且跨层相似度达到了近0.5左右。

下图4代表的是同1个block内部不同head的的注意力图相似性，发现块内不同头之间的相似性都低于30%，它们呈现出足够的多样性。

Attention collapse这个问题究竟是如何影响ViT模型的性能的呢？

如图5所示的黑色线代表最后一个特征图与前面每个特征图的相似度；红色线代表相邻的层中具有相似的attention map的比例。从大概20层之后，这个相似度变得非常高，使得学习到的特征图和attention map都差不多了。这一观察表明，attention collapse这个问题是ViT不可扩展问题的原因。

作者又对注意力坍塌这个问题做了进一步分析：如下图6所示的shared blocks这个参数代表把模型最后的这么多层的  和  设置为一模一样的值以后得到的Acc。比如说第2行的这个11，就代表把一个24层的ViT模型的最后11层的  和  设置为一模一样的值，使得最后11层的attention map也是一样的。发现掉点并不严重，这表明最后深层的attention map其实本来就很接近，即：注意力坍塌这个问题确实是存在的，使得简单地加深ViT模型是无效的。

为了解决这个问题，作者提出了2种解决办法：

第1种是增加self-attention模块的embedding dimension，这样做的目的是为了增加每个token的表达能力。这样一来，生成的注意力图可以更加多样化，每个块的attention map之间的相似性可以减少。作者设计了4种不同的embedding dimension，分别是256,384,512,768。如下图7所示，随着embedding dimension的增长，相似的block的数量在下降，同时模型的Acc在上升。换句话讲，注意力坍塌的问题得以缓解，这证明了本文的论点，即：注意力坍塌问题是影响ViT模型scaling的原因之一。

尽管增加embedding dimension是有效的，但增加embedding dimension也会显著增加计算成本，带来的性能改进往往会减少，且需要更大的数据量来训练，增加了过拟合的风险。

鉴于此，作者提出了第2种方法。第2种是通过一个改进的模块，叫做Re-attention。 其想法来自于作者发现同一个层的不同head之间的相似度比较小，所以说来自同一自注意力层的不同head关注输入token 的不同方面。那如果可以把不同的head的信息结合起来，利用它们再生出一个attention map，是不是能够避免注意力坍塌问题呢？

所以，Re-attention以一种可学习的方式交换来自不同attention heads的信息来再生注意力图。具体来讲，作者把得到的attention map 经过一步Linear Transformation  ，得到再生的attention map，写成公式就是：

注意这个可学习的  是作用在head这个维度上的。Norm取得是BatchNorm而不是传统的LayerNorm。

这样做维度好处是充分利用了不同head之间的信息，因为这些信息的相似度低，所以可以有效地缓解注意力坍塌的问题，而且很容易实现，也不用像增加embedding diemnsion一样增加那么多的参数量和计算量。

下图9为原始ViT模型和DeepViT的Feature map的跨层余弦相似性，可以看出，用Re-attention取代原始的Self-attention可以显著降低特征Feature map的相似性。

如下图10所示为普通的self-attention与Re-attention的attention map的可视化结果。

普通的self-attention的attention map的特点是：只学习到一些局部的patch之间的relationship，而且在网络的深层时attention map很接近。

Re-attention的attention map的特点是：学习到更大范围内的patch之间的relationship，而且在网络的深层时attention map有差别。

Re-attention使得网络深层的attention map保持了多样性，并与相邻块具有更小的相似性。

Experiments：

优化器： AdamW

初始学习率： 5e-4

学习率衰减策略： cosine

batch size： 256

模型配置如下图11所示：

作者进行的一个很重要的是研究是分析Re-attention的作用。就是想看看加了Re-attention以后，模型会不会随着层数的加深而达到饱和，会不会随着层数的加深而相似度逐渐升高？

结果如下图12所示，传统的ViT模型当层数分别加深为16,24,32层时，相似的块的数量分别是5,11,16，也就是说不相似的只有前11层，且Acc不会继续增长。但是使用了Re-attention以后，没有出现相似的块，换句话说，Re-attention使得更深的模型的每个block都是有贡献的。这使得随着模型的加深，Acc也会逐渐增高，说明Re-attention可以解决注意力坍塌的问题。

作者也尝试了一种解决注意力坍塌问题的办法，就是通过给self-attention加一个temperature  ：

的作用是什么呢？

的减小可以让attention map的分布变得更加的sharp，使得attention map的数值拉开，在实际实现的时候作者使得  的值随着层数的加深而逐渐衰减，或者干脆把  设置为可学习的参数。

作者也尝试了一种解决注意力坍塌问题的办法，就是随机drop掉attention map中的一些值。因为对于不同的block来说，会随机drop掉不同的值，故attention maps之间的相似性可以得到减少。

这2种做法的结果如下图13中的绿线所示和图14所示，无论是加上temperature参数，还是采用drop attention的操作，相似的block的数量仍然很多，以及不同block之间的相似性仍然很大。所以，加上这个不断衰减的  参数对缓解注意力坍塌问题来讲只起到非常小的作用。

采用drop attention确实可以在一定程度上减弱不同blocks之间的相似性，但是attention map之间相似性的减弱来源于mask位置的不同，这对feature map影响不大。也就是说，深层的，不同层的feature map还是非常相似的。

最后作者可视化了最后一层的Transformation matrix  的样子，如图15所示，  为每个head分配了不同的权重，使得再生过程改变了不同的head的权重。比如说第31层是一个  ，第32层是一个不同的  ，凭借这一点来降低block之间的attention map的相似性。

与SOTA的对比：

作者提出了两种尺度的模型，DeepViT-S 和 DeepViT-L 层数分别是16层和32层。那么加深ViT模型会使得参数量急剧上升，为了保持和ViT的参数量大致相当，作者把embedded dimension做小了一点，ViT-B 和ViT-L的embedded dimension分别是768和1024，而DeepViT-S 和 DeepViT-L 的 embedded dimension分别是396和408。结果如下：

相比于ViT-B 和ViT-L，DeepViT-S 和 DeepViT-L 以更小的参数量获得了更好的性能。

总结：

本文提出了一种加深ViT模型的方法。所基于的论点是作者发现随着模型的加深，以一个32层模型为例，深层的 (17-32层)的不同层之间的attention map的相似度很高，然而不同的head之间的相似度却很低。基于此，作者通过一个转移矩阵，把不同的head的信息给结合起来了。借此，使得深层的block里面，不同的层训练出来的这个转移矩阵是不同的，达到异化attention map的目的。所以，DeepViT能够有效缓解注意力坍塌，使得ViT模型的性能可以随着层数的加深而增加。

18 CaiT：Going deeper with Image Transformers

论文名称：Going deeper with Image Transformers

论文地址：

Going deeper with Image Transformersa

https://arxiv.org/abs/2103.17239

• 18.1 CaiT原理分析：

18.1.1 优秀前作DeiT

CaiT和DeiT一样都是来自Facebook的同一团队，一作都是Hugo Touvron。

关于DeiT的论文，论文和代码的详细解读请读者们参考下文：

Training data-efficient image transformers & distillation through attention

https://arxiv.org/abs/2012.12877

科技猛兽：Vision Transformer 超详细解读 (原理分析+代码解读) (三)

https://zhuanlan.zhihu.com/p/349315675

在这里简单重复下DeiT和ViT的不同之处：

DeiT可以通过一组很优秀的超参数使得在不改变任何架构的前提下实现了ViT的涨点。

当然，DeiT也可以改变架构，DeiT在训练时可加入distillation token，如下图18所示，如果你觉得这部分我讲的有点简单了，请参考原文或上面博客链接。

ViT的输出是一个softmax，它代表着预测结果属于各个类别的概率的分布。ViT的做法是直接将这个softmax与GT label取  。

而在DeiT中，除了这个 以外，还要加上一个蒸馏损失：

蒸馏分两种，一种是软蒸馏（soft distillation），另一种是硬蒸馏（hard distillation）。先讲讲软蒸馏，如下式所示，右半部分， 和 分别是student model和teacher model的输出，  表示  散度， 表示softmax函数， 和  是超参数。

硬蒸馏如下式所示，  表示交叉熵。

简而言之，蒸馏的含义就是：学生网络的输出  与真实标签取  ，接着如果是硬蒸馏，就再与教师网络的标签取  。如果是软蒸馏，就再与教师网络的softmax输出结果取  。

值得注意的是，硬标签也可以通过标签平滑技术 (Label smoothing) 转换成软标签，其中真值对应的标签被认为具有  的概率，剩余的  由剩余的类别共享。  是一个超参数，这里取0.1。

训练示意图如下图19所示：

在训练过程中，class token来自ViT，保持不变。只是右侧加上了distillation token。

左侧class token的输出通过一个MLP head得到predicted label，与真实标签GT label取  。

右侧distillation token的输出通过一个MLP head得到predicted label，与教师模型的输出Teacher predicted label取软蒸馏损失 $\mathcal{L}\mathrm{global}  \mathcal{L}\mathrm{global}^\mathrm{hard Distill} $ 。

以上就是DeiT相对于ViT的变化，可以概括为2点：

• 一个distillation token配上蒸馏损失。
• 一组优秀的超参数。

18.1.2 正式介绍CaiT

Motivation

这个工作的Motivation也很直接，就是为了去优化Deeper Transformer。

本文所基于的核心论点是：网络架构 (architecture) 和优化 (architecture) 是相互影响，相互作用的。

换句话讲，一个网络优化的结果怎样，与其架构是密切相关的。架构能影响优化结果，而优化方式也可以决定架构。

作者先以著名的ResNet为例，其表达式可以写作：

其中，  一般是残差，即Identity操作，而  代表这个block的具体操作。

ResNet的设计就是上述核心论点的最好证明。kaiming大佬也在论文中这样讲到：

Residual networks do not offer a better representational power. They achieve better performance because they are easier to train.

也就是说，ResNet不具备更优秀的表征能力，只是使网络变得更易训练罢了。再次证明了网络架构 (architecture) 和优化 (architecture) 是相互影响，相互作用的。

再回到vision transformer上面来，其每个block的表达式可以写作：

现在的问题就是：如何对这个架构进行更好地归一化 (Normalize)，训练权重 (Weigh) 和初始化 (Initialize)？

CaiT给出的答案，相对DeiT来讲，也可以概括为2点：

• LayerScale： 使Deep Vision Transformer易于收敛，并能提高精度。
• class-attention layers：高效的处理class token的方式。

和DeiT一样，依然是2个改进，很简单但是相当有用。

最重要的是，ResNet，DeiT，CaiT的这种“小改进”是真的能够给vision transformer带来性能提升的，非常solid的工作，代码开源完整，follower很多。而且这些改进你都可以很快地在你自己的工程中复现并且收获满意的结果。

具体方法1：使Deep Vision Transformer易于收敛，并能提高精度的LayerScale

LayerScale这个操作的目的也很直接，就是为了使得vision transformer在训练时更稳定。作者首先对比了如下图20所示的4种不同的transformer blocks的正则化策略，看看它们哪个有助于提升优化的效果。

图20 (a) 中的操作：是ViT/DeiT中使用的pre-norm结构，即先进行Layer Normalization，再进行Self-attention或者FFN。其结果直接与这个block的输入相加。

图20 (b) 中的操作：引入了一个可学习的参数  ，作用在residual block的输出，并取消了Layer Normalization的操作。

ReZero方法把  初始化为0，Fixup方法把  初始化为1，但是作者在实验中发现这些方法都不会使训练收敛。

图20 (c) 中的操作是把(a) 和 (b)的操作结合在一起了，发现(c)的操作是可行的。把  初始化为一个比较小的数值，这种办法在网络较深时也可以收敛。

图20 (d) 中的操作就是本文提出的LayerScale操作了。

具体做法是保持Layer Normalization，并对Self-attention或者FFN的输出乘以一个对角矩阵。目的是给不同的channel乘以不同的  值。比如在channel这个维度，一共有  个channel，为它们分别乘以  以后再与原来的  相加。

式中的  和  都是可学习的参数。在18层之前，它们初始化为0.1。若网络更深，则在24层之前初始化为  。若网络更深，则在之后更深的网络中初始化为  。这样做的目的是希望使得每个block在一开始的时候更接近Identity mapping，在训练的过程中逐渐地学习到模块自身的function。作者通过实验证明，以这种方式训练网络更容易，而且这样做比 (c)的方式更加灵活。

LayerScale这种方式不会影响网络的表达能力，因为它也可以看成是融入了之前的Self-attention或者FFN模块中。

具体方法2：class-attention layers：高效的处理class token的方式

下面介绍详细的CaiT的结构：

下图21中最左侧是ViT的结构。作者先指出了ViT的一个问题，即：ViT在优化参数时，其实是要class token及其后续层的对应的class embedding同时做到2件事情，即：

1. 引导attention过程，帮助得到attention map。

2. 这些token还要输入到classifier中，完成分类任务。

作者认为class embedding要同时做到这2件事其实有点矛盾，参数会朝着矛盾的方向优化。所以作者想能不能把class token往后放，如下图middle所示。为啥要往后放呢？主要是想让class token不参与前面的Self-attention，使得参数的优化方向尽量不产生矛盾。

如果把这种思想再进一步，就得到了CaiT的结构，如图21右侧图所示。

• self-attention与ViT一致，只是不再使用class token了。
• 在网络的最后2层，self-attention变成了一个叫class-attention的东西。它依然是attention的功能，只是只更新最左侧的class embedding，而不更新其它的  个patch embedding (被freeze住了)。换句话说，class-attention这个模块只有patch embedding传递给class embedding的信息，而没有class embedding传递给patch embedding的信息。而且只有class embedding会传给FFN。

仔细观察下这个class-attention层 (CA layer)的特点：patch embedding不更新，只更新class embedding，我们可以思考下作者为啥要这样设计，下面是我个人的理解：

首先你想完成分类任务，就得靠class embedding。class embedding后面一定要接MLP head完成分类任务。那么class embedding就得获得这个图片的全部信息。

而图片的信息都在patch embedding里面，所以我们需要一个layer来从patch embedding里面抽取信息给class embedding。所以这个CA层必须具有这个功能，即：

从patch embedding里面抽取信息给class embedding。

然后，按照上文的分析，为了使得参数的优化方向尽量不产生矛盾，class embedding就只干好分类的事就行，尽量别参与attention过程。你的任务就是分类的，不要参与其他活动比如attention过程，这样使得class embedding的角色更专一。所以：

class embedding的信息不能给patch embedding。

所以根据提炼出的这的2个要求，CA层的设计就容易理解了。

下面是把CA层用数学公式表达出来。

假设为了有  个head和  个patch，embedding dimension为  。

，计算CA层的  ：

式中  。

式中  。它加在了  上面得到这个CA层的输出。

实际上CA layer使用2层就够了，前面依旧是12个正常的Transformer Block。

计算复杂度也不高，因为普通的self-attention计算  ，而CA计算  ，复杂度由  变成了  。

Experiments

实验的超参数设计方面，作者在DeiT的基础上做了一些微小的变动。为了加速训练并优化显存占用，作者使用了Fairscale library，fp16精度。

作者观测到依旧使用DeiT的超参数在训练深层的ViT模型时会出现不收敛的现象，很难训练，尤其是在仅仅使用ImageNet数据集的时候。

实验1：使训练更稳定的方法

• 调整不同深度的drop-rate：

就是使得每一层的drop-rate随着层数的加深而线性增加。但是作者实验证明这种方法并没有太大的作用。最多训到18层深的ViT，若想把模型变得更深，则需要加大dropout-rate，如下图22左侧所示。实际实现时的不同层的dropout-rate设置为：

• Normalization可以使深层ViT收敛：

这个地方1的实验还是在展开图20的 (b) 和 (c)，(b)代表使用了Fixup 和 T-Fixup等技术，(c)代表作者又加上了Layer Normalization。作者发现只使用Fixup 和 T-Fixup等技术，模型无法收敛，只有加上了Layer Normalization以后才会收敛，甚至会提升模型性能，如下图22中间侧所示。

结果显示，Fixup 和 T-Fixup等技术在模型比较浅时与LayerScale效果相当，但是它们实现更复杂，因为需要为不同类型的layer设置不同的初始化方法。

• LayerScale的作用：

作者训练了一个36层的模型，并计算了每一层的  ，这个值越大，代表残差块的作用越大，代表模型距离Identity越远，其结果如下图23所示。图23中的上面2个图代表Self-attention，图23中的下面2个图代表FFN。橘色的图代表不使用LayerScale的结果，而蓝色的图代表使用了LayerScale的结果。实验发现使用了LayerScale之后这个比例变得更加uniform了，这个好处也是模型更易训练的原因之一。

实验2：Class-attention layers的作用

如下图24所示为CA层在不同设置下的性能对比。结论是使用12层Self-attention和2层CA可以得到最佳的性能。