ICCV 2023 | MIT韩松团队提出基于轻量级多尺度注意力的EfficientViT - 极术社区

导读

本文旨在解决在边缘设备上部署最先进的语义分割模型所面临的计算成本过高的问题。作者指出，先前的语义分割模型通常依赖于自注意力机制、计算密集型的大卷积核或复杂的拓扑结构，以获得良好的性能，但这些方法不适用于边缘设备。为此，论文提出了一种新的语义分割模型家族EfficientViT，旨在实现在边缘设备上的高效语义分割。

EfficientViT的核心是一种新颖的轻量级多尺度注意力模块，这个模块还使用小卷积核对附近的Token进行聚合，生成多尺度token，并在这些多尺度token上执行ReLU-based全局注意力，将计算复杂度从二次降低到线性，同时保留了同等的特征提取能力，并能够很好的将全局感受野与多尺度学习结合起来。

最终，该模型在性能和硬件效率之间取得了良好的平衡，为在边缘设备上部署语义分割应用提供了一种可行的解决方案。EfficientViT 相对于先前的语义分割模型在速度和性能上都有显著的优势，使其成为实际应用的有力选择。

方法

Lightweight Multi-Scale Attention

轻量级多尺度注意力模块旨在在边缘设备上进行语义分割时实现性能和效率的平衡。同时，在性能方面，全局感受野和多尺度学习对语义分割任务非常重要，以提高模型的性能。不同于先前的多尺度注意力机制模块，本文探索的关键点在于如何仅依赖对硬件友好的算子来实现同等的全局感受野和多尺度学习，这无疑就点名了 Transformer 原生的 MHSA，计算成本高，效率低。

如图2右侧所示，展示了 LMSA 的结构，其核心思想是通过降低一定的模型容量来显著提高计算效率，这里面涉及到一个关键组件便是用于实现全局感受野的轻量级ReLU-based Attention。

ReLU-based Attention

众所周知，为了实现全局感受野，传统方法通常使用heavy的自注意力机制，但该机制计算复杂度高，不适用于边缘设备。因此，本文提出使用轻量级的ReLU-based全局注意力来代替复杂度高的自注意力机制。这是一种线性注意力模块，之前在其他领域有应用，但在语义分割中尚未使用。

通过使用ReLU-based全局注意力，可以实现全局感受野并将计算复杂度从二次降低到线性，而不改变原有的特征提取能力。此外，ReLU-based全局注意力不涉及像softmax等硬件效率低下的操作，因此在硬件上更加高效。

Multi-Scale Token

单独使用ReLU-based注意力的容量有限，因此为了增强多尺度学习能力，论文提出了从附近的 Q/K/V token 中汇聚上下文信息以获取多尺度token的方法。这个信息聚合过程对于每个头部中的每个Q、K和V都是独立的。为了避免对硬件效率的影响，作者使用小卷积核来执行信息聚合。

此外，为了进一步在GPU上高效执行这些操作，本文基于组卷积范式，将所有的深度可分离卷积（DWConv）融合为单个DWConv，即将所有的 1x1 卷积融合为单个 1x1 组卷积。在完成多尺度标记生成后，对它们进行全局注意力操作以提取多尺度的全局特征。最后，我们便能够将来自不同尺度的特征沿着头部维度进行连接，然后输入到最终的线性投影层以融合这些特征。

EfficientViT Architecture

如前所述，EfficientViT是一种基于轻量级多尺度注意力（Lightweight Multi-Scale Attention，MSA）模块构建的模型家族，专为在边缘设备上进行语义分割任务而设计，框架如下：

如图所示，EfficientViT采用了标准的backbone-head（编码器-解码器）架构设计，其核心构建块是EfficientViT模块，如图2左侧所示。这个模块包括两个关键组件：

轻量级多尺度注意力模块（Lightweight MSA Module）：这个模块用于提取上下文信息，以帮助模型理解图像中的全局信息。
MBConv：用于提取局部信息，有助于捕获图像中的局部特征。

下面简单介绍下其网络结构：

骨干网络：EfficientViT 的骨干也遵循标准的设计，包括steam外接四个stage，其中特征图大小逐渐减小，通道数逐渐增加。其中，EfficientViT 模块被插入到第三和第四阶段。为了进行下采样，使用了带有步长为2的MBConv。
头部：P2、P3和P4分别代表第二、第三和第四阶段的输出，形成一个特征图金字塔。为了简化和提高效率，使用了 1x1 卷积和标准的上采样操作（例如双线性/双三次上采样）来匹配它们的空间和通道大小，然后通过加法操作来融合它们。由于 EfficientViT 的骨干已经具备了强大的上下文信息提取能力，头部设计相对简单，包括几个 MBConv 块和输出层（即预测和上采样）。在实验中，研究人员发现这种简化的头部设计足以实现SOTA性能，这得益于轻量级MSA模块的作用。

此外，为了满足不同的效率要求，同样提供了缩放机制，详细配置见下表：

实验

首先，如下表所示，在ImageNet图像分类任务上，EfficientViT的骨干网络表现综合性能还行。特别是EfficientViT-B3在 ImageNet 上取得了84.2的Top-1准确度，相较于 EfficientNet-B6 提供了0.2 的准确度提升，并且速度快了 7.9 倍。与 ConvNeXt 相比，性能相当情况下推理速度也挺溜。

其次，在Cityscapes语义分割数据集上，EfficientViT相对于先前的SOTA语义分割模型也取得了显著的效率改进，而不会牺牲性能。特别地，与SegFormer相比，EfficientViT实现了高达13倍的MACs（浮点运算次数）节省和15倍的延迟减少，同时提供更高的mIoU（平均交并比）。与SegNeXt相比，EfficientViT 在移动设备上实现了高达9.3倍的速度提升，同时保持更高的mIoU。哪怕是在具有相似计算成本的情况下，EfficientViT-B3 相对于SegFormer-B1取得了+4.7的mIoU增益。

在ADE20K语义分割数据集上，EfficientViT同样实现了显著的效率改进。例如，EfficientViT-B1相对于SegFormer-B1取得了+0.5的mIoU增益，同时减少了5.9倍的MACs和6.5倍的延迟。EfficientViT-B2相对于SegNeXt-S取得了+0.8的mIoU增益，同时减少了2.7倍的MACs和5.2倍的延迟。