DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification 论文阅读笔记

简介

文章利用Transformer的token非结构化序列的特点，设计了一种Token稀疏化剪枝方法，通过删除信息量不大的tokens降低计算量，加速Transformer。该文章为Transformer加速开辟了新的方向。

原文链接：​arxiv.org/abs/2106.02034

代码：​github.com/raoyongming/DynamicViT

1. Introduction

CNN类架构为计算机视觉带来了非常大的进步。最近Transformer开始在目标检测和分类等任务当中代替CNN。正如CNN模型过去几年发生的一样，Transformer模型也需要加速以适应实时应用。

CNN的一个常见加速方式是剪去不重要的Filter。然而Transformer对输入的处理方式——将输入图像分为多个相互独立的Patches，提供了另一种正交的方法使用信息稀疏性进行加速：通过剪枝掉输入序列当中对输出贡献不大的tokens。这个只在诸如Transformer这样的self-attention能够处理任意长度序列作为输入的模型才可能实现，因为非结构化的输入并不影响self-attention模块。相比而言，在卷积当中扔掉部分像素很难实现真正加速，因为非结构化邻域会导致卷积难以并行计算。CNN的结构化降采样能够有效提升模型效率，我们希望能够为Transformer探索非结构化的降采样，以进一步利用self-attention优势。

文中提出一种轻量化的预测模型用于决定剪去哪些tokens，称为DynamicViT。对于每个输入序列，使用二值化masking决定哪些tokens是没什么信息而需要丢弃的。该模块被添加到Transformer的多个层当中，从而采用逐层稀疏化，逐渐增加被剪去的Token数量。一旦一个Token被剪除，其不再会参与前馈过程。该轻量化模块增加的计算量可以忽略不计，考虑到其节省的计算量时更是划算。

两个Strategy：

1. Gumbel-Softmax用于克服从特定分布当中采样的无差异问题，从而使得End to End训练成为可能。

2. 关于如何使用学习到的二值化决定剪去不重要的Tokens。考虑到每个输入样本都有不同的mask，直接移除信息量不大的Tokens将使得训练过程无法并行计算。且这个会阻碍预测模块的反向传播——即使某个Token被移除，预测模块也需要计算其需要保留下来的概率。另外，直接将移除的Token置为0也不是明智的选择，因为其最终还是会参与attention 矩阵的计算。因此本文提出了attention mask策略，基于得到的mask扔掉剪除的token和其它token之间的联系，从而克服上述困难。另外为了约束特定层后被扔掉的Token比例，修改了原始的Transformer训练目标函数。在Inference状态，由于无需保留可微分的梯度信息，可以直接扔掉被剪去的tokens。

本文方法在ImageNet上进行测试，基准模型采用DeiT和LV-ViT。结果：通过剪去66%的输入to