[Transformer]SPViT:Pruning Self-attentions into Convolutional Layers in Single Path

Paper
Code

Abstract

ViT在各类计算机视觉任务中取得了令人印象深刻的性能，但是基于MHSA计算也有两大局限：一是庞大的资源消耗，二十缺乏局部归纳偏置。一种统一的解决方案是通过NAS的剪枝方法进行搜索，用类似卷积的归纳偏置来替换一些MSA层。但是如果将MSA和卷积作为单独的候选路径进行训练，又会导致搜索成本大大提升，优化过程也十分困难。


因此本文提出一种新的MSA与卷积之间权重共享的方案，将搜索问题转化为寻找每个MSA层中的参数自己。本文提出的这种SPViT（Single Path Vision Transformer）剪枝的方法使得能够满足一定的效率约束情况下快速的将预训练好的ViT修剪成更加精确、紧凑的混合模型，显著降低了搜索成本。本文在两类代表性的ViT模型上进行了广泛的实验，实验结果表明本文方法能够很好的权衡精度和效率。

Section I Introduction

ViT在分类、检测、识别等诸多领域引起了人们极大的研究兴趣，但是MHSA的两大局限限制了Vit的进一步应用、开发和部署，尤其是对于长序列的建模。因此也有诸多工作致力于高效ViT的研究，比如利用剪枝来修剪并不十分中间的一部分组件；另一问题是MSA缺少局部依赖关系，为了解决这一问题，现有的方法通常是插入卷积层来引入局部哦安置，比如在FFN内部或者encoder之前或者MSA层之前插入卷积。


为了同时解决这两个问题，一种解决方案就是使用NAS自动搜索，将一些MSA通过剪枝替换为卷积操作，这样可以移除一些计算昂贵的MSA层提升模型效率，同时又引入了局部归纳偏置。

因此近期也有一些One-shot的方法提出将MSA，卷积包含到搜索空间中，但是他们是单独的路径进行训练，参见Fig 1(a)，这种多路径的方法会大大提升搜索成本，因为每次搜索是独立更新每条路径的。
本文建议将搜索看做一个自己选择问题，其中卷积运算的输出直接从MSA中间结果索引得到，这样可以简化优化过程同时减少搜索成本。


为了解决上述问题，本文提出一种新的单路径ViT剪枝方法——SPViT，通过对预训练好的ViT剪枝获得高性能更精简的紧凑模型。受到MSA层可以关注到局部区域这一实验现象，本文提出一种MSA与卷积之间权重共享的方案；此外还可以使用预训练好的MSA参数来对卷积核进行初始化，在搜索过程中对MSA的中间结果进行索引获得卷积运算的输出。
通过将所有的候选操作都封装在一个MAS层中进行，这样形成了一个单路径搜索方法，参考Fig 1(b),可以通过剪枝的方法将MSA剪枝成卷积，因此不需要想多路径方法一样在不同的候选操作中进行选择，而是将搜索问题转化为在MSA中选择自己的问题。这样就可以大大降低搜索成本，同时优化搜索过程。
除此之外考虑到FFN层也占据很大一部分的计算复杂度，比如对17.5G乘加运算中占据了11,1G,b本文提出一种细粒度的MLP扩展方法来搜索MLP的拓展比例，即SPViT会基于门控机制确定每个FFN hidden dim的权重，然后修剪不那么重要的维度。
本文的贡献总结如下：
（1）本文提出一种MSA与卷积权重共享的方案，这样可以将所有候选操作封装到一层MSA中，这样只需要single-path就会更新全部，将搜索问题转化为寻找MSA参数的罪域子集组合问题，这大大降低了搜索成本；

（2）基于single-path方案，本文提出的SPViT会自动对一些不重要的MSA剪枝成卷积操作，并且对MLP层的hidden dim也进行搜索；

（3）在CIFAR-10和ImageNet-1k上的实验结果表明SPViT可以显著降低搜索成本，性能优于基线方法，同时还对一些搜索结果进行分析，得到一些Vit设计的经验性结果。

Section II Related Work

Pruning Transformers

为了降低Transformer的计算成本，常使用剪枝。目前对Transformer的剪枝大致有两类，分别是模块剪枝和token剪枝。模块剪枝主要调整head number,linear projection的通道数，weight等，近期Zhu等人则提出可以动态计算、修剪不那么重要的token特征；Chen等人则基于彩票算法来修剪Transformer参数。token的修剪则是会动态计算、修剪 不那么重要的token使得模型满足一些FLOPs的要求；但也因为token的减少，如果要将其应用到一些密集的预测任务，如分割任务，就十分有挑战性。

本文也是兼职，但与之前的工作区别在于专注于剪枝没那么重要的全局的MSA，将MSA修剪为卷积，这样可以以较低的搜索成本快速部署搜索Transformer模型。


Convolutional vs self-attention layers



本文的另一方面工作则是探究卷积核SA，因为卷积更适合提取局部纹理特征，SA更适合建模全局相关性，提取物体整体形状，为了获得二者优点，许多工作尝试将二者单独处理组成混合模型，比如听过向ViT中插入卷积层，或者在CNN之外堆叠SA层；另一类思路则是探究两种操作的相关性，比如SA结合可学习的位置编码可以表征任意卷积层，ConViT则是将软性的归纳偏执作为注意力分组的一部分，本文的工作也属于后面这一类。



本文提出的权重共享表明MSA的子集可以表示卷积运算，这样就能将候选操作集合在MSA中，从而大大降低搜索成本。

Section III Method

Part 1 Weight-sharing between MSA and convolutional operations

权重共享指的是不同操作之间共享参数子集，比如在候选卷积操作之间共享权重可以大大降低搜索成本，本文使用了一种新的MSA与卷积共享权重的方法。
回顾卷积和MSA
卷积是CNN的基础模块，标准的卷积层会将感受野内的信息进行聚合，最终的输出可以表示为：

而MSA则是Transformer的基础模块，但是其感受野是整个输入序列也就是其感受野是全局的，MSA计算过程可以表示为：

MSA计算复杂度为：

因此当输入分辨率很高的时候计算复杂度会成指数增长；

但是卷积的计算复杂度一直都是：


权重共享策略
接下来叙述MSA如何与卷积进行参数共享，通过逐步消除MSA独有的参数，得到最后二者共享的参数。

首先，因为卷积只处理感受野内的局部信息，因此再用MSA表示卷积时需要将非局部部分的注意力分数置零，将局部区域的分数固定为1：



因此MSA计算式子变为：

注意到MSA(😆,并且需要将注意力限制在kxk大小的局部窗口内。
接下来需要将MSA参数转换为表示卷积核的参数，一种方法是建立一个映射关系，选定head子集中的一些head分配到局部窗口的特性位置，表示为：

即每个head只关注局部窗口中的某个位置，类似DeiT-Ti中的方案，这种方法适合一些大型网络。

此外也有研究表明在MSA中不同head倾向于关注不同的区域，因此如何定义head子集到局部位置的额映射是很重要的，本文使用可学习的参数z并且经过softmax激活后作为映射的结果。



本文相信这样处理使得head集合中不同head会关注局部窗口中一些预定义的位置，softmax则是确保head集合保持原来的比例。
完成了这一步就可以将MSA的参数表示卷积核，卷积核权重就是子集中每一个headsoftmax后的结果。


同时处于计算复杂度的考虑，本文选择ResNet中的Bottleneck模块并进行了一定的修改：

注意到卷积运算的Cout=Ch，主要还是为了计算效率。然后继续添加BN和ReLU，得到最终的BConv计算表达式：

Remark


注意到通过权重共享可以使得MSA建模全局的同时处理局部特征。因为局部区域确实是全局的一部分，那么MSA自然有处理局部区域的能力；其次ViT的一些head往往会聚焦到局部区域，那么使用这些头表示卷积操作我们就能得到合理的输出。但是不同的ViT模型，表示局部区域的最佳映射可能会不一样，因此本文的SPViT会动态的选择。

Part 2 Single Path vision transformer pruning

2.1 Search Formulation
Searching for MSA and Convolutional layers




Step 1:
首先看将MSA修剪为卷积的最佳映射位置。



具体来说就是将ViTblock中用一个UMSA层，通过前一节描述的权重共享方案同时获得MSA的输出和卷积操作的输出，然后再学习门控机制配置具体的映射方案。




为此在UMSA中可以通过调控或者缩放V来获得、调控卷积输出，而不用将卷积操作作为一个单独的路径。流程参考Fig 2(a).






这样就以较低的计算成本同时获取了MSA和卷积的输出。根据计算成本从低到高排序，就可以获得一个由BCOnv和MSA表示的一个集合P。



Step 2：定义一系列二进制门操作来代表选择操作，这一分布使用梯度自动优化。



某一个门服从伯努利分布，经过sigmoid激活后代表选择第p个操作的可能性。


为了提升搜索效率，本文设置每层只选择一种操作，因此每一层UMSA最多只有一个gate打开：

最终UMSA层的输出可以表示为：


其中p表示第p个操作的输出，然后每一个UMSA层后跟着残差连接和LN，像ViT一样。如果所有门都是关闭的则使用shortcut操作。

因此UMSA搜索空间总共有三种类型的候选操作，分别是：Shortcut,不同卷积核大小的BConv，MSA。


Searching for fine-grained MLP expansion ratios


接下来说说细粒度MLP扩展率的搜索，代表的是FFN层的hidden dimension，之前的工作一般都比较粗粒度的手工设置一系列组合。但本文会在预训练的模型对应的expansion ratio-α上进一步对MLP Expansion ratio进行搜索。


UFFN层的输出可以表示为：





与前一小节一样，依旧通过二进制门来寻找更细粒度的expansion ratio，来定义每一个gate输出与否，这样可以删除一些不重要的隐藏层维度。



Searching objective



搜索目标就是获得满足效率约束的理想模型，因此使用计算复杂度辅助损失函数。


具体来说就是定义了一个查找表，表中包含了所有候选操作和模型的计算复杂度，Lcomp定义为：

其中F（X）是当前网络的计算复杂度，Fhat为目标复杂度。
因此整体的损失函数是交叉熵损失和复杂度损失二者的结合，λ表示二者的相对权重。
2.2 Fine tunning

微调过程中，因为二进制规定了每一层最多只有2个门打开，因此UMSA每一层最多选择一种操作，输出可以表示为：

如果选择的操作是BConv则可以将参数组装为卷积核，进行硬件部署：

类似的对于UFFN层也有一个表示选定的hidden dim的集合，微调就从这个集合里选。

Section IV Experiments

对比网络有DeiT,ViT，Swin
实验设置
基本与DeiT和Swin的训练保持一致





Part 1 Main results

Compared methods




本文使用两中剪枝算法来进行比较，一是基于敏感性的剪枝，会迭代的修剪MSA和FFN层中不重要的head或者channel，二是STEP(基于评估其的修剪)，这样可以保留重要性分数最高的attention head和隐藏层维度。





Table 1,2展示了SPViT与不同基线模型在两个任务上的性能，从Table 1可以看到通过将DeiT修剪为更为紧凑的模型，FLOPs下降了20%，同时精度也十分具有可比性。而SPViT-DeiT-B节省了33.5%de FLOPs,知道了0.2%的点。


需要注意的是SBP是从头开始修剪DeiT，SPViT是在预训练模型上进行的修剪，但是SBP的训练了600个epoch而本文只训练了130个epoch。



从Table 2中可以看到当修剪更加紧凑的模型时，SPViT在修剪使得FLOPs下降20%时依旧取得了很高的性能。

Part 2 Observations on searched architectures

本文还可视化了SPViT搜索过程中的一些结果，本文分别可视化了FLOPs在3.9G,2G,2.4G下SPViT-DeiT-S搜索到的结构，精度分别是79.3%，78.9%等精度。

Locality is encouraged in shallow blocks


参考Fig 3(a)-(f)的结果，可以看到当SPViT-DeiT-S和Ti会将浅层MSA修剪为卷积或skip connection，这与之前的一些工作观察到的结果是一致的，说明浅层MSA包含的全局性有一定冗余性；还可以看到最后两个block依旧是MSA层说明了在较深层使用全局性是很有必要的。


Last few FFN layers have more redundancy in standard ViTs


预训练模型中所有模块的MLP expansion ratio设置为同样大小，但从搜索结果可以看出，SPViT倾向于修剪后面模块的HIDDEN DIM.因为隐藏层dim越大表明容量越大，因此本文认为最后几个FFN层比其它层的容量要小，因为他们在原始VIT中需要处理的注意力更少。

**Shallower FFN layers within each stage require higher capacity in hierarchical ViTs.**




与原始ViT的区别在于Swin模型有层次性结构，分了不同的stage每个stage有诸多模块。在每个stage开始会进行token merging会被合并成更高维度。



参考Fig 3(d)-(f)可以看到浅层会保留更多的隐藏层维度，本文推测如果要建模特征空间的变化需要ViT更高的模型容量。









Part 3 Ablation Atudies

Single-path vs multi-path





Table 3展示了使用单路径搜索和多路径搜索的对比结果，可以看到在相同的计算复杂度下单路径搜索比多路径搜索性能更高，参数更少，搜索成本更低，比如搜索到的SPViT-DeiT-S比对应的多路径搜索节骨高0.3%,并降低了50%的搜索成本。







Effect of the search strategy







为了证明本文搜索策略的有效性，本文在CIFAR-100上与随机搜索的结果进行了对比，本文在搜索空间中随机选了10种计算复杂度接近的结构，然后进行微调，对比结果可以看到在FLOPs和参数接近的情况下SPViT搜索结果比随机搜索高出1.3%和0.6%,证明了本文搜策略的有效性。









Prunning proportions for MSA and FFN layers









Table5展示了MSA和FFN层中剪枝的比例，可以看到SPViT-DeiT-B对MSA层剪枝比例较高，但SPViT-Swin-B对FFN层有较高的剪枝比例，本文的提议是在一定效率的约束条件下，针对不同模型灵活的调整选择合适的剪枝比例。

Section V Conclusion and future work

本文提出一张新的MSA与卷积之间权重共享的方案，允许将卷积作为MSA的参数子集，可以同时优化两种操作。基于权重共享的方案本提出了SPViT在降低搜索成本的前提下也加入了归纳偏置，还对隐藏层维度进行了搜索，针对搜索结果得到了一些观察结果。










Limitations and social impacts










SPViT的一大局限是目前将MSA简直到卷积的粒度过于粗糙，可能会导致识别能力性能大大下降，本文未来将探索head pruning的方式。另一种未来的方向是探索将权重共享方案扩展到其他卷积变体，比如深度卷积、组卷积等。
社会影响是虽然本文的剪枝可以将密集的模型修剪为精简模型，但是训练过程仍然有大量的碳排放，尤其是在获得预训练模型这一过程。

Appendix

实验实施细节

ImageNet-1K:


batch_size 1024 在8块V100上



CIFAR-100：batch_size 84 单块V100




Fig 4展示了搜索到的SPViT-DeiT-B和SPViT-Swin-B结构，与前文的观测结果是一致的，比如前面浅层MSA会被修剪为卷积，FFN层后面几层的隐藏层维度比较低。

此外本文还可视化了中间层的attention map，可以看到深层次和浅层次关注的模式还是有很大不同的，可以看到MSA头部主要关注局部区域，而深层模块更加关注全局信息，因此注意力的概率图很好的证明了本文的思路，就是用SPViT浅层MSA修剪成卷积，并且保持深层MSA不变。