DVT:动态Transformer用于高效图像识别任务
Paper Code
Abstract
ViT在大规模图像识别任务中取得了巨大的成功,通过将2D图像切分成固定数量的图像切片,将每一个patch视作token。通常一张图片表征的token数越多,识别精度越好,但是随之而来的是计算成本的指数上升。为了有效的权衡二者,通常将token的数目设置为16x16或14x14… ?本文认为,每张图都有自己的特征,在理想的情况下其输入token是独一无二的。同时本文还观察到对于一些较为简单的图像,仅使用4x4个token就可以进行准确预测,只有一部分较难的图像需要更精细的特征表述。 ??受此启发,本文提出动态Transformer,会自动为输入图像分配合适数量的token,这是通过级联不同token数目的Transformer实现的,他们会自适应的依次激活,一旦预测结果达到一定的置信度,就停止继续激活。? ??本文进一步设计了有效的特征复用和关系复用的组件添加在DVT中,这样可以减少冗余计算。在ImageNet,CIFAR-10,CIFAR-100数据集上的实验结果表明本文提出的DVT在理论计算效率、实际推理速度等方面均优于基线模型。
Section I Introduction
NLP领域中的Transformer在图像识别领域也大放异彩,尤其是在ImageNet benchmark上性能已经达到了SOTA,同时展现出了高度的可扩展性和灵活性。通过将图像拆分成固定数量的token,然后将其结合位置嵌入作为Transformer的输入。一般来说token数目越多识别精度越高,但这也导致计算成本指数上升。目前token的经验值为16x16或14x14. 本文认为,给所有图像使用同样数目的token并不是最优的选择,因为不同的图像之间存在巨大的差异(背景、物体尺寸等),因此需要根据不同的输入给定token数目,这一问题也关乎着模型的计算复杂度。
Table 1就展示了T2T-Vit-12模型使用不同token数目对精度和计算量的影响,可以看到使用14x14token只比4x4精度提升了15.9%但是计算成本增加了8.5x,换言之这些计算资源主要就浪费在了许多简单的图像上,并不需要14x14tokens.? 基于以上实验现象,本文提出了动态Transformer,可以根据图像内容动态分配token数目。具体来说设置了token数目逐渐增多的Transformer级联,他们在训练时会从token较少的依次激活,当有足够高的预测置信度后就停止推理过程,通过调控token数目就会使得计算资源动态分配给简单或难的图像上。?
本文还进一步提出了特征复用和关系复用机制来减少冗余计算,特征复用使得下游模型可以基于之前提取的深层次特征之上进行训练;关系复用使得模型可以利用上游任务中得到的注意力关系,学习更加精确的注意力图谱。详情参见Fig 1。? DVT是一个通用的框架,当前SOTA的一些backbone,比如ViT,DeiT,T2T-ViT都可以直接使用,主要是为了提升计算效率;这种动态分配计算资源的方式非常适合在资源有限或资源波动的情况下,而这两种情况在现实世界十分常见,比如搜索引擎和移动应用程序。 ?本文在ImageNet,CIFAR数据集上验证了DVT的性能,与其他基准模型相比,有效的提升了计算效率。比如与T2T-ViT相比,计算成本降低了1.6-3.6倍,精度也没有下降。
Section II Related Work
Vision Transformer? ?受Transformer在NLP领域中的启发,ViT已经被用于图像识别任务,虽然ViT本身目前还不能与CNN在benchmark上相比较,但是在更大规模数据集上预训练之后可以获得有益的性能;DeiT则是研究了ViT的训练策略,并提出使用知识蒸馏,超越了ResNet的性能,随后又提出了一系列ViT的变体,如DeepViT,CPVT,LocalViT,CrossViT等。? ?另一种优化思路是尝试将CNN的归纳偏置集成到Transformer中,也有尝试将ViT用于其他视觉任务,如目标检测、语义分割等。当前主流的方法都是使用固定数目的token。 ?据我们所知,本文是第一个考虑基于输入来决定token数目的工作。 Efficient deep network? 网络的计算效率在实际应用场景中至关重要,关乎着实际的功耗、碳排放或延迟时间。在降低CNN计算出成本方面前人已经做了诸多探索。但在高效ViT方便仍有很大的探索空间,T2T-ViT提出一个轻量级的Tokens-2-token模块,精度可以匹配MobileNet;LeViT则通过引入卷积层来加速Transformer的推理;Swin Transformer则是用过基于移动窗口的方法来提升效率。 本文则是通过动态调整每个样本对应的网络结构来提升效率。 ?Dynamic Models? ?设计动态网络是一种有效的深度学习算法,MSDNet机器变体通过设计一个多分类器CNN网络,对简单的样本进行早退;还有的网络会跳过冗余层或者冗余的通道。但它们都是 基于CNN设计的本文则是为Vision Transforme量身定制的。? ?值得注意的是,虽然本文看起来有点类似多出口网络,但是还是有一定区别的。比如MSDNet和RANet是通过自适应调整网络深度获得效率提升的,即在单个网络中使用更少的网络层次来处理简单样本;而本文是通过级联多个Transformer,测试时是从多个网络获得推理结果;DVT并没有使用密集连接或者多尺度网络,调整的是token数目。并且上述网络都是纯CNN网络,而本文是基于Transformer的。
Section III Dynamic Vision Transformer
ViT会将2D图像转成1Dtoken,为了更好的识别较难的图像,需要token数目非常大,从而使得计算成本指数增长。因此本文提出使用多个Transformer,它们各自具有不同数目的token,它们会逐渐被激活直到得到一个置信度较高的预测结果。这样不同图像分配的计算资源是不一样的,从而提高了整体的计算效率。同时,如果所有Transformer都是单独训练的,那么下游Transformer计算时上游Transformer计算的结果将会被丢弃,这样也会导致效率低下,因此本文引入了特征复用和关系复用机制。
Part 1 Overview
Inference. ??首先从DVT的推理流程开始介绍。示意图参见Fig 2.对于每一个测试样本,首先用一组数量较少的token进行粗略表示,比如直接通过patch获得,然后经过一个Vision Transformer得到初始的预测结果;这时计算效率还比较高效;随后使用一定的评价指标来评估预测结果的置信度,确定结果是否可靠。如果结果足够可靠就会提前终止。 ??
以输入图像为条件获得适当的token数目,本文会级联不同token数目的Transformer,它们会被依次激活,直到获得一个足够可靠的预测结果。 ??Transformer之间引入了特征复用和关系复用。?如果预测结果不满足终止标准,原始图像会被切成更多patch,从而用更多token进行计算。 ??细节:token嵌入的维度不变,变得是token的数量,这样可以获得更细粒度的特征表示。同时会激活另一个结构一样的Transformer,但是他们权重不同。 ??这样就可以在保证效率的前提下更好的权衡困难样本的学习;同时为了进一步提升效率,新的Transformer还可以利用原先模型学到的特征和关系。上述的过程会重复进行,直至最终退出。 ???Training? ???训练时训练DVT在每个出口的预测结果,目标函数是:
x,y分别代表训练集的某个样本及对应标签,损失函数计算的是交叉熵损失函数。?Transformer backbone?DVT是一种通用且灵活的框架,可以使用当前主流的Vision Transformer作为backbone,如ViT,DeiT,T2T-ViT等。
Part 2 Feature reuse and Relationship reuse
本文的另一挑战是如何提高DVT的计算复用结果,因为一旦下游Transformer被激活,上游Transformer的计算结果就被丢弃了,这无疑是十分低效的,因为上游模型也是针对当前物体进行训练的,也提取到了有效的特征,虽然精度差了一点,因此本文提出特征复用和关系复用来重用上有计算结果,从而在引入最小的计算成本基础上获得性能提升。? ?
Feature reuse ?DVT中不同的Transformer都是为了提取特征表述,因此下游Transformer应该是基于上游的Transformer做进一步学习,而不是从头提取特征,因为上游计算结果对推断和下游模型都有贡献。 ?因此本文的feature reuse允许利用上游Transformer最后一层中输出的token(Zup)来学习下游模型中的layer-wise Embedding(El). ?
Fl指的是LN-MLP等一系列操作,目的是引入非线性,增加更灵活的变换。?随后图像token会会reshape到原始图像中对应的位置,在进行上采样、展平来匹配下游Transformer的token数量。 ??Fl中映射的维度一般设置会比较小,也是考虑到计算的效率。?最后上游Transformer计算得到的嵌入El会作为先验知识加入到下游模型中,可以看到会将El与Zl级联,变化只是将LN+MLP层中的维度由之前的D换成现在的D+D’.?El被称之为上下文嵌入(context embedding)。 ???细节:不会复用classification token和zeropad;?这种特征复用方式可以认为是增大了模型的深度。 ???
Relationship Reuse? Vision Transformer的一大优势是其自注意力模块可以集合整张图的信息,有效的建模了长程依赖关系。通常每一层Transformer layer需要学习一组注意力图来描述token之间的关系,因此下游模型还可以利用上有模型学习到的attention map的信息。 ?对于给定的输入zl,SA计算公式如下:
其中Al就是一组注意力图的数值。为了进行关系复用,本文首先将上游模型学习到的注意力图拼接起来,Nup代表token的数目,Natt-up代表注意力图=NhL,Nh表示注意力的头数,L则是层数。这样下游模型就可以同时利用自身的token和上游的注意力图进一步进行学习。 ?
这样下游模型的每一个SA模块都可以访问上游的注意力图学到的关系,另一方面,因为新生成的注意力图数值和上游的注意力图在是混合后一起进行计算的,因此他们之间的相对重要性可以通过调整注意力图的数值来自动学习。
rl并不是简单的上采样,详情参见Fig 5.每一块是HxW token对应的注意力图,现在H=W=2,需要先将所有token对应的att map reshape成HxW大小,再将他们上采样到H’xW’大小,再讲他们组合成HW x H’W’大小,完成上采样操作。
Part 3 Adaptive Inference
前面提到过,DVT会根据输入样本的易难程度逐渐增加token的数目,即不同难度的输入图像处理起来的计算复杂度也不尽相同。即在进行softmax获得最终预测结果时,会与设定的一个阈值进行比较,只有当预测超过阈值推断才会结束,同时将预测结果输出;否则输入图像将会被切分成更多token,并激活下游Transformer进行计算。因此最后一个Transformer输出的地方,阈值设定的是0.?由于上述算式不可微,因此需要通过[48]的遗传算法进行求解。
Section IV Experiments
Dataset: ?ImageNet,CIFAR-10/100 ??training:testing = 50000:10000? ??Backbone:? ??T2T-ViT,DeiT?Part 1 ??
Main Results ???Fig 6,7展示的是ImageNet上的分类结果。可以看到与原始的T2T-ViT相比,DVT精度达到了82.3%,计算量减少了3.6x。此外DVT仅通过调整输出阈值就可以灵活的获得曲线上的所有点。 ???
Practical efficiency of DVT? 本文还在单张2080Ti上测试了DVT的推理速度,Table 2是对比结果。可以看到在相同的推理速度下更小的模型精度提高了3.7-6.8%;在大型模型上性能不掉的前提下推理速度提升了1.4-1.7倍。?T able 3展示的是在CIFAR数据集上的对比结果。比T2T-ViT的计算量减少了3-9倍。 ?Table 4展示的则是与SOTA框架的对比结果.
Table 5则是与其他多出口网络模型对比的结果,可以看到明显优于其他基线网络。? Table6则是消融实验,验证各个部件的有效性。可以看到feature reuse和relationaship reuse都可以显著提升第二第三出口的精度,最多提升6%。但会给第一出口精度带来略微下降,主要是因为上有模型对下游模型做出的妥协。 ?Table 7展示的是不同复用设置对精度的影响?分别重用:上游对应层、classification token,只对第一层复用,移除LN层。 ??Table 8则是关系复用的结果。?分别重用:上游对应层注意力、只用最后一层注意力、MLP层替换成为线性层,采用原始的上采样操作。结果表明,最好使用所有层的注意力、以及不要使用原始的上采样。? ??
Table 9对比的是不同退出方法,分别是随即退出、基于熵、基于置信度。 可视化结果
Fig 9展示了第一级到第三级才预测正确的图像,可以看到分别对应简单图像和复杂图像。可以看到简单的样本一般姿态比较清晰、并且具有足够大的分辨率;而复杂的样本通常包含复杂的场景、姿势并不典型,需要更多的token进行表示。 ?Fig 10展示的是不同计算成本下不同级别出口的退出数量,可以看到随着后面级别输出图像的增多,DVT精度显著提高。
Section V Conclusion
本文的DVT致力于为每一类图像分配最优的token数目,通过逐渐激活不同token数目的Transformer来动态分配计算资源,本文还引入了feature reuse和relationship reuse来进一步提升计算效率。大量的实验表明,本文的DVT对计算效率有显著提升。
|