训练Backbone你还用EMA？ViT训练的大杀器EWA升级来袭-轻识

点击下方卡片，关注「集智书童」公众号

结构重参化是一种适用于卷积神经网络（CNNs）的通用训练方案，它在不增加推理成本的情况下实现了性能提升。随着视觉Transformer（ViTs）在各种视觉任务中逐渐超越CNNs，人们可能会提出一个问题：是否存在一种专门针对ViTs的训练方案，也能在不增加推理成本的情况下实现性能提升？

最近，由于多个分支结构可以有效地扩展Transformer的容量，Mixture-of-Experts（MoE）引起了越来越多的关注。考虑到MoE也可以被看作是一个多分支结构，我们是否可以利用MoE来实现类似于结构重参化的ViT训练方案？

在本文中，作者肯定地回答了这些问题，并提出了一种新的ViT通用训练策略。具体来说，作者将ViT的一些前馈网络（FFNs）替换为经过特殊设计的、更高效的MoE，这些MoE通过随机均匀分区将Token分配给 Expert，并在每次迭代结束时对这些MoE进行 Experts Weights Averaging（EWA）。训练后，通过对 Expert进行平均，将每个MoE转化为一个FFN，将模型转换回原始的ViT以进行推理。

作者还提供了理论分析，以展示为什么以及如何运作。对各种2D和3D视觉任务、ViT体系结构和数据集进行的综合实验验证了所提出的训练方案的有效性和通用性。此外，当对ViTs进行微调时，作者的训练方案也可以应用于提高性能。

最后但同样重要的是，所提出的EWA技术可以在各种2D视觉小数据集和3D视觉任务中显着提高朴素MoE的效果。

1、简介

近年来，通用训练方案在深度学习中发挥了关键作用。这些方法在不修改模型架构的情况下增强了现有模型，并且不需要对特定架构做出任何假设。因此，它们可以应用于各种模型，并提供均匀的好处。

对于CNN架构，广泛使用的通用训练方案是结构重参化（structural re-parameterization）（图1(b)）。这种方法通过在训练过程中使用多分支结构，将CNN模型的训练和推理阶段解耦，并在推理过程中等效地转换回原始的单分支结构。结构重参化可以在不增加推理成本的情况下实现性能提升。

然而，在训练过程中使用的多分支结构可能会降低训练速度（如Tab. 1的第2列所示，其训练延迟高达普通训练的2.18倍）。此外，结构重参化仅适用于CNN，而不适用于其他像ViTs（如Tab. 1的第6列所示）这样的最近更强大的架构。

对于Transformer架构，最近的通用训练方法涉及将前馈网络（FFN）层替换为稀疏的 Expert混合（MoE）（图1(a)）。MoE是一种新颖的网络组件，包括多个具有唯一权重和可训练Routing网络的 Expert（FFN）。

在训练和推理阶段，Routing网络为每个输入选择一种稀疏的 Expert组合，通过稀疏计算实现了Transformer模型容量的有效扩展。在自然语言处理领域，许多稀疏的MoE模型已经证明它们的性能优于密集模型。

然而，MoE仍然具有一些限制，限制了它在Vision Transformers（ViTs）中的广泛使用。

首先，V-MoE和SwinV2-MoE是在大规模数据集（JFT-300M，ImageNet-22K）上进行预训练的，而当在ImageNet-1K上进行预训练时，V-MoE的性能低于ViT。此外，还没有研究研究MoE Vision Transformer在小数据集（如CIFAR-100和Tiny-ImageNet）上的情况。

其次，当前的MoE ViTs专注于2D视觉任务，没有用于3D视觉任务的MoE ViTs（见Tab. 1的第4行）。

第三，MoE本身的前kRouting机制并不高效和有效。如Tab. 1的第4行所示，由于前kRouting需要额外的计算来为每个输入做出Routing决策，训练和推理的延迟都增加到了1.12倍，而训练和推理的参数都大幅增加到了3.25倍，给模型带来了负担。

此外，前kRouting容易导致 Expert之间的负载不平衡，因此许多研究致力于通过各种设计来解决这个问题，如添加辅助负载平衡损失，在全局范围内做出Routing决策，或重新设计Routing算法。

因此，以下问题产生了：

是否有一种新的训练方法，可以在训练阶段高效地训练稀疏激活的 Expert，而在推理阶段可以将稀疏模型转化为稠密模型而不需要额外的计算成本？是否可以进一步将MoE的有效性扩展到2D视觉数据集和3D视觉任务中？

在本文中设计了一种名为"Experts Weights Averaging"（EWA）训练的新型通用训练策略，用于ViTs。本文的方法实现了性能提升，同时不增加推理延迟和参数，如Tab. 1的第2行和最后一行所示。

具体而言，作者在ViT的训练和推理阶段进行了解耦，如图1(c)所示。在训练过程中，用经过特殊设计的、更高效的MoEs替换了普通ViT中的一些FFN，这些MoEs通过随机均匀划分（RUP）将Token分配给 Expert，不需要额外的参数、特殊设计或辅助损失。

此外，每次迭代结束时，对每个MoE执行"Experts Weights Averaging"（EWA）。在推理过程中，通过简单地将每个MoE的 Expert平均为单个FFN，可以将稀疏的MoE模型转换为普通的ViT模型，而不会有任何性能损失。

作者对各种2D和3D视觉任务、ViT架构和数据集进行了全面的实验，以验证所提出的训练方案的有效性和泛化性。例如，作者提出的训练方法在图像分类任务上实现了1.72%的性能提升，在点云语义分割任务上实现了1.74%的平均IoU提升。

此外，EWA训练还可以用于微调预训练的ViT模型。本文的训练方法将CIFAR100上的ViT-B微调结果从90.71%提高到91.42%（↑ 0.71%）。最后，作者发现EWA训练中的"Experts Weights Averaging"关键技术可以显著提高2D小型视觉数据集和3D视觉任务中原始MoE的效果。

主要贡献总结如下：

提出了一种新的通用训练方案，用于ViTs，实现了性能提升，同时不增加推理延迟和参数。在训练过程中，将ViT架构中的某些FFN替换为使用随机均匀划分的MoE，并在每次模型权重更新后对这些MoE层执行"Experts Weights Averaging"。训练结束后，通过将模型中的每个MoE层的 Expert进行平均，将模型转换回原始的ViT架构，用于推理。
验证了EWA训练可以为不同的2D/3D视觉任务和数据集上的各种Vision Transformer架构带来统一的性能提升。此外，作者发现EWA训练还可以用于微调预训练的ViT模型，实现更进一步的性能提升。作者提供了理论分析，展示了为什么以及如何有效。
探索了在小型2D视觉数据集和3D视觉任务中原始MoE的有效性。作者发现使用原始MoE的性能低于密集模型，但所提出的"Experts Weights Averaging"技术可以显著提高原始MoE的效果。

2、相关工作

2.1、MoE

稀疏门控混合 Expert模型（Mixture-of-Experts，MoE）在训练时间、模型容量和性能方面显示出显著的改进。此外，当MoE被引入基于 Transformer的语言模型中时，它受到越来越多的关注。迄今为止，许多工作都改进了MoE的关键组件，即Routing网络。

一方面，对于传统的可学习Routing网络，Switch Transformer简化了Routing，仅为每个Token选择前1个 Expert，Base Layer通过处理线性分配问题来执行Routing，Expert Choice将前k个TokenRouting到每个 Expert。

另一方面，Hash Layer为TokenRouting设计了确定性哈希，THOR随机激活每个输入的 Expert。除了Routing网络的改进，许多工作还探索了MoE在现代硬件、扩展性和应用方面的实现。

然而，大多数MoE工作都是基于 Transformer的语言模型，计算机视觉领域中只有少数MoE Vision Transformer的工作存在。

Riquelme等人通过将ViT的一些FFN层替换为MoE层，对JFT-300M或JFT-3B进行预训练，取得了在图像分类方面的巨大成功。Hwang等人在Swin-transformer V2架构中添加MoE层，以进行ImageNet-22K的预训练，用于图像分类和目标检测任务。Xue等人通过将FFN替换为MoE并在 Transformer块之间共享MoE层，提出了一种更加参数高效和有效的框架WideNet，在ImageNet-1K预训练中实现了性能提升。Komatsuzaki等人使用预训练的密集 Transformer检查点来启动MoE模型的训练，并在JFT-300M预训练和ImageNet-1K微调方面取得了有希望的结果。Xue等人通过从预训练的MoE模型中进行知识蒸馏，来学习一个密集模型，并在ImageNet-1K上进行了验证。

在本文中提出了一种新的通用训练方法，以提高ViT的性能，而不增加推理成本，其目标和实现方法与MoE不同。此外，以前的MoE-ViT研究主要集中在更大的2D视觉数据集上，在较小的2D视觉数据集和3D视觉方面进行了较少的探索。

本文表明，在ViT中添加原始的MoE会降低在小型2D视觉数据集和3D视觉上的性能，并引入了一种新颖的"Early EMA"训练方法，可以极大地提高它们的性能。

2.2、结构重参化

结构重参数化的本质是在训练过程中使用多分支结构代替卷积或全连接层，以增强模型。训练完成后，使用等效的参数转换或特定算法融合多个分支以进行推理。典型的例子有ACNet、DBB和RepVGG。

具体而言，ACNet在训练过程中将常规的K × K卷积替换为具有K × 1、1 × K和K × K卷积并行结构的ACBlocks。训练完成后，ACBlocks等效地转换为单个K × K卷积层。DBB在训练过程中使用1 × 1卷积、K × K卷积、1 × 1-K × K卷积和1 × 1卷积-平均池化的多分支结构来替代原始的K × K卷积，推理过程中等效地转换为单个K × K卷积。RepVGG在训练过程中在每个3 × 3卷积层上添加并行的1 × 1卷积分支和恒等分支，在训练后等效地转换为单个3 × 3卷积层。

由于能够在不增加推理成本的情况下提高性能，结构重参化在各种计算机视觉任务中受到了越来越多的关注和应用。例如，ExpandNets使用它来设计紧凑型模型，RepNAS将其用于神经架构搜索，ResRep将其与剪枝相结合。然而，以前对于结构重参化的研究主要集中在CNN架构上，这使得它们与ViT架构不兼容。

本文引入了一种新颖的方法，将ViT架构的训练和推理阶段解耦。在训练过程中，使用高效的MoE稀疏计算，采用随机均匀划分和 Experts Weights Averaging机制。在推理过程中，将模型转换回原始的ViT架构，从而在不增加推理参数和延迟的情况下实现性能改进。这种方法是首次将MoE稀疏计算应用于改善ViT架构性能的方法。

2.3、Weight Averaging

权重平均是深度学习中常用的技术，在各种任务中得到了广泛应用。例如，在自监督和半监督学习中，有方法使用学生模型的指数移动平均（EMA）权重作为教师，为学生提供更平滑的目标。同样，在在线知识蒸馏中，PCL将每个分支的EMA权重作为对等的平均教师，协同地在分支之间传递知识。

此外，许多研究使用权重平均来增强模型的泛化能力。例如，随机权重平均（SWA）对SGD优化轨迹上的多个点的权重进行平均，实现了比传统训练更好的泛化能力。Wortsman等人对原始的零样本模型和微调模型进行平均，从而在分布内和分布外的泛化中都取得了改进。Model Soup在微调期间对具有相同预训练初始化但不同超参数配置的模型进行平均，以获得更好的模型。DiWA对权重平均在OOD方面的成功提供了理论分析，并提出通过对多个独立训练的模型的权重进行平均来减少预测之间的协方差，从而提高OOD的泛化性能。

与上述在模型级别上进行权重平均以获得单个模型的研究不同，本文的方法在MoE的训练过程中在 Expert级别上执行权重平均。通过将其他 Expert的权重平均到每个 Expert上，获得了多个 Expert。

2.4、Dropout and Other variants for ViT

Dropout是一种常见的技术，通过在训练过程中随机地丢弃神经元及其对应的连接来防止过拟合，从而提高网络的泛化能力。对于ViTs，大多数方法在自注意力或FFN层中使用固定的丢弃率进行dropout。与此同时，随机深度（Stochastic Depth）通常用于ViTs，随机地丢弃一些块，并对更深的块应用较高的丢弃率以提高泛化性能。最近，一种针对ViT的新型dropout变体被引入，称为dropkey。Dropkey通过在注意力计算过程中随机丢弃一些Key为每个注意力块分配自适应的操作符。这种策略约束了注意力分布，使其更加平滑，更有效地减轻了过拟合问题。此外，dropkey会随着块数的增加逐渐降低丢弃率，避免了在训练过程中使用恒定丢弃率的传统dropout方法可能导致的不稳定性。

在本文的方法中，对于每个MoE层，在训练过程中随机且均匀地将所有Token分配给所有 Expert。这等效于在每个 Expert处随机丢弃大量的Token。这样的操作可能有助于防止过拟合，同时确保训练延迟几乎与原始训练相当。

3、本文方法

3.1、PRELIMINARY

1、MoE

如图1(a)所示，标准的MoE层包括一组N个 Expert（即N个FFN），表示为，以及具有权重的Routing网络G。给定输入示例，训练时和推断时的MoE层的输出可以写成：

其中，表示到第个 Expert的Routing分数，代表第个 Expert的输出，表示选择前个 Expert，并将其他 Expert的设置为0。通常情况下，，这意味着是一个稀疏的N维向量。当时，不需要计算。

2、结构重参化

如图1(b)所示，结构重参化在训练期间使用多分支CNN结构替代常规卷积层。用表示总共的M个分支（通常是具有兼容尺寸的M个不同运算符）。对于给定的输入，训练时结构的输出可以表示为：

训练完成后，这些M个分支将被等效地转换为单个卷积层进行推断。对于推断时的输入，输出变为。

3.2、 A New General Training Scheme

1、动机

本论文旨在设计一种新的通用训练方案，用于ViT，在不增加推断成本的情况下实现性能提升。如公式3所示，经典的结构重参化仅适用于由卷积组成的CNN，由于多分支结构，它显著增加了训练成本。如公式2所示，最近的MoE由N个 Expert和一个Routing网络组成。

首先，Routing网络引入了需要学习的附加参数，并容易导致负载不均衡问题。
其次，N个 Expert大大增加了部署模型的负担。
第三，Routing分数在将MoE转换为FFN时成为瓶颈。

考虑到这些问题，首先提出了一种特别设计的、更高效的MoE，用于替换ViT的一些FFN，通过随机均匀划分将Token分配给 Expert。然后，提出了一种简单但有效的技术，称为 Experts Weights Averaging，它在每个训练迭代结束时将其他 Expert的权重平均到每个 Expert。因此，在训练完成后，可以将每个MoE转换为单个FFN进行推断。

2、EWA Training

Overall

在将EWA训练应用于从头开始训练ViT模型时，首先通过用基于随机均匀划分的MoE层替换一些FFN来创建一个随机初始化的MoE ViT。在训练过程中，在每次权重更新后对每个MoE层进行 Experts Weights Averaging，共享率为，并且会在训练时期从0线性增加到超参数共享率。在训练完成后，通过对 Expert进行平均将每个MoE转换为一个FFN。

MoE layer with Random Uniform Partition

考虑一个由个 Expert（即N个相同的FFN）组成的MoE层，表示为，假设有T个输入，其中可以被整除。具有随机均匀划分（RUP）的MoE层的输出可以表示为：

如第4式所示，个输入被随机均匀地划分为部分，其中表示分配给第个 Expert的个输入。进一步地，如第5式所示，对于中的每个输入，其输出为。

Experts Weights Averaging

给定个 Expert的权重和共享率，对每个MoE层执行的Experts Weights Averaging（EWA）可以表述为：

其中表示第个 Expert的新权重。简而言之，将其他 Expert的权重平均分配给每个 Expert，从而获得多个 Expert。

Convert MoE into FFN

训练结束后，每个MoE层将通过简单地平均 Expert的权重来转换为一个FFN层。对于由N个 Expert 组成的MoE层，相应的推断时FFN可以表示为：

通过这种方式，训练时的MoE模型将在推断时转换为原始的ViT模型。

3、EWA Fine-tuning

在将EWA训练应用于微调预训练的ViT模型时，与从头开始训练ViT的唯一区别是创建的MoE ViT模型的权重初始化。不同于随机初始化，使用给定的ViT模型检查点来初始化MoE ViT。在用MoE层替换某些FFN层时，特定MoE层的每个 Expert被初始化为预训练ViT模型中对应FFN层的副本。对于保持不变的其他层，它们的权重直接继承自原始的ViT检查点。

4、理论分析

本节从理论上分析了EWA训练为何有效以及其工作原理。为了方便分析，关注一个MoE层和m个训练步骤。假设当前训练步骤为，在根据公式6计算后，经过 Experts Weights Averaging化后的第个 Expert的新权重可以表示为：

接下来，开始第个训练步骤，MoE的每个 Expert的权重更新为，其中，表示学习率。进一步进行 Experts Weights Averaging化，如公式9所示。根据公式8和，可以将公式9进一步重写为公式10。

类似地，可以得到，和。通过数学归纳，得到公式12。

根据上述推导，有两个发现：

观察公式9和公式11的第一个项，进行 Experts Weights Averaging化时会出现逐层的权重衰减；
观察公式10和公式12的最后一个项，会在EWA训练的训练迭代过程中持续聚合多 Expert历史指数平均权重。

5、实验

5.1、EWA Training

1、在各种架构和数据集上的2D分类

如表2所示，与普通训练相比，EWA训练可以为不同的ViT架构在不同的数据集上带来一致的性能提升。对于标准的ViT-S，EWA训练可以在CIFAR-100和Tiny-ImageNet数据集上分别获得{1.72％，2.09％}的显著提升。

即使对于专门设计的ViT架构及其增强版本，添加EWA训练方案仍然可以提高它们的性能，这进一步证实了其有效性和泛化性。例如，增强版Swin-T在CIFAR-100和Tiny-ImageNet数据集上的普通训练已经取得了相对较高的top-1准确率，分别为80.31％和64.34％。即使对于每个数据集，它仍然可以从EWA训练中获得{0.21％，0.79％}的性能提升。

2、3D视觉的任务

如表3所示，提出的EWA训练方案在点云分类和语义分割任务中获得了一致的性能提升。值得注意的是，本文的方法在分割任务中实现了0.52％的整体准确度（OA）提升，1.40％的平均准确度（mAcc）提升和1.74％的平均交并比（mIoU）提升。

5.2、EWA Fine-tuning

如表4所示，展示了标准微调和EWA微调之间的性能比较。尽管标准微调的性能已经相当高，但EWA微调仍然可以带来进一步的改进。具体而言，在经过EWA微调后，ViT-S和ViT-B在{Food-101，CIFAR-100}数据集上分别获得了{0.38％，0.45％}和{0.54％，0.71％}的性能提升。

5.3、消融研究

Share schedule and Share rate

在上述内容中，默认情况下共享计划设置为线性增加，在这里将其与常数共享计划进行比较。对于共享率，在{0.1，0.2，0.3，0.4，0.5}范围内进行调整。使用不同的共享计划和共享率在CIFAR100上训练ViTS，并进行EWA训练。

如图2所示，不同的共享计划和共享率的EWA训练始终可以优于香草（即基准）训练，其中线性增加的共享计划在大多数情况下表现更好。

5.4、更多讨论和细节

更有趣的是，作者发现在小型2D视觉数据集和3D视觉任务上，标准的MoE相对于其原始的密集对应模型表现较差，而EWA技术可以无缝地应用于MoE并帮助其取得更好的学习效果。

1、naive MoE

按照在CIFAR-100、Tiny-ImageNet、ModelNet40和S3DIS上训练ViT-S的设置，进一步训练了V-MoE-S。V-MoE-S将ViT-S的部分FFN替换为naive top-1RoutingMoE，每隔一个块进行一次。将 Expert数量设置为4，负载平衡损失λ为0.01，2D数据集的容量比C为1.05，3D数据集的C为1.2。如表5所示，在使用naive训练时，naive V-MoE-S在不同数据集上的性能较ViT-S有所下降。

2、naive MoE + Early EWA

当将Experts Weights Averaging（EWA）引入naive MoE时，情况发生了变化。考虑到在整个训练过程中始终使用EWA最终会导致各个 Expert之间的权重几乎相同，只在训练的早期阶段对naive MoE执行EWA。默认情况下，使用恒定的共享计划，并且只在总训练轮数的前半部分执行早期EWA。通过早期EWA训练，V-MoE-S的性能得到了极大的改善，在CIFAR-100和Tiny-ImageNet上分别获得了{2.14%，3.63%}的top-1准确率提升。

此外，早期EWA训练还可以提高V-MoE-S在各种3D视觉任务上的性能，ModelNet40分类任务上准确率提升了0.36%，S3DIS Area5分割任务上OA提升了0.30%，mAcc提升了1.40%，mIoU提升了1.31%。

7、总结

总之，本文为ViTs提出的训练方案在不增加推理延迟和参数的情况下实现了性能提升。通过在训练过程中设计高效的MoEs和EWA，并在推理过程中通过对 Expert的平均来将每个MoE转化为FFN，将ViTs的训练和推理阶段解耦。在各种2D和3D视觉任务、ViT架构和数据集上进行的全面实验表明了其有效性和通用性。

理论分析进一步支持了本文的方法，作者的训练方案还可以应用于fine-tuning ViTs，并在各种视觉任务中改善naive MoE的效果。