迈向更高效通用的加速之路：谷歌提出视觉和多任务MoE方法

大数据文摘授权转载自智源社区

作者：梁子、戴一鸣

审校：戴一鸣

近日，MoE家族又添两个新成员：V-MoE和Task MoE，由谷歌提出，能够显著提升计算机视觉和多任务学习模型的计算效率，进一步扩大模型的规模。

MoE是近年来快速发展的一种模型加速方法，研究者认为，如同人类有不同领域的专家一样，模型在面对不同的数据时也可以由更为“擅长”的“专家”（特定的网络结构）来处理。

与其对每个输入都要训练整个模型，不如在训练过程中根据输入的不同，仅触发并更新该模型中的某一部分，以降低模型总体的训练代价，学习出更大参数量的模型；并在推理时，可以针对不同的下游任务，将大模型拆解得到一个体积更小，但是效果更好的子模型，以实现更加高效的推理。

在这种思路的指导下，一系列混合专家模型（Mixture of Expert，MoE）方法应运而生。依托MoE方法，Switch Transformer等模型参数提升至了万亿级别。而谷歌近日提出的新型MoE方法，让这一方法逐步扩展到更多领域，成为一种更为通用，且相比知识蒸馏等方法更为高效的加速策略。

V-MoE

1.基本原理

Vision Transformers（ViT）是CV领域中比较有代表性的模型之一，该模型通过将图像均匀切分成若干个片段（Patch）来模拟Transformer的输入。V-MoE（Vision MoE）则是在ViT的基础上，通过将Transformer模块的内部结构改为MoE形式而得到的变体，是一种稀疏混合专家模型。

V-MoE的模型结构和MoE Transformer的结构类似，如下图所示。研究者将ViT模块中的feedforward层（FFN）替换为多个FFN的混合集成（其中每一个FFN为一个Expert）。之后，对于每一个模块的输入，ViT通过学习路由单元（Router）来选择性地激活若干个Expert，并为每一个Expert赋予其权重。

在V-MoE中，每张图片中不同的片段（即不同的Token）将被Router分配到不同的几个FFN中。一般而言，Router会为每个Token选择Top-K个FFN（K一般取1或2），而在每一个ViT模块中大概会有E个FFN（在本论文中E取值为32）。通过这种方式，V-MoE模型的大小可以得到提升，而且在模型规模提升的情况下，每一个Token对应的计算量都不会发生太大变动。下图展现了V-MoE encoder模块的更多细节。

2.实际应用

然而，在实际应用中，受限于硬件，MoE系统一般不会动态地调整Buffer，而是人为地设置一个固定大小的Buffer容量，超过该容量的Token将会被截断。为加快V-MoEs的推理时间，研究者认为，可以为这个Buffer 容量设置一个较小的阈值（甚至可以比一般待处理的Token数目更小）。这些FFN将会强制地被要求跳过一部分Token的处理。

在V-MoE之前，已有的工作大多是随意地对Token进行选择，而V-MoE则通过学习，对这些Token进行排序，优先处理那些更加重要的Token。通过这种方式，V-MoEs可以在降低计算代价的同时，保证模型仍然具有较好的预测效果。

下图展现了V-MoE在设置不同buffer容量下的精度情况。可以发现，即使V-MoE所处理Token只有原来的20%，模型的效果也比单个的ViT（图中的横线）稍好。而通过选择哪些Token更加重要（即BPR），可以在FFN的Capacity较低的情况下一样获得较好的效果。

最后，谷歌的研究者还对MoE的机理进行了探索实验。如下图所示，左图和右图分别是在V-MoE中底层模块和顶层（靠近输出）的可视化结果。横轴为每个Expert的id，纵轴为每个类别的id，颜色越深代表其关联性更强。可以发现，对于高层次的MoE，Class与Expert表现出了较强的对应关系。

3.实验结果

研究者在JFT-300M上进行了预训练，并在ImageNet-5shot（每一类有5个样本）上进行了测试，效果如下：

从图中可以看出，V-MoE在拥有较好的模型效果的同时，也拥有更快的训练速度。

多任务MoE

1.基本原理

Task MoE（多任务MoE）也是在MoE Transformer的基础上进行改进的。MoE Transformer是以Token为单位，而Task MoE则是以训练任务为基本单位。相比于过去所有神经元都需要进行计算的语言模型， Task MoE通过一个路由单元（Router Neural Networks），将不同类型输入中的所有Token分发到只和该类型输入相关的子神经网络中。因此，对于输入而言，Task MoE中仅有一部分神经元被激活 。路由单元在这里是起到了Gate Function的作用，可用如下公式表示：

以机器翻译任务为例，对于不同的翻译任务，每个任务都有一个One-hot编码，某一门特定语言的所有Token都会被路由到相同的子网络。这一点和Token MoE等采用了Router模块的模型有些不同，后者是使用Router将不同的Token映射到不同子网络中。

Task MoE在推理时相比Token MoE具有更多优势。Token MoE需要为每一个Token判断对应的Expert，因此在训练时和推理时都需要维护完整的一个Transformer模型，计算代价很高。

而对于Task MoE，Router所处理的对象是Task-level的，可以在模型推理时将整个模型拆分开，为每一个任务仅生成一个子网络，并删掉不会被使用的Experts。通过这种方式，Task MoE可以通过训练一个大型的多任务的模型。而在推理时，则针对每一个任务将模型拆分出一个对应的子网络，以提升推理速度。

与知识蒸馏等算法不同，上述提取子网络的提取过程不需要进行额外的训练或其他处理。下图展现了Task MoE的训练过程以及为每个任务提取子网络的流程。

2.实验结果

为了展示Task MoE在推理时的优势，研究者设计了在不同的Batch size下模型的吞吐率（即每秒所生成的Token个数）对比试验。该实验的Baseline是Token MoE（533M）和常规的全激活的Transformer（143M）。

由于提取了子网络，Task MoE的模型参数两只有Token MoE的七分之一，该参数级别完全可以运行在单个的TPUv3上，而后者则需要64块相同硬件的支持。通过实验结果可以发现，Task MoE的吞吐量达到了Token MoE的两倍。并且，Token MoE还要花大量的推理时间（25%）在设备间通信上，而Task MoE则几乎没有。

由于在推理时只需要保留任务对应的子网络，Task MoE的参数量比Token MoE小很多，还可以在单GPU上运行，吞吐量上也远远超过常规的MoE模型。

此外，作为一种可行的降低模型参数的方法，论文还将Task MoE和知识蒸馏方法进行了比较。如下图所示，相比于参数量不相上下的蒸馏模型，Task MoE同样可以表现出较高的性能。

其他MoE算法

1.Mixture of Experts（MoE）

MoE算法提出于2017年，其基本思想是在训练推理过程中，仅激活一部分神经网络。选择更为合适的神经网络结构，进行学习。通过在神经网络中引入多个Expert网络（Expert Network），可以降低需要激活的神经元数量，提升模型的计算效率和规模。

以基于RNN的MoE结构为例。对于当前模块的输入x（在自然语言处理中，这样的一个输入多是当前输入Token的一个表示向量），MoE会维护多个Expert（正常情况下模型只有一个Expert，MoE在此处进行了集成），并通过一个门控单元（Gating Network）来选择当前的输入需要激活哪几个Expert，并决定每一个Expert的权重。

每一个输入的Token都仅会激活很多Expert中的一个——这一部分也被认为是和当前输入相关的部分。在MoE中，最重要的结构是门控单元和Expert。一般而言，此处的Expert和常规的神经网络没有差别，而门控单元主要有两种设计：

1. Softmax门控，即对输入x映射到n维后使用Softmax来获取Top-k个Expert：

2. Noisy Top-k 门控，即在Softmax门控的基础上，加入Noise来实现负载均衡（即各个Expert的weight相差不要太大），并只维持top-k来确保所选择Expert的稀疏性：

2.MoE Transformer

MoE Transformer是在MoE的基础上把Backbone从RNN换成了Transformer。对于每一个Transformer模块，MoE Transformer选择对其中的FFN进行扩充和集成，如下图所示：

MoE Transformer与Transformer唯一的区别是将部分Transformer模块中的FFN扩充成了多个，如下公式所述：

在公式中，wi和wo分别对应的是FFN中第一层和第二层的权重矩阵。而x_s则是Token s的表示向量。

3.FastMoE

MoE和MoE Transformer是由谷歌提出的，使用MoE的目的是为了构建更大规模的语言模型，因此MoE算法和代码等实现和谷歌已有的软硬件深度绑定。

为了进一步推动MoE系列技术的发展，让广大研究者可以使用MoE算法进行大模型研发，智源研究院、清华大学等机构的研究者推出了FastMoE。这是一种基于PyTorch，开源的MoE系统。该系统简洁易用，支持多种Transformer变体模型，只需要几行代码即可在PyTorch中实现构建和调用。FastMoE主要流程如下:

伪代码实现如下：

FastMoE的处理流程如下：

总结

以上工作便是MoE通过集成和路由，实现模型加速方法的介绍。通过谷歌的实验可以发现，混合结构确实可以降低模型训练上的负担，从而使得训练比常规模型更大的稀疏模型成为可能，而通过这种方式，模型也可以获得更好的效果。

除此之外，Task MoE也表明了MoE系列模型在多任务学习上的优势。当然，该模型所涉及到的任务仍然局限在机器翻译这样的背景下，如果能通过多种下游任务的联合训练学习一个比单任务微调更好的MoE模型，又可以通过使用子网络分别地为每一个下游任务瘦身，那么MoE系列模型则有更大的优势。

参考文献

[1] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, Jeff Dean:  Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. ICLR (Poster) 2017

[2] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, Zhifeng Chen:  GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. ICLR 2021

[3] Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, Neil Houlsby:  Scaling Vision with Sparse Mixture of Experts. CoRR abs/2106.05974 (2021)

[4] Sneha Kudugunta, Yanping Huang, Ankur Bapna, Maxim Krikun, Dmitry Lepikhin, Minh-Thang Luong, Orhan Firat:  Beyond Distillation: Task-level Mixture-of-Experts for Efficient Inference. EMNLP (Findings) 2021: 3577-3599

[4] Jiaao He, Jiezhong Qiu, Aohan Zeng, Zhilin Yang, Jidong Zhai, Jie Tang:  Fast MoE: A Fast Mixture-of-Expert Training System. CoRR abs/2103.13262 (2021)

[5] Scaling Vision with Sparse Mixture of Experts. Carlos Riquelme

https://ai.googleblog.com/2022/01/scaling-vision-with-sparse-mixture-of.html (2022)

[6]Learning to Route by Task for Efficient Inference. Sneha Kudugunta https://ai.googleblog.com/2022/01/learning-to-route-by-task-for-efficient.html (2022)