FastViT：一种使用结构重新参数化的快速混合视觉变换器

摘要

论文：https://arxiv.org/pdf/2303.14189v1.pdf

代码链接：https://github.com/apple/ml-fastvit

最近，Transformer和卷积设计的融合使得模型的准确性和效率得到了稳步提高。在这项工作中，我们引入了FastViT，这是一种混合视觉Transformer架构，实现了最新的延迟-准确率权衡。为此，我们引入了一种新型的令牌混合操作符RepMixer，这是FastViT的一个构建块，它使用结构重参数化来降低网络中的跳连接，从而降低内存访问成本。我们进一步应用训练时过参数化和大核卷积来提高准确性，并从经验上证明这些选择对延迟的影响最小。我们的模型比CMT(最近最先进的混合Transformer架构)快3.5倍，比EfficientNet快4.9倍，比移动设备上的ConvNeXt快1.9倍，在ImageNet数据集上具有相同的精度。在相似的延迟下，我们的模型在ImageNet上获得的Top-1精度比MobileOne高4.2%。我们的模型在几个任务上始终优于竞争对手的架构——图像分类、检测、分割和3D网格回归，在移动设备和桌面GPU上的延迟都有显著改善。此外，我们的模型对分布外样本和损坏具有很高的鲁棒性，优于其他竞争的鲁棒模型。

1、简介

视觉Transformer[13]在图像分类、检测和分割等多项任务上取得了最先进的性能[34]。然而，这些模型传统上在计算上是昂贵的。最近的研究[28,38,41,57,66]提出了降低视觉Transformer的计算和内存要求的方法。最近的混合架构[9,16,41,62]有效地结合了卷积架构和Transformer的优势，构建了在广泛的计算机视觉任务上具有高度竞争力的架构。我们的目标是构建一个模型，实现最先进的延迟-精度权衡。

最近的视觉和混合Transformer模型[16,41,42,52]遵循Metaformer[65]架构，该架构由具有跳过连接的令牌混合器以及具有另一个跳过连接的馈送前馈网络（FFN）。这些跳过连接占用了由于内存访问成本增加而导致的显著延迟开销[12,55]。为了解决这种延迟开销，我们引入了RepMixer，这是一种完全可重新参数化的令牌混合器，使用结构重参数化来去除跳过连接。RepMixer块还使用深度卷积进行类似ConvMixer[54]的信息空间混合。然而，关键的区别在于我们的模块可以在推理时进行重参数化以去除任何分支。

为了进一步提高延迟、FLOPs和参数计数，我们将所有密集的k×k卷积替换为其分解版本，即深度卷积后接点对点卷积。这是一种由高效架构[24,25,46]使用的常见方法，用于提高效率指标，但是，如表1所示，简单地使用这种方法会损害性能。为了增加这些层的容量，我们使用线性训练时过参数化，如[10-12,17,55]中引入的。这些额外的分支仅在训练期间引入，并在推理时进行重参数化。

此外，我们在网络中使用大核卷积。这是因为尽管基于自注意力的令牌混合在获得具有竞争力的准确性方面非常有效，但它们在延迟方面效率低下[38]。因此，我们在馈送前馈网络（FFN）和补丁嵌入层中引入了大核卷积。这些变化对模型的整体延迟影响最小，同时提高了性能。

因此，我们介绍了基于三个关键设计原则的FastViT - i)使用RepMixer块来删除跳过连接，ii)使用线性列车时间过参数化来提高准确性，iii)在早期阶段使用大卷积核来替代自关注层。

FastViT相比其他混合视觉Transformer架构在延迟方面取得了显着的改进，同时保持在图像分类、目标检测、语义分割和3D手部网格估计等几个任务上的准确性。我们在iPhone 12 Pro设备和NVIDIA RTX-2080Ti台式GPU上部署了最新的最先进架构，并进行了全面的分析。

在图1中，我们展示了在ImageNet Top-1精度为83.9%的情况下，我们的模型在iPhone 12 Pro移动设备上比EfficientNetB5[49]快4.9倍，比EfficientNetV2-S[50]快1.6倍，比CMT-S[16]快3.5倍，比ConvNeXtB[35]快1.9倍。在ImageNet Top-1精度为84.9%的情况下，我们的模型在桌面GPU上的速度与NFNetF1[1]一样快，而在移动设备上的速度要小66.7%，使用的FLOPs减少50.1%，速度提高42.8%。在iPhone 12 Pro移动设备上的延迟为0.8ms时，我们的模型在ImageNet上的Top-1精度比MobileOne-S0高4.2%。对于使用Mask-RCNN[19]头部的MS COCO对象检测和实例分割，我们的模型达到了与CMTS[16]相当的性能，同时减少了4.3倍的骨干延迟。对于ADE20K上的语义分割，我们的模型比PoolFormer-M36[65]提高了5.2%，同时在iPhone 12 Pro移动设备上的骨干延迟降低了1.5倍。在3D手部网格估计任务上，我们的模型在GPU上进行基准测试时比MobileHand[15]快1.9倍，比最近最先进的MobRecon[4]快2.8倍。

除了准确性度量外，我们还研究了我们的模型对损坏和分布外样本的鲁棒性，这些样本并不总是与准确性很好地相关。例如，PVT[58]在ImageNet数据集上实现了高度竞争力的性能，但根据Mao等人的报道[37]，其对损坏和分布外样本的鲁棒性非常差。在现实世界应用中，在这种情况下使用鲁棒性模型可以显着提高用户体验。我们在流行的基准测试中展示了我们的架构的鲁棒性，并表明我们的模型对损坏和分布外样本具有高度鲁棒性，同时显着快于竞争的鲁棒模型。总之，我们的贡献如下：

• 我们介绍了FastViT，这是一种混合视觉变换器，使用结构性的重新参数化来降低内存访问成本并提高容量，实现了最先进的精度-延迟权衡。
• 我们表明，我们的模型在两个广泛使用的平台（移动设备和桌面GPU）上在延迟方面是最快的。
• 我们表明，我们的模型适用于多种任务-图像分类、目标检测、语义分割和3D手部网格回归。
• 我们表明，我们的模型对于 损坏和 OOD 样本具有鲁棒性，并且比竞争的鲁棒模型要快得多。

2、相关工作

在过去的十年中，卷积神经网络一直是视觉模型的标准架构[12,20,24,25,35,43,46,49,50,55,64]。最近，Transformer在计算机视觉任务上取得了巨大的成功[13,34,41,42,52,53,58,63]。与卷积层不同，视觉转换器中的自注意层通过建模远程依赖关系来提供全局上下文。不幸的是，这种全局作用域通常需要付出很高的计算代价[38]。像[28,38,41,57]这样的工作解决了与自关注层相关的计算成本的方法。在我们的工作中，我们探索了一种有效的自关注层替代方案，以降低延迟。

混合视觉Transformer。为了在保持准确性的同时设计高效的网络，最近的工作引入了混合架构，结合了卷积和Transformer设计，以有效地捕获本地和全局信息。一些设计用卷积层(convolutional layers)[63]代替patchify stem[13]，引入早期卷积阶段(convolutional stages)[8,41]，或者通过窗口注意(windotional attention)进行隐式杂交(implicit hybridize)[6,34]。最近的研究构建了明确的混合结构，以便在令牌(或补丁)之间更好地交换信息[9,16,62]。在大多数混合架构中，令牌混合器主要是基于自关注的。最近，MetaFormer[65]引入了池化，这是一种简单而有效的代币混合候选方法。

结构重参数化。最近的工作[12，55]表明了对跳连接进行重参数化以降低内存访问成本的益处。在我们的模型中，我们引入了一个可重参数化的新颖架构组件RepMixer，可以在推理时进行重参数化。为了提高效率，像[24,25,36,46,69]这样的工作使用深度卷积或分组卷积以及1×1点状卷积引入了k×k卷积的分解。虽然这种方法在提高模型的整体效率方面非常有效，但较低的参数数量可能导致容量下降。最近，线性训练时间过参数化在[10，11，55]中引入以提高此类模型的容量。我们在模型中使用分解的k×k卷积，并使用线性训练时间过参数化来提高这些层的容量。据我们所知，在任何先前的混合变换器架构中都没有尝试使用结构重参数化来删除跳连接和线性过参数化。

3、体系结构

3.1、概述

FastViT是一种混合Transformer，有四个不同的级，在不同的规模下运行，如图2所示。我们在表2中详细介绍了所有FastViT变体。

FastViT使用RepMixer，这是一个令牌混合器，可以重新参数化跳过连接，这有助于降低内存访问成本(参见图2)。为了进一步提高效率和性能，我们将主干和补丁嵌入层中常见的密集k×k卷积替换为使用训练时间过参数化的因式卷积(参见图2a)。

自注意力[13]令牌混合器在计算上非常昂贵，特别是在更高分辨率的情况下[38,42]。虽然[16,41]探索了自注意力层的有效版本，但我们使用大核卷积作为提高网络结构早期阶段感受野的替代方案（参见图2c）。

我们在表1中分析了基于PoolFormer[65]基线设计FastViT时所做的各种架构选择，并在下面详细介绍了我们的方法。

3.2、FastViT

3.2.1、重新参数化跳过连接

RepMixer。卷积混合最早是在ConvMixer中引入的[54]。对于输入张量X，层中的混合块实现为:
$$Y=\mathrm{BN}(\sigma (\mathrm{DWConv}(X)))+X\text{\hspace{1em}(1)}$$

其中$\sigma$是非线性激活函数，BN是批标准化[26]层，DWConv是深度卷积层。虽然这个块被证明是有效的，但在RepMixer中，我们简单地重新排列操作并去掉非线性激活函数，如下所示，

$$Y=\mathrm{DWConv}(\mathrm{BN}(X)+X\text{\hspace{1em}(2)}$$

我们设计的主要优点在于，它可以如上所示和图2d所示在推理时间被重参数化为单个深度卷积层。
$$Y=\mathrm{DWConv}(X)\text{\hspace{1em}(3)}$$

位置编码。我们使用条件位置编码[6,7]，它是动态生成的，并取决于输入令牌的局部邻域。这些编码是由深度卷积运算符产生的，并添加到补丁嵌入中。注意这组操作中缺乏非线性，因此该块被重新参数化，如图2a所示。

实证分析。为了验证重新参数化跳过连接的好处，我们在MetaFormer S12架构中消融使用最有效（就延迟而言）的令牌混合器之一，即Pooling和RepMixer。这两个被消融的模型都有\sim1.8 G FLOPs。我们在iPhone 12 Pro移动设备上对各种输入分辨率从224×224到1024×1024的时间模型。从图3可以看出，RepMixer显著优于Pooling，特别是在较高的分辨率下。在384×384下，使用RepMixer可将延迟降低25.1％，在类似1024×1024的大分辨率下，延迟显著降低43.9％。

3.2.2、线性训练时间过参数化

为了进一步提高效率（参数数、FLOPs和延迟），我们将所有的密集$\mathrm{k}\times \mathrm{k}$卷积替换为其分解版本，即$\mathrm{k}\times \mathrm{k}$深度卷积后接$1\times 1$逐点卷积。然而，分解后的参数数量降低可能会降低模型的能力。为了提高分解层的能力，我们进行了如MobileOne[55]中所述的线性训练时间过参数化。stem、patch嵌入和投影层中的MobileOne风格的过参数化有助于提高性能。

然而，由于添加的分支带来的计算开销，训练时间过参数化导致训练时间增加。在我们的体系结构中，我们只对那些用上面描述的分解形式替换密集$\mathrm{k}\times \mathrm{k}$卷积的层进行过参数化。这些层位于卷积stem、patch嵌入和投影层中。这些层的计算成本低于网络的其他部分，因此对这些层进行过参数化不会显著增加训练时间。例如，FastViT-SA12需要比在相同设置下训练未经过参数化的变体多6.7%的时间，而FastViT-SA36需要多4.4%的时间，如第4.1节所述。

3.2.3、大核卷积

RepMixer的感受野是局部的，相对于自注意力令牌混合器。然而，基于自注意力的令牌混合器在计算上是昂贵的。在不使用自注意力的情况下提高早期阶段的感受野的一种计算效率高的方法是引入深度大核卷积。我们在FFN和patch嵌入层中引入深度大核卷积。从表4中我们可以注意到，使用深度大核卷积的变体可以与使用自注意力层的变体相竞争，同时适度增加延迟。当我们比较V5与V3时，模型大小增加了11.2％，而延迟增加了2.3倍，但Top-1准确率仅增加了0.4％。V2比V4大20％，且延迟比V4高7.1％，但它在ImageNet上的Top-1准确率与V4相似。补充材料中提供了对内核大小和延迟的进一步消融分析。在表1中，我们对FFN和patch嵌入层中的大核卷积进行了消融。总体而言，大核卷积为FastViT-S12带来了0.9％的Top-1准确率提升。

我们的FFN和补丁嵌入层的架构如图2所示。我们的FFN块的结构类似于ConvNeXt[35]块，但有几个关键的区别，见图2c。我们使用批归一化而不是层归一化，因为它可以在推理时与前一层融合。此外，它不需要额外的重塑操作来为LayerNorm获得适当的张量布局，就像在ConvNeXt块的原始实现中那样。

除了增加感受野之外，如[59]中所述，大核卷积还有助于提高模型的鲁棒性，并且如[37]中观察到的那样，卷积-FFN块通常比vanilla-FFN块更稳健。因此，引入大核卷积是提高模型性能和鲁棒性的有效方法。

4、实验

4.1、图像分类

我们记录了广泛使用的ImageNet-1K数据集[45]的结果，该数据集包含约130万张训练图像和50万张验证图像。我们遵循[52,65]中规定的训练配方，即使用AdamW优化器对模型进行300次训练，权重衰减为0.05，峰值学习率为10−3，总批大小为1024。预热周期数设置为5，余弦调度用于衰减学习率。我们的实现使用timm库[60]，所有模型都在8个NVIDIA A100 gpu上进行训练。有关所有变体使用的超参数的详细信息，请参阅补充材料。对于384×384输入大小，我们对模型进行了30个epoch的微调，权重衰减为$10^{-8}$，学习率为$5\times 10^{-5}$，批大小为512[35]。为了测量延迟，我们使用对应于各自方法的输入大小。对于iPhone延迟测量，我们使用Core ML Tools (v6.0)导出模型，并在带有iOS 16的iPhone12 Pro Max上运行它，并将所有模型的批量大小设置为1。我们遵循与[55]相同的方案。对于GPU延迟测量，我们将跟踪模型导出为TensorRT (v8.0.1.6)格式，并在批处理大小为8的NVIDIA RTX-2080Ti上运行。我们报告100次运行的中位数延迟估计。

与SOTA模型的比较。在表5中，我们将我们的模型与ImageNet-1k数据集上最新的最先进的模型进行了比较。为了公平比较，我们对官方实现的ConvNeXt[35]进行了修改，避免了昂贵的重塑操作，详见补充材料。我们无法可靠地导出LITv2[41]，因为这两个库对可变形卷积的支持都很差。与两种不同的计算结构(即桌面级GPU和移动设备)上的最新最先进模型相比，我们的模型获得了最佳的精度-延迟权衡。我们的模型在参数计数和FLOPs方面都比LITv2[41]有所改进，在最高精度为84.9%的情况下，FastViT-MA36比LITv2- b小49.3%，消耗的FLOPs减少55.4%。在iPhone 12 Pro上，fastvi - s12比MobileOne-S4[55]快26.3%，在GPU上快26.9%。fastviti - ma36的准确率为83.9%，排名第一，比iPhone 12 Pro上优化后的ConvNeXt-B模型快1.9倍，在GPU上快2.0倍。FastViT-MA36在GPU上的准确率为84.9%，与NFNet-F1[1]一样快，但在移动设备上的速度要小66.7%，FLOPs要少50.1%，速度要快42.8%。

知识蒸馏。我们在表6中记录了使用蒸馏目标进行训练时模型的性能。我们遵循DeiT[52]中描述的设置，以RegNet16GF[43]作为教师模型。根据DeiT[52]，我们使用硬蒸馏，其中将教师的硬决策设置为真标签。我们的模型被训练了300次。与[30,52,53]不同，我们没有为蒸馏引入额外的封头。FastViT优于最近最先进的模型EfficientFormer[30]。fastvitv - sa24的性能与EfficientFormer-L7相似，但参数减少3.8倍，FLOPs减少2.7倍，延迟降低2.7倍。

4.2、鲁棒性评价

我们在以下基准上评估了我们的模型的分布外鲁棒性:(i) ImageNet-A[23]，一个包含自然发生的示例的数据集，被ResNets错误分类;(ii) ImageNet- r[21]，包含具有不同纹理和局部图像统计的ImageNet对象类的自然再现的数据集;(iii) ImageNet- sketch[56]，包含所有ImageNet类的黑白草图的数据集，使用google图像查询获得;(iv) ImageNet- c[22]，一个由应用于ImageNet测试集的算法生成的损坏(模糊，噪声)组成的数据集。我们使用[37]提供的实现来评估我们的模型。我们的模型比最近的视觉Transformer更鲁棒，同时比纯基于cnn的模型更快，后者鲁棒性较低，如表7所示。架构选择，如在FFN中使用大的内核卷积和补丁层与自关注层相结合，有助于提高我们模型的鲁棒性，如第3.2.3节所讨论的。表7中比较的所有模型仅在ImageNet-1k数据集上使用类似的训练食谱进行训练。从表7中可以看出，我们的模型与RVT和ConvNeXt竞争非常激烈，事实上，FastViT-M36具有更好的清洁精度，对损坏的鲁棒性更好，并且与ConvNeXt- s具有相似的分布外鲁棒性，后者比我们的模型多610万参数，FLOPs多10%。

4.3、3D Hand网格估计

近期关于实时3D手部网格估计的研究在CNN的基础上引入了复杂的网格回归层。这些骨干网络通常属于ResNet或MobileNet家族，但METRO和MeshGraphormer除外，它们使用HRNet[47]进行特征提取。虽然大多数硬件设备对从2D CNN中提取特征进行了高度优化，但这些方法中使用的复杂网格回归头并不是这样。在我们的方法中，我们用简单的回归模块替换了复杂的网格回归头，该模块回归MANO模型[44]的弱透视相机、姿势和形状参数。我们认为，使用学习底层图像良好表示的特征提取骨干可以减轻网格回归的学习挑战。而其他实时方法则使用复杂的网格回归层来补偿弱特征提取骨干，我们使用更好的特征提取骨干和一个简单的网格回归层。我们将我们的方法与FreiHand[73]数据集上的其他已发表的方法进行比较。为了公平比较，我们引用了仅使用FreiHand数据集进行训练的方法的结果，因为有些方法使用额外的姿势数据集进行预训练、训练或微调。我们只使用ImageNet-1k数据集进行预训练，然后使用[31]中描述的实验设置专门在FreiHand数据集上进行训练。详情请参见补充资料。从表8中可以看出，在实时方法中，我们的方法在所有关节和顶点误差相关指标上都优于其他方法，同时比MobileHand[15]快1.9倍，比最新的MobRecon快2.8倍。

4.4、语义分割和目标检测

对于语义分割，我们在ADE20k上验证了模型的性能[70]。该数据集包含20K张训练图像和2K张验证图像，包含150个语义类别。我们使用语义FPN[27]解码器训练语义分割模型。带有语义FPN头部的模型使用与[65]相同的设置。所有模型都使用相应图像分类模型的预训练权值进行初始化。对512×512图像作物的FLOPs和骨干延迟进行了估计。由于输入图像的分辨率更高，在表9和表10中，GPU延迟估计为批处理大小为2。在表9中，我们将我们的模型与最近的作品进行了比较。FastViT-MA36模型的mIoU比PoolFormer-M36高5.2%，后者在桌面GPU和移动设备上具有更高的FLOPs、参数计数和延迟。

我们在MS-COCO[33]数据集上进行了目标检测训练，该数据集包含80个类别，包含118K张训练图像和5K张验证图像。在表10中，我们将我们的模型与最近的工作进行了比较。所有模型都按照[65]中的1x调度进行训练，并使用Mask-RCNN[19]头部。所有模型都使用其相应图像分类模型的预训练权重进行初始化。我们表明，我们的模型在多个延迟制度下均达到了最佳性能。FastViT-MA36模型在桌面GPU和移动设备上的性能与CMT-S相似，但分别快2.4倍和4.3倍。

5、结论

我们提出了一种通用的混合视觉变换器，它在多个计算结构上都非常高效：移动设备和桌面级GPU。通过结构性的重新参数化，我们的模型减少了内存访问成本，从而显著提高了运行时间，尤其是在更高的分辨率下。此外，我们提出了进一步的架构变化，以提高ImageNet分类任务和其他下游任务(如对象检测，语义分割和3D手部网格估计)的性能。我们的经验表明，我们的主干对分布外样本具有高度鲁棒性，同时比竞争的鲁棒模型要快得多。