EfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention
1、
从三个角度探讨如何提高vision transformers的效率:内存访问、计算冗余和参数使用。
2.1. Memory Efficiency
红色字体表示操作所花费的时间主要由内存访问决定,而用于计算的时间要小得多。
vision transformers中内存不高效的操作:reshaping, element-wise addition, and normalization
本文通过通过减少内存低效层来节省内存访问成本。
存储低效的层MHSA 层比FFN层多。大多数模型使用相同数量的这两层,无法实现最佳效率。于是作者探索了不同比例MHSA 层比FFN层设置。20%-40% MHSA 层效果比较好。
结果表明,适当降低MHSA层利用率可以在提高模型性能的同时提高内存效率。
2.2. Computation Efficiency
注意力计算非常耗费资源,并且有一些其实不重要。于是,作者探索了如何减少冗余注意力计算。测量每个头部和每个块内的剩余头部的最大余弦相似性。结果如下:
在注意力头之间存在较高相似性,尤其最后几个block。这一现象表明,许多头部学习相同完整特征的相似投影,并产生计算冗余。为了明确地鼓励头部学习不同的模式,我们应用了一种直观的解决方案,只给每个头部提供完整特征的一部分。我们用改进的MHSA训练缩减模型的变体,并计算相似性,如图4所示。这表明在不同的头部中使用不同的通道分割特征,而不是像MHSA那样对所有头部使用相同的完整特征,可以有效地减少注意力计算冗余。
2.3. Parameter Efficiency
典型的ViT主要继承了NLP变换器的设计策略,例如,使用Q、K、V投影的等效宽度,逐级增加头,并将FFN中的膨胀比设置为4。在轻量级模型中这些组成部件应该被重新精细设计。作者采用Taylor structured pruning -Taylor结构修剪来自动找到Swin-T和DeiT-T中的重要成分,并探索参数分配的基本原理。修剪方法在一定的资源约束下去除不重要的通道,并保留最关键的通道以最好地保持准确性。它使用梯度和权重的乘积作为信道重要性,这近似于去除信道时的损耗波动。
图5 表明:1) 前两个阶段保留了更多的维度,而最后一个阶段保留的维度要少得多;2) Q、K和FFN的尺寸在很大程度上被修剪,而V的维度几乎被保留下来,并且仅在最后几个块处减小。这些现象表明:1)典型的通道配置,在每个阶段后将通道加倍,或对所有块使用等效通道,可能会在最后几个块中产生大量冗余;2) 当它们具有相同的维度时,Q、K中的冗余度远大于V。V更喜欢相对较大的通道,接近输入嵌入维度。
- Efficient Vision Transformer
3.1. EfficientViT Building Blocks
Sandwich Layout:
self-attention layers减少,FFN layers增加。额外每个FFN之前使用深度卷积(DWConv)增加token interaction。引入局部结构信息的归纳偏差,提高模型的性能。
Cascaded Group Attention:
提出级联组注意力(CGA),它向每个头部提供完整特征的不同分割,从而明确地分解头部之间的注意力计算。
把heads分开分别计算注意力,再合起来,具体看图6.c
将每个头部的输出添加到后续头部,以逐步细化特征表示:
Parameter Reallocation:
Q和K投影设置了小通道尺寸。对于V投影,允许它具有与输入嵌入相同的维度。由于其参数冗余,FFN中的膨胀比也从4降低到2。
3.2. EfficientViT Network Architectures
每个阶段堆叠所提出的Ef ficientViT构建块,并且在每个子采样层,令牌的数量减少4倍(分辨率的2倍子采样)。为了实现高效的二次采样,提出了一种高效的ViT二次采样块,它也具有三明治布局,只是自注意层被倒置的残差块取代,以减少二次采样过程中的信息损失。在整个模型中采用BatchNorm(BN)而不是Layer Norm(LN),因为BN可以折叠到前面的卷积或线性层中,这是比LN的运行时优势。我们还使用ReLU[54]作为激活函数,因为常用的GELU或HardSwish要慢得多,而且有时不能很好地得到某些推理部署平台的支持。
- Experiments