EfficientViT(2023CVPR)：具有级联组注意力的内存高效视觉Transformer!

EfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention

EfficientViT: 具有级联组注意力的内存高效视觉Transformer

万文长字，请耐心观看~

论文地址： https://arxiv.org/abs/2305.07027

代码地址： Cream/EfficientViT at main · microsoft/Cream · GitHub

1、前言

Vision Transformers，简称ViTs，因为其强大的模型能力而取得了巨大的成功。然而，这些模型的卓越性能伴随着高昂的计算成本，这使得它们不适合实时应用。在本文中，作者提出了一种新型的高速视觉Transformer，名为EfficientViT。研究发现，现有的变换器模型的速度通常受到内存操作效率低下的限制，尤其是多头自注意力（MHSA）中的张量重塑和逐元素函数。因此，作者设计了一种新的模块，采用三明治布局，即在高效的前馈网络（FFN）层之间使用单个内存绑定的MHSA层，能在提高内存效率的同时增强了通道间的通信。

此外，作者发现不同头（heads）之间的注意力图（attention maps）有很高的相似性，这导致了计算冗余。为了解决这个问题，作者提出了一种级联组注意力（cascaded group attention）模块，它将完整的特征（full feature）分割（splits）成不同的部分提供给注意力头，这不仅节省了计算成本，还提高了注意力的多样性。综合实验表明，EfficientViT 在速度和准确性之间取得了良好的平衡，超过了现有的高效模型。例如，EfficientViT-M5 在准确性上比 MobileNetV3-Large 高出1.9%，同时在Nvidia V100 GPU和Intel Xeon CPU上的吞吐量分别提高了40.4%和45.2%。与最近的高效模型MobileViT-XXS相比，EfficientViT-M2在准确性上提高了1.8%，运行速度在GPU/CPU上分别快了5.8倍/3.7倍，并且在转换为ONNX格式时快了7.4倍。

视觉变换器（Vision Transformers，简称ViTs）因其强大的模型性能和优越的表现在计算机视觉领域引起了一场革命。然而，这种不断提高的准确率是以模型尺寸和计算开销的增加为代价的。例如，SwinV2模型使用了30亿（3.0B）个参数，而V-MoE (Vision Mixture of Experts)模型更是使用了147亿（14.7B）个参数，在ImageNet数据集上达到最先进的性能。

近期有一些研究致力于设计轻量级且高效的视觉Transformer模型。然而，这些方法大多数都集中在减少模型参数或者计算量（Flops），这些只是速度的间接度量，并不能真实反映模型实际的推理吞吐量。例如，MobileViT-XS模型尽管只有700M Flops的计算量，但在Nvidia V100 GPU上运行的速度却比计算量为1,220M Flops的DeiT-T模型慢很多。尽管这些方法在减少Flops或参数的同时取得了良好的性能，但与标准的同构或层次化的Transformer模型相比（例如DeiT和SwinT），它们在实际的墙上时钟时间（wall-clock time，即实际经过的时间）上并没有显著的速度提升，因此并没有得到广泛的采用。

为了解决这个问题，本文探索了如何使视觉Transformer运行得更快，寻找设计高效Transformer架构的原则。基于流行的视觉Transformer模型DeiT 和SwinT，作者系统地分析了影响模型推理速度的三个主要因素，包括内存访问、计算冗余和参数使用。特别是，作者发现Transformer模型的速度通常受到内存访问限制。换句话说，内存访问的延迟阻碍了GPU/CPU计算能力的充分利用，从而对Transformer模型的运行速度产生了严重的负面影响。大多数效率低下的内存操作是多头自注意力（MHSA）中频繁的张量重塑和逐元素函数。作者观察到，通过适当调整 MHSA 和 FFN 层之间的比例，可以在不影响性能的情况下显著减少内存访问时间。此外，作者还发现一些注意力头倾向于学习类似的线性投影，导致注意力图冗余；分析表明，通过向每个注意力头提供不同的特征来显式分解每个头的计算可以缓解这个问题，同时提高计算效率。此外，现有的轻量模型经常忽略不同模块中的参数分配，因为它们主要遵循标准 Transformer 的配置。为了提高参数效率，作者使用结构化剪枝（structured pruning）来识别最重要的网络组件，并总结模型加速参数重新分配的经验指导。

红色文本表示内存限制操作(超过50%)，即操作所花费的时间主要由内存访问决定，而Multiplication乘法计算所花费的时间要少得多。

基于之前的分析和发现，作者提出了一个新的Transformer模型家族，名为EfficientViT。

首先，作者设计了一个具有三明治布局（Sandwich Layout）的新模块来构建模型。这种布局在前馈网络（FFN）层之间应用单个受内存限制的多头自注意力（MHSA）层，减少了 MHSA 中内存限制操作带来的时间成本，并应用更多的 FFN 层来允许不同通道之间的通信，从而提高内存效率。

接着，作者提出了一个新的注意力模块，称为级联组注意力（Cascaded Group Attention, CGA）。这个模块的核心思想是增强输入到注意力头（heads）的特征多样性；与以前所有注意力头使用相同特征的自注意力机制不同，CGA为每个注意力头提供不同的输入分割，并将输出特征跨注意力头进行级联。这种模块不仅减少了多头注意力中的计算冗余，而且还通过增加网络深度来提升模型的容量。

最后，作者通过扩展关键网络组件（如值投影value projections）的通道宽度，同时缩小那些重要性较低的组件（如FFN中的隐藏维度），来重新分配参数，最终提高了模型的参数效率。

总结：EfficientViT模型通过这三个关键的设计改进来实现更高的效率和性能：

内存效率：通过三明治布局减少内存操作的时间消耗。

计算效率：通过CGA模块减少计算冗余。

参数效率：通过重新分配参数来优化模型的关键组件。

2、Going Faster with Vision Transformers：使用更快的Vision Transformers

在本节中，作者将从内存访问、计算冗余和参数使用三个角度探讨如何提高视觉Transformer的效率。通过实证研究来识别潜在的速度瓶颈，并总结有用的设计指南。

2.1、Memory Efficiency

内存访问开销是影响模型速度的一个关键因素。Transformer模型中的许多操作，如频繁的重塑（reshaping）、逐元素加法（element-wise addition）和归一化（normalization），都是内存效率低下的，需要跨不同的内存单元之间进行耗时的访问。

尽管当前有些方法试图通过简化标准softmax自注意力的计算来解决这个问题，例如稀疏注意力（sparse attention）和低秩近似（low-rank approximation），但这些方法通常以牺牲准确性和有限的加速为代价。

在这项工作中，作者通过减少内存效率低下的层来节省内存访问成本。最近的研究表明，内存效率低下的操作主要位于 MHSA 而不是 FFN 层。然而，大多数现有的视觉Transformer模型都使用相同数量的MHSA和FFN层，这可能无法实现最优的效率。

作者探索了在具有快速推理能力的小模型中MHSA和FFN层的最优分配。具体来说，作者将Swin-T和DeiT-T模型缩小到几个子网络，这些子网络的推理吞吐量提高了1.25倍和1.5倍，并比较了不同比例的MHSA层的子网络性能。如下图所示，具有 20%-40% MHSA 层的子网络往往会获得更好的精度，这样的比率比采用 50% MHSA 层的典型 ViT 小得多。

横轴（MHSA Proportion）：表示MHSA层在模型中所占的比例。纵轴（Top-1）：表示模型的top-1准确率百分比。

左侧（Swin-T作为基线），右侧（DeiT-T作为基线）

1.25×/1.5×分别表示将基线模型进行缩小从而加速1.25/1.5倍。

此外，作者还测量了内存绑定操作的时间消耗，以比较内存访问效率，包括重塑、逐元素加法、复制和标准化。在具有 20% MHSA 层的 Swin-T-1.25× 中，内存限制操作减少到总运行时间的 44.26%，该观察结果还推广到 DeiT 和具有 1.5 倍加速的较小模型。事实证明，适当降低 MHSA 层利用率可以在提高模型性能的同时提高内存效率。

总结：这一节主要强调了内存访问开销对模型推理速度的影响，并提出了通过优化模型结构中的MHSA和FFN层的比例来提高效率的方法。

作者通过实验验证了减少MHSA层的比例并增加FFN层可以提高模型的推理吞吐量，从而实现更快的推理速度。

2.2、Computation Efficiency

MHSA将输入序列嵌入到多个子空间（也就是“头”）中，并分别为每个头计算注意力图。这种方法已被证明可以有效地提高模型性能。

尽管MHSA有效，但计算注意力图的成本很高。有研究表明，并非所有的注意力图都是至关重要的，意味着许多计算可能是冗余的。

为了节省计算成本，作者探索了如何在较小的视觉 Transformer 模型中减少冗余的注意力计算。作者训练了宽度缩小的 Swin-T 和 DeiT-T模型，这些模型的推理速度提高了1.25倍。在训练过程中，测量每个块内每个头与其他头之间的最大余弦相似度。

下图显示了在不同块（blocks）中每个头（head）的平均最大余弦相似度，用来衡量多头自注意力（MHSA）中不同头之间计算冗余。作者观察到每一块的注意力头之间存在高度的相似性，特别是在网络的最后几个块中，这表明许多注意力头学习了相似的特征投影，导致了计算上的冗余；使用特征分割的变体模型（深蓝色线）显示出较低的余弦相似度，这表明通过让每个头只处理部分特征，可以有效地降低头之间的相似度，减少计算冗余。

蓝色线（Blue lines）：表示原始缩小规模的模型（Swin-T-1.25x和DeiT-T-1.25x），这些模型的每个头都使用完整特征。

深蓝色线（Darkblue lines）：表示变体模型，这些模型修改了MHSA，使得每个头只接收完整特征的一个分割部分。

为了鼓励每个注意力头学习不同的模式，作者采用了一种直观的解决方案：只向每个注意力头提供完整特征的一部分。这类似于在卷积神经网络中使用的组卷积的思想。作者训练了修改后的MHSA的缩小模型变体，并计算了注意力相似度。结果表明，与MHSA使用相同的完整特征对所有头进行计算不同，使用不同通道的特征分割可以有效减少注意力计算冗余。

总结：这一节强调了在设计视觉Transformer模型时，通过优化MHSA部分，模型可以减少不必要的计算，提高计算效率，同时保持或甚至提高性能。

2.3、Parameter Efficiency

经典的ViT主要继承了NLP Transformer 的设计策略，例如，使用Q、K、V投影的等效宽度、增加每阶段的注意力头数以及将 FFN 中的扩展比设置为4。

受先前研究的启发，作者采用了泰勒结构化剪枝（Taylor structured pruning）方法来自动识别Swin-T和DeiT-T模型中的重要部分。

下图绘制了剩余输出通道与输入通道之间的比率，并且还给出了未剪枝模型中的原始比率以供参考。

QKV ratio(origin)：表示剪枝前，Query（Q）、Key（K）和Value（V）的输出通道与输入嵌入的比例。

QK ratio(pruned)：表示剪枝后，Q、K的输出通道与输入嵌入的比例。

V ratio(pruned)：表示剪枝后，仅Value（V）的输出通道与输入嵌入的比例。

FFN ratio(origin)：表示剪枝前，前馈网络（Feed Forward Network，FFN）的输出通道与输入嵌入的比例。

FFN ratio(pruned)：表示剪枝后，FFN的输出通道与输入嵌入的比例。

SwinT剪枝带来的影响：

（1）Baseline accuracy：剪枝前的准确率为79.1%；Pruned accuracy：剪枝后的准确率为76.5%，准确率有所降低。

（2）剪枝后，Q、K、V 和 FNN 的输出通道比例都有所减少，这表明通过剪枝移除了一些不重要的通道，减少了模型的复杂度。

什么是结构化剪枝？

结构化剪枝是一种网络优化技术，它通过移除不重要的通道（即神经网络中的过滤器或神经元），在保持模型性能的同时减少模型的复杂度和计算成本。它使用梯度和权重的乘积作为通道重要性的度量，这可以用来近似估计当该通道（或参数）被移除时，模型的性能下降程度。如果某个参数的权重和梯度乘积较大，那么移除这个参数可能会导致损失函数的值有较大增加，即模型性能下降较多。因此，通过使用权重和梯度的乘积来评估参数的重要性，可以在剪枝过程中优先保留那些对模型性能影响最大的参数，而移除那些影响较小的参数，以此来平衡模型的计算效率和性能。

3、Efficient Vision Transformer：高效的视觉Transformer

基于上述分析，在本节中，作者提出了一种新型的快速推理分层模型，名为 EfficientViT。架构概览如下图所示。

(a) EfficientViT的架构； (b) 三明治布局块（Sandwich Layout block）； (c) 级联组注意力（Cascaded Group Attention）

3.1、EfficientViT Building Blocks

作者提出了一种新的高效视觉Transformer模块，如上图 (b) 所示。它由内存高效的三明治布局、级联的群体注意力模块和参数重新分配策略组成，分别致力于提高内存、计算和参数方面的模型效率。

3.1.1、Sandwich Layout三明治布局

它减少了内存效率低下的自注意力层SA，增加了更高效的前馈网络FFN层，以促进不同特征通道之间的通信。具体来说，它在FFN层 ΦFi 之间应用单个自注意力层 ΦAi 来进行空间信息的混合，即先FFN->Attention->FFN，公式表示为：

其中 Xi 是第 i 个块的完整输入特征（full input feature），该块在单个自注意力层之前和之后使用 N个FFN 将 Xi 转换为 Xi+1，注意有N个FFN，但是只有1个Attention层。这种设计减少了模型中自注意力层带来的内存时间消耗，并应用更多的FFN层以允许不同特征通道之间高效地通信。作者还在每个FFN之前使用深度卷积（DWConv）引入局部结构信息的归纳偏置，增强模型能力。

3.1.2、Cascaded Group Attention级联组注意力CGA

MHSA 中的注意头冗余，导致计算效率低下；受到高效 CNN 中的组卷积的启发，作者为视觉Transformer提出了一个名为级联组注意力（CGA）的新注意力模块，它为每个注意力头提供完整特征的不同分割，从而分解各个头的注意力计算。形式上，这种注意力可以表述为：

其中Xi 是第 i 个块的完整输入特征，第 j 个头计算 Xij 上的自注意力，Xij 是输入特征 Xi 的第 j 个分割，即 Xi = [Xi1，Xi2，...，Xih] 且 1 ≤ j ≤ h，h 是头的总数。WQij 、WKij 和 WVij 是将输入特征映射到不同子空间的投影层，将各个头计算后得到的分割特征在维度上进行拼接，WPi 是将拼接特征投影回与输入特征维度一致的线性层，简而言之，WPi 是一个调整特征维度的线性层。

值得注意的是，如上图（c）的黄色箭头所示，级联操作会将每个头的输出添加到后续头的输入中，逐步细化特征表示；例如head1计算完自己对应的分割特征后，会将输出的特征添加到head2的输入特征中。

Xij' 是第 j 个头的输入分割特征 Xij 和第（j-1）个头的输出特征 Xij~ 之和。

此外，在Query投影后应用一个token交互层（使用深度卷积），假设 Q 是Query的原始输出，经过深度卷积处理后的Query表示为 Q'，公式为：Q' =DepthwiseConv(Q)。经过Query后的Q本身就包含全局信息，再经过深度卷积，获取到细致的局部信息，使得注意力机制可以更全面地理解输入特征。

总结：这种级联设计有两个优点。首先，为每个头提供不同的特征分割可以提高注意力图的多样性。与组卷积类似，级联组注意力可以将Flops和参数节省 h 倍，因为QKV层中的输入和输出通道减少了 h 倍。其次，级联注意力头可以增加网络深度，从而进一步提升模型容量，而无需引入任何额外参数。由于每个头中的注意力图计算使用较小的 QK 通道维度，因此它只会产生较小的延迟开销。

3.1.3、Parameter Reallocation参数重新分配

为了提高参数效率，通过扩大关键模块的通道宽度，同时缩小不重要模块的通道宽度，来重新分配网络中的参数。具体来说，为所有阶段的每个头中的 Q 和 K 投影设置较小的通道尺寸；对于 V 投影，允许它具有与输入特征相同的维度。FFN中的扩展比（expansion ratio）也从4减少到2。经过以上参数重分配，重要模块具有更多数量的通道来学习高维空间中的表示，这防止了特征信息的丢失；同时去除了不重要模块中的冗余参数，加快推理速度，提高模型效率。

3.2、EfficientViT Network Architectures：EfficientViT 网络架构

EfficientViT 的整体架构如下图所示。具体来说，作者引入了重叠补丁嵌入（overlapping patch embedding），将 16×16 补丁嵌入到具有 C1 维度的token中，从而增强了模型在低级视觉表示学习中的能力。

该架构包含三个阶段stage。每个阶段都堆叠若干个 EfficientViT Block，并且每个下采样层（subsampling）的 token 数量减少 4 倍（长和宽各缩小两倍）。为了实现高效的下采样，作者提出了一种 EfficientViT subsample block，它也具有三明治布局，只不过将自注意力层替换为反向残差块（inverted residual block）以减少下采样期间的信息损失。值得注意的是，作者采用 BatchNorm (BN) 而不是 LayerNorm (LN)，因为 BN 可以折叠到前面的卷积层或线性层中，这比 LN 具有运行时优势；还使用 ReLU 作为激活函数，因为常用的 GELU 或 HardSwish 速度慢得多，而且有时某些推理部署平台无法很好地支持。

作者构建了具有六种不同宽度和深度比例的模型系列，并为每个阶段设置了不同数量的头部。在早期阶段使用的block数比后期阶段使用的block数量少，类似于 MobileNetV3 和 LeViT，因为在早期阶段需要处理较大分辨率的特征，block数量过多会更耗时。作者用一个小因子factor（≤ 2）增加阶段stage的宽度，以减少后期阶段的冗余，如第 2 节中分析的那样。架构细节如下表所示：

Ci、Li、Hi 分别指第 i 个阶段的宽度（通道数）、深度（block块数）和注意力头数。

4、Experiments

4.1、Implementation Details

作者在 ImageNet-1K 上进行图像分类实验。使用 AdamW 优化器和余弦学习率调度器在 8 个 Nvidia V100 GPU 上从头开始训练 300 个epoch。总批量大小设置为 2,048；输入图像被调整大小并随机裁剪为 224×224；初始学习率为 1×10−3，权重衰减为 2.5×10−2；使用与DeiT 相同的数据增强，包括混合Mixup、自动增强auto-augmentation和随机擦除random erasing。

4.2、Results on ImageNet

作者将 EfficientViT 与 ImageNet 上流行的高效 CNN 和 ViT 模型进行比较，并在下表中报告结果。结果表明，在大多数情况下，作者提出的EfficientViT 在不同的评估设置中实现了最佳的准确性和速度权衡。

（1）与高效 CNN 的比较。首先将 EfficientViT 与普通 CNN 模型（例如 MobileNet、EfficientNet）进行比较。

具体来说，与 MobileNetV2 1.0× 相比，EfficientViT-M3 获得了 1.4% 的 top-1 准确率，同时在 V100 GPU 和 Intel CPU 上的运行速度分别快了 2.5× 和 3.0×。与最先进的 MobileNetV3-Large 相比，EfficientViT-M5 的准确度提高了 1.9%，但运行速度更快。此外，EfficientViT-M5 的精度与搜索模型 EfficientNet-B0 相当，同时在 V100 GPU/Intel CPU 上运行速度快 2.3 倍/1.9 倍。尽管EfficientViT 使用了更多参数，但它减少了影响推理速度的内存效率低下的操作，并实现了更高的吞吐量。

（2）与高效 ViT 的比较。将EfficientViT 与最新的高效视觉Transformer进行了比较。当在 ImageNet-1K 上获得相似的性能时，作者提出的EfficientViT-M4 在测试的 CPU 和 GPU 设备上的运行速度比最近的 EdgeViT-XXS 快 4.4 倍和 3.0 倍。与最先进的 MobileViTV2-0.5 相比，EfficientViT-M2 实现了稍微更好的性能和更高的吞吐量，在 GPU 和 CPU 设备上测试的吞吐量分别提高了 3.4 倍和 3.5 倍。与 Swin-T 相比，EfficientViT-M5 的精度低 4.1%，但在 Intel CPU 上速度快 12.3 倍，证明了所提出设计的效率。

（3）使用更高分辨率进行微调。最近关于 ViT 的研究表明，更高分辨率的微调可以进一步提高模型的能力。将最大模型 EfficientViT-M5 微调到更高分辨率。EfficientViT-M5↑384 在 V100 GPU 上的吞吐量达到 3,986 张图像/秒，top-1 精度达到 79.8%，而 EfficientViT-M5↑512 进一步将 top-1 精度提高到 80.8%，展示了处理更高分辨率图像的效率以及良好的模型能力。

4.4、Ablation Study

（1）三明治布局块的影响。我们首先提出了一项消融研究，通过用原始的 Swin 块替换三明治布局块来验证所提出的三明治布局设计的效率。 Top-1 精度下降了 3.0%，这验证了应用更多 FFN 而不是受内存限制的 MHSA 对于小模型来说更有效。此外，为了分析 self-attention 前后 FFN 数量 N 的影响，我们将数量从 1 更改为 2 和 3，块数量相应减少以保持相似的吞吐量。如表所示，由于缺乏远程空间关系，进一步增加FFN的数量并不有效，并且N=1实现了最佳效率。

（2）级联群体注意力的影响。我们提出了CGA来提高MHSA的计算效率。如表所示，用 MHSA 替换 CGA 会使准确度降低 1.1%，ONNX 速度降低 5.9%，这表明解决头冗余问题可以提高模型效率。对于没有级联操作的模型，其性能与 MHSA 相当，但比 CGA 差，证明了增强每个头的特征表示的功效。

（3）参数重新分配的影响。与没有 QKV 通道维度重新分配或 FFN 比率降低的模型相比，我们的 EfficientViTM4 的 top-1 精度分别提高了 1.4%/1.5%，GPU 吞吐量提高了 4.9%/3.8%，表明参数重新分配的有效性。

（4）其他组件的影响。我们消除了使用 DWConv 进行令牌交互、标准化层和激活函数的影响，如上表 1 所示。使用 DWConv，精度提高了 1.4%，且延迟开销较小，证明了引入局部结构信息的有效性。用 LN 替换 BN 会使精度降低 0.9%，GPU 速度降低 2.9%。使用 HardSwish 代替 ReLU 可提高 0.9% 的准确度，但会导致 ONNX 速度大幅下降 20.0%。激活函数是逐元素操作，在 GPU/CPU 上占用大量处理时间，因此使用 ReLU 代替更复杂的激活函数效率更高。

此外，我们研究了每个头中 QK 维度的选择以及 V 维度与输入嵌入的比率，如下图所示。结果表明，随着QK维度从4增加到16，性能逐渐提高，进一步增加会带来较差的性能。当 V 维与输入嵌入之间的比率从0.4增加到1.0时，性能从70.3%提高到71.3%；当进一步放大该比率至1.2时，仅获得0.1%的提升。因此，将 V 的通道设置为接近输入嵌入可以实现最佳参数效率。