清华提出DynamicViT：动态Token稀疏化的高效视觉Transformer

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

作者：赵文亮   编辑：Amusi  

https://zhuanlan.zhihu.com/p/379054782

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

论文地址: https://arxiv.org/abs/2106.02034

代码仓库链接: 

https://github.com/raoyongming/DynamicViT

项目主页: 

https://dynamicvit.ivg-research.xyz/

预训练模型和训练/测试代码均已开源，欢迎使用

简介

自从去年 ViT 的提出，许多工作开始将 transformer 应用到各种视觉任务中。限制视觉transformer 的一个最大的瓶颈在于自注意力机制的巨大计算量（与token数的平方成正比）。然而，图片中所有的token对分类是同等重要的吗？为了探究这个问题，我们首先利用 transformer 可视化工具 可视化出图片中每个 token 对输出结果的影响：

可以看到，模型的预测结果只与图片中的少部分的token有关。也就是说，我们可以动态去除掉一些重要性较低的token，而不会对模型的准确率带来较大的影响。之前也有一些工作关注在以更低的复杂度去计算或近似自注意力机制（例如 Performer等），但是这些方法中仍然对所有的token计算FFN（Feedforward Network）。相比之下，如果我们可以去除掉大部分不重要的 token，那么自注意力和FFN的计算量都会随之减少。

基于上面的发现，我们提出了动态 token 稀疏化的方法。动态 token 稀疏化可以看成一种下采样的方式，但是每次下采样中保留哪些token是由当前的输入来动态确定的。作为对比，CNN 中的下采样方式是预先定义好的结构化下采样，而动态 token 稀疏化是非结构化的。需要指出的是，由于CNN需要时刻保持feature map的空间结构，动态 token 稀疏化方法在CNN上并不能实现加速。而在 transformer 中，自注意力的计算无需考虑token之间的空间关系，所以我们可以通过动态token稀疏化直接去除不重要的token。

方法

下面介绍我们的整体框架，包括一个主干网络和多个预测模块。其中，主干网络可以是各种视觉 transformer 模型，例如 DeiT，LV-ViT；预测模块用来动态预测token 稀疏化策略。

需要注意的是，我们的 token 稀疏化是以层次化的方式进行的。例如，假设 transformer 中有 12 个 block，我们会在第 4 个，第7个，第10个 block 之前分别实施一次稀疏化操作。

预测模块

为了减小额外的计算开销，我们预测模块的具体实现方式是几层轻量化的MLP，其输出结果表示了每个 token 被保留下来的概率。

注意力 masking 策略

在训练的时候，我们利用 Gumbel-Softmax 的技术将上面的概率 π变成二值化的mask。Gumbel-Softmax 既可以保证期望不变，又可以保证二值化的过程是可导的。为了实现并行化，我们在训练时并不会将无用的 token 剪掉，而是提出了一个注意力 masking 策略 来显式切断这些token和其他token之间的联系。

这种注意力 masking 策略可以保证最终的预测结果只与保留下来的 token 有关。在测试时，我们直接对每个token的保留概率排序，每次保留固定比例的token即可。例如，假设每次的保留比例为p=70%，经过三次 token 稀疏化操作后只有p^3=34.3%的token被保留下来。

实验结果

我们首先以 DeiT-S，LV-ViT-S，LV-ViT-M 三种模型作为主干网络分别用 0.7~1.0 的保留比例进行 token 稀疏化。可以看到我们的方法可以减少 31%~37% 的 FLOPs，并实现43%~54% 的实际加速，同时分类准确率几乎不受影响。

通过动态 token 稀疏化，我们可以实现非常有竞争力的准确率/复杂度 tardeoff。

我们的模型超过了包括EfficientNet-B5 在内的许多SOTA的CNN 模型。这个现象从下面的图中也可以很明显地看到。

此外，我们还发现我们的方法比直接在 channel 维度缩放具有更好的性能。下图中，我们的DynamicViT-256/0.7 在相近的计算量下的ImageNet分类准确率比 DeiT-Ti 高了4% 左右。这个现象也说明了我们的方法提供了一个新的模型缩放的方向。

为了验证动态token稀疏化的有效性，我们对比了结构化下采样（与CNN中相同）、使用静态稀疏化策略（即策略与输入无关），发现我们的动态 token 稀疏化的方法能在相同 FLOPs 下取得最高的准确率。

我们将我们的动态token稀疏化的过程可视化出来，其中白色的位置表示该位置的token已经被剪掉。可以看到，在经过3个stage过后，大部分背景的token都被逐渐剪掉，模型关注在与分类有关的物体上面。这一可视化也说明我们的方法具有很好的可解释性。

DyVIT论文PDF和代码下载

后台回复：DyViT，即可下载论文PDF和代码

CVPR和Transformer资料下载

后台回复：CVPR2021，即可下载CVPR 2021论文和代码开源的论文合集
后台回复：Transformer综述，即可下载最新的两篇Transformer综述PDF
CVer-Transformer交流群成立
扫码添加CVer助手，可申请加入CVer-Transformer 微信交流群，方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch和TensorFlow等群。
一定要备注：研究方向+地点+学校/公司+昵称（如Transformer+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群
▲长按加小助手微信，进交流群▲点击上方卡片，关注CVer公众号
整理不易，请给CVer点赞和在看