还在用全部token训练ViT？清华&UCLA提出token的动态稀疏化采样，降低inference时的计算量

详细信息如下：

• 论文链接：https://arxiv.org/abs/2106.02034

• 项目链接：https://github.com/raoyongming/DynamicViT

导言：

在Vision Transformer中，Attention其实是比较稀疏的。作者观察到，在视觉Transformer中的最终预测仅基于信息最丰富的token的一个子集，这足以进行准确的图像识别。基于这一观察结果，作者提出了一个动态token稀疏化框架，基于输入逐步、动态地删除冗余的token。

具体来说，作者设计了一个轻量级的预测模块，来估计给定特征的每个token的重要性分数。该模块被添加到不同的层中，分层地删除冗余token。为了以端到端的方式优化预测模块，作者提出了一种attention masking策略，通过阻止一个token与其他token的交互来修剪token。

得益于 self-attention的机制，非结构化的稀疏token对硬件是友好的，这使得本文的框架很容易实现实际的加速。通过分层修剪66%的输入token，本文的方法大大减少了31%∼37%FLOPs，提高了40%以上的吞吐量，而各种视觉Transformer的精度下降在0.5%以内。

      01      

Motivation

这些年来，CNN架构的发展带来了计算机视觉带来的巨大进步。最近的一些工作开始取代CNN，通过使用Transformer来完成许多视觉任务，如目标检测和图像分类。就像过去几年对CNN的架构所做的那样，加速Transformer的模型，使它们更适合实时应用，也是非常有意义的。

CNN网络加速的一个常见做法是修剪不那么重要的滤波器。而视觉Transformer及其变体处理输入的方式是将输入图像分割成多个独立的patch，这提供了新的的方法来引入加速的稀疏性。也就是，可以删除输入实例中不那么重要的token，因为许多token对最终预测的贡献很小。

当然，这只适用于类似Transformer的模型，因为其中的Self-Attention模块可以以可变长度的token序列作为输入，而非结构化修剪的输入不会影响Self-Attention模块。但是，丢弃某一部分像素并不能加速卷积操作，因为卷积所使用的非结构化邻域将使它很难通过并行计算来加速。在本文中，作者探索视觉Transformer的非结构化和依赖数据的降采样策略，以进一步利用Self-Attention的优势。

如上图所示，CNN模型通常利用结构化降采样策略来构建层次架构，如(a)所示。(b)中的非结构化和数据依赖的降采样方法可以更好地利用输入数据的稀疏性。由于自注意力操作的性质，非结构化的token集也很容易通过并行计算来加速。(c)可视化了每个空间位置对DeiT-S模型中最终预测的影响。这些结果表明，视觉Transformer的最终预测仅基于信息最丰富的token的一个子集，因此大量的token可以在不损害性能的情况下被删除。

在这项工作中，作者使用了一个轻量级的预测模块来确定哪些token将被动态地修剪，称为DynamicViT。对于每个输入实例，预测模块生成一个定制的二进制决策mask，以决定哪些token没有信息，需要丢弃。该模块被添加到视觉Transformer中，随着每个预测模块后逐步增加修剪token的数量，可以分层进行稀疏化。一旦token在某一层之后被修剪，它将永远不会在FFN过程中使用。与通过消除无信息token而节省的计算开销相比，这个轻量级模块引入的额外计算开销相当小。

为了能够与视觉Transformer主干一起以端到端方式联合优化，作者采取了两种专门的策略。

第一个方法是采用Gumbel-Softmax 来克服从分布中采样的不可微问题，从而进行端到端训练。

第二个方法是关于如何应用这个学习到的二进制决策mask来修剪不必要的token。考虑到二进制决策mask中的零元素数量是不同的，在训练过程中直接消除每个输入实例的无信息token将使并行计算是不可能的。此外，这也将阻碍预测模块的反向传播，因为即使某个token最终被消除，预测模块需要计算是否保留该token的概率分布。此外，直接将废弃的token设置为零向量也不行，因为零向量仍然会影响注意力矩阵的计算。因此，作者提出了一种称为注意力mask 的策略，即基于二进制决策mask，将从丢弃的token到注意力矩阵中所有其他token的连接删除。

此外，作者还修改了视觉Transformer的原始训练目标函数，通过添加一个项来限制特定层后的修剪token的比例。在推理阶段，可以直接在每个输入实例的特定层后放弃固定数量的token，因为这时已经不再需要考虑该操作是否可微，这将大大加快推理速度。

作者在 DeiT和 LV-ViT模型上基于ImageNet数据集上做了实验，结果表明通过分层修剪66%的输入token，可以大大减少31%∼37% GFLOPs和40%的吞吐量，精度下降在0.5%以内。

      02      

方法

2.1 Overview

模型总体框架如上图所示。DynamicViT由一个普通的视觉Transformer作为主干和几个预测模块组成。预测模块负责生成删除/保留token的概率。token稀疏化通过在整个网络的某些位置分层执行。例如，给定一个12层Transformer，作者可以在第4、第7和第9个block之前进行token稀疏化。在训练过程中，由于注意力mask策略，预测模块和主干网络可以进行端到端的优化。在推理过程中，只需要根据预定义的剪枝比和预测模块计算的分数来选择信息最多的token，就可以降低计算量和推理时间。

2.2 Hierarchical Token Sparsification with Prediction Modules

DynamicViT的一个重要特征是，token稀疏化是分层执行的，也就是说随着计算的进行，网络逐步删除无信息的token。为了实现这一点，作者维护了一个二进制决策掩码来指示是删除或保留每个token，其中N=HW为patch的数量。

首先将决策掩码中的所有元素初始化为1，并逐步更新掩码。预测模块以当前决策和token 作为输入。首先使用一个MLP来映射token到一个新的维度：

类似地，可以通过以下方法来计算一个全局特征：

Agg是聚合所有现有token信息的函数，可以用AvgPool实现：

局部特征对某个token的信息进行编码，而全局特征包含整个图像的上下文，因此它们都是有信息的。因此，作者结合局部和全局特征，获得局部-全局embedding，并将它们提供给另一个MLP，以预测删除/保留token的概率：

其中，π，表示丢弃第i个token的概率，π，是保留它的概率。

然后，通过从π中采样来生成当前的决策D，并通过更新：

一旦一个token被删除，它将永远不会被使用。

2.3 End-to-end Optimization with Attention Masking

虽然本文的目标是执行token稀疏化，但作者发现在训练过程中实现并不简单。首先，从π得到二进制决策掩码D的采样是不可微的，这阻碍了端到端的训练。为了克服这一问题，作者使用Gumbel-Softmax技术从概率π中采样：

第二个困难是在训练中修剪token时，决策mask通常是非结构化的，不同样本的mask包含不同数量的1。因此，简单地丢弃的token就会导致Batch内样本的token数量不一致，这使得计算难以并行化。因此，必须保持token的数量不变，同时减少修剪后的token和其他token之间的相互作用。作者发现，仅仅根据二进制掩码将需要丢弃的token置为零是不可行的。

因为如果直接将Attention置为0，在Softmax的时候还是会有影响。为此，作者设计了一种注意力mask的策略，它可以完全消除被丢弃的token的影响。具体来说，可以通过以下方法来计算注意矩阵：

（这一步之所以不用-inf直接对Attention Map进行mask，而是在Softmax中加入了乘上mask map的操作，是因为这里作者还需要训练预测mask的结构。如果直接对Attention Map用-inf来mask，这一步是不可导的，就不能进行端到端的训练了。 ）

2.4 Training and Inference

DynamicViT的训练包括训练预测模块，使它们能够产生有利的决策，并微调主干网络，使其适应token稀疏化。交叉熵损失函数如下：

其中为预测的概率，为ground truth。

为了最小化token稀疏化对性能的影响，作者使用原始的主干网络作为教师模型，并希望DynamicViT的行为尽可能接近教师模型。具体来说，作者从两个方面来考虑这种约束。首先，将剩余的token接近教师模型的token，这可以看作是一种自蒸馏：

其中，和 分别表示DynamicViT和教师模型的最后一个block之后的第i个token。，是在第s个稀疏化阶段的第b个样本的决策mask。

其次，通过KL散度来最小化DynamicViT与其老师之间的预测的差异：

其中，是教师模型的预测结果。

最后，作者还约束了保留的token与预定义值的比值。给定S阶段的一组目标比率ρρ（），，ρ，作者使用了MSE损失来监督预测模块：

最后总的损失函数为：

在预测的时候，作者就是根据比例来挑选需要丢弃的token，用剩下的token的预测：

通过这种方式，DynamicViT在运行时动态地删除信息较少的token，从而可以降低推理过程中的计算成本。

      03      

实验

3.1 Main results

DynamicViT可以将计算成本降低31%∼37%，并在运行时加速推理提高43%∼54%，而性能的影响可以忽略（−0.2%∼−0.5%）。

3.2 Comparisons with the-state-of-the-arts

我们还可以观察到，DynamicViT比LV-ViT系列可以实现更好的trade-off。

3.3 Analysis

DynamicViT for model scaling

通过将DynamicViT应用于DeiT-256，可以获得了一个与DeiT-Ti具有相当的计算复杂度的模型，但ImageNet的Top-1准确率高了4.3%左右。

Visualizations

上图展示了被丢弃的token分布，可以看到，DynamicViT可以逐渐删除无信息标记。

作者使用DynamicViT生成ImageNet验证集中所有图像的决策，并计算所有三个阶段中每个token的保持概率。中间的token往往被保留下来，因为在大多数图像中，物体位于中心。

Comparisons of different sparsification strategy

上表展示了不同采样策略的结果，可以看出，动态稀疏化采样的效果是比较好的。

      04      

总结

在这项工作中，作者利用输入图像中patch信息的稀疏性，开辟了一条加速视觉Transformer的新路径。对于每个输入实例，DynamicViT模型根据从轻量级预测模块输出的定制的二进制决策mask，以动态的方式修剪不太重要的token，它融合了token中包含的局部和全局信息。作者将预测模块添加到多层，以便以分层方式执行token剪枝。为了实现端到端的训练，作者还是用了Gumbel-Softmax和Attention masking。在推理阶段，DynamicViT可以通过逐步修剪66%的输入token来大大提高效率，而不同Transformer主干网络的精度下降小于0.5%。

作者介绍

研究领域：FightingCV公众号运营者，研究方向为多模态内容理解，专注于解决视觉模态和语言模态相结合的任务，促进Vision-Language模型的实地应用。

知乎/公众号：FightingCV

END,入群👇备注：TFM