论文:https://arxiv.org/abs/2205.03436

 

导读

在计算机视觉领域,基于Self-attention的Transformer模型已经成为CNN之外的热门架构。尽管Transformer模型的变种具有越来越高的识别精度,但由于Self-attention的二次复杂度,现有的ViT在计算和模型大小方面都有较高的要求。此前,CNN的一些成功的设计选择(例如,卷积和分层结构)已经被引入到最近的ViT中,但它们仍然不足以满足移动设备有限的计算资源需求。这促使研究人员基于MobileNet-v2开发了轻型模型MobileViT,但MobileViT与MobileNet-v2仍然存在性能差距。在这项工作中,作者进一步推进这一研究方向,引入了EdgeViTs这一新的轻量级ViTs家族,也是首次使基于Self-attention的视觉模型在准确性和设备效率之间的权衡中达到最佳轻量级CNN的性能。作者在图像分类、目标检测和语义分割方面进行了大量实验,验证了EdgeViTs在移动硬件上的准确性-效率平衡能力,几乎在所有情况下都实现了对其他ViT的超越,并可以达到最高效CNN的性能。

 

贡献

Transformer模型的建模能力自提出依赖就受到了诸多研究人员的一致拥护,现有研究同样证实其性能与模型深度同样可以呈现出正比关系。为了在保证精度的同时降低计算量,早期研究提出了以下三种思路:1)使用分层架构,其空间分辨率(即令牌序列长度)在各个阶段中逐渐下采样,2)使用控制输入标记序列长度和参数共享的局部分组自注意机制,3)对注意力机制汇总的key、value进行子采样。然而,这些进步仍然不足以满足移动平台和边缘平台的设计要求和限制。

本文以计算量小的ViT为目标,考量了MobileViT中的一系列设计需求:

  • 推理效率,应当满足移动平台的低算力;
  • 模型大小,对边缘设备的存储空间友好可接受;
  • 部署方便,可以在多种运算框架下进行训练和推理。

作为在这个方向上的初步尝试,MobileViTs仍然落后于 CNN 同行。 此外,它的评估协议将模型大小(即参数数量)作为竞争对手的选择标准(即,仅将模型的准确性与相似的参数数量进行比较),但这并不能充分体现模型的对比结果。

本文通过引入一个基于Self-attention和卷积的最优集成的高成本的local-global-local(LGL)信息交换瓶颈,提出了EdgeViTs这一新的轻量级ViTs家族。对于移动设备专用的评估,不依赖于不准确的proxies,如FLOPs的数量或参数,而是采用了一种直接关注设备延迟和能源效率的实用方法。本文首次使基于Self-attention的视觉模型在准确性和设备效率之间的权衡中达到最佳轻量级CNN的性能。

 

 

方法

为了设计适用于移动/边缘设备的轻量级ViT,作者采用了最近ViT变体中使用的分层金字塔结构(图2(a))。Pyramid Transformer模型通常在不同阶段降低了空间分辨率同时也扩展了通道维度。每个阶段由多个基于Transformer Block处理相同形状的张量,类似ResNet的层次设计结构。

基于Transformer Block严重依赖于具有二次复杂度的Self-attention操作,其复杂度与视觉特征的空间分辨率呈2次关系。通过逐步聚集空间Token,Pyramid Transformer可能比各向同性模型(ViT)更有效。

在这项工作中,作者深入到Transformer Block,并引入了一个比较划算的Bottlneck,Local-Global-Local(LGL)(图2(b))。LGL通过一个稀疏注意力模块进一步减少了Self-attention的开销(图2(c)),实现了更好的准确性-延迟平衡。

 

Local-Global-Local bottleneck

Self-attention已被证明是非常有效的学习全局信息或长距离空间依赖性的方法,这是视觉识别的关键。另一方面,由于图像具有高度的空间冗余(例如,附近的Patch在语义上是相似的),将注意力集中到所有的空间Patch上,即使是在一个下采样的特征映射中,也是低效的。

因此,与以前在每个空间位置执行Self-attention的Transformer Block相比,LGL Bottleneck只对输入Token的子集计算Self-attention,但支持完整的空间交互,如在标准的Multi-Head Self-attention(MHSA)中。既会减少Token的作用域,同时也保留建模全局和局部上下文的底层信息流。

为了实现这一点,作者将Self-attention分解为连续的模块,处理不同范围内的空间Token(图2(b))。

这里引入了3种有效的操作:

  • Local aggregation:仅集成来自局部近似Token信号的局部聚合
  • Global sparse attention:建模一组代表性Token之间的长期关系,其中每个Token都被视为一个局部窗口的代表;
  • Local propagation:将委托学习到的全局上下文信息扩散到具有相同窗口的非代表Token。

将这些结合起来,LGL Bottleneck就能够以低计算成本在同一特征映射中的任何一对Token之间进行信息交换。

 

Local aggregation

对于每个Token,利用Depth-wise和Point-wise卷积在大小为k×k的局部窗口中聚合信息(图3(a))。

 

Global sparse attention

对均匀分布在空间中的稀疏代表性Token集进行采样,每个r×r窗口有一个代表性Token。这里,r表示子样本率。然后,只对这些被选择的Token应用Self-attention(图3(b))。这与所有现有的ViTs不同,在那里,所有的空间Token都作为Self-attention计算中的query被涉及到。

 

Local propagation

通过转置卷积将代表性Token中编码的全局上下文信息传播到它们的相邻的Token中(图3(c))。最终,LGL bottleneck可以表达为:

这里,x表示输入张量。Norm是Layer Normalization操作。LocalAgg表示局部聚合算子,FFN是一个双层感知器。GlobalSparseAttn是全局稀疏Self-attention。LocalProp是局部传播运算符。为简单起见,这里省略了位置编码。注意,所有这些操作符都可以通过在标准深度学习平台上的常用和高度优化的操作来实现。因此,LGL bottleneck对于实现是友好的。

 

实验

LGL bottleneck与最近的PVTs和Twins-SVTs模型有一个相似的目标,这些模型试图减少Self-attention开销。然而,它们在核心设计上有所不同。PVTs执行Self-attention,其中Key和Value的数量通过strided-convolutions减少,而Query的数量保持不变。换句话说,PVTs仍然在每个网格位置上执行Self-attention。

在这项工作中,作者质疑位置级Self-attention的必要性,并探索由LGL bottleneck所支持的信息交换在多大程度上可以近似于标准的MHSA。Twins-SVTs结合了Local-Window Self-attention和PVTs的Global Pooled Attention。这不同于LGL bottleneck的混合设计,LGL bottleneck同时使用分布在一系列局部-全局-局部操作中的Self-attention操作和卷积操作。

如实验所示(表2和表3)所示,LGL bottleneck的设计在模型性能和计算开销(如延迟、能量消耗等)之间实现了更好的权衡。

在目标检测和语义分割任务上,EdgeViT均可以在保证轻量级要求的前提下实现一流的性能。

 

参考

  • EdgeViTs: Competing Light-weight CNNs on Mobile Devices with Vision Transformers
  • https://mp.weixin.qq.com/s/7JYRbOXsQtecCnVEg6hQxA
  • https://mp.weixin.qq.com/s?__biz=MzU5OTA2Mjk5Mw==&mid=2247494785&idx=1&sn=1a1cd39a39e1ba7d5160b90f6792a10a&chksm=feb80c3fc9cf85298e397999e2c5e841b667e43915b447db5c2593d9f9edf9633b4f7f633c47&scene=21#wechat_redirect
  • https://zhuanlan.zhihu.com/p/395228490

内容中包含的图片若涉及版权问题,请及时与我们联系删除