ElasticViT: Conflict-aware Supernet Training for Deploying Fast Vision Transformer on Diverse Mobile

论文：ElasticViT

    Vison transformer在各种计算机视觉任务中取得了显着的成功（优点不赘述），然而，成功的代价是巨大的——ViT 重量级且推理延迟成本高，这给将 ViT 引入资源有限的移动设备带来了巨大的挑战。设计准确且低延迟的 ViT 成为一个重要但具有挑战性的问题。
    Vit的思想是把图片分割成小块，然后将这些小块作为一个线性的embedding作为transformer的输入，处理方式与NLP中的token相同，用监督训练的方式进行图像分类。在实验过程中，由于vit缺少CNN隐藏层中的bias，所以当在数据量不足的情况下进行训练时，不能很好地泛化，效果不如CNN。不过在训练大规模数据时，vit的效果会反超CNN
    首先，需要把图片输入进网络，和传统的卷积神经网络输入图片不同的是，这里的图片需要分为一个个patch，如图中就是分成了9个patch。每个patch的大小是可以指定的，比如1 16×16等等。然后把每个patch输入到embedding成，也就是Linear Projection of Flattened Patches，通过该层以后，可以得到一系列向量（token），9个patch都会得到它们对应的向量，然后在所有的向量之前加入一个用于分类的向量*，它的维度和其他9个向量一致。此外，还需要加入位置信息，也就是图中所示的0~9。然后把所有的token输入Transformer Encoder中，然后把TransFormer Encoder重复堆叠L次，再将用于分类的token的输出输入MLP Head，然后得到最终分类的结果。

    为了克服在移动设备上部署 ViT 的挑战，最近的进展采用了两阶段 NAS 流程。而它的前身是One-Shot NAS，就是这个BigNas，包含搜索阶段和评估阶段，。。。。。而两阶段NAS包含训练阶段和搜索阶段；在第一阶段，超网通过权重共享机制训练，旨在探索并优化一个广泛的架构空间。这一阶段的目标是找到一组权重，这组权重在整个搜索空间中的所有可能架构上都表现得相对优秀。因此，当超网训练得当，它学习到的权重可以很好地泛化到搜索空间中的各种子网络架构上。
    在第二阶段，通过搜索算法从超网中选择出性能最佳的子网络架构。由于这些子网络是从同一个超网中派生出来的，它们可以直接利用超网中已经优化好的权重，这意味着这些子网无需从头开始训练即可直接部署使用。
    OFA 是一种更通用的方法（允许动态提取满足各种设备的特定性能标准（例如延迟、能源效率）的子网），用于跨各种硬件平台优化 DNN，强调解耦训练和搜索效率。相比之下，ViT 的两级 NAS 专门针对 Vision Transformers 的独特挑战

    两阶段 NAS 的成功很大程度上依赖于第一阶段超网训练的质量。然而，对于移动设备来说训练高质量的 ViT 超网极具挑战性，由于移动设备的多样性：移动应用程序必须支持具有不同计算能力的各种移动电话，从最新的高端设备到 CPU 速度慢得多的旧设备。例如，Google Pixel 1 的运行速度比 Pixel 6 慢 4 倍。因此，超网必须覆盖从弱设备的微小尺寸（< 100M FLOPs）到强设备的大尺寸的 ViT。然而，与两级 NAS 中的典型搜索空间相比，同时包含小型和大型 ViT 会导致搜索空间大得多。而且在这样的搜索空间上训练超网会因大小差异很大的子网引起的优化干扰而导致性能下降。

    所以为了有效地服务强移动设备和弱移动设备。elasticVIT建立在针对移动设备优化的单个非常大的搜索空间上，包含各种视觉转换器，大小范围从 37M 到 3G FLOP。这个搜索空间比典型的两级 NAS 搜索空间大 10^7 倍，使我们能够适应具有各种资源限制的各种移动设备。论文首先调查在过大的搜索空间上训练 ViT 超网时性能不佳的根本原因。我们发现主要原因是先前的超网训练方法依赖于均匀采样，它可以轻松地从我们的搜索空间中采样具有截然不同的模型大小（例如，50M 与 1G FLOP）的子网。这会导致子网梯度之间的冲突并带来优化挑战。我们做了两个关键的观察：（i）两个子网之间的梯度冲突随着它们之间的 FLOPs 差异而增加；（ii）如果它们是好的子网，相同大小的子网之间的梯度冲突可以显着减少。
    在论文中提到的"1G FLOPs"是指每秒10亿次浮点运算。这是一个衡量计算性能的指标，特别是在评估深度学习模型或其他计算密集型任务时。FLOPs常用来描述模型的大小和计算复杂度，即模型在处理单个输入时需要进行多少次浮点运算。具体到这篇论文中，提及超过1G FLOPs的视觉变换器（ViT）模型时，作者指的是那些在进行一次前向传播时，需要执行超过10亿次浮点运算的模型。这样的模型通常具有较高的计算需求，可能不适合直接部署到计算资源受限的移动设备上，因为这会导致延迟过高，影响用户体验。

     在设计训练方法之前，论文给出一些列观点，虽然许多工作旨在设计具有高精度和小 FLOP 的 ViT 模型，但观察到小 FLOP 的模型在移动设备上可能具有高延迟。例如，NASViT-A0 的 FLOPs 比 MobileNetV3 少，但在 Google Pixel 4 上运行速度慢 2 倍。MobileFormer [9] 只有 52M FLOPs，但运行速度慢 5.5 倍。这是因为这些模型结合了有效但不适合移动设备的操作，减少flop次数。
    解决：
    基于对移动设备友好的CNN-ViT架构构建搜索空间
    在CNN阶段，使用了MobileNetv2和MobileNetv3模块
    在ViT阶段，则对NASViT的注意力机制进行了关键修改以提高效率，包括去除了延迟较高的talking heads操作，使用Hswish激活函数代替Gelu，并且决定不使用Swin中的shifted window注意力机制
    CNN阶段：
    使用了MobileNetv2和MobileNetv3模块：这两种模块是由Google开发的，旨在移动设备上提供高效的卷积神经网络解决方案。MobileNetv2和MobileNetv3都使用了深度可分离卷积，这种卷积方式可以大幅减少模型的参数数量和计算成本，同时保持相对较高的精度。这使得它们非常适合用于资源受限的环境。
    ViT阶段：

1. Talking Heads Attention
   Talking heads attention 是一个变体的注意力机制，它在自注意力（self-attention）的计算中引入了附加的线性变换，这些变换作用于注意力矩阵的行和列。这样可以增加模型的表达能力，改进注意力的分配，从而可能提高模型在特定任务上的表现。然而，这种增强的注意力机制需要额外的计算资源，导致计算延迟增加。在资源受限的移动设备上，这种延迟的增加可能不可接受，因此在设计轻量级和低延迟的ViT模型时，通常会省略Talking heads attention。
2. GELU激活函数
   Gaussian Error Linear Unit（GELU）是一种常用的激活函数，尤其在Transformer架构中广泛应用。GELU相比于其他激活函数如ReLU或Swish，其计算涉及更多的数学运算（例如乘法和指数运算），这可能在硬件上导致更高的计算成本和更长的处理时间。为了优化ViT模型的延迟和计算效率，可以选择使用Hswish激活函数替代GELU。Hswish提供类似的非线性处理能力，但计算更简单，因此更适合在计算资源有限的设备上使用。
3. Swin Transformer的Shifted Window Attention
   Swin Transformer引入了一种新颖的shifted window机制来计算自注意力。这种方法通过限制每个自注意力操作在一个局部窗口内进行，然后在不同层之间交错这些窗口，从而有效减少了计算量。尽管这种方法能显著减少参数数量和计算成本，但它仍然相对复杂，尤其是在输入尺寸较小时可能不会带来明显的性能提升。因此，对于旨在进一步降低延迟的轻量级ViT模型，可能选择不采用这种机制。

     搜索空间大小：作者提出的搜索空间非常大，它比传统两阶段NAS中的搜索空间大了10^7倍，从而能够覆盖广泛的ViT配置，适应各种移动设备的需求。
     ViT配置范围：这个搜索空间包括了从非常小（<100M FLOPs）到非常大的ViT模型配置，以满足不同计算能力的移动设备。例如，性能高的设备可以运行高达500M FLOPs的模型来满足约30ms的延迟限制，而性能较弱的设备（如Pixel 1）则需要运行小于100M FLOPs的模型来满足同样的延迟限制
     Stage: 表示超级网络中的不同阶段。不同的阶段代表模型中不同的层或结构组件。
     Depths: 表示在每个阶段可以采用的层的数量范围。例如，在MBv2 block阶段，深度可以在1到2之间变化。
     Channels: 代表各个阶段所使用的通道数的范围，以及在搜索过程中可以调整的步长。例如，在Conv 3x3阶段，通道数可以在16到24之间变化，步长为8。
     Kernel size (V scale): 描述卷积层中卷积核的尺寸。在ViT结构中，卷积核的尺寸也可能与所谓的"V scale"相关联，后者可能是指与ViT模型中注意力机制相关的某种尺度调整。
     Expansion ratio: 这通常指的是逐点卷积（pointwise convolution）后跟随的线性瓶颈层的维度扩张比率。例如，如果扩张比率为3，这意味着该层的输出维度将是输入维度的三倍。
     Stride: 在卷积操作中，步长是过滤器（或称卷积核）移动的距离。步长为2意味着过滤器在每个方向上每次移动两个像素点。
MBPool: 是一个池化层，可能是MobileNet架构中使用的特定类型的池化。1984可能表示该层使用的特征维度。
     Input resolution: 提供了超级网络可以处理的不同输入图像分辨率。分辨率的大小可以直接影响模型的复杂性和计算需求。

    超网训练与标准单网络训练的不同之处在于，所有子网共享其公共部分的权重。共享权重可能会收到冲突的梯度，从而导致不同的优化方向，从而降低最终的精度。已经提出了几种技术来缓解这个问题，并证明在典型搜索空间上训练高质量超网方面取得了成功。其中，三明治规则对于保证性能下限和上限至关重要。具体来说，它在每次迭代中对最小子网、最大子网和两个随机子网进行采样。然而，当在表 1 中的巨大搜索空间上训练 ViT 超网时，我们观察到使用以前的最佳实践时准确率显着下降。具体来说，我们使用三明治规则并遵循 NASViT [15] 中相同的训练收据。图(a) 比较了通过继承超网权重在 ImageNet 上进行再训练所实现的 20 个随机子网的准确性。与重新训练相比，源自超级网的模型的准确率下降了高达 8%。
    超级网络结构: 与典型的视觉变换器（ViT）超级网络相比，作者的超级网络包含了许多小型子网络，且子网络间的大小差异很大。
梯度相似性分析: 作者从超级网络中随机抽取了多个子网络（这些子网络的浮点运算次数FLOPs范围从50到1200），并计算了在相同训练数据批次下，每对子网络共享权重     梯度的余弦相似性。余弦相似性较低意味着梯度差异较大，这增加了训练超级网络的难度。
    FLOPs差异与梯度冲突: 观察#1指出，子网络间共享权重的梯度相似性与它们的FLOPs差异负相关。即如果两个子网络的FLOPs相近，它们的梯度相似性会比较高；如果FLOPs差异很大，它们的梯度相似性可能接近于0。
    子网络质量的影响: 超级网络中子网络的质量也可能影响梯度的相似性。如果在采样和训练过程中选中了一个表现不佳的子网络，它可能会扰乱表现好的子网络的权重。
    优质子网络的选择: 为了验证子网络质量对梯度相似性的影响，作者随机抽取了50个同FLOPs水平的子网络，并计算了最优子网络和随机抽取子网络之间的梯度余弦相似性。观察#2表明，如果抽取的是质量好的子网络，即使FLOPs水平相同，它们之间的梯度相似性也可以显著提高。

    受到观察结果的启发，我们提出了两种方法来解决梯度冲突问题。 （i）我们引入复杂性感知采样来限制相邻训练步骤之间的 FLOP 差异，因为较大的差异会导致梯度冲突问题（第 4.2 节）。 (ii) 由于“好的子网”可以相互增强，从而进一步减少梯度冲突，我们提出性能感知采样，以确保训练过程可以尽可能多地采样好的子网（第 4.3 节）。
    超网目标: 两阶段NAS的目的是在给定的搜索空间A中训练一个共享权重的超级网络，并联合优化所有子网络的参数。
    三明治规则: 由于不可行同时训练搜索空间A中的所有子网络，因此通常采用三明治规则来抽样子网络。这个规则在每个训练步骤中采样两种类型的子网络：
    (i) 通过先验分布Γ从搜索空间采样的子网络集合，作为期望项的近似。
    (ii) 固定的子网络集合S，通常包括最大和最小的子网络。
优化问题: 超级网络的训练可以表述为一个优化问题，其中w代表所有子网络共享的权重，f(·)表示神经网络，L是训练集D上的损失函数，ws是子网络s的独有权重。
子网络类型: 任何搜索空间中都有两种类型的子网络，随机子网络，它们的各个维度参数在最大和最小设置之间采样；以及大小最小和最大的子网络，其维度参数分别是最小和最大设置。
    实施效果: 在Table 1中描述的针对移动设备优化的超级网络训练过程中，使用三明治规则可能会由于两个原因导致梯度冲突：
    首先，它总是采样最小（37M FLOPs）和最大（3191 MFLOPs）的子网络，以及两个随机子网络，这会导致FLOPs的显著差异，从而使梯度相似性接近于0。
其次，两个随机采样的子网络的大小和质量无法保证，这加剧了梯度冲突的问题。
    总的来说，这段话解释了在训练超级网络时，尽管三明治规则被广泛采用，但在某些情况下，特别是在搜索空间非常大时，这一规则可能会导致梯度冲突，降低了训练效果。因此，可能需要其他策略来优化超级网络的训练过程，特别是在处理极端大小差异的子网络时。

    问题：在极大的搜索空间中，随机采样子网络可能导致在同一训练步骤内或不同训练步骤之间的FLOPs差异很大。
    相邻步骤采样：为了解决这一问题，提出了一种相邻步骤采样策略，该策略限制了每个训练步骤中M个随机子网络的FLOPs，使得同一训练步骤中的子网络具有相同级别的FLOPs。
    复杂度水平：定义了一系列逐渐增加的复杂度水平C1​,…,CK​（例如，100, 200, …, 800 MFLOPs），这覆盖了从小到大的ViT模型范围。
    采样策略：如果在步骤t−1中采样了复杂度级别C(t−1)i​，那么在步骤t中将从C(t)j​的复杂度级别采样M个子网络。为了满足相邻步骤采样的要求，这两个步骤之间的FLOPs距离必须满足以下条件：
                        g(s(t);Cj(t)​)=∣Cj(t)​−Ci(t−1)​∣≤Z,
    其中Z是一个正整数复杂度阈值，用于控制两个步骤之间的采样距离。当Z趋向于正无穷大时，这等同于随机采样。
    具体实施：在实际操作中，步骤t的复杂度级别要么等于步骤t−1，要么与之相邻。这提供了三个选项：减少一个FLOPs级别、保持当前FLOPs级别或增加一个FLOPs级别。

    删除最大子网络主要是因为在超级网络的搜索空间中，最大的子网络因其巨大的FLOPs（在本例中为3191 MFLOPs）而自然地与其他的复杂度级别存在很大的差异。这种巨大的差异会导致梯度方向发生冲突，因为最大子网络需要的梯度更新可能会对其他较小的子网络不利。因此，移除最大子网络可以减少这种梯度冲突，从而有助于稳定整个超级网络的训练过程。
    相对而言，最小的子网络通常不会引起这样的问题，原因包括：
    影响力: 最小子网络的影响通常较小，因为它们涉及的权重和特征较少，梯度更新的影响范围有限。
    梯度稀释: 较小子网络对共享权重的更新影响较弱，因为它们的梯度相对于较大子网络可能会被“稀释”。
    训练频率: 最小子网络更频繁地被训练，因为在随机采样中它们更容易被选中（特别是当使用均匀分布时），这可能使得它们的权重更新对整个网络的影响更加平衡。
    硬件资源适应性: 在实际应用中，最小子网络往往是为了适应资源受限的设备而设计的。在优化超级网络时，保留最小子网络可以确保在资源有限的环境中还能获得合理的性能。

    最小子网络的影响: 由于最小子网络只有37 MFLOPs，它在较大复杂度范围内（例如Ci ≥ 500 MFLOPs）采样时也会引入每一步的大FLOPs差异。不同于最大子网络，最小子网络决定了整个搜索空间的性能下限，因此不能简单地被移除。
    多层次最小子网络（HSS）策略: 为了解决不能移除最小子网络的问题，作者提出不在每一步都采样37 MFLOPs的最小子网络，而是从预定义的多层次最小子网络集合中选择一个与当前复杂度水平最接近的最小子网络。
    HSS集合的定义: HSS集合包括了N = 3个预定义的复杂度不同的子网络，这些子网络逻辑上将整个搜索空间划分为N个层次的子空间。
    采样策略: 在采样复杂度水平C(t)时，选择HSS集合中与其复杂度最接近的子网络sn​作为最小子网络。
    实验选择: 在实验中，作者选择了37 MFLOPs的原始最小子网络（min1），以及新增的160 MFLOPs（min2）和280 MFLOPs（min3）的子网络作为第二和第三小的子网络。
    HSS集合的讨论: HSS集合有两个优点：
    (i) 它允许整个空间统一共享权重，使小子网络能够从大子网络中获益。研究已经证明，将小模型整合到更大的模型中可以显著提高小子网络的性能。
    (ii) 它不依赖于任何启发式空间重新设计或强烈的维度之间的关联假设（例如在其他研究中使用的线性相关性），这增强了该方法的普适性。

    该方法不需要基于特定的预设规则（启发式规则）来调整或优化搜索空间的结构，也不需要假定搜索空间中不同参数（或维度）之间存在某种固定的、明显的相关性。这里的几个关键点包括：
    启发式空间重新设计：
    启发式设计通常指的是基于特定经验或直觉的规则来调整搜索空间的结构。例如，可能会基于先前的研究或某些特定的观察来选择包括或排除某些特定的网络配置。
    在此上下文中，方法不依赖于这种类型的设计，意味着它采用了一种更通用或灵活的方法来处理搜索空间，而不是基于特定的、可能具有主观性的经验规则。
    强烈的维度之间的关联假设：
    在许多机器学习模型和算法设计中，经常需要假设输入特征（或参数维度）之间的某种关联性，如线性关系、统计依赖等。
    不依赖于这种关联假设意味着模型的设计和优化不需要预先设定参数间复杂的相互作用模式。这种做法增加了方法的普适性，因为它不局限于特定的关系模型或数据分布特性，使得方法在更广泛的场景和配置下都可能有效。

    我们观察到性能最好的子网可以进一步缓解梯度冲突问题。受此启发，我们引入了性能感知采样来以可能更高的精度对子网进行采样
    具体来说，对于 U 中 M = 3 个相同大小的子网（如方程 4 中定义），我们的目标是从有利于良好子网的新分布中采样，而不是从具有随机性能的均匀分布中采样。为了实现这一目标，我们提出了一种探索和开发策略，该策略基于质量感知记忆库和路径偏好规则构建新的分布。质量感知记忆库用于以 q 的概率探索    历史良好的子网，而偏好规则则以 1 − q 的概率探索宽度更宽、深度更浅的子网
    性能感知采样：选择可能具有更高准确性的子网进行采样，而不是随机性能的均匀采样。
    探索与利用策略：
    质量感知的记忆库（Quality-aware Memory Bank）：存储历史上表现良好的子网。这些“好”的子网是通过比较当前小批次（min-batch）上的交叉熵损失来确定的。记忆库允许以概率 q 利用这些历史表现良好的子网。
    路径偏好规则（Path Preference Rule）：偏好选择宽度更宽、深度更浅的子网进行探索，以概率 1−q 进行。
    混合分布采样：新的子网分布是基于动态记忆库和未探索子网的混合分布，旨在平衡探索新的可能性和利用已知的良好表现。

    质量感知的记忆库（Quality-aware Memory Bank）：这些“好”的子网是通过比较当前小批次（min-batch）上的交叉熵损失来确定的。记忆库允许以概率 q 利用这些历史表现良好的子网。

    在训练初期，q 的值较小（例如，0.2），这意味着均匀采样在训练中占主导地位，以探索有前途的子网。
    随着训练的进行，记忆库逐渐填满，q 的值也逐渐增加，以便利用这些记忆中的良好子网。

    宽度和深度的定义：这里的宽度 (t)width​）和深度 (t)depth​）分别指子网在Transformer阶段的通道数（或宽度）和层数（或深度）。
    与锚定模型的比较：通过与一个锚定模型（sac​）进行比较，来量化子网的FLOPs分布在深度和宽度维度上的偏好。这通过生成两个新的子网（一个调整为锚定模型的宽度，另一个调整为锚定模型的深度），并计算它们与锚定模型的FLOPs差异来实现。
    如果Φa​（宽度引起的FLOPs变化）大于或等于Φb​（深度引起的FLOPs变化），则认为子网络s(t)在Transformer阶段具有更宽的宽度和更浅的深度。这样的子网络将按照路径偏好规则进行训练。否则，将重新采样一个新的子网络并重复上述步骤。
    锚定模型的质量：锚定模型的质量对比较的有效性有重要影响。假设存储库中捕获了这些偏好，因此选择存储库中损失最小的子网络作为锚定模型。