开发: C++知识库 Java知识库 JavaScript Python PHP知识库人工智能区块链大数据移动开发嵌入式开发工具数据结构与算法开发测试游戏开发网络协议系统运维
教程: HTML教程 CSS教程 JavaScript教程 Go语言教程 JQuery教程 VUE教程 VUE3教程 Bootstrap教程 SQL数据库教程 C语言教程 C++教程 Java教程 Python教程 Python3教程 C#教程
数码: 电脑笔记本显卡显示器固态硬盘硬盘耳机手机 iphone vivo oppo 小米华为单反装机图拉丁

-> 人工智能 -> DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification 论文阅读笔记 -> 正文阅读

[人工智能]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%的输入tokens能够降低模型31%~37%的GFLOPs，效率提升超过40%，所有Transformer模型精度降低在0.5%以内。该工作探索了采用空间稀疏性对Transformer进行加速的可行性，为未来Transformer类型的模型加速开辟了新的研究方向。

2. Related works

Vision transformers. Transformer模型首先在NLP当中被广泛研究。其提供了使用Self-attention代替RNN及其变种的可能。最近研究表明Transformer也是CNN很好的替代者，其在诸如图像分类、目标检测、实例分割、3D分析的视觉任务当中都有着很好的效果。DETR是首个将Transformer用于图像领域的工作，该工作将目标检测看作一个集合预测问题，采用编码器-解码器生成一系列的Bounding Boxes。ViT是第一个直接在非重叠图像paches上将Transformer用于图像分类的工作。相比于CNN，ViT通过大规模预训练实现了更好的效果。DeiT提出了许多训练技术，使我们可以只在ImageNet1K上训练无卷积Transformer，实现了比ViT更好的性能。LV-ViT通过引入一种新的称为Token labeling的训练目标，进一步改进了性能。ViT及其后续技术将输入图像分割成多个独立的图像块并进行变换，将这些图像块转换为Tokens以供进一步处理。这使得为Transformer类的模型引入空间稀疏性成为可能。

Model acceleration. 模型加速技术对于深度模型部署到边缘设备上具有重要意义。有许多技术可以用来提高深度模型的推理速度，包括量化、剪枝、低秩分解，知识蒸馏等。也有许多工作专门工作旨在提高Transformer模型的推理速度。例如，TinyBERT提出了一种蒸馏方法来加速Transformer的推理。Star-Transformer通过将全连接层替换为星形拓扑，将与输入序列成平方关系的空间和时间复杂性降低到线性。上述工作集中在NLP领域，很少有工作探讨利用视觉任务特性加速Vision Transformer。另外，Transformer和CNN之间的差异也使得探索采用诸如Filter剪枝等CNN加速方法以外的方法成为可能。我们的方法旨在通过探索信息丰富的图像块分布的稀疏性，修剪不重要的Tokens而非修剪神经元。。

3. Dynamic Vision Transformers

3.1 Overview

DynamicViT的总体架构如图2所示。DynamicViT由一个正常的Transformer骨干网络和若干个预测模块组成。骨干网络可采用各种类型的Vision Transformer，如ViT、DeiT、LV-ViT等等。预测模块负责生成丢弃/保留Tokens的概率。对各个位置的Tokens稀疏化在整个网络上逐层进行。例如，给定12层Transformer，可以在第4、7、10层之前进行稀疏化。在训练过程中，由于我们新设计的attention masking策略，骨干网络当中的预测模块通过端到端训练方法进行优化。测试过程当中，我们只需根据预定义的剪枝率和预测模块计算的分数，选择信息量最大的tokens。

3.2? Hierarchical Token Sparsification with Prediction Modules

标记稀疏化分层执行是DynamicViT的一个重要特征，例如随着计算的进行，逐渐丢弃缺乏信息的tokens。为了实现这一点，我们保留了一个二进制决策mask? ${?{\hat D}} \in {\{ 0,1\} ^N}$ 。表示丢弃或保留每个token，其中N=HW是嵌入的patch数量。首先将所有mask初始化为1，并逐渐更新mask。预测模块将当前决策 ${?{\hat D}}$ 与tokens? $x \in {\mathbb{R}^{N \times C}}\$ ?作为输入，首先使用MLP处理这些tokens。?

${z^{?{\text{local}}}} = \operatorname{MLP} (x) \in {\mathbb{R}^{N \times C'}}$

其中 $C'$ 可以是一个低一些的维度，我们在实现中使用C'=C/2。相似的，可以计算全局特征：

${z^{?{\text{global}}}} = \operatorname{Agg} (\operatorname{MLP} (x)) \in {\mathbb{R}^{C'}}$

其中Agg是一个聚合所有现有tokens的函数，可以简单地采用平均池化：

${Agg} (u,\hat D) = \frac{\sum_{i=1}^{N}\hat D_i u_i}{\sum_{i=1}^{N}\hat D_i}, u\in\mathbb{R}^{c'}$

局部特征编码特定token的信息，而全局特征包含整个图像的上下文信息。因此两者都有很大的信息量。因此我们结合局部与全局特征获得局部-全局嵌入，并将其输入另一个MLP来预测丢弃/保留tokens的概率。

${z_i} = [{z^{?{\text{local}}}},{z^{?{\text{global}}}}],1 \leqslant i \leqslant N\$

$\pi = \operatorname{Softmax} (\operatorname{MLP} (z)) \in {\mathbb{R}^{N \times 2}}$

其中 ${\pi _{i,0}}$ 表示丢弃第i个token的概率， ${\pi _{i,1}}$ 表示保留它的概率。

随后通过从 $\pi$ 当中采样获得当前决策D，并更新 $\hat D$ 。

$\hat D \leftarrow \hat D \odot D$

其中 $\odot$ 表示Hadamard乘法，表示一旦某个token被丢弃，其在后续层当中再也不会被使用。

3.3 End-to-end Optimization with Attention Masking

虽然本文目标是将tokens稀疏化，但实践当中对其进行训练并不简单。首先，在 $\pi$ 当中获取二进制采样结果D是不可微的，这就阻碍了端到端的训练。为了克服这一点，我们应用了Gumbel-Softmax技术从概率 $\pi$ 当中抽样:

${G_{ij}} = \left\{ \begin{gathered} 1,i = j \\ {?{\hat D}_j},i \ne j \\ \end{gathered} \right.$

其中使用下标1，因为D表示保留tokens的mask。Gumbel-Softmax的输出是一个一次张量，其期望正好等于 $\pi$ 。同时，Gumbel-Softmax是可微分的，因此能够支持端到端训练。

第二个阻碍出现在试图在训练当中剪去token时。决策mask? $\hat D$ ?通常是非结构化的，不同的样本所包含的mask为1的数量不相同。简单地丢弃? $\hat D_i$ ?为0的tokens会导致一个batch当中每个样本tokens数量不一致，使得计算并行化变得困难。因此我们必须保持token数量不变，同时剪断被丢弃的tokens和其它tokens之间的相互关系。我们还发现仅仅采用二进制mask? $\hat D$ ?清零被丢弃的tokens是不可行的，因为在self-attention矩阵的计算当中：

$G_i_j=1$ $A = \operatorname{Softmax} \left( {\frac{?{Q{K^T}}}{?{\sqrt C }}} \right)$

被清零的tokens仍然通过softmax运算影响着其它tokens。为此我们设计了一种称为注意力mask的策略完全剔除被丢弃tokens的影响。attention矩阵的计算如下：

$P = Q{K^T}/\sqrt c$

${G_{ij}} = \left\{ \begin{gathered}1,i = j \\ {?{\hat D}_j},i \ne j \\ \end{gathered} \right.$

${A_{ij}} = \frac{?{\exp ({P_{ij}}){G_{ij}}}}{?{\sum\limits_{k = 1}^N {\exp ({P_{ik}}){G_{ik}}} }}$

上式构建了一个图，其中 $G_i_j=1$ 表示第j个token将对第i个token的更新做出贡献。注意我们显式得为每个token 添加一个自循环以提高数值稳定性。也很容易证明自循环不会影响结果：如果 $\hat D_j = 0$ ，第j个token将不会为任何除自己以外的其它token做贡献。最后一个式子计算加权Attention矩阵A，其等价于训练过程中只考虑了被保留的tokens但训练当中有一个固定N×N尺寸的Attention矩阵。