Vision Transformer在CV任务中的速度如何保证

我的PhD课题是Vision Transformer的efficiency问题，这个topic范围很广，可以是inference efficiency, 力争speed and accuracy trade-off。也可以是training efficiency，让模型更容易训练，如训练时间少，需要的硬件资源不用很多。也可以是energy efficiency，减少模型部署时在硬件上的功耗。当然也有data efficiency，怎么用更少的数据在短时间训练出较好的模型。

本文主要讲inference time的效率问题，简单说就是如何让模型更快，同时性能不掉太多甚至反升。从个人角度来看，如果性能的增长是以巨大的速度牺牲为代价，那么模型的意义便少了许多。

背景

  2021开年后，Transformer在CV任务中大放光彩，Google的初版Vision Transformer（ViT）虽然方法简单，但通过大量的实验证明了Transformer在计算机视觉（CV）任务上，特别是image classification上的可能性。论文展示，ViT经过large-scale pretrain之后在ImageNet上可以达到和CNN媲美的性能。

然而，CV的任务可不止image classification，给一个图片预测一个label就结束了。对于大量dense prediction的任务而言， 比如object detection和segmentation，Transformer都有其局限性，而其中最关键的有两点，

1. 缺少金字塔特征（Pyramid Feature Map)。层级特征对物体的检测很重要，特别是小物体，因为feature map的scale如果比较小的话，小物体相当于直接被忽视了。常见的CNN，如ResNet将网络划分四个stage，每个stage处理一种scale的feature map, 如1/4，1/8， 1/16，1/32。这样出来的层级特征可以送入FPN进行处理，然后作为输入给detection或者segmentation的head。而ViT只处理1/16这个scale下的feature map，直观上来看是对dense prediction任务不利的。  
2. Multi-head Self-attention（MSA）计算开销问题。MSA是一个quadtic complexity的模块，小图还好，但是遇到分辨率大的时候，无论是显存占用还是计算速度都会大打折扣。然而像detection和segmentation这种任务，图片一般都比较大，如512x512。按照ViT的patch embedding划分，图片会分成 （512/16）x (512/16) = 1024。此时ViT-Base的单层MSA会消耗4G FLOPs。然而整个ResNet-50也才消耗4G FLOPs. 因此MSA的计算开销限制了ViT在高分辨率图像上的效率。  

对于第一个问题，一个自然的想法就是给ViT引入pyramid feature map, 这也是ICCV 2021积极涌现的idea，如Swin, PVT。然而，引入结构上的hierarchy造成了前期feature map的scale比较大，此时MSA的计算开销上涨。这使得efficient attention在ViT上成为了必需品，正如Swin用到的local window attention和PVT的spatial reduction attention，当然后面有更多的architecture，更多的attention，在此就不一一举例了。

自此，ViT适配下游任务已经基本没有了大的障碍。 大致结构上和CNN保持一致，如划分stage，逐渐对feature map做downsampling。得益于attention的结构, ViT在各个task上性能都保证的不错。那接下来的主要问题就是对ViT在CV任务中进行优化。

什么最影响ViT的速度？

影响一个模型速度的东西无非是两个大方面，一是结构，二是处理的数据规模，剩下基本是工程上的优化。从个人观点来看，对于ViT来说，影响结构效率的有两个，

1. Efficient General Architecture. 整体结构上的设计需要变得有效并且scalable
2. Efficient Multi-head Self-attention. 注意力机制必须足够efficient

对于处理的数据规模问题，这个CNN也有，但有一个明显区别于CNN的是，ViT对input是permutation-invariant的，只要加入了positional encoding就可以无视2D spatial order。所以针对ViT而言，减少input tokens是可以做推理加速的，如本人在ICCV 2021上的工作HVT: Scalable Vision Transformers with Hierarchical Pooling。

结构上怎么设计最合理？ - LITv1

HVT之后，我们曾经的一个想法是进一步对token做一个dynamic的选取，因为直觉上来讲dynamic要比unifrom的pooling要好，了解这一领域的同学可能有点耳熟，没错，这个思路就是DynamicViT。

然而，实验过程中的一个发现，使得我们展现给大家的是LIT, Less is More: Pay Less Attention in Vision Transformers (AAAI 2022)。

和DynamicViT利用pretrained weights和plain ViT不一样，我们直接从hierarchical Vision Transformer入手并且train from scratch，如PVT，用几层FC对前两个stage中每个block的input tokens做一个dynamic的选取，使得前期MSA只需要处理1/4, 1/8 scale下选取的部分token。这个思路听起来感觉没什么问题，可视化效果也确实不错，如下面是去年基于PVT-S做sparse选取token做的一个可视化：

虽然但是，当我们对比random sampling的baseline时，发现性能差距并不大，一番分析之后，我们直接训练了一个完整的基于原始MSA的PVT-Small，通过可视化attention head, 我们发现前期的head只能关注于很小的local区域，尤其是第一个stage的MSA，attention weight基本都在关注自己，我们进一步在论文给出了理论上的解释。

基于这个发现，我们做了一个在当时来看（2021.05）比较大胆的决定，直接拿掉前两个stage的所有MSA layer，这样前期的Transformer block中只剩下了FFN， 而后期的两个stage直接采用最原始的MSA。这使得我们模型的整体结构看起来非常简单，如下面所示。

我们把这种设计原则在Swin Transformer和PVT上做了验证，结果显示这种方案不仅不降性能，还因为避开了前期MSA的计算开销，减小了理论复杂度，大幅度加快了推理速度。

  另外我们在LIT中提出了一个Deformable Token Merging module (DTM)，使得feature map在下采样时可以动态的根据前景信息对token做选取。经过可视化发现学习到的offset十分精准地捕捉到了前景物体信息。相关可视化代码已公开在GitHub。  

阶段总结：LIT在结构上主要陈述了一个事情，backbone前期采用MSA的效益并不高。即使是擅长global attention的MSA也依然重点关注local区域，何况前面的feature map比较大，MSA处理高分辨图像的开销很大，即使采用efficient attention也难以避开这部分计算。

所以我们看到后来的Uniformer, MobileViT，EfficientFormer都是这种前面local后面global的结构。想要快，前面还是得尽量避开MSA。

Architecture和Attention一起优化 - LITv2

在前期的工作中，我们定下了ViT的Architecture设计原则，但第一版LIT有两个问题影响实际推理速度，

1. RPE虽然能涨点，但对于LIT来说，由于后面需要动态根据input大小做插值，所以对速度影响很大。
   
2. 后两个stage的原始MSA针对高分图像还是会有很大的计算开销。
   

  基于此，我们提出了LITv2：Fast Vision Transformers with HiLo Attention (NeurIPS 2022), 主要对LIT做了两点改进，  

1. 去掉了relative positional encoding (RPE), FFN中加入depthwise convolution来引入conditional positional embedding。通过引入depthwise conv不仅能引入positional encoding, 还能同时增大网络前期的receptive field。（因为LIT前面两个stage只有FFN）
2. 换上了全新的attention方法，HiLo。
   

Attention还能更快

直接影响ViT在高分辨率图像上的效率的还是attention。从21年到现在，这一领域已经针对ViT提出了不少attention的变种，如SRA（ICCV-21），local window(ICCV-21), focal attention(NeurIPS-21), QuadTree (ICLR-22)。和MSA相比，这些方法都能保证更优的理论复杂度。但是我们注意到FLOPs只能反映理论复杂度，而真正想在CV任务中beat过纯CNN的话，除了性能，实测速度也十分关键，因为理论复杂度并不能反应GPU和CPU上的实际速度。影响实测throughput/latency的还有memory access cost和其他难以被统计进FLOPs的因素，如for-loop。

以实测速度为核心，我们在LITv2中提出了HiLo。HiLo的设计理念很简单: 在信号处理中，高频关注local detail, 因此就用部分head做local window attention去捕获，而低频关注global structure，对local detail并不太在乎，所以另一部分head只需处理average pool的低频信号，这样高低频同时捕捉，实现单层MSA同时抓取local和global信息。相对应的branch，我将它命名为Hi-Fi和Lo-Fi，如下图所示。

为了保证实际速度，HiLo避开了一些不能算到FLOPs但损害速度的运算，如overlapping window和interpolation。基于此，HiLo在实测速度，memory占用，以及性能上都超越了SOTA方法。

在CPU和GPU上，单层HiLo的推理速度非常快。

又快，又好，又容易训练的LITv2

LITv2共推出了三个变种，small, medium, base, 为了公平对比，三个模型的宽度和深度和Swin Transformer一样，超参seed一点没调。LITv2整体在ImageNet-1K上效果占优，同时降低了训练显存要求，推理速度更快。

由于资源实在有限，对比Swin-Large和ImageNet-22k上的setting确实没法跑。。。。

另外我们在多个GPU平台上测试了推理速度（来自reviewer的意见）。对比ResNet-50和其他11个最近的ViT模型，相似参数量下，LITv2-S在速度和理论复杂度上实现了全面超越，同时ImageNet的Top-1 accuracy在第一梯队。

未来还能做什么

Vision Transformer的故事还在继续，今年NeurIPS投稿也能看到大家陆续也在关注实测速度了，如TRT ViT，EfficientFormer, 但实际上ViT的工业落地相比CNN可能还是会有障碍。如量化上，ViT还是在低bit上还有很大提升空间。除此之外，我们组也在探索Transformer的其他efficiency问题，比如

ViT的Training efficiency还有很大的空间。ViT普遍难训练，需要的显存大，训练时间长。特别是backbone的公认setting是1024的batch size + 8 GPUs，使得很多小组连ViT跑都跑不起来（8卡 32GB V100不是所有人都有的），针对显存消耗问题，我们提出了一个针对Transformer的memory-efficient的训练框架: Mesa: A Memory-saving Training Framework for Transformers. 和普通训练相比，Mesa可以做到显存节省一半，同时与checkpointing和gradient accumulation等技术方案不冲突。针对训练时间长的问题，组里有同学在CVPR 2022提出了一个progressive learning的方法，大幅度缩减原先的训练epochs数，性能甚至更好，参考Automated Progressive Learning for Efficient Training of Vision Transformers。

Transformer的能源消耗也是一个问题，特别是当下很多大组在跑以Transformer为backbone的大模型，训练时所造成的电力损耗，碳排放对Green AI这一长远目标不利。针对Energy efficiency，我们组在NeurIPS 2022最新的工作 EcoFormer: Energy-Saving Attention with Linear Complexity 提出了一个Transformer二值化的全新方案，全新设计的EcoFormer是一个general的efficient attention, 性能和功耗上都优于一众线性复杂度的Linformer, Performer等，同时在45nm CMOS microcontroller上能耗更低，速度更快。

至于CNN和Transformer哪个好的问题，这个答案其实不必再过多讨论了，因为两者互补。通过合理地配置CNN和MSA在Backbone中的位置，可以让网络得到双倍的快乐。更不用说Convolution和MSA两者之间本身存在一种联系，比如我们组的另一篇工作：Pruning Self-attentions into Convolutional Layers in Single Path， 巧妙地通过参数共享将计算复杂度较高的attention layer剪成更为efficient的convolutional layer，工程细节可以参考GitHub。

个人主页：Zizheng Pan

实验室主页：Zhuang Intelligent Processing Lab (ZIP Lab)

文中如有错误，欢迎指出，同时欢迎各位进行学术交流~