[Transformer]Evo-ViT：Slow-Fast Token Evolution for Dynamic Vision Transformer

Paper
Code

Abstract

Vision Transformer是目前的研究热点，但是其高昂的计算成本仍然是一个严重的问题，ViT的计算复杂度会随着输入序列的长度指数增长，因此目前主流的做法是减少token的数量，如使用结构化的空间压缩、搭建金字塔模型减少空间分辨率，或者动态删除一些冗余的token。

基于token剪枝目前存在两方面的局限性：一是剪枝会使得空间信息不完全，与结构化的空间压缩方法不兼容，但结构化的空间信息是当前深-窄Transformer的主流；二是这种方法通常需要耗时很久的预训练过程。
为了解决上述限制并且扩展token剪枝的应用场景。

本文提出一种简单但高效的slow-fast token进化修剪的方法，叫做Evo-ViT。具体来说就是基于原始的全局的cls 注意力来选择非结构化的instance级别的token，然后本文提出使用不同的计算路径来进行slow-fast token的更新，这样即保持了空间信息也加速了ViT的训练过程。
实验结果限制本文的方法可以有效降低Transformer的计算成本，比如增加了DeiT-S 60%的吞吐量，在ImageNet-1K上值减少了0.4%的top-1精度，在精度和效率上都优于当前的token剪枝方法。

Section I Introduction

ViT在计算机视觉领域大放异彩，因为它具有CNN缺少的建模长程依赖的能力，但是代价是计算复杂度的爆炸，因为图像中傲寒大量低级纹理和无信息的背景，因此图像token的密集计算是十分低效的。
当前解决ViT计算低效主流有两大方法，参见Fig1.


一是基于一定的局部空间先验知识进行结构化的压缩，Fig 1中间的那条路径，比如进行线性映射、卷积映射、或者基于窗口的计算，但是依旧糊将对象和背景赋予同样的优先级来进行修剪；二是token的剪枝，会以一种非结构化的方式将无用的token进行丢弃，主要基于预训练网络的重构来识别和去除冗余的标记，训练好的剪枝mask对所有instance都是固定的。尽管可以加速训练，但是由于其空间结构和信息流都是不完整的，不利于结构化的压缩Transformer，并且无法从头训练。



本文提出的方法参见Fig 1的第三条路径，从训练一开始就使用这种数据依赖的方式来动态处理，同时这是一种结构化压缩方法。本文会对一些对最终预测无用的token进行标记，这些无用的token在建模长程依赖时只是起占位的作用，只能增加计算成本。
本文会以一种非结构化、动态的方式来区分每个instance的有用token和无用token，然后使用不同的路径来更新这两种类型的token,而不是在预训练的网络中进行冗余修剪。因为保留了占位token，因此保持了完整的空间结构和信息流，这种方式使得本文的方法从训练的一开始就能成为大多数平面结构或者深层结构通用的插件。
这种自激励的动态token标记的fast-slow计算方法称之为Evo-ViT.
self-motivated表示Transformer可以自然的区分每个instance的信息token和占位token,这样可以洞悉token之间的依赖关系。为了不失一般性，本文采用DeiT作为基准框架。通过对DeiT进行可视化本文发现cls token包含更多的信息，并具有系数的注意力响应，因此使用cls token可以确定哪些是信息token哪些是占位token,这样可以保证Transformer浅层网络完整的信息流，在确定了两类token后会以一种slow-fast的方式进行更新。
具体来说，占位token会被精炼为一个具有代表性的token，会随着完整的Transformer以一种缓慢且复杂的方式进行更新，那些信息token则以一种快速的方式进行更新，来获得更具代表性的特征。
本文会在扁平的DeiT和深层且窄的LeViT上测试Evo-ViT的有效性，本文的Evo-ViT中的self-motivated能够提升40%-60%的吞吐率。

Section II Related Work

Vision Transformer

在ViT的一系列优化方法中，DeiT借助调整训练策略和知识蒸馏有效解决了ViT总数据利用低效的问题，为了实现更好的精度-速度权衡还有研究通过下采样、局部平均池化、卷积采样等来减少中间层的token数目。

Redundancy Reduction

计算成本一方面来自MHSA的二次计算复杂度另一方面FFN也增加了hidden dim的维度。目前的加速方法主要由：稀疏注意力、知识蒸馏。稀疏注意力由低秩分解、局部计算等，还有的工作提出在注意力矩阵中增加以了可以进化的全局注意力，基于此本文也提出进化的全局注意力来进行每一层的token selection。

近期相关的一项研究是Patch Sliming提出的一种自顶向下的patch瘦身方法，会计算特征重构误差来进行patch的瘦身，这样使得修剪后的网络依旧能以较低的计算承恩保证原有的性能。另外Dynamic ViT还提出一个轻量级的预测模块来估计每个token的重要性分数，这样可以非结构化的修剪每个token。与他们的区别在于本文提出的是保留占位token这样可以保证完整的空间结构，获得更高的性能；只不过以不同的计算路径来更新信息token和占位token；并且完整的信息流使得我们可以加速Transforemr的训练。

Section III Preliminaries

ViT会将图像展平成一系列的patch并且使用一个维阿的cls token作为全局的图像表示。此外由于Transformer中的自注意力是位置无关的因此ViT还是用位置嵌入，然后送入Transformer进行处理，最后cls token用来分类。


Transformer由一系列encoder堆叠而成，每个encoder包含MSA模块和FFN模块，FFN包含两次线性变换+激活函数，并且encoder模块还是用了残差连接，最后进行层归一化。整个处理流程表述为：

其中注意力模块会使用三个可学习的线性映射，WQ,WK,WV来计算注意力：

MSA是SA的延伸，会将Q,K,V复制n次，并行的计算注意力，然后级联作为最终的注意力输出。CLS token是Transformer区别于CNN的一点，cls token作为所有token的总结，会计算cls token与所有patch token之间的相似性分数：



计算复杂度分析
ViT中MSA和FFN的计算复杂度为：O(4NC2+2N2C)和O(8NC2)

本文可实现修剪η%的标记，至少可以在FFN和MSA模块中减少η%的FLOPs并且可以迁移到下游应用。

Section IV Methodology

Overview
本文致力于解决Transformer建模低效的问题，Evo-ViT的流程参见Fig 3，主要包含两部分：结构化token选择slow-fast token更新模块。
在token selection中或通过进化的全局class attention来分类informative token和placeholder token，二者会以不同的方式进行更新，比如占位token会汇总为一个代表token单独进行更新；informative token则是通过MSA-FFN进行更新，用以建模长程依赖和学习丰富的特征。
首先本文会详细介绍提出的结构化token保留、选择模块；然后再介绍如何进行slow-fast更新。

Part 1 Structure perserving token selection

本文建议保留所有token，但是会区分信息token和占位token，从而保证完整的信息流。因为在Vision Transformer中浅层和中间层修剪token后果十分严重，尤其是在训练开始阶段，会带来严重的信息丢失。
本文使用中心对齐（CKA）相似性来评测中间层token与最终cls token之间的相似性，如Fig 4(a)所示可以看到占位token和信息token与cls token的相似性还是有很大区别的，但是在网络浅层二者之间的区别还不明显，此时进行修剪会带来信息损失；
其次tokentoken query在网络浅层的相关性也比较低，本文评估了DeiT-S不同网络深度中token的皮尔森相关系数来显示冗余性，参见Fig 4(b)可以看到浅层相关性更低、方差更大，也表明了浅层特征的冗余性。
注意力权重是最流行的一种解释模型决策过程的一种方式，本文提出的evolved global class attention可以动态的区分信息token和占位token，因为信息token与cls token有更高的CKA相关性，占位token的相关性则较低。

本文会选择注意力分数的top-k个token，其他token被认为是包含较少信息的占位token，但是本文并不进行token的修剪，而是被保留这些占位token。本文为了更好的捕获不同层token之间的低层信息，使用了一个global class attention来增强跨层之间的信息流动，具体就是使用残差连接和一些正则化操作，这样可以增加训练稳定性和效率，只更新信息token部分的global class attention分数：

Part 2 Slow-fast updating

当确定好信息token和占位token后就已slow-fast的方式来分别更新，而不是直接将占位token进行删除。信息token会通过MSA和FFN进行细致的更新，占位token则被总结为一个总结性的token进行粗略的快速更新。



对于N个token，首先将他们分成k个信息token和N-k个占位token，其中占位token会被总结为一个representative token:
agg代表一个聚合函数，比如加权和或者转置线性投影等，本文使用的是基于计算出的全局注意力分数的加权和。




然后信息token和总结后的representative token会送入MSA和FFN模块，进行以下操作：


这样信息token和representative token会以一种缓慢而复杂的方式更新。
而占位token则是通过残差进行快速更新：


并且可以通过残差连接保证slow更新和fast更新在相同的数量级内。

Part 3 Training Strategies

Layer-to-stage training schedule


本文的token selection会随着训练进行越来越稳定，参考Fig 5可以发现随着训练进行不同层选的token越来越一致，说明Transformer倾向于利用计算资源使用信息token，因此本文提出一种分层到不同stage的训练策略进一步提升效率。

即在前200个epoch逐层进行token selection 在最后100个epoch则是在每个stage开始进行token selection，然后在每一层进行slow update，比如对于DeiT-T设置每四层为一个stage.


Assisted CLS token loss


为了训练更稳定本文在训练期间使用基于cls标记和最终的平均池化特征来计算损失函数：

Section V Experiments

本文会在ImageNet-1K上进行广泛的实验，基于DeiT和LeViT进行，token selection的比例设置为0.5，所有的模型会训练300个epoch.

Part 1 Comparisons with existing pruning methods.

Table 1展示了本文与其他token pruning的对比结果，可以看到以DeiT为例本文的方法可以将吞吐率提升60%而没有带来很明显的精度下降。




Comparisons with state-of-the-art ViTs

由于保留了占位token因此本文比其他方法拥有更多的空间结构，因此本文将Evo-ViT进一步应用到LeViT上，这是一个更加深且窄的网络，结果如Table 2所示，可以发现比DeiT精度下降的更大，原因是LevIt深层token较少，冗余的比较少，但是本文对于LeViT依旧有加速的效果，说明了本文方法对密集输入的有效性。

Part 2 Ablation Analysis

Effectiveness of each module


Table 3展示了消融实验的结果，分别对比：

原始的token选择：只根据每一层原始的类别注意力删除token;

结构化保留：保留占位token但不进行更新；

快速更新：快速更新保留的占位token;

layer-to-stage:本文提出的方法。


DeiT上的结果表明本文的方法可以有效的保持完整的信息流，进一步提高了选择的有效性；global class attention则增强了不同层选择的一致性，并取得了很好的性能；
快速更新策略对DeiT的影响小于对LeViT的影响，本文认为基于结构保留和全局分类注意力对DeiT已经达到了饱和，对LeViT仍有一定改进空间。



LeViT使用空间池化来减少token的数目，这使得每个阶段进行非结构化的token裁剪更加困难，通过fast update可以收集一些占位token的额外信息，来进一步增强特征表述；本文还评估了layer-to-stage的策略，发现能在保证准确性的同时进一步加速推理速度。




Different Token Selection Strategy




本文还对比了不同的token选择方法，结果参见Table 4.可以看到本文的evolved global class attention优于其他方法，虽然从Fig 2中可以看出深层次注意力更关注物体上但是最终精度却与有所下降，可能是因为网络需要一些背景信息来辅助分类，而在整个训练过程中限制所有层只关注对象导致对背景信息的保留不够。




Visualization

Fig 5可视化了token selection的结果，可以看到Evo-ViT针对的主要是对象而不是背景，也说明了这种方法可以有效的辨别信息token和占位token，不同层选择的结果趋于一致，证明了本文layer-to-stage的有效性。




另一个有趣的发现是本文的结构保留策略使得浅层遗漏的一些token在深层次中被检索到了，比如以棒球图像为例，随着网络成熟的加深，球棒的token被标记出来了。
Consistent Keeping Ratio

本文还在每一层设置了不同的标记保持比，来找到最优的加速结果，结果参见Fig 6，在所有层中保持token数目抑制然后调整每一层的keep ratio，发现在所有层中保持一致的keeping ratio效果是最好的，最终本文设置keeping ratio = 0.5.


Section VI Conclusions

本文提出一种自己离的slow-fast标记更新方式（Evo-ViT）来加速Transformer，有效避免了当前通过剪枝token的局限性。在广泛的实验上证明Evo-ViT可以显著加速各种Transformer并且没有带来明显的性能下降。

未来本文将探究将Evo-ViT迁移到下游任务的性能，如目标检测和实例分割。