DynamicViT：动态Token稀疏化实现高效的视觉Transformer系列

作者 | 科技猛兽  编辑 | 极市平台

点击下方卡片，关注“自动驾驶之心”公众号

ADAS巨卷干货，即可获取

点击进入→自动驾驶之心技术交流群

后台回复【数据集下载】获取Nuscenes、KITTI、COCO等计算机视觉近30种数据集！

导读

 

来自清华大学与UCLA的工作，提出了一种基于动态token稀疏化的高效视觉transformer，通过分层剪枝66%的输入tokens，可减少31%~37%的FLOPs，并将模型运行速度提高了40%以上，保证精度下降在0.5%以内，可应用于各种视觉 transformer模型中。

本文目录

50 DynamicViT：动态 Token 稀疏化实现高效的视觉 Transformer
(来自清华大学，周杰，鲁继文团队，UCLA)
50.1 DynamicViT 论文解读
50.1.1 背景和动机
50.1.2 DynamicViT 方法简介
50.1.3 通过预测模块进行 Token 的分层稀疏化
50.1.4 通过注意力掩码进行端到端的训练
50.1.5 Dynamic ViT 的训练和推理
50.1.6 实验结果

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型，现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制，不采用 RNN 的顺序结构，使得模型可以并行化训练，而且能够拥有全局信息。

这篇文章观察到视觉 Transformer 模型的最终预测结果往往是由一小部分信息量比较大的 tokens 决定的，而大部分的 tokens 其实是冗余的。因此，基于这一观察，本文提出了一种 token 稀疏化的框架，使得视觉 Transformer 模型在前向传播的过程中根据输入图片，动态地识别出冗余的 tokens 并剪掉。通过逐渐地剪掉约 66% 的输入图片的 tokens，DynamicViT 可以在将精度掉点控制在 0.5% 以内的前提下，节约 31%∼37% 的计算量，并提高 40% 以上的吞吐量。

50 DynamicViT：动态 Token 稀疏化实现高效的视觉 Transformer

论文名称：DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification

论文地址：

https://arxiv.org/pdf/2106.02034.pdf

50.1.1 背景和动机

计算机视觉领域的巨大进步和发展部分原因是以视觉 Transformer 架构为代表的通用视觉模型的发展带来的。视觉 Transformer 已经在目标检测，语义分割等诸多下游任务中取得了广泛的应用。就像针对 CNN 的研究那样，领域也同样致力于 Transformer 类模型加速算法的研究，希望这类模型能够更好地应用在实时移动设备上。之前领域大量研究致力于 CNN 的加速，比如剪掉 CNN 中不太重要的权重。在视觉 Transformer 模型中，我们有一种新的思路。因为 Transformer 将输入图像分割成多个独立的小块 (叫做 Patch 或者 token)，这为我们提供了另一种引入稀疏性的正交方法，即我们可以剪掉不重要的 tokens，因为这些冗余的 tokens 对最终模型预测结果的影响很小。

50.1.2 DynamicViT 方法简介

DynamicViT 通过一个极其轻量化的预测模块，来动态地决定哪个 token 应该被剪掉，这个预测模块将被加到 Transformer 模型的多个层里面。对于每个输入图片，预测模块会产生一个针此输入的二进制决策掩码，以决定哪些 tokens 的信息量比较小，需要被剪掉，哪些 tokens 的信息量较大，需要被保留。

如下图1所示，预测模块添加到 Transformer 模型的多个层里面，这样一张输入图片可以在通过模型的过程中，逐渐地减小 tokens 的数量，也就是以一种分层的方式进行稀疏化。一旦一个 token 在某一层之后被剪掉，它就永远不会在模型的前向传播中使用。

虽然在模型中引入了这些额外的预测模块，但是它们所带来的额外计算开销与剪掉冗余 tokens 带来的计算开销的节省相比是相当小的。

预测模块的权重可以随着 Transformer 骨干模型的权重端到端地训练，为此，作者采取了两种专门的策略：

1. 采用 Gumbel-Softmax[1] 来克服从分布中采样的不可微问题，从而实现端到端训练。

2. 关于如何根据学习到的二进制决策掩码来剪掉冗余的 tokens：对于每个输入图片而言，其二进制决策掩码中0值的位置肯定是不一样的。因此，在训练过程中直接消除每个输入图片的冗余的 tokens 将无法实现并行计算，也无法计算预测模块的反向传播。而且，直接将被剪掉的 tokens 设置为零向量也不可以，因为零向量仍然会影响 Attention 矩阵的计算。

因此，作者提出了一种称为注意力掩码的策略，基于二进制决策掩码把注意力矩阵中被剪掉的 tokens 与其他 tokens 的 attention 丢弃。作者还改了原始的训练目标函数，增加了一项 loss 来限制剪掉 tokens 的比例。

在推断阶段，针对不同的每个输入图片，可以将冗余的 tokens 剪掉，不再需要考虑操作是否可微，这将大大加快模型的推理速度。

图1：本文方法框架：所提出的预测模块被插入到连续的 Transformer 层中间，根据前一层输出的特征选择性地剪掉信息量较少的 token。这样一来，后续的层需要处理的 token 就减少了

50.1.3 通过预测模块进行 Token 的分层稀疏化

这里 "分层 (Hierarchical)" 的含义是：随着模型的层数由浅到深，token 的数量逐渐减小，即随着计算的进行，逐渐丢弃无信息的 tokens。举个例子：前几层是1000个，中间几层是500个，最后几层只有200个。这个过程是 "分层 (Hierarchical)" 的。

为了实现这个过程, 作者维持了一个二值的 decision mask 向量：, 来指示这  个 tokens 是丟弃还是保留, 其中  是图片的 Patch 数。将决策掩码  中的所有元素初始化为1, 并逐步更新掩码。预测模块的输入是当前的决策掩码  和特征  。

预测模块的具体做法是：

a) 首先用一个 MLP 映射输入 tokens：

式中， 。

b) 计算全局的特征：

式中， \text{Agg}\text{Agg} 是聚集全聚信息的，比如可以是 average pooling：

局部特征编码了某个 token 的信息，而全局特征包含了整个图像的上下文，因此两者都包含了信息。

c) 因此，作者结合局部和全局特征来获得局部-全局嵌入，并将它们输入到另一个 MLP，以预测丢弃/保留 tokens 的概率:

式中,  表示丟弃第  个 token 的概率,  表示保留第  个 token 的概率。然后, 可以通过从  采样来生成当前决策 , 并通过:

更新决策向量  。式中,  代表 Hadamard 积, 代表一旦某个 token 的决策掩码被置为了0, 它就再也不会被使用了。

50.1.4 通过注意力掩码进行端到端的训练

现在我们通过预测模块得到了一个概率向量 , 然后通过从  采样来生成当前决策  。再根据当前决策 D 来执行 tokens 的稀疏化过程。但作者发现在训练期间实际实现这个过程并不容易。首先, 从概率向量  到二进制决策掩码  的采样过程是不可微的，这会阻碍端到端的训练。为了克服这个问题, 作者使用 Gumbel-Softmax 技术对概率向量  进行采样:

这里作者使用了下标 1，是因为  代表保留的 token。Gumbel-Softmax 的输出是一个独热编码，它的期望恰好等于概率向量 。同时，Gumbel-Softmax 是可微的，这使得端到端的训练成为可能。

第二个难点是如何根据学习到的二进制决策掩码来剪掉冗余的 tokens：对于每个输入图片而言，其二进制决策掩码中0值的位置肯定是不一样的，不同样本的掩码包含不同数量的1值。因此，在训练过程中直接消除每个输入图片的冗余的 tokens (即简单地丢弃 将无法实现并行计算，也无法计算预测模块的反向传播。因此，我们必须保持 tokens 的数量不变，同时切断被剪掉的 tokens 与其他 tokens 之间的交互。而且，直接将被剪掉的 tokens 设置为零向量也不可以，因为零向量仍然会影响 Attention 矩阵的计算：

通过 Softmax 操作，值为0的 tokens 仍会影响其他 tokens。为此，作者设计了一种称为注意力掩码的策略，可以完全消除掉一些 tokens 的影响。具体来说，通过以下方式计算注意力矩阵  ：

这个式子意味着如果第  个 token 被剪掉了, 那么 , 则 , 这意味着它不会再对注意力矩阵产生任何的影响。这样一来, 注意力矩阵  的维度还是 , 但是它只考虑了保留的 tokens。

50.1.5 Dynamic ViT 的训练和推理

DynamicViT 的训练包括训练预测模块，使它们能够产生有利的决策掩码  ，并微调主干网络的权重，使其适应 token 的稀疏化。假设一个 Batch 的样本数为 B 。

DynamicViT 的首先是传统的 Cross-Entropy Loss：

式中,  是 DynamicViT 的预测结果,  是标签。

为了最小化 token 稀疏化对性能的影响，作者使用未稀疏化的原始模型作为教师模型，并希望稀疏化之后的模型的行为尽可能接近教师模型。具体而言，作者从两个方面考虑这个约束

其一，使 DynamicViT 的最后剩余的 tokens 接近教师模型的 tokens，这可以被看作是一种 Self-distillation：

其中,  和  分别代表 DynamicViT 和教师模型的第  个 token,  指的是第  个样本在第  个 stage 的决策值。

其二，通过 logit 蒸馏，即 KL 散度最小化 DynamicViT 和教师模型之间的预测差异：

最后, 作者希望将保留的 token 的比率限制为一个预定义的值。比如一共有  个 stages, 那么每个 stage 的保留 token 的比例可以预设为 , 作者通过 MSE Loss 来监督预测模块:

总的目标函数为：

其中，在所有实验中设置。

DynamicViT 的推理过程就是给定目标 token 的稀疏度 , 可以通过由预测模块产生的概率向量 , 直接丢弃信息较少的 tokens, 使得在第  个 stage 保留  个tokens。比如对于第  个 stages, 对这些 tokens 的概率值进行排序:

根据排序之后的结果取  为保留的 tokens, 并丟弃掉其他的 tokens。这样, DynamicViT 可以在推理时剪掉信息量较少的 tokens, 从而减少推理过程中的计算开销。

50.1.6 实验结果

所有实验中设置稀疏化的 stage 数为 3 , 即通过3个阶段逐渐进行 token 的稀疏化, 3个阶段的稀疏率分别是  。在训练 DynamicViT 时, 作者遵循 DeiT 中使用的大多数训练策略。使用预训练的 DeiT 权重来做初始化, 并微调整个模型 30 Epochs。预测模块的学习率设置为 , backbone 的学习率设置为  。

ImageNet-1K 实验结果

在几种不同的视觉 Transformer 模型的结果如下图1所示。DynamicViT 适用于不同类型的 Transformer 模型中，作者汇报了 ImageNet-1K 精度，计算量和吞吐量 (吞吐量是在单个 NVIDIA RTX 3090 GPU 上测量的，Batch size 固定为32)。DynamicViT 可以减少 31%~37% 的计算成本，并在运行时加速 43%~54% 的推理速度，而性能影响可以忽略 (0.2%~0.5%)。

图2：ImageNet-1K 实验结果

与 SOTA 模型对比

如下图3所示是 DynamicViT 与其他 SOTA 的分类模型对比的结果。可以观察到，DynamicViT 表现出良好的复杂性/精度权衡。值得注意的是 DynamicViT-LV-M/0.7 击败了先进的 CNN 模型 EfficientNet-B5 和 NFNet-F0。

图3：与 SOTA 模型的对比

DynamicViT 用于模型的缩放

在视觉 Transformer 中，改变通道的数量是最常用的缩放模型的方法，DynamicViT 提供了另一个强大的工具来执行 token 的稀疏化。首先，作者训练了几个 Embedding dimension 从192 (DeiT-Ti) 到 384 (DeiT-S) 变化的 DeiT 模型。然后基于  训练 DynamicViT。可以发现，在进行 token 稀疏化之后，DynamicViT-384/0.7 的复杂度降低到与 DeiT-320 相似；DynamicViT-320/0.7 的复杂度降低到与 DeiT-256 相似。

这样一来，通过将 DynamicViT 应用于 DeiT-256，就可以获得一个计算复杂度与 DeiT-Ti 相当的模型，但ImageNet Top-1 的精度却高出约 4.3%。

图4：动态 token 稀疏化方法与模型宽度缩放的比较：前者可以得到与后者复杂度相当但精度更高的模型

可视化结果

为了进一步观察 DynamicViT 的行为，作者在下图5可视化了稀疏化的过程。图5展示了原始输入图像和3个阶段后的稀疏化结果，其中掩码表示相应的 tokens 被剪掉了。可以观察到每张图片都通过分层的 token 稀疏化，DynamicViT 都会逐渐丢弃掉信息量相对较少的 tokens，最终聚焦于图像中的物体。这一现象还说明，DynamicViT 可以有更好的可解释性，即：它可以定位图像中对分类结果贡献最大的重要部分。

图5：逐渐稀疏化 tokens 的可视化结果。每张图片都通过分层的 token 稀疏化，DynamicViT 都会逐渐丢弃掉信息量相对较少的 tokens，最终聚焦于图像中的物体。这一现象还说明 DynamicViT 有更好的可解释性

除了以上不同样本的可视化结果之外，作者还对稀疏化 token 决策的统计特征感兴趣。也就是说，DynamicViT 究竟从数据集学习了什么样一般的模式？作者使用 DynamicViT 为 ImageNet 验证集中的所有图像生成决策，并计算所有3个阶段中每个 token 的保留概率，如图6所示。作者将结果通过 Pooling 操作变成维度是 7×7，使得它们更容易可视化。作者发现图像中间的 tokens 往往被保留，这也是合理的。因为在大多数图像中，物体一般位于图片的中心。

图6：每个阶段 token 被保留的概率

不同稀疏化策略的比较

如下图7所示，为了讨论动态稀疏化是否优于其他的稀疏化策略，作者进行了消融实验，结果如下图7所示。一种方法是结构化下采样 (structural downsampling)，这种方法作者使用了一个 2×2 大小的 average pooling，这种方法与 DynamicViT 有类似的 FLOPs。另一种稀疏化的策略是静态 token 稀疏化 (static token sparsification)，即稀疏化 token 与输入无关。通过实验可以发现，虽然这三种策略的计算复杂度相似，但 DynamicViT 的精度是最高的。

图7：不同稀疏化策略的比较结果

总结

DynamicViT 是一种动态 token 稀疏化的框架，使得视觉 Transformer 模型在前向传播的过程中根据输入图片，动态地识别出冗余的 tokens 并剪掉。通过逐渐地剪掉约 66% 的输入图片的 tokens，DynamicViT 可以在将精度掉点控制在 0.5% 以内的前提下，节约 31%∼37% 的计算量，并提高 40% 以上的吞吐量。

参考

1. ^Categorical Reparameterization with Gumbel-Softmax

【自动驾驶之心】全栈技术交流群

自动驾驶之心是首个自动驾驶开发者社区，聚焦目标检测、语义分割、全景分割、实例分割、关键点检测、车道线、目标跟踪、3D目标检测、多传感器融合、SLAM、光流估计、轨迹预测、高精地图、规划控制、AI模型部署落地等方向；

加入我们：自动驾驶之心技术交流群汇总！

自动驾驶之心【知识星球】

想要了解更多自动驾驶感知（分类、检测、分割、关键点、车道线、3D目标检测、多传感器融合、目标跟踪、光流估计、轨迹预测）、自动驾驶定位建图（SLAM、高精地图）、自动驾驶规划控制、领域技术方案、AI模型部署落地实战、行业动态、岗位发布，欢迎扫描下方二维码，加入自动驾驶之心知识星球（三天内无条件退款），日常分享论文+代码，这里汇聚行业和学术界大佬，前沿技术方向尽在掌握中，期待交流！