【导读】 由于随着ViT中的token数量的增长，会导致计算成本呈平方级急剧增加！近期，清华黄高团队提出了自适应序列长度的DVT方案，本篇文章，我们将介绍另一种方法。6月3日，清华鲁继文教授团队提出了对Token进行动态稀疏化处理的方案，提出了一种动态Token稀疏化框架。看来最近针对Token上的优化将会出来一系列工作。我们不妨一起来期待一下！话不多说，下面来看看鲁继文老师DynamicViT这项工作吧！
链接：https://arxiv.org/abs/2106.02034
代码地址：https://github.com/raoyongming/DynamicViT
摘要： 在视觉Transformer中，注意力是稀疏的。因此，视觉Transformer的最终预测往往仅基于最具有信息量的token。所以，只要有效提取到这些最具信息量的token，那么我们就可以实现精准的图像识别任务。基于此，我们提出了一种动态Token稀疏化框架，该框架可以根据输入信息，逐步动态地裁减冗余token。具体地说，我们设计了一个轻量级预测模块，它可以对当前给定特征信息中每个token的重要性进行分数预估。同时该模块可以添加到不同的层，从而可以对每层冗余的token进行丢弃。此外，为了实现端到端的训练方式，我们进一步对该预测模块进行优化，提出了一种注意屏蔽策略。该策略可以阻止需要修剪的token与其它token之间的信息交互。另一方面，得益于自注意力自身的优点，非结构化稀疏token仍然是硬件友好型的，这使得我们的框架很容易在实际推理中进行模型加速。最终，通过对66%的输入token进行分层修剪，我们的方法大大减少了31% -37% FLOPs，提高了40%以上的吞吐量，与此同时，我们对各种视觉Transformer的精度损失都控制在0.5%以内。在ImageNet数据集上，与当前最好的CNN模型以及ViT相比，由于DynamicViT模型配备了动态token稀疏化框架，因此，它在复杂性/准确性之间实现了形成一种非常好的平衡。
简介 ：过去针对CNN架构的加速方法，现在也可以应用到ViT模型中。首先，由于ViT及其变体对输入图像进行处理的方式,都是将输入图像分割成多个独立的小块，这种分割方式，为我们提供了另一种正交的方法来进行token稀疏性处理。我们可以删除输入实例中不那么重要的token，因为许多token对最终预测的贡献很小。但这种方法仅适用于ViT模型，因为自注意力模块可以将可变长度的token序列作为输入，对输入进行非结构化的剪枝并不会影响自注意力模块。由于CNN的级联方式是一种下采样结构，该方式大大提高了模型的效率。因此，我们希望通过探索ViT的非结构化和依赖数据的下采样策略，从而进一步利用自我注意力的优势(我们的实验也表明，非结构化的稀疏可以比结构化降采样的方法，更能提高ViT模型的性能)。
我们的方法的基本思想如下图所示。
具体而言，我们使用一个轻量级预测模块以一种动态的方式选择需要修剪的token，这种方法我们称为DynamicViT。特别是，对于每个输入实例，预测模块将生成一个特定的二进制决策掩码，来确定哪些token是无信息的，需要裁剪掉的。该模块被添加到ViT的多个层中，这样我们可以在每个预测模块之后，逐渐增加剪枝token的数量时，继而实现分层地稀疏化处理。一旦一个token在某个层之后被删除，它将永远不会在后续过程中使用。这个轻量级模块引入的额外计算开销非常小，为了消除裁剪token所带来的计算开销。我们将该预测模块与ViT主干网络进行端到端共同优化。基于此，我们采取了两项策略。第一种是采用Gumbel-Softmax来克服分布采样不可微的问题，从而实现端到端训练。第二个是关于如何应用这个学习到的二进制决策掩码来删除不必要的token。由于二进制决策掩码中每个实例的零元素数量不同，在训练过程中通过直接消除每个输入实例的无信息token，并实现并行计算，几乎是不可能的。此外，这也会阻碍预测模块的反向传播，因为，预测模块需要计算token最终被消除后是否保留概率分布。而且，直接将裁剪的token设为零向量也是不明智的，因为零向量仍然会影响注意力矩阵的计算。所以，我们又提出了一种称为注意屏蔽的策略，即基于二进制决策掩码，将注意力矩阵中需要丢弃的token与所有其它token之间的连接删除掉。通过这种方式，我们可以解决上述困难。此外，我们还修改了ViT的原始训练目标，通过增加一个正则项来约束某一层后剪枝token的比例。在推理阶段，我们可以直接在每个输入实例的特定层之后放弃固定数量的token，因为我们不再需要考虑该操作是否可微，这将大大加快推理效率。最后，我们使用DeiT和LV-ViT作为骨干网络，在ImageNet上验证了我们这种方法的有效性。实验结果证明了在速度和精度之间的完美权衡下。通过分层修剪66%的输入token，我们可以大大减少31% - 37% GFLOPs，提高吞吐量超过40%，而准确性的下降控制在0.5%以内，针对所有不同的ViT模型。我们的DynamicViT演示了利用空间稀疏性对基于ViT变体进行加速的可能性。我们希望，这种尝试可以为后面的基于ViT模型的加速的工作提供一种新的思路。
Method
DynamicViT的总体框架如图2所示。DynamicViT由一个普通的ViT作为backbone和几个预测模块组成。backbone可以采用ViT， DeiT， LV-ViT等等。预测模块负责生成丢弃/保存token的概率。Token稀疏化是在整个网络的特定位置分层执行的。例如，给定一个12层的变压器，我们可以在第4、7、9块之前进行token稀疏化。在训练过程中，由于我们新设计的注意掩蔽策略，预测模块和骨干网络可以同时进行端到端优化。在推理过程中，我们只需要根据预定义的修剪比率和预测模块计算的分数选择信息量最大的标记。
预测模块中的分层Token稀疏化
在这里，主要通过一个二进制决策掩码 来确定每个token是保留还是丢弃。首先将决策掩码中的所有元素初始化为1，并逐步更新掩码。预测模块以当前决策 和token 作为输入。我们首先使用MLP对token进行变换:
其中 可以是一个更小的维度，在我们的实现中，使用 。类似地，我们可以计算一个全局特征:
其中Agg是一个集合所有保留的token信息的函数，可以通过一个简单地平均池化来实现:
因为，局部特征是对特定标记的信息进行编码，而全局特征则包含整个图像的上下文信息，两者都是有用的。因此，我们通过结合局部和全局特征来获得局部-全局嵌入，并将它们放入到另一个MLP中来预测丢弃/保留token的概率:
式中 表示丢弃第i个token的概率， 表示保留它的概率。然后，我们可以通过从 中取样产生当前的决策D，并更新
其中 是Hadamard Product，这表明一旦一个令牌被丢弃，它将永远不会被使用。
端到端优化
虽然我们的目标是将token稀疏化，但我们发现在训练过程中实现它并不简单。首先，从 中采样得到的二值决策掩码 是不可微的，这就阻碍了端到端的训练。为了克服这个问题，我们应用Gumbel-Softmax技术从概率 中取样:
这里我们使用索引“1”，因为D表示保留的token的决策掩码。Gumbel-Softmax的输出是一个one-hot 向量，其期望正好等于 。同时，Gumbel-Softmax具有可微性，这就可以实现端到端训练。第二个问题，当我们试图在训练期间修剪token时。决策掩码 通常是非结构化的, 因此,不同样本的掩码包含不同的"1"。采取简单地丢弃 的令牌就会导致一批样本的token数量不一致，这使得并行化计算会十分困难。因此，我们必须保持令牌的数量不变，同时减少修剪后的token与其它token之间的交互。我们还发现，仅仅用二进制掩码去掉令牌是不可行的，因为在计算自注意力矩阵时：
零向量的token仍然会通过Softmax操作影响其它token。为此，我们设计了一种名为注意力屏蔽的策略，可以完全消除丢弃token带来的的影响。具体来说，我们通过以下方法来计算注意力矩阵:
这里，我们构造了一个图，其中 表示第 个token将有助于第 个token的更新。请注意，我们显式地为每个token添加了一个自循环，以提高数值稳定性。也很容易证明自循环不会影响结果:如果 ，第 个token除了它自己之外不会对任何token做出贡献。这里计算的是掩体注意力矩阵 ，它只考虑保留token的注意力矩阵，但在训练过程中其大小保持不变 。
训练与推理
DynamicViT的训练包括训练预测模块，使其能够产生有利的决策，并对主干进行微调，使其适应令牌稀疏化。假设我们处理的是一小批B样本，我们采用标准交叉熵损失:
表示DynamicVit预测值， 代表真实值。为了尽量减少token稀疏化对性能的影响，我们使用原来的骨干网络作为Teacher model，并希望我们的DynamicViT的行为尽可能接近这个model。具体来说，我们主要从两个方面来考虑这个约束。首先，我们使DynamicViT的最后剩余token数量接近Teacher mode，这可以看作是一种自我蒸馏:
第二，我们通过KL散度来最小化我们的DynamicViT和它的teacher模型之间的预测的差异:
最后，我们希望将保留的token的比率设置为一个预定义值。给定S阶段的一组目标比率 ，我们利用MSE损失来对预测模块进行有监督训练:
整个训练损失函数如下：
在实验中，我们设置
在推断期间，给定目标比率 ，我们可以通过预测模块产生的概率直接丢弃信息较少的token，这样只有确切 的token会在 阶段被保留。形式上，对于 阶段，我们让
为保持概率 排序的索引,我们将需要保留的token索引所在位置的标记为 ，而丢弃其他标记。通过这种方式，我们的DynamicViT在运行时动态地删除信息较少的标记，从而可以减少推断期间的计算成本。动态调整 的结果如下：
实验结果
DynamicViT的最大的一个优点是，它可以应用于广泛的ViT架构中，同时减少计算复杂度和实现较少的性能损失。我们将DynamicViT模型与目前最好的图像分类模型在FLOPs和参数数量进行了比较。并且在ImageNet数据集比较了Top-1 acc（%），如下所示：
我们的DynamicViT-LV-M/0.7 打败了EfficientNet-B5和NFNet-F0，如图3所示，我们画出了他们的FLOPS-accuracy变化曲线图,图4代表的是调整模型width的变化情况:
最后，我们对逐层稀疏化处理token的结果进行了可视化
总结
在本项工作中，我们利用输入图像中信息分块的稀疏性，为ViT模型加速开辟了一条新的路线。对于每个输入实例，我们的DynamicViT模型可以根据来自轻量级预测模块的定制二进制决策掩码输出概率结果，以动态的方式删除重要性较低的token，同时该模块也融合了token中包含的局部和全局信息。预测模块可以通过添加到多个层中，从而实现按分层的方式进行token剪枝操作。此外，我们还将Gumbel-Softmax和注意力屏蔽技术与预测模块一起用于ViT模型的端到端训练。在推理阶段，我们的方法可以通过逐步修剪66%的输入token来大大提高效率，应用在不同的ViT骨架网络中，准确率下降也严格控制在0.5%以内。本文中，我们重点研究了图像分类任务。后续，也可以将我们的方法扩展到其它场景，如视频分类和密集预测任务。大家可以去尝试一下哟！
重磅！DLer-Transformer&CV交流群已成立！
欢迎各位Cver加入Transformer微信交流群，本群旨在交流模型压缩/量化/剪枝、Transformer、多模态学习、自监督学习、无监督学习、元学习等内容。欢迎对这些研究方向感兴趣的小伙伴加群一起交流学习！
加群请备注：研究方向+学校/公司+昵称（如Transformer+上交+小明）
???? 长按识别添加，邀请您进群！
扫一扫
843
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
