引言: 近年来，视觉Transformer（Vision Transformer，简称ViT）在计算机视觉任务中的应用日益广泛，从图像分类到对象识别等，均显示出优越的性能。然而，ViT模型也面临一些挑战，特别是在模型尺寸和推理效率方面。为解决这些问题，最新，在国际计算机视觉与模式识别会议（ICCV）上提出的ElasticViT方法可以很好的解决该问题。ElasticViT主要利用了神经架构搜索（NAS）技术，通过一个两阶段的搜索方法来优化ViT的架构，使之既高效又节省资源。特别注意的是，ElasticViT模型能够在不牺牲准确性的前提下，显著降低模型的运行延时和资源消耗，也为ViT模型在移动设备上的广泛应用提供了一种新的可能性。让我们一起来看一看吧。
论文地址：https://openaccess.thecvf.com/content/ICCV2023/papers/Tang_ElasticViT_Conflict-aware_Supernet_Training_for_Deploying_Fast_Vision_Transformer_on_ICCV_2023_paper.pdf
- 研究背景
1.1 视觉Transformer
视觉 Transformer 是将自然语言处理（NLP）领域中成功的 Transformer 模型扩展到了计算机视觉领域，专门用于处理视觉任务，如图像分类、目标检测和语义分割等。与传统的卷积神经网络（CNN）不同，ViT将输入图像分割为图块，并将它们作为序列输入Transformer模型。并且ViT不直接使用卷积操作，而是使用自注意力机制来捕捉图像中的全局和局部关系。这种自注意力机制允许模型在处理图像时同时考虑到所有像素之间的相互作用，能够更好地捕捉长距离依赖关系，从而取得了在图像分类等任务上的优异性能。
图1：ViT模型的基本架构
然而，尽管ViT在图像分类等任务上表现出色，但其推理效率却面临挑战。主要原因在于ViT采用的自注意力机制导致了较高的计算复杂度。自注意力机制需要在每个位置上计算所有位置之间的关联程度，这导致了计算量随着输入序列长度的增加而呈平方级增长。对于图像而言，输入序列的长度往往非常大，因为图像被切分成了大量的图块，每个图块都要与其他所有图块进行关联计算，这使得计算复杂度极高。
因此，研究者们提出了许多方法来设计高效的ViT。他们使用不同的方式，如新的架构或模块[4]，更好的注意力操作[5]和混合CNN-Transformer[6]。混合模型通常通过引入特殊操作，在小尺寸下表现良好。例如，MobileFormer采用了带有双向桥的平行CNN-Transformer结构。然而，尽管FLOPs减少了，但由于双向桥接等对移动不友好的操作，这些ViT仍然具有很高的延迟。
1.2 神经架构搜索
神经架构搜索（Neural Architecture Search，NAS）是一种自动化机器学习技术，其旨在通过自动搜索神经网络结构的空间，找到最优的网络设计，以解决特定任务。传统的深度学习模型设计通常依赖于人工经验和专业知识，但随着深度学习模型的复杂性不断增加，手动设计变得越来越困难。NAS的出现填补了这一空白，使得神经网络的设计变得更加自动化和智能化。最近，一些研究应用NAS来设计改进的ViT，如Autoformer、ViTAS和ViT-ResNet。这些方法的重点是搜索超过1G FLOPs的模型。对于小ViT、UniNet和NASViT在小FLOPs下搜索混合CNN-ViT，取得了很好的结果。然而，这些NAS主要针对FLOPs进行优化，而没有考虑不同移动设备的效率，导致性能不理想。
1.3 两阶段NAS
在NAS的发展过程中，研究者们逐渐意识到一次性搜索整个网络结构空间是非常耗时且计算资源密集的。为了提高搜索效率，逐步发展出了两阶段NAS方法。两阶段NAS将训练和搜索解耦。通过训练一个超大网络结构空间的超网，然后使用典型的搜索算法在各种资源约束下寻找最佳子网。搜索的子网可以直接继承超级网络权重进行部署，达到与从头开始重新训练的子网相当的准确性。两阶段NAS的成功很大程度上依赖于第一个阶段超网训练的质量，因为搜索子网时需要继承超网的权重来评估子网的优劣，一个高质量的超网能够更好地区分子网的性能，从而搜索到更优秀的网络架构。
超级网络训练与标准单网训练的不同之处在于，所有子网对其公共部分共享权值。共享权值可能会收到冲突的梯度，导致优化方向不同，从而降低最终的精度。然而，现有的研究采用统一采样来对子网进行采样，这可能导致在更大的搜索空间中对大小差异很大的子网进行采样，从而导致梯度冲突和性能下降。因此，训练一个高质量的ViT超网是极具挑战性的。
2. ElasticViT
近期，Tang等人提出了一个新的两阶段神经结构搜索方法对于轻量级ViT——（ElasticViT: Conflict-aware Supernet Training for Deploying Fast Vision Transformer on Diverse Mobile Devices）[1]，该方法减少了超网训练时产生的冲突，提高了超网训练权重的质量，并应用于搜索适配移动设备的轻量级ViT，取得了不错的效果。
ElasticViT方法旨在设计精确的ViT网络，可以在移动设备上实现低延迟。具体而言，这是一种两阶段的NAS方法，它在一个非常大的搜索空间上训练一个高质量的ViT超级网络，以覆盖广泛的移动设备，然后搜索一个最佳子网以进行直接部署。此外，对超网训练问题的根本原因进行探究，并针对观察结论提出了两种方法来解决梯度冲突问题：复杂度感知采样和性能感知采样。
2.1 两个观察
观察一：FLOPs差异与梯度冲突的关系
这个观察指出，两个子网络之间共享权重的梯度相似度与它们的FLOPs差异呈负相关。具有类似FLOPs的子网络之间会有更高的梯度相似度，而如果两个子网络之间的FLOPs差异较大，则梯度相似度接近于0。这表明，当训练的子网络大小差异较大时，可能会出现更大程度的梯度冲突，这可能导致超级网络的性能下降。
图2：不同FLOPs模型的梯度余弦相似度
观察二：子网质量对梯度相似度的影响
这个观察指出，同等大小的子网络之间的梯度相似度可以得到显著改善，如果它们是高质量的子网络。如果采样和训练了一个质量较差的子网络，它会干扰好的子网络的权重，导致共享权重的梯度相似度下降。这表明，训练高质量的子网络是至关重要的，以确保超级网络能够获得良好的性能并避免梯度冲突的影响。
图3：在相同的FLOPs水平下，良好模型比随机抽样模型具有更相似的梯度。
2.2 两个关键技术
2.2.1 复杂度感知采样
由于在搜索空间A中训练所有子网络是不可行的，通常采用广泛使用的“三明治规则”来实现。其超网的训练可被描述为以下优化问题：
其中w是所有子网络的共享权重，f(·)表示神经网络，L是训练集D上的损失函数， 是子网络s的专用权重。三明治规则通常通过从均匀分布Γ中随机采样M = 2个子网络来近似方程中的第一个项。在第二个项中，通常包括最大子网络 和最小子网络。然而，在针对移动设备专门设计的超级网络训练中，应用三明治规则可能会导致梯度冲突。这主要是由于两个原因：首先，三明治规则总是采样最小、最大和两个随机子网络，这导致了显著的FLOPs差异，并经常导致梯度相似度接近于0；其次，无法保证随机采样的两个子网络的大小和质量，进一步加剧了问题。
因此，ElasticViT方法引入了复杂度感知采样技术。其定义了一组递增的复杂度级别（例如，100,…,800MFLOPs）来覆盖从小到大的ViT模型。假设步骤t-1采样第i个复杂度等级 ，步骤t采样 第j个复杂度等级下的M个子网。为了满足相邻步采样，这两步之间的FLOPs距离必须满足以下条件:
其中 是控制相邻两步采样距离的复杂度阈值。当 时，相当于随机抽样。ElasticViT方法希望步骤t的复杂度等于或接近步骤t-1。因此，为 提供了三个选项: , , ，分别表示将FLOPs水平降低一个，保持当前FLOPs水平，或将FLOPs水平提高一个的选择。
通过复杂度感知采样，可以将超网训练重新表述为以下公式：
其中t为当前训练步骤，U为包含M=3个均匀子网的随机子网集，其中每个子网对于步骤t的FLOPs级别为
ElasticViT方法删除了对最大子网的采样，因为它大的FLOPs差异会导致梯度冲突。此外，还使用多个分层最小子网（例如37M, 160M, 280MFLOPs），替换三明治规则中采样的最小子网。σ(·)从HSS中选择最近的最小子网。如果s_n是距离 最近的最小子网，，否则为0。
图4：冲突感知超级网络训练的整体结构
综上，ElasticViT方法通过删除最大子网采样、替换最小采样子网、限制FLOPs差异的随机子网采样，来限制相邻训练步骤之间的模型FLOPs差异，提高了采样模型梯度相似度，从而提高了超网的训练效果。
图5：ElasticViT的搜索空间
2.2.2 性能感知采样
鉴于之前的观察2所述，同FLOPs条件下性能优异的子网可以减轻梯度冲突问题。ElasticViT方法引入了性能感知采样方法，以便从具有更高准确度潜力的子网中进行采样。具体而言，提出了一个探索和开发策略，该策略基于性能感知记忆库和路径偏好规则构建了一个新的分布。性能感知记忆库用于存储历史上优秀的子网，并以一定概率利用这些子网。子网的期望为：
其中表示搜索空间A中FLOPs等级为的所有子网，为未探索子网的均匀分布。q的值在训练早期较小，便于使均匀采样为主导，后期随着记忆库被填充，q值也逐渐增加，以利用历史的优秀子网。
最近有研究表明[2,3]，在纯ViT模型中，与CNN模型相比，后期Transformer阶段倾向于更宽的通道而不是更多层。因此，当采样未探索的子网时，ElasticViT方法将这种偏好纳入了超网采样中，以过滤掉一些劣质子网。该工作提出了一种方法，选择记忆库中损失最小的子网作为锚模型，通过将子网的FLOPs分布与锚模型进行比较来量化子网的深度和宽度维度。如果一个采样的子网被认为在ViT阶段具有更宽的宽度和更浅的深度，则符合路径偏好规则，将对其进行训练。否则，将重新对子网络进行采样并重复上述步骤。
3. 实验结果
在实验中，将该工作应用提出的技术来训练600个epoch的超大ViT超级网络。训练中使用的复杂性级别设置为{100,200,300,400,500,700,900,1200}MFLOPs，这适用于移动状态下的ViT。其他训练设置和超参数遵循现有的最佳实践。评估了训练好的ViT超级网络在四种不同资源水平的手机上的有效性:弱(Pixel1)，中性(Pixel4, Xiaomi11)和强(Pixel6)。对于每个设备，设置了一系列延迟约束，并使用nn-Meter构建了一个延迟预测器，以实现高效搜索。我们使用OFA中的进化搜索方法为每个延迟约束搜索5k个子网。
实验结果显示，ElasticViT在搜索轻量级ViT模型方法取得了显著成效，搜索到范围37到800 MFLOPs的轻量级低延迟ViT模型，使得准确的ViT模型可以应用于各种移动设备上。ElasticViT模型在准确性和延迟方面表现优于现有的ViT和CNN基线模型。
表1：ImageNet-1K上的ElasticViT性能与最先进的高效CNN和ViT模型的比较
为了进一步验证ElasticViT的有效性，该工作进行了相应的消融实验。以传统的三明治规则训练的超网为基线，逐一加入改进的算法来增强超网训练。表2说明了ElasticViT在一个巨大的ViT搜索空间上增强超级网络训练的有效性，相邻步采样和多层最小子网策略都有效地控制了训练子网之间的FLOPs差异，从而提高了top-1精度指标。此外，使用性能感知采样来训练良好的子网，进一步的提高了最佳搜索ViT的准确度。
表2：:ImageNet上的消融研究结果
通过与 OFA 模型延迟-准确率方面对比，实验结果（如图5）显示，无论是在性能较弱还是性能较强的移动设备上，本文算法搜索到的ViT 模型都一致地超过了 OFA 模型，满足了各种延迟约束。值得注意的是，本文算法搜索到的 ViT 模型是首个在移动设备上实现实时推理延迟和模型精度协同优化的 ViT模型。这些结果证明了所提方法的有效性，表明其具备在移动边缘设备上搜索到高效、高性能的模型的潜力。
图5 ElasticViT 模型与 OFA 模型延迟-准确率方面的对比实验
4. 总结
本文探讨了将神经架构搜索技术应用于移动设备的ViT模型的策略，并重点介绍了ElasticViT方法。ElasticViT的主要目标是设计出精确的ViT网络，以实现在移动设备上低延迟的推理过程。该工作深入分析了超网训练中梯度冲突的根本原因，并针对这一问题引入了两项关键技术：限制采样子网之间的FLOPs差异和采样潜在的优秀子网来解决梯度冲突问题，大大提高了超级网络的训练质量。实验证明，ElasticViT搜索到的ViT模型在ImageNet数据集上优于现有技术的高效CNN和ViT。这项研究为在移动设备上部署准确性高、延迟低的ViT模型提供了有力的方法和技术支持。
