论文

 

导读

ViT 的最新进展在视觉识别任务中取得了出色的表现。卷积神经网络 (CNN) 利用空间归纳偏差来学习视觉表示,但这些网络是空间局部的。ViTs 可以通过其self-attention机制学习全局表示,但它们通常是heavy-weight的,不适合移动设备。

在本文中提出了Cross Feature Attention(XFA) 以降低 Transformer 的计算成本,并结合高效的mobile CNNs形成一种新颖的高效轻量级 CNN-ViT 混合模型 XFormer,可作为通用主干学习全局和局部表示。实验结果表明,XFormer 在不同的任务和数据集上优于众多基于 CNN 和 ViT 的模型。

在 ImageNet-1K 数据集上,XFormer 使用 550 万个参数实现了 78.5% 的 top-1 准确率,在相似数量的参数下,比 EfficientNet-B0(基于 CNN)和 DeiT(基于 ViT)的准确率分别提高了 2.2% 和 6.3%。XFormer的模型在转移到目标检测和语义分割任务时也表现良好。在 MS COCO 数据集上,XFormer 在 YOLOv3 框架中超过 MobileNetV2 10.5 AP(22.7 → 33.2 AP),只有 6.3M 参数和 3.8G FLOPs。在 Cityscapes 数据集上,只有一个简单的 all-MLP 解码器,XFormer 实现了 78.5 的 mIoU 和 15.3 的 FPS,超过了最先进的轻量级分割网络。

 

贡献

Transformer是一种流行的神经网络结构,它使用高度表达性的注意力机制计算网络输出。它们起源于自然语言处理(NLP)社区,已被证明能有效地解决NLP中的广泛问题,如机器翻译、表征学习和问答。

最近,视觉Transformer在视觉界越来越受欢迎,并已成功应用于广泛的视觉应用,如图像分类、目标检测、图像生成和语义分割。为了推进视觉Transformer在工业界的应用,轻量级Transformer逐渐受到了诸多研究人员的关注,然而,该领域的研究仍处于初期阶段。例如,当模型参数量固定在600万时,DeiT在 ImageNet 分类任务上的准确率比 MobileNetV3 低 3%,并且 FLOP 多 6 倍。因此,轻量级 CNN 已作为许多移动视觉任务的通用主干,而当前大多数基于 ViT 的模型重量级大,难以优化和缩减,应用到下游任务时需要参数量昂贵的解码器。 最近的工作已经证明了为特定任务(如语义分割)设计轻量级 ViT 模型的可能性,但基于 ViT 的模型远没有被认为是移动设备的实用选择。 因此,设计一个高效且轻量级的 ViT 模型至关重要。

目前已经提出了几种方法 [8, 42, 54, 9, 32] 来改进 ViT。 一个简单的想法是通过在早期阶段使用卷积穿插到Transformer模块中来组合混合模块。 这些方法 [8, 42] 旨在通过引入空间局部偏差来降低整体模型的计算成本并使优化更容易。 另一种方法是优化转换器中的自注意力操作。 原始的self-attention在token数量上具有二次计算复杂度,这给下游任务的转移造成了巨大的瓶颈。 以前的工作采用分解[9]、低秩近似[54]或散列[32]等方法来降低复杂性,但一般的权衡是性能显着下降。

本文中提出了Cross Feature Attention(XFA) 以降低 Transformer 的计算成本,并结合高效的mobile CNNs形成一种新颖的高效轻量级 CNN-ViT 混合模型 XFormer,可作为通用主干学习全局和局部表示。实验结果表明,XFormer 在不同的任务和数据集上优于众多基于 CNN 和 ViT 的模型。在 ImageNet-1K 数据集上,XFormer 使用 550 万个参数实现了 78.5% 的 top-1 准确率,在相似数量的参数下,比 EfficientNet-B0(基于 CNN)和 DeiT(基于 ViT)的准确率分别提高了 2.2% 和 6.3%。XFormer的模型在转移到目标检测和语义分割任务时也表现良好。在 MS COCO 数据集上,XFormer 在 YOLOv3 框架中超过 MobileNetV2 10.5 AP(22.7 → 33.2 AP),只有 6.3M 参数和 3.8G FLOPs。在 Cityscapes 数据集上,只有一个简单的 all-MLP 解码器,XFormer 实现了 78.5 的 mIoU 和 15.3 的 FPS,超过了最先进的轻量级分割网络

 

方法

一个标准的 ViT 模型首先使用patch size  h×w 将输入 图像X转化为一系列flattened patches  Xf∈RN*PC,其中P=hw和N=HW/P表示token数。然后将Xf投影到固定的 D 维空间XD∈RN*D并使用一堆transformer blocks来学习inter-patch表示。由于忽略了空间归纳偏差,ViT 通常需要更多参数来学习视觉信息。此外,transformers 中 self-attention 的昂贵计算导致优化此类模型的瓶颈。

图1 XFormer

在本节提出了 XFormer,这是高效、轻量级的 CNN-ViT 框架,以解决 ViT 中的上述问题。首先介绍了一种提高自注意力效率的新方法,然后说明了新的 CNN-ViT 混合模型的架构设计。

1.1 Cross Feature Attention (XFA)

1、Attention Overview

Transformer 的主要计算瓶颈之一在于 self-attention 。在最初的 self-attention 过程中, XD首先用于通过线性投影生成query Q、key K 和value V。它们都具有相同的维度(N × D),其中 N 是图像token数,每个维度为 D。然后计算注意力分数为:

其中 σ 是 softmax 操作,dhead 是head维度。计算注意力分数的计算复杂度为O(N2D)。 self-attention 的二次复杂性导致了巨大的计算瓶颈,这使得 ViT 模型难以在移动设备上按比例缩小。

2、Efficient Attention

为了解决 self-attention 中的二次复杂性问题,作者提出了一种新的注意力模块结构,称为Cross Feature Attention(XFA)。在之前的工作之后,首先沿特征维度 D 对query Q 和key K 应用 L2 归一化:

QKT在原始 self-attention 中的直觉是定位应该关注的重要图像块。

但是直接计算会导致不必要的冗余和计算开销。相反,作者为 K 构建了2个中间分数:查询上下文分数 Kc和查询特征分数Kf。使用2个卷积核矩阵Wc和Wf 沿token维度 N 计算 ,沿特征维度 D 计算 。借助卷积滤波器中间分数向量可以表示计算注意力图的更紧凑的表示,同时也降低了计算成本:

最后,将Cross Feature Attention(XFA)定义为:

其中 λ 是一个温度参数,用于动态调整不同transformer层中的比例因子,从而提高训练稳定性。注意到归一化将注意力值限制在一定范围内,因此放弃了冗余且昂贵的 softmax 操作。与原始的具有二次复杂度的 self-attention 不同,XFA模块将计算成本从O(N2D)降低到O(ND2).

3、Comparison with Self-attention

本文提出的 XFA 模块和原始注意力之间的主要区别是:

  • XFA通过构建中间查询上下文和特征分数,沿特征维度D计算注意力图,大大降低了计算成本;
  • XFA使用可学习的温度缩放参数来调整归一化,并且不受 softmax 操作的影响。

本文方法为二次复杂度问题提供了解决方案,并且对于资源受限的设备更有效且对移动设备更友好。

1.2 Building XFormer

1、MobileNetV3 Block

MobileNetV2 首先引入了inverted residual和linear bottleneck以构建更高效的层结构。MobileNetV3 随后添加了squeeze excitation (SE) 模块以处理更大的表示特征。最近的工作证明了通过在 ViT 的早期阶段结合卷积层来提高 ViT 性能的合法性。受这种直觉的启发,继续为轻量级模型探索这种 CNN-ViT 混合设计。MobileNetV3 块内的操作可以表述为:

其中xin是前一层的输入特征, xout是 MV3 Block 的输出特征,SE 是squeeze excitation模块。Convdw表示depth-wise卷积操作, Convpw表示point-wise卷积操作。

2、XF Block

利用提出的Cross Feature Attention模块介绍了 XF Block,一个精心设计的轻量级 transformer module。XF Block内的操作可以表述为:

其中xin是前一层的输入特征, xmid是 XFA 模块的特征,  xout是 XF Block 的输出特征。LN 表示层归一化操作,MLP 为全连接层。

3、Patch Size Choice

对于较大的模型(ViT-Large),基于 ViT 的模型通常采用 8×8、16×16 甚至 32×32 的Patch Size。具有较大Patch Size的优点之一在于,对于分类等图像级任务,ViT 可以有效地提取图像块信息,而不会增加过多的计算开销。

最近的工作表明,当迁移到语义和实例分割等下游任务时,更小的Patch Size更受青睐,因为它可以增强 Transformer 学习更好的像素级信息的能力,这通常会带来更好的性能。

此外,随着Patch Size的减小,token数 N 会大得多。线性复杂度 XFA 模块可以避免潜在的计算瓶颈。在网络设计中,每个 XF 块的Patch Size设置为 2×2。

4、XFormer

在 MobileNetV3 和 XF Block 的基础上提出了 XFormer,这是一种 CNN-ViT 混合轻量级模型,由堆叠的 MobileNetV3 Block 和 XF Block 组成,用于学习全局和局部上下文信息。与之前设计高效 CNN 的工作一样,本文的网络由Stemin块组成,用于提取原始图像特征,在特征由 CNN 和transformer blocks处理后,使用 Stemout、全局池化和全连接层来产生最终的 logit 预测。

在主要处理块中,XFormer 有5个阶段。前2个阶段仅包含 MobileNetV3 Block (MV3) ,因为卷积块比全 ViT 模型更能提取重要的图像级特征表示,并随后帮助转换器块看得更清楚。最后三个阶段中的每一个都包括一个 MV3 块和几个 XF 块。结合来自 MV3 块的局部归纳偏差和来自 Transformer Block 的全局信息,网络可以学习更全面的特征表示,可以轻松地转移到不同的下游任务。

先前关于 ViT 的工作表明,应在更深的Transformer layers中使用更大的 MLP 比率,并且 Q-K-V 维度应相对小于嵌入维度,以便在性能和模型大小之间进行更好的权衡。作者遵循这些建议并相应地设计轻量级模型。对于三个不同阶段的 XF Block ,它们的 MLP 比率、嵌入维度和 Q-K-V 维度分别设置为(2、2、3)、(144、192、240)和(96、96、96)。在全连接层中将 ReLU 替换为 GELU 激活;在所有其他层中使用 SiLU。对于 MV3 Block,扩展比均设置为 4。规格如表 1 所示。

5、Model efficiency

模型的总参数大小只有 550 万。与类似大小的基于 ViT 的模型相比,本文的模型可以更有效地处理高分辨率图像并避免潜在的内存瓶颈(见表 2)。

例如,当输入分辨率为 1024×1024 时,与使用原始自注意力的 MobileViT 相比,XFormer 的推理速度提高了近 2 倍,GPU 内存使用量减少了 32%。本文的模型可以轻松处理高分辨率吞吐量,而不会出现内存瓶颈。最重要的是,XFormer 提供了比比较模型更好的精度(参见表 3),在模型大小和性能之间实现了很好的平衡。

 

实验

表3显示,XFromer在不同模型大小上的性能优于其他轻量级基于cnn的模型。x前者将mobileNetV2提高了5.2%,MobileNetV3提高了3.3%,效率net-b0提高了2.2%。与重量CNN模型相比,我们的模型比ResNet-50[24]和DenseNet[31]具有更少的参数。

表4显示,在YOLOv3框架内,Xfrome在目标检测方面始终优于以前的移动方法。YOLOv3的性能比MobileNetV2基线大大提高了46%(22.7→33.2AP),比最先进的检测模型更准确,如CenterNet[18](+7.9AP),YOLOX[22](+4.7AP),同时保持相对轻(6.4M参数,3.8GFLOPs)。在推理速度方面,特别是对于高分辨率输入,XFromer计算速度比基于XFA的检测网络快,如YOLOS[21](+17FPS)和MobileViT[42](+3FPS),这表明由于我们的XFA模块,我们的方法确实计算效率高。

表5显示,XFromer在城市景观数据集上取得了与最先进的移动语义分割方法竞争的结果。与基于cnn的模型相比,我们的方法的性能优于MobileNetV2,不需要大量的解码器。此外,使用MobileViT提高了3.8%(75.6→78.5mIoU),也优于最先进的轻量级LVT[59](+3.6mIoU)和[58](+2.0mIoU)。最重要的是,Xfrome比之前的ViT骨干实现了更快的推理速度,验证了我们的模型在处理高分辨率输入方面的效率。这表明我们提出的XFromer模型可以作为语义分割任务的有效骨干。定性结果见图3。

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