DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification

thalch

已于 2025-01-02 10:11:41 修改

阅读量97

点赞数

文章标签：深度学习人工智能计算机视觉

于 2023-03-27 10:16:47 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_41878880/article/details/129789333

版权

视觉转换器的注意力是稀疏的。我们观察到视觉变形器的最终预测仅基于信息量最大的标记的子集，这对于准确的图像识别是足够的。在此基础上，我们提出了一种动态令牌稀疏化框架，根据输入逐步动态地修剪冗余令牌。具体来说，我们设计了一个轻量级的预测模块，在给定当前特征的情况下估计每个令牌的重要性得分。

该模块被添加到不同的层，以分层地删除冗余令牌。为了以端到端方式优化预测模块，我们提出了一种注意力屏蔽策略，通过阻止一个令牌与其他令牌的交互来区分地修剪一个令牌。得益于自关注的性质，非结构化稀疏令牌仍然是硬件友好的，这使得我们的框架很容易实现实际的加速。通过对66%的输入标记进行分层修剪，我们的方法大大降低了31% ~ 37%的FLOPs，提高了40%以上的吞吐量，而对各种视觉转换器的精度下降在0.5%以内。DynamicViT模型配备了动态令牌稀疏化框架，与ImageNet上最先进的cnn和视觉转换器相比，它可以实现非常有竞争力的复杂性/准确性权衡。

近年来，CNNtype体系结构的发展给计算机视觉带来了巨大的进步[12,18]。最近的一些工作开始使用transformer取代CNN，用于许多视觉任务，如物体检测[36,20]和分类[25]。就像过去几年对cnn类型的体系结构所做的一样，还需要加速类似变压器的模型，使其更适合实时应用程序。

加速cnn类型网络的一种常见做法是修剪不太重要的过滤器。视觉转换器及其变体处理输入的方式，即将输入图像分割成多个独立的补丁，为我们提供了另一种引入加速度稀疏性的正交方法。也就是说，我们可以修剪输入实例中不太重要的标记，因为许多标记对最终预测的贡献很小。这仅适用于类变压器模型，自注意模块可以将变长令牌序列作为输入，非结构化修剪后的输入不会影响自注意模块，而去掉某一部分像素并不能真正加速卷积运算，因为卷积使用的非结构化邻域会使并行计算难以加速。由于具有结构下采样的cnn的层次结构提高了模型在各种视觉任务中的效率，我们希望探索非结构化和数据依赖的视觉变压器的下采样策略，进一步利用自注意的优势(我们的实验也表明，与结构下采样相比，非结构化稀疏可以导致视觉变压器更好的性能)。我们的方法的基本思想如图1所示。

在这项工作中，我们建议使用一个轻量级的预测模块来确定哪些令牌需要以动态方式修剪，称为DynamicViT。特别是，对于每个输入实例，预测模块都会生成一个定制的二进制决策掩码，以决定哪些令牌是没有信息的，需要放弃。该模块被添加到视觉转换器的多个层中，这样当我们在每个预测模块之后逐渐增加修剪令牌的数量时，可以以分层的方式进行稀疏化。一旦标记在某一层之后被修剪，它将永远不会在前馈过程中使用。这个轻量级模块引入的额外计算开销非常小，特别是考虑到消除无信息标记所节省的计算开销。

该预测模块可以与视觉转换器骨干网端到端共同优化。为此，采用了两种专门的策略。第一种是采用Gumbel-Softmax[15]，克服了从分布中采样的不可微问题，从而可以进行端到端训练。第二部分是关于如何应用这个学习到的二进制决策掩码来修剪不必要的标记。考虑到二进制决策掩码中的零元素数量对于每个实例都是不同的，在训练过程中直接消除每个输入实例的无信息令牌将使并行计算成为不可能。此外，这也会阻碍预测模块的反向传播，它需要计算是否保留令牌的概率分布，即使它最终被消除。此外，直接将废弃的代币设为零向量也是不明智的，因为零向量仍然会影响注意力矩阵的计算。因此，我们提出了一种称为注意力屏蔽的策略，其中我们基于二进制决策掩码将从废弃标记到注意力矩阵中所有其他标记的连接删除。这样做，我们就能克服上述困难。我们还修改了视觉转换器的原始训练目标，在某一层之后增加了一项来约束修剪标记的比例。在推理阶段，我们可以在每个输入实例的某些层之后直接放弃固定数量的令牌，因为我们不再需要考虑操作是否可微，这将大大加快推理的速度。

我们使用DeiT[25]和LV-ViT[16]作为主干，说明了我们的方法在ImageNet上的有效性。实验结果证明了速度和准确性之间的竞争权衡。特别是，通过对66%的输入标记进行分层修剪，我们可以大大降低31% ~ 37%的GFLOPs，并将吞吐量提高40%以上，而对于所有不同的视觉转换器，精度下降在0.5%以内。我们的DynamicViT演示了利用空间稀疏性来加速类变压器模型的可能性。我们期望我们的尝试能为今后类似变压器模型的加速工作开辟一条新的道路。

视觉变形金刚。变压器模型首先在NLP社区[26]中被广泛研究。证明了用自我注意代替递归神经网络及其变体的可能性。最近的进展表明，变压器的变体也可以成为cnn的竞争替代品，并在图像分类[8,25,20,35,23]、物体检测[2]、语义分割[34,5]和3D分析[31,33]等不同视觉任务上取得了有前景的结果。DETR[2]是第一个将变压器模型应用于视觉任务的工作。它将目标检测任务制定为集合预测问题，并在变压器中遵循编码器-解码器设计生成边界框序列。ViT[8]是第一个将变压器架构直接应用在非重叠图像补丁上进行图像分类任务的工作，整个框架不包含卷积操作。与cnn类型的模型相比，ViT在大规模的预训练下可以获得更好的性能。如果架构可以在没有任何预先培训的情况下达到最先进的水平，那就更好了。DeiT[25]提出了许多训练技术，使我们可以只在ImageNet1K[7]上训练无卷积变压器，并获得比ViT更好的性能。LV-ViT[16]通过引入一种称为令牌标记的新训练目标进一步提高了性能。ViT及其后续技术都将输入图像分割为多个独立的图像补丁，并将这些图像补丁转换为令牌进行进一步处理。这使得在这些类变压器模型中考虑空间维度上的稀疏性是可行的。

加速度模型。模型加速技术对于在边缘设备上部署深度模型非常重要。有很多技术可以用来加快深度模型的推理速度，包括量化[9,27]、剪枝[13,22]、低秩因子分解[30]、知识蒸馏[14,19]等。也有许多工作旨在加快变压器模型的推理速度。例如TinyBERT[17]提出了一种蒸馏方法来加速变压器的推断。星形变换器[10]通过用星形拓扑取代全连接结构，将二次空间和时间复杂度降低为线性。然而，这些工作都集中在NLP任务上，很少有工作探索利用视觉任务的特性来加速视觉转换器的可能性。此外，Transformer和CNN特性之间的差异也使得我们可以采用另一种方式进行加速，而不是使用CNN加速的方法，如filter pruning[13]，从深度模型中去除非关键或冗余的神经元。我们的方法旨在通过利用信息图像补丁的稀疏性来修剪不太重要的标记，而不是神经元。

我们的DynamicViT的总体框架如图2所示。我们的DynamicViT由一个普通的视觉转换器作为骨干和几个预测模块组成。骨干网可实现多种视觉转换器(如ViT[8]、DeiT[25]、LV-ViT[16])。

预测模块负责生成丢弃/保留令牌的概率。令牌稀疏化在整个网络的特定位置按层次执行。例如，给定一个12层的变压器，我们可以在第4、7和10块之前进行令牌稀疏化。在训练过程中，由于我们新设计的注意力掩蔽策略，预测模块和骨干网可以以端到端的方式进行优化。在推理过程中，我们只需要根据预定义的剪枝比例和预测模块计算出的分数来选择信息量最大的令牌
在这里插入图片描述

我们的DynamicViT的一个重要特征是标记稀疏化是分层执行的，也就是说，随着计算的进行，我们逐渐删除无信息的标记。为了实现这一点，我们维护了一个二进制决策掩码D∈{0,1}N来表示是删除还是保留每个令牌，其中N = HW是补丁嵌入的数量s2。我们将决策掩码中的所有元素初始化为1，并逐步更新掩码。预测模块将当前决策D和令牌x∈RN×C作为输入。我们首先使用MLP投射令牌:
在这里插入图片描述
其中C0可以是更小的维度，我们在实现中使用C0 = C/2。类似地，我们可以通过

其中Agg是聚合所有现有令牌信息的函数，可以简单地实现为平均池:

局部特征编码某个标记的信息，而全局特征包含整个图像的上下文，因此两者都是有信息的。因此，我们结合局部和全局特征来获得局部-全局嵌入，并将它们提供给另一个MLP来预测丢弃/保留令牌的概率:
在这里插入图片描述
其中πi,0表示丢弃第i个令牌的概率，πi,1是保留它的概率。

然后，我们可以通过从π取样生成当前决策D，并通过
在这里插入图片描述
是Hadamard产品，表示一旦一个令牌被丢弃，它就永远不会被使用。

虽然我们的目标是执行令牌稀疏化，但我们发现在训练期间在实践中实现它并不简单。首先，从π采样得到二进制决策掩码D是不可微的，这阻碍了端到端训练。为了克服这个问题，我们应用Gumbel-Softmax技术[15]从概率π中取样:
在这里插入图片描述
这里我们使用索引“1”，因为D代表保留的令牌的掩码。GumbelSoftmax的输出是一个单热张量，其期望正好等于π。同时，Gumbel-Softmax是可微的，这使得端到端训练成为可能。

第二个障碍是当我们尝试在训练期间修剪代币时。决策掩码D通常是非结构化的，不同样本的掩码包含不同数量的1。因此，简单地丢弃在 Di = 0 `处的令牌将导致批处理中样本的令牌数量不一致，这使得计算难以并行化。因此，我们必须保持令牌的数量不变，同时减少修剪后的令牌与其他令牌之间的交互。我们还发现，仅仅使用二进制掩码D将要丢弃的令牌归零是不可行的，因为在自注意矩阵[26]的计算中
在这里插入图片描述
已归零的令牌仍将通过Softmax操作影响其他令牌。为此，我们设计了一种称为注意力掩蔽的策略，可以完全消除掉的代币的影响。具体地说，我们用

由式(10)我们构造了一个图，其中Gij = 1表示第j个令牌将有助于第i个令牌的更新。注意，我们显式地为每个令牌添加了一个自循环，以提高数值稳定性。也很容易表明自循环不影响结果:如果mj = 0，第j个令牌将不会贡献给除它自己以外的任何令牌。式(11)计算掩码注意矩阵~ A，它等价于仅考虑保留的标记计算的注意矩阵，但在训练过程中具有恒定的形状N × N。–类似于mask矩阵

损失包括三个方面：
1、真值损失
2、蒸馏损失：蒸馏损失的理解是要对网络的裁剪加上损失引导，由于没有真值，也就是没有明确的该保留哪些token的真实值，所以使用蒸馏。教师-学生模型类似于半监督？但在点云配准里，可以使用重叠损失引导采样。
3、数量限制

我们现在描述了DynamicViT的培训目标。DynamicViT的训练包括训练预测模块，使它们能够产生有利的决策，并对骨干进行微调，使其适应令牌稀疏化。假设我们处理的是一小批B样本，我们采用标准交叉熵损失:
在这里插入图片描述
其中y是DynamicViT的预测(在softmax之后)，¯y是ground truth。

为了尽量减少我们的令牌稀疏化对性能的影响，我们使用原始骨干网作为教师模型，并希望我们的DynamicViT的行为尽可能接近教师模型。
具体来说，我们从两个方面考虑这个约束。首先，我们使DynamicViT最后剩下的标记接近教师模型的标记，这可以被视为一种自我蒸馏:
在这里插入图片描述
其中ti和t0i分别表示DynamicViT和教师模型的最后一个块之后的第i个标记。Db,s为第s个稀疏阶段第b个样品的决策掩码。
其次，我们通过KL散度最小化我们的DynamicViT和它的老师之间的预测差异:

其中y0是教师模型的预测。

最后，我们希望将保留的令牌的比例限制为一个预定义的值。给定一组S阶段的目标比率ρ = [ρ(1)，…， ρ(S)]，我们利用MSE损失来监督预测模块:
在这里插入图片描述

在推理过程中，给定目标比率ρ，我们可以通过预测模块产生的概率直接丢弃信息量较小的令牌，这样只有精确的ms = bρsNc令牌被保留在s阶段。形式上，对于第s阶段，令
如果指数按保持概率π *，1排序，则我们可以保留指数在Is1:ms中的记号，而丢弃其他记号。通过这种方式，我们的DynamicViT在运行时动态删除信息量较小的令牌，从而可以减少推理期间的计算成本。