加速DeiT-S 60%+吞吐量！腾讯优图提出高性能Transformer加速方法

©PaperWeekly 原创 · 作者 | 小马

单位 | FightingCV公众号运营者

研究方向 | 计算机视觉

写在前面

视觉 Transformer（ViT）最近引起了非常大的热度，但巨大的计算成本仍然是一个严重的问题。由于 ViT 的计算复杂度相对于输入序列长度是呈二次关系的，因此减少计算的主流范式是减少 token 的数量。

现有的设计包括结构化空间压缩，使用渐进收缩金字塔来减少大型特征映射的计算，以及动态删除冗余 token 的非结构化 token 剪枝。然而，现有的 token 剪枝存在两方面的局限性：1）由剪枝引起的不完全空间结构与 Transformer 常用的结构化空间压缩不兼容 ；2）需要一个耗时的预训练阶段 。

为了解决限制和扩展 token 剪枝的适用场景，作者提出了 Evo-ViT，这是一种用于视觉 Transformer 的自激励的慢-快 token 进化方法（self-motivated slow-fast token evolution approach） 。

具体来说，作者通过利用视觉 Transformer 固有的全局类注意来进行非结构化的实例级 token 选择。然后，作者用不同的计算路径来更新所选的信息 token 和非信息 token，即慢-快更新（slow-fast updating） 。由于慢-快更新机制维持了空间结构和信息流，Evo-ViT 可以从一开始就加速 Transformer 的训练过程。

实验结果表明，该方法在保持图像分类性能的同时，显著降低了视觉 Transformer 的计算开销。本文的方法将 DEIT-S 加速了 60% 以上的吞吐量，而只牺牲了 ImageNet-1K 上 0.4% 的 Top-1 精度，在精度和效率上都优于当前的 token 剪枝方法 。

论文标题：

Evo-ViT: Slow-Fast Token Evolution for Dynamic Vision Transformer

论文地址：

https://arxiv.org/abs/2108.01390

代码链接：

https://github.com/YifanXu74/Evo-ViT

Motivation

近年来，视觉 Transformer 在各种计算机视觉任务上表现出了强大的性能，如图像分类，物体检测，实例分割。将 Transformer 引入计算机视觉的原因在于卷积神经网络（CNN）所缺乏的特性，特别是建模长距离依赖关系的特性。然而，由于图像包含大面积的低级纹理和无信息区域，因此对图像 token 之间的长期依赖关系进行密集建模的计算效率是很低的 。

现有的方法遵循两个方向来解决 ViT 中 token 之间的长期依赖建模的低效问题，如上图所示。第一个方向（如第二条路径所示），是基于局部空间先验进行结构化压缩 ，如局部线性投影，Shift Windows 等。大多数具有深窄结构的 Transformer 都在这 条pipeline 内。然而，结构化压缩模型以相同的优先级处理信息对象 token 和非信息背景 token 。

因此，第二个方向是 token 修剪 ，以非结构化的方式识别和删除非信息 token（如上图的第一条路径所示）。这种方法通常是根据预训练模型的误差，来去除冗余的 token。也有方法提出了通过分层的方式去除冗余 token 来加速预训练的 Transformer 网络，并通过自蒸馏探索了一种数据端降采样策略。尽管有显著的加速，但这些非结构化 token 剪枝方法由于其空间结构和信息流不完整，存在两方面的限制：不适用于结构化的压缩 Transformer，也不能从头开始训练 。

在本文中，第三条路径如图所示。作者提出从训练过程的一开始，以动态数据依赖的方式来处理低效问题。作者认为无信息的 token，它们对最终预测的贡献很小，但计算占位符 token（placeholder token）的远程依赖时，依旧会带来计算成本。

不同于直接减少局部空间冗余的结构化压缩，作者提出以非结构化和动态的方式，将信息 token 与占位符 token区分开来，并使用不同的计算路径更新两种类型的 token 。不同于在预训练的网络中搜索冗余和剪枝，作者通过保留占位符 token，可以保持完整的空间结构和信息流。通过这种方式，本文的方法可以从训练的一开始就插入到各种 ViT 的结构中，进行加速 。

具体地说，本文提出了一种动态的自激励 Evo-ViT 进化方法。作者认为，由于 Transformer 的类 token（class token）可以深入了解图像 token 之间的全局依赖关系，并学习分类，因此它自然能够区分信息 token 记和占位符 token，这都是自激励的 。

以上图中的 DeiT 为例，作者发现 DeiT-T 的 class token 估计了每个 token 对于依赖关系建模和最终分类的重要性。特别是在更深的层（例如，layer10），class token 通常会用更高的 attention weight 增强信息性 token，这与可解释的可视化结果非常相似。在浅层（例如第 5 层），class token 的注意力相对分散，但主要集中在信息区域。

因此，利用 class token 确定信息性 token 和占位符 token 的优势，就可以保留占位符 token 确保 Transformer 浅层中的完整信息流动，以确保建模精度。确定这两种 token 后，将以“慢-快”的方式更新它们。具体而言，占位符 token 被总结为一个代表性 token，该 token 通过 Transformer 编码器与信息性 token 以缓慢而精细的方式进行更新进化。然后，利用进化的代表性 token 快速更新占位符 token 以获得更多的代表性特征。

方法

3.1. Preliminaries

ViT 提出了一种简单的 token 化策略，该策略通过将图像分为多个 patch 并将每个 patch 线性投影到 embedding 中来处理图像。额外的类 token（CLS）被添加到序列中，并用作全局图像表示。此外，由于 transformer 的编码器中的自注意是位置不可知的，并且视觉应用对位置信息的要求很高，因此 ViT 将位置 embedding 添加到每个 token 中，包括 CLS token。之后，所有 token 通过堆叠的 transformer 编码器传递，CLS token 用于最终分类。

transformer 由一系列堆叠的编码器组成，其中每个编码器由两个模块组成，即多头自注意（MSA）模块和 feed forward network（FFN）模块。FFN 模块包含两个带有激活函数的线性变换。MSA 和 FFN 模块都使用残差连接和 LayerNorm（LN）。给定 ViT 的输入 ，第 k 个编码器的处理可以表示为：

自注意（SA）模块使用三个可学习的线性映射，将输入序列投影到 query、key、value 向量（即 ）中。然后，通过以下方法计算序列中所有 value 的加权和：

值得注意的是，ViT 与 CNN 的一个非常不同的设计是 CLS token。CLS token 与每个编码器上的 patch token 相互作用，并总结了最终用于所有 patch token 的表示。将 CLS token 和 patch token 之间的相似性分数表示为类注意力 ，计算表示为：

其中，  是 CLS token 的 query 向量。

在 ViT 中，MSA 和 FFN 模块的计算成本分别为 和 。对于剪枝方法，通过剪枝 η% 的 token，至少可以减少 FFN 和 MSA 模块中 η% 的 FLOPS。本文的方法可以在达到相同的效率的同时，更适用于从头开始训练的模型（而非对预训练好的模型进行剪枝）。

3.2. Methodology

3.2.1. Overview

在本文中，作者的目标是从 Transformer 训练过程开始，解决每个输入实例的低效建模问题。如上图所示，Evo-ViT 的 pipeline 主要包含两个部分：结构保留 token 选择模块（Structure preserving token selection） 和慢-快 token 更新模块（slow-fast token updating module） 。在结构保留 token 选择模块中，信息 token 和占位符 token 由全局类注意力决定，并以不同的方式进行更新。

3.2.2. Structure preserving token selection

在这项工作中，作者保留所有的 token，并动态区分信息 token 和占位符 token，以实现完整的信息流。原因是修剪视觉 Transformer 的浅层和中间层的 token 并不简单，特别是在训练过程的开始。

作者用层间和层内的方式来解释这个问题。首先，浅层和中间层通常呈现出快速增长的特征表示能力。修剪 token 会带来严重的信息丢失。作者使用中心核对齐（CKA）相似度测量每层中间 token 特征和动态 CLS token 的相似性。

结果如上图（a）所示，当模型越深时，DeiT-T 的 token特征不断更新，最终的 CLS token 特征与浅层的 token 特征有很大的不同。这表明，浅层或中间层的特征表示编码不足，使得 token 剪枝相当困难 。第二，token 在浅层之间的相关性较低。作者评估了 DeiT-S 模型中不同 patch token query 中关于网络深度的皮尔逊相关系数（PCC），以显示冗余。如上图（b）所示，浅层的方差较大，证明了区分浅层特征冗余的难度较大 。

注意力权重是简单的方法来解释一个模型的决策，并获得关于 token 中信息传播的 insights。作者发现，全局类注意能够成为一个简单的度量，以帮助动态区分视觉 Transformer 中的信息 token 和占位符 token。在上图中（a），信息 token 与最终的 CLS token 具有较高的 CKA 相关性，而占位符 token 具有较低的 CKA 相关性。

根据类注意 ，作者选择了 Top-k 个类注意力得分的 token 作为信息 token。其余 N−k 个 token 为包含较少信息的占位符 token。与 token 修剪不同，占位符 token 被保留并快速更新，而不是删除 。

为了更好地捕获不同层中 token 之间的底层信息，作者提出了一种全局类注意。具体来说，作者设计了不同层次的类注意之间的残差连接，以促进具有正则化效应的注意信息流。表示如下：

其中 是第 k 层的全局类注意， 是第 k 层的类注意。

3.2.3. Slow-fast token updating

信息 token 和占位符 token 由全局类的注意决定，接着作者提出以一种“慢-快”的方式更新 token，而不是直接删除占位符 token。信息 token 通过 MSA 和 FFN 模块更新，而占位 token 总结为一个代表 token 并进行更新。

对于 N 个 patch token ，作者首先通过上述 token 选择策略将它们分成 k 个信息 token 和 N−k 个占位符 token 。然后，将占位符 token 聚合为一个具有代表 token ，如下所示：

然后，将信息性 token 和代表性 token 分别输入 MSA 和 FFN 模块，更新过程计算如下：

因此，信息性 token 和代表性 token 以一种缓慢而复杂的方式进行更新。

最后，通过 残差的方式快速更新占位符 token ：

因此，这里的占位符 token 通过代表性 token 进行快速而简单的更新。

实验

4.1. Main Results

4.1.1 Comparisons with existing pruning methods

上表显示了本文方法和现有剪枝方法的性能和吞吐量对比。

4.1.2 Comparisons with state-of-the-art ViT models

上表展示了本文方法和 SOTA 的视觉 Transformer 的对比，可以看出，本文的方法在保证性能的情况下，能够在很大程度上进行加速。

4.2. Ablation Analysis

4.2.1 Effectiveness of each module

上表显示了不同更新方式的性能对比，可以看出本文提出“慢-快”更新方式能够具备性能上的优势。

4.2.2 Different Token Selection Strategy

此外，作者还对比了其他 token 选择方法，来验证本文方法的有效性。可以看出，全局类注意在准确性和效率上都优于其他选择策略和常见的采样方法。

4.2.3 Visualization

上图展示了 token 选择的可视化结果。结果表明，一些信息 token，在早期被确定为占位符 token。随着训练 epoch 的增加，本文的方法在 token 选择方面逐渐趋于稳定。

4.2.4 Consistent keeping ratio

作者还设置了各层不同的 token 保持比，以研究 Evo-ViT 的最佳加速结构。当所有层的保持比设置为 0.5 时，能够获得最佳的计算量-准确率的 trade-off。

总结

在这项工作中，作者提出一种自激励的慢-快 token 进化（Evo-ViT）方法，来研究视觉 Transformer 的效率。作者提出了一个保留 token 选择和具有“慢-快”更新策略的结构，以充分利用完整的空间结构和信息流。对 DeiT 和 LeViT 的实验表明，所提出的 Evo-ViT 方法显著加速了各种 Transformer 结构，同时保持了相似的分类性能，特别是在密集输入的情况下。

更多阅读

#投 稿 通 道#

 让你的文字被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注 

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[e-mail] 

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

·