本文将介绍 MoE 的构建模块、训练方法以及在使用它们进行推理时需要考虑的权衡因素。
与稠密模型相比,预训练速度更快; 与参数数量相同的模型相比,推理速度更快; 因为要把所有专家模型都加载在内存中,所以需要大量显存; 在微调方面面临许多挑战,但最近在 MoE 指令微调方面的工作很有希望能解决这些问题。
使用稀疏的 MoE 层代替稠密的前馈网络(FFN)层。MoE 层中有一定数量(如 8 个)的 「专家」,每个专家都是一个神经网络。实际上,专家可以是 FFN,也可以是更复杂的网络,甚至是 MoE 本身,这样就会形成有多层 MoE 的 MoE。 使用门控网络或者路由来决定将哪个 token 发送给哪个专家。例如,在下图中,「More」被发送给第二个专家,而 「Parameters」被发送到第一个专家。如何将 token 通过路由发送给专家是使用 MoE 时需要计划的重点之一,这是因为路由同样由学习到的参数组成,并与网络的其他部分同时进行预训练。
训练:MoE 能够大大提高预训练的计算效率,但在微调过程中难以实现泛化,从而导致过拟合。 推理:虽然 MoE 可能有很多参数,但在推理过程中只使用其中的一部分。与参数数量相同的稠密模型相比,推理速度要快得多。然而,所有参数都需要加载到 RAM 中,因此对内存的要求很高。例如,给定一个像 Mixtral 8x7B 这样的 MoE,需要有足够的 显存 来容纳一个 47B 参数的稠密模型。为什么是 47B 参数而不是 8 x 7B = 56B?这是因为在 MoE 模型中,只有 FFN 层被视为独立的专家网络,其余模型的参数都是共享的。同时,假设每个 token 只传入两个专家网络,那么推理速度(FLOPs)就像使用 12B 模型(而不是 14B 模型),因为它进行的是 2x7B 的矩阵乘法运算,同时有些层是共享的(后文将详细介绍)。
专家网络组件化:在传统的 MoE 中,整个系统由一个门控网络和多个专家网络组成。在 SVM、高斯过程和其他方法中,研究者们对作为整体模型的 MoE 进行了探索。Eigen、Ranzato 和 Ilya 的研究,将 MoE 作为更深层网络的组成部分进行了探索。MoE 能够作为多层网络中的组成部分,使得模型既庞大又高效成为可能。 条件计算:传统网络每一层都会对所有输入数据进行处理。接着,Yoshua Bengio 研究了根据输入的 token 动态激活或停用组件网络的方法。
随机路由:在 top-2 的设置中,研究者总是选择排名第一的专家,同时第二位专家以与其权重成正比的概率被选中。 专家容量:还可以设定一个阈值,即一个专家可以处理多少 token。如果两位专家的处理能力都已达到阈值,则 token 会被视为溢出,并通过剩余连接发送到下一层(或在其他项目中完全丢弃)。这一概念将成为 MoEs 最重要的概念之一。为什么需要设置专家容量呢?因为所有张量形状都是在编译时静态确定的,但无法提前知道每个专家将获得多少 token 作为输入,所以需要确定容量。
减少了路由的计算量 每个专家的 batch size 至少减半 通信成本降低 模型的质量仍能得到保证
单任务微调 多任务指令微调 多任务指令调整后进行单任务微调
数据并行:相同的权重在所有内核上复制,数据在内核上分割。 模型并行:模型在不同内核间分区,数据在不同内核间复制。 模型和数据并行:在不同内核间划分模型和数据。请注意,不同的内核处理不同批次的数据。 专家并行:将专家放在不同的工作站上。如果与数据并行相结合,则每个内核都有不同的专家,数据在所有内核之间进行分割
Switch Transformer 的作者进行过蒸馏实验。将 MoE 蒸馏到稠密模型中,可以保留 30-40% 的稀疏化增益。因此,蒸馏可以带来更快的预处理速度以及模型更小的好处。 创新的路由算法:将完整的句子或任务路由至专家网络,使用用于提取的子网络来提供服务。 MoE 聚合:这种技术可以合并专家权重,从而减少推理时使用的参数。
Megablocks: https://github.com/stanford-futuredata/megablocks Fairseq: https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm OpenMoE: https://github.com/XueFuzhao/OpenMoE
Switch Transformer (Google):从 8 到 2048 个专家的基于 T5 的模拟引擎集合。最大的模型有 1.6 万亿个参数。 NLLB MoE (Meta):NLLB 翻译模型的 MoE 变体。 OpenMoE:基于 Llama 的 MoE 。 Mixtral 8x7B (Mistral):一种高质量的翻译模型,其性能优于 Llama 2 70B,推理速度也更快。
将 Mixtral 提炼为稠密模型 探索专家模型的合并技术,分析其对推理时间的影响 Mixtral 的量化技术