Peng Jin1,2,3∗, Bo Zhu4, Li Yuan1,2,3 , Shuicheng Yan4
1北京大学深圳研究生院电子与计算机工程学院,中国深圳
2鹏城实验室,中国深圳
3北京大学深圳研究生院AI for Science (AI4S)-Preferred Program,中国
4昆仑2050研究院 & Skywork AI,新加坡
jp21@stu.pku.edu.cn, yuanli-ece@pku.edu.cn
代码: https://github.com/SkyworkAI/MoE-plus-plus

A1 主要贡献

本文旨在同时提升混合专家(Mixture-of-Experts, MoE)方法的有效性和效率。为实现此目标,本文提出了MoE++,一个集成了前馈网络(FFN)和零计算专家的通用异构MoE框架。

核心问题与研究目标
- 问题: 现有的大多数MoE方法(如【12, GLaM: Efficient scaling of language models with mixture-of-experts, 2022, ICML】,【13, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022, Journal of Machine Learning Research】,【28, BASE layers: Simplifying training of large, sparse models, 2021, ICML】,【42, Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale, 2022, ICML】)为所有令牌(token)激活固定数量的FFN专家。然而,并非所有令牌在语言任务中都具有相同的预测难度。例如,简单的令牌(如逗号)可能只需要一个专家,而与现有专家匹配度差的令牌甚至可能从完全跳过MoE层中受益。这种僵化的混合机制导致训练和推理效率低下,并最终导致模型性能次优。
- 目标: 提出一种灵活的计算分配机制,能够根据令牌的难易程度动态分配计算资源,从而在降低计算开销的同时提升模型性能。

创新点
1. 零计算专家 (Zero-computation experts):首次为MoE架构提出零计算专家的概念。具体引入了三种类型:
* 零专家 (zero expert):对应于丢弃(discard)操作。
* 复制专家 (copy expert):对应于跳过(skip)操作。
* 常数专家 (constant expert):对应于替换(replace)操作。
通过引入这些专家,MoE++相比于香草(vanilla)MoE,在计算开销更低的情况下拥有更高的拟合能力。

图1:香草MoE与我们提出的MoE++架构的高层比较。子图(a)展示了使用Top-2路由策略的标准MoE层,而子图(b)展示了在MoE++中集成零计算专家。值得注意的是,这些零计算专家需要的参数几乎可以忽略不计,确保MoE++的总参数数量与香草MoE保持在同一水平。
图1:香草MoE与我们提出的MoE++架构的高层比较。子图(a)展示了使用Top-2路由策略的标准MoE层,而子图(b)展示了在MoE++中集成零计算专家。值得注意的是,这些零计算专家需要的参数几乎可以忽略不计,确保MoE++的总参数数量与香草MoE保持在同一水平。

  1. 门控残差 (Gating residuals):将前一层的门控分数(gating scores)融入当前层的专家选择中。这种设计使得每个令牌在选择当前MoE++层中的合适专家时,能够考虑到其在前一层所采取的路径。

  2. 灵活的计算分配 (Flexible computation allocation):MoE++通过为简单令牌分配更少的FFN专家来优化计算分配,从而让更多的FFN专家可以专注于处理具有挑战性的令牌。这一机制不仅提升了整体性能,还带来了高达2倍于香草MoE方法的专家前向吞吐量。

核心优势
- 低计算开销: 与香草MoE中对所有令牌采用统一混合机制不同,MoE++允许每个令牌与动态数量的FFN交互,或通过常数向量进行调整,甚至完全跳过MoE层。
- 高性能: 通过让简单令牌使用更少的FFN专家,MoE++使得更多的专家可以专注于处理困难的令牌,从而释放出比香草MoE更大的性能潜力。
- 部署友好: 由于零计算专家的参数可以忽略不计,可以将所有零计算专家部署在每个GPU上,消除了分布式FFN专家带来的显著通信开销和专家负载不平衡问题。

A2 方法细节

一个标准的混合专家(MoE)层由N个专家网络 $E = \{E1, E2, ..., EN \}$ 和一个激活Top-K个专家的路由器G组成。通常,激活的专家数量K是固定的,并且远小于专家总数N。形式上,给定输入令牌x,MoE层的输出令牌y是K个激活专家的输出的加权和:

公式1
公式1

香草MoE的构成。一个香草MoE层通常由多个结构相同的专家组成,其中每个专家都是一个标准的前馈网络(FFN)。此外,路由器通常实现为一个可训练的权重矩阵。形式上,香草MoE层中的专家和路由器可以定义为:

公式2
公式2

其中 $W \in R^{N \times D}$ 是可训练的权重矩阵,D表示模型的隐藏大小。由于在香草MoE层中为简单和困难的令牌都激活了固定数量的FFN,因此在处理简单令牌时,香草MoE层可能存在训练和推理效率低下的问题。

MoE++框架概述。我们提出的MoE++是一个通用的、异构的MoE框架,它集成了FFN和零计算专家。此外,MoE++通过门控残差增强了路由器,允许每个令牌在当前MoE++层中选择合适的专家时,考虑其先前的路径。另外,为了有效地训练异构专家结构,我们引入了异构负载均衡损失和异构专家容量分配策略。所提出的MoE++的核心组件如图2所示。

图2:MoE++的核心组件。子图(a)展示了FFN专家和零计算专家的架构,而子图(b)展示了门控残差,它允许每个token在选择合适的专家时考虑其先前的路径。
图2:MoE++的核心组件。子图(a)展示了FFN专家和零计算专家的架构,而子图(b)展示了门控残差,它允许每个token在选择合适的专家时考虑其先前的路径。

3.1 零计算专家

零计算专家的设计原则。在MoE++中,重新设计的专家架构应满足特定标准:(i)它应该尽可能精简,以高效处理简单令牌;(ii)为确保与香草MoE的公平比较,新专家引入的参数数量应几乎可以忽略不计。在这些原则的指导下,我们引入了零计算专家,每个专家只执行最基本的操作。具体来说,我们提出了三种类型的零计算专家:零专家、复制专家和常数专家,它们分别对应于丢弃、跳过和替换操作。

零专家(Zero Experts)。最简单的零计算是丢弃当前输入。给定输入令牌x,零专家的输出为0,公式如下:

公式3
公式3

本质上,零专家的存在可以将Top-2 MoE++层降级为Top-1 MoE++层。具体来说,如果一个零专家被激活,其零输出使得Top-2 MoE++层的输出等同于单独另一个专家的输出。因此,在MoE++中,零专家的引入为同时处理简单和困难的令牌增加了灵活性。

复制专家(Copy Experts)。受残差网络(【17, Kaiming He, et al. Deep residual learning for image recognition. 2016, CVPR】)的启发,我们提出了复制专家,其输出等于输入,因此等同于一个快捷连接(shortcut):

公式4
公式4

直观上,复制专家提供了跳过当前MoE++层的选项。具体来说,如果输入令牌与现有专家匹配不佳,它可能会从绕过MoE++层中受益。

常数专家(Constant Experts)。零专家和复制专家都不包含可训练的参数,因此它们在处理输入令牌方面的灵活性有限。为了解决这个限制,我们提出了常数专家,它用一个可训练的向量v替换输入令牌x。然而,完全替换会丢失所有输入信息。因此,我们使用一个可训练的权重矩阵Wc来动态预测替换的比例。形式上,常数专家的输出可以定义为:

公式5
公式5

其中 $Wc \in R^{2 \times D}$ 是可训练的权重矩阵,D是输入令牌x的隐藏大小。

零计算专家的组合效益。通过为简单令牌分配更少的FFN专家,并为困难令牌投入更多的FFN专家,MoE++优化了计算分配。因此,MoE++以比香草MoE更少的计算量实现了更好的性能。此外,MoE++显著拓宽了子网络的范围。例如,将一个FFN专家与一个常数专家结合,相当于使用一个可训练的向量来调整FFN专家的输出。同样,零专家与复制专家的组合允许输入令牌完全绕过当前层。

3.2 路径感知路由器

路径感知路由器的重要性。由于MoE++包含异构专家,路由器的设计变得比香草MoE更为关键。为此,我们提出了路径感知路由器,它在选择合适的专家时会考虑前一层所采取的路径。

表1:MoE++与MoE的复杂度比较。超参数 $\tau$ 控制零计算专家和FFN专家之间的令牌分配比例。

表1
表1

门控残差机制。类似于残差网络(【17, Kaiming He, et al. Deep residual learning for image recognition. 2016, CVPR】),我们将前一层的路由分数加到当前层预测的路由分数上。具体来说,给定第j层的输入令牌xj和N个专家,我们使用一个可训练的变换矩阵 $W^j_g \in R^{N \times N}$ 将前一层得分整合到当前层:

公式6
公式6

其中 $W^j \in R^{N \times D}$ 是可训练的权重矩阵,D是隐藏大小。这些门控残差有效地在MoE++层之间建立了连接,从而确保了稳定的路由。

3.3 预训练期间的负载均衡

预训练中的负载均衡问题。直接训练MoE模型通常会导致大多数令牌被分配到少数几个专家,而其他专家则训练不足(【45, Noam Shazeer, et al. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. 2017, arXiv preprint arXiv:1701.06538】)。遵循先前的工作(【27, Dmitry Lepikhin, et al. Gshard: Scaling giant models with conditional computation and automatic sharding. 2021, ICLR】、【56, Fuzhao Xue, et al. Openmoe: An early effort on open mixture-of-experts language models. 2024, arXiv preprint arXiv:2402.01739】、【11, Damai Dai, et al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. 2024, arXiv preprint arXiv:2401.06066】、【51, Tianwen Wei, et al. Skywork-moe: A deep dive into training techniques for mixture-of-experts language models. 2024, arXiv preprint arXiv:2406.06563】),我们应用负载均衡损失和专家容量来确保预训练期间的负载均衡。

异构负载均衡损失。在香草MoE方法中,每个专家都是一个标准的前馈网络(FFN),因此所有专家被分配相同数量的令牌。然而,在我们提出的MoE++中,零计算专家和FFN专家的架构和参数数量差异显著,因此为这两种类型的专家分配相同数量的令牌是次优的。为此,我们引入一个超参数 $\tau$ 来控制零计算专家和FFN专家之间的令牌分配比例。具体来说,给定第t个输入令牌xt,异构负载均衡损失Lb的公式为:

公式7
公式7

其中T表示令牌的数量,N是专家的数量,1(∗)表示指示函数。较小的超参数 $\tau$ 意味着更多的令牌被分配给零计算专家。相比之下,较大的 $\tau$ 意味着更少的令牌被分配给零计算专家。

异构专家容量分配。专家容量被提出来通过限制路由到每个专家的最大令牌数来缓解严重的负载不平衡。由于MoE++为不同类型的专家分配不同数量的令牌,我们也为每种类型的专家设计了不同的专家容量。对于一个具有NFFN个FFN专家和NZC个零计算专家的MoE++模型,专家总数为N = NFFN + NZC。给定超参数 $\tau$ ,专家容量定义为:

公式8
公式8

其中 $\gamma$ 是预设的容量因子,T是令牌的数量。同样,较小的超参数 $\tau$ 意味着更多的容量被分配给零计算专家。对于两种类型的专家,如果一个专家未被充分利用,其未使用的容量将用填充令牌填充。一旦一个专家达到其容量,任何额外分配给该专家的令牌都将被丢弃,这意味着额外的令牌将直接传递到后续的Transformer(【50, Ashish Vaswani, et al. Attention is all you need. 2017, NeurIPS】)块。

总训练目标。最后,总的训练损失是交叉熵损失Lce和异构负载均衡损失Lb的加权和:

公式9
公式9

其中 $\beta$ 是一个权衡超参数,用于降低路由崩溃的风险。

3.4 效率分析

效率分析。值得注意的是,零计算专家处理一个令牌所需的计算和通信成本可以忽略不计。如表1所示,对于一个具有NFFN个FFN专家、NZC个零计算专家和超参数 $\tau$ 的MoE++模型,其计算复杂度仅为具有相同参数数量的MoE模型的一部分。

A4 实验环境

模型设置:
- 模型架构: 本文使用Megatron(【47, Mohammad Shoeybi, et al. Megatron-lm: Training multi-billion parameter language models using model parallelism. 2019, arXiv preprint arXiv:1909.08053】)作为训练框架,并在一系列规模上(从0.6B到7B)构建了MoE++模型及其对应的香草MoE基线。具体模型配置如表2所示,例如,“MoE++ 0.6B/(8+4)E”代表一个约0.6B参数的MoE++模型,包含8个FFN专家和4个零计算专家(1个零专家、1个复制专家、2个常数专家)。

表2:MoE++和香草MoE模型的尺寸和架构。“0.2B/0.6B”表示一个约0.6B参数的架构,在推理期间每个token激活0.2B参数。“1/1/2”表示一个带有一个零专家、一个复制专家和两个常数专家的MoE++模型。

表2
表2

硬件配置:
- GPU: 训练在一个包含4个节点、32个A100 GPU的集群上进行。

软件配置:
- 训练框架: Megatron。
- 评估工具: lm-evaluation-harness(【15, Leo Gao, et al. A framework for few-shot language model evaluation, 2024, URL https://zenodo.org/records/12608602】)。
- 分词器 (Tokenizer): 使用LLaMA2的分词器,词汇量为65,536。

数据集:
- MoE++仅在公共数据集上进行训练,包括从RedPajama(【9, Together Computer. Redpajama: an open dataset for training large language models, 2023a, URL https://github.com/togethercomputer/RedPajama-Data】)、Dolma(【48, Luca Soldaini, et al. Dolma: An Open Corpus of Three Trillion Tokens for Language Model Pretraining Research. 2024, arXiv preprint】)和Pile(【14, Leo Gao, et al. The Pile: An 800gb dataset of diverse text for language modeling. 2020, arXiv preprint arXiv:2101.00027】)中按不同概率采样的数据。详细的采样比例和训练策略见附录B.1。

A4 实验结果

与香草MoE的比较
- 实验内容: 在0.6B到7B等不同模型规模上,比较MoE++和香草MoE模型在9个简单基准测试上的性能,所有模型均使用100B tokens进行训练。
- 实验结果: 如表3所示,MoE++在所有规模上都持续优于香草MoE模型。
- 分析结论: MoE++不仅性能更优,而且专家前向吞吐量比同等规模的香草MoE模型高出15%~111%。这证明了MoE++在开发先进高效语言模型方面的坚实基础。
表3:MoE++与香草MoE模型的比较。下表中列出的所有MoE++和香草MoE模型的训练预算均为100B tokens。

表3
表3

与等效激活参数模型的比较
- 实验内容: 将MoE++模型的训练预算扩展到1T tokens,并与具有相当激活参数量的密集模型和其他MoE模型进行比较。
- 实验结果: 如表4所示,MoE++模型的性能与激活参数量是其2到3倍的密集模型相当。特别地,MoE++优于更大且使用更多tokens(1.1T)训练的OpenMoE-8B/32E模型,而其计算成本仅为后者的约57%。
- 分析结论: MoE++是训练大语言模型的一个极具前景的解决方案。
表4:与等效激活参数的LLMs的比较。“§”表示该模型不是从头开始训练,而是通过剪枝并使用LLaMA2-7B的权重进行持续微调。因此,为公平比较,我们将Sheared-LLaMA和LLaMA-MoE灰显。OpenMoE-8B/32E的结果来自其论文,因此只有部分任务结果可用。

表4
表4

消融分析
- 超参数 $\tau$ 的影响:
- 结果: 表3的对比显示,较小的 $\tau$ 值会为零计算专家分配更多tokens,从而带来更高的吞吐量;而较大的 $\tau$ 值则通常带来更好的性能。
- 结论: 为平衡计算成本和性能,默认将 $\tau$ 设置为0.75。
- 零计算专家的影响:
- 结果: 表5的消融研究表明,常数专家对模型的提升最大,这归因于其处理tokens的灵活性。包含所有三种零计算专家的完整模型性能最佳。
- 结论: 三种零计算专家共同为语言模型带来了益处。
表5:关于“MoE++ 1B/(16+4)E”模型中每种零计算专家影响的消融研究。所有模型均以100B tokens的预算进行训练,超参数τ设置为0.75。

表5
表5

- 门控残差的影响:
- 结果: 表6显示,移除门控残差会导致性能下降。
- 结论: 简单的门控残差能有效地在MoE++层之间建立连接,从而确保路由的稳定性。
表6:关于“MoE++ 1B/(16+4)E”模型中门控残差的消融研究。所有模型均以100B tokens的预算进行训练,超参数τ设置为0.75。
表6
表6

- 常数专家数量 $n_{const}$ 的影响:
- 结果: 如图3所示,随着常数专家数量的增加,模型平均性能先升后降。
- 结论: 性能下降的原因是常数专家的增加会减少其他类型专家的容量。论文提出一个自适应确定 $n_{const}$ 的公式:$n_{const} = \lfloor N_{FFN}/4 \rfloor$。
图3: 常数专家数量n_const的影响。该图显示了“MoE++ 1B/(16+4)E”模型在不同n_const下的平均性能。所有模型都以100B tokens的预算进行训练,超参数τ设为0.75。
图3: 常数专家数量n_const的影响。该图显示了“MoE++ 1B/(16+4)E”模型在不同n_const下的平均性能。所有模型都以100B tokens的预算进行训练,超参数τ设为0.75。

定性分析
- 任务级别的专家负载分布 (图4):
- (i) FFN专家的激活数量因任务内容和复杂度而异,而非任务整体难度(例如ARC-C比ARC-E激活更多FFN专家)。
- (ii) 零专家是平均激活次数最多的专家,且在较简单的任务中激活次数更多。
- (iii) 不同主题的任务表现出显著不同的专家分配模式。

图4:任务级别的专家负载分布可视化。结果来自“MoE++ 7B/(16+4)E”模型的第12层,超参数τ设置为0.75。
图4:任务级别的专家负载分布可视化。结果来自“MoE++ 7B/(16+4)E”模型的第12层,超参数τ设置为0.75。

- 令牌级别的FFN专家激活数量 (图5):
- (i) 动词(如"touch")倾向于激活大量FFN专家(平均接近上限2)。
- (ii) 名词通常激活中等数量的FFN专家(1.5到1.7之间)。
- (iii) 语义简单的令牌(如词根"pper")激活的FFN专家较少(低于1.5)。
- 结论: 这证实了MoE++的核心思想:为简单令牌使用更少的FFN专家,为困难令牌保留更多计算资源。
图5:令牌级别上每个token激活的FFN专家数量的可视化。结果来自“MoE++ 7B/(16+4)E”模型,超参数τ设置为0.75。我们在超过60,000个token上进行评估,并对所有MoE++层的结果取平均。分词器通常会将一个词分成多个部分,导致像“icken”这样在语义上无意义的token。
图5:令牌级别上每个token激活的FFN专家数量的可视化。结果来自“MoE++ 7B/(16+4)E”模型,超参数τ设置为0.75。我们在超过60,000个token上进行评估,并对所有MoE++层的结果取平均。分词器通常会将一个词分成多个部分,导致像“icken”这样在语义上无意义的token。

- 门控残差对路由分数的影响 (图6):
- 结果: 门控残差有效地降低了路由分数的方差,而没有改变其均值和取值范围。
- 结论: 门控残差有助于MoE++中异构专家架构的稳定路由。
图6:门控残差对路由分数影响的可视化。我们展示了“MoE++ 1B/(16+4)E”模型在WinoGrande基准测试上的最高和次高分数。所有模型都以100B tokens的预算进行训练,超参数τ设置为0.75。
图6:门控残差对路由分数影响的可视化。我们展示了“MoE++ 1B/(16+4)E”模型在WinoGrande基准测试上的最高和次高分数。所有模型都以100B tokens的预算进行训练,超参数τ设置为0.75。

A5 结论

本文介绍了MoE++,一个集成了FFN和零计算专家的通用异构MoE框架。与香草MoE方法对所有令牌使用固定的混合机制不同,MoE++通过为简单令牌分配更少的FFN专家来优化计算分配,从而使更多FFN专家能够专注于处理困难的令牌。因此,MoE++实现了比香草MoE更低的计算开销和更好的性能。此外,由于零计算专家无需跨GPU部署,MoE++具有高度的部署友好性。大量的实验结果表明,MoE++不仅持续优于香草MoE方法,而且其专家前向吞吐量达到了同等规模香草MoE模型的约1.1至2.1倍。值得注意的是,MoE++是一个通用框架,可以与任何MoE方法集成,以同时提升模型吞吐量和性能。我们相信MoE++为开发先进高效的MoE相关模型奠定了坚实的基础。

A6 附录

A 附加讨论

A.1 专家架构

MoE专家架构的演进。MoE模型中的专家通常与密集模型中使用的标准前馈网络(FFN)相同。最近,有一些工作致力于改进专家架构。DeepseekMoE(【11, Damai Dai, et al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models. 2024, arXiv preprint arXiv:2401.06066】)和XMoE(【57, Yuanhang Yang, et al. Enhancing efficiency in sparse models with sparser selection. 2024, arXiv preprint arXiv:2403.18926】)将密集模型中的FFN拆分为更小的FFN,从而在减小每个专家大小的同时增加了激活专家的数量。PEER(【18, Xu Owen He. Mixture of a million experts. 2024, arXiv preprint arXiv:2407.04153】)和MH-MoE(【53, Xun Wu, et al. Multi-head mixture-of-experts. 2024, arXiv preprint arXiv:2404.15045】)更进一步,不仅减小了专家的大小,还将输入令牌拆分为更小的单元。尽管这些方法取得了一些进展,但现有MoE模型中的专家结构仍然主要基于FFN,对非FFN或非参数专家的探索甚少。据我们所知,我们是第一个为异构MoE架构提出零计算专家的。

A.2 局限性与未来工作

局限性与未来工作。在本节中,我们阐述了我们工作的局限性,并概述了未来研究的途径。

  • 不同层间的异构MoE++。MoE++在单个MoE层内实现了异构专家。此外,如附录D所示,我们观察到,与中间层相比,浅层和最终层的专家分配模式在不同任务之间变化更为显著。这表明模型主要通过这些层来适应任务。未来的工作可以探索设计跨不同层的异构MoE++配置,以进一步增强模型对广泛任务的适应性。

  • 将MoE++与其他模块结合。当前的MoE++方法是作为Transformer中FFN层的替代品。未来的工作可以探索集成其他模块,例如将注意力层与我们的MoE++方法相结合。

  • 大语言模型的脆弱性。我们工作的重点是构建先进高效的混合专家大语言模型(LLM),因此也继承了LLM常见的脆弱性。

  • 幻觉。LLM中的幻觉仍然是一个重大的未解决挑战。这些虚幻的回答可能导致在开放式对话中出现无根据的断言,解决这个问题可以极大地加速该领域的进展。有关大型LLM常见弱点的更深入分析,请参阅(【5, Tom Brown, et al. Language models are few-shot learners. 2020, NeurIPS】、【40, Jack W Rae, et al. Scaling language models: Methods, analysis & insights from training gopher. 2021, arXiv preprint arXiv:2112.11446】)。
  • 长序列处理。基于Transformer的语言模型在面对比训练时看到的序列长得多的测试序列时,通常难以泛化。这一限制在多轮对话中尤为突出,模型可能会丢失先前的上下文,导致回答错误。
  • 提示敏感性。上下文学习对演示的各个方面(如提示格式)表现出令人不安的敏感性(【61, Zihao Zhao, et al. Calibrate before use: Improving few-shot performance of language models. 2021, ICML】)。值得注意的是,提示格式的变化可能导致完全矛盾的输出。解决这个问题可以显著加速该领域的进展。

  • 更多模态。语言仅是交流的一个方面。视觉和音频信息有助于增强我们对世界的理解(【24, Peng Jin, et al. Chat-univi: Unified visual representation empowers large language models with image and video understanding. 2024, CVPR】、【23, Peng Jin, et al. Video-text as game players: Hierarchical banzhaf interaction for cross-modal representation learning. 2023, CVPR】、【22, Peng Jin, et al. Expectation-maximization contrastive learning for compact video-and-language representations. 2022, NeurIPS】)。未来的工作可以探索替代模态,如视觉和音频输入。多种模态的结合有望拓宽模型可以解决的任务范围,并有可能通过利用这些不同模态之间的协同作用来增强其性能(【24, Peng Jin, et al. Chat-univi: Unified visual representation empowers large language models with image and video understanding. 2024, CVPR】)。

  • 更多参数。由于计算限制,我们实验中MoE++模型的最大参数数量限制在7B。然而,我们的MoE++方法具有很高的泛化性,可以在未来的研究中扩展到更大的模型。

B 实现细节

B.1 数据细节

数据详情。与先前的工作一致,我们使用LLaMA2的分词器,其中包含65,536个词汇令牌。值得注意的是,MoE++完全在公共数据集上训练,使其适用于学术研究环境。具体来说,我们根据不同的采样概率从以下数据集中采样:
- The RedPajama(【9, Together Computer. Redpajama: an open dataset for training large language models, 2023a, URL https://github.com/togethercomputer/RedPajama-Data】)包括来自七个领域的训练数据:CommonCrawl、C4、Github、Wikipedia、Books、ArXiv和StackExchange。
- The Dolma(【48, Luca Soldaini, et al. Dolma: An Open Corpus of Three Trillion Tokens for Language Model Pretraining Research. 2024, arXiv preprint】),一个大型且多样化的开放英语文本语料库,包含从七个来源采样的3万亿个令牌,包括来自Common Crawl的网页、来自The Stack的代码、来自C4的精选网络数据(【41, Colin Raffel, et al. Exploring the limits of transfer learning with a unified text-to-text transformer. 2020, Journal of machine learning research】)、来自Reddit的社交媒体对话、来自PeS2o的学术论文、来自Project Gutenberg的公共领域书籍,以及来自Wikipedia和Wikibooks的综合内容。
- The Pile(【14, Leo Gao, et al. The Pile: An 800gb dataset of diverse text for language modeling. 2020, arXiv preprint arXiv:2101.00027】),一个用于训练大型语言模型的开源英语文本语料库,包括22个多样化的公开可用数据集,如Wikipedia、NIH ExPorter、ArXiv、Books3、BookCorpus2、OpenSubtitles、YoutubeSubtitles和Enron Emails。

采样策略。表A显示了不同开源数据集的详细采样比例。我们发现在训练后期增加高质量数据(如Books和Wikipedia)的比例可以显著提升模型性能。因此,对于表4中的“MoE++ 7B/(16+4)E”模型,我们在最后的100B令牌中增加了高质量数据的比例。具体来说,该模型在前900B令牌使用策略1进行训练,在最后100B令牌使用策略2,总训练预算为1T令牌。相比之下,为简单起见,表3中的所有MoE++和MoE模型都使用策略1进行训练,预算为100B令牌。
表A:不同开源数据集的采样比例。表3中的所有MoE++和MoE模型都使用策略1进行训练,预算为100B tokens。相比之下,对于表4中的“MoE++ 7B/(16+4)E”模型,前900B tokens应用策略1,最后100B tokens应用策略2,总训练预算为1T tokens。

表A
表A

B.2 训练超参数

训练超参数。表B显示了详细的训练超参数。具体来说,MoE++的超参数是根据密集Transformer语言模型的常规实践选择的。我们用MoE++层替换了Transformer中所有的FFN层,并为每一层设置Top-K为2,这导致计算量大约是密集模型的两倍。异构负载均衡损失的权重β设置为0.01,专家容量因子γ设置为1.1。MoE++使用AdamW优化器(【35, Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. 2017, arXiv preprint arXiv:1711.05101】)进行训练。训练期间,应用了0.1的权重衰减和1.0的梯度裁剪。表3中的所有MoE++(除了具有8路流水线并行的“MoE++ 7B/(16+4)E”)和MoE模型都使用策略1进行训练,最大学习率为5e-4,批量大小为400万个令牌,序列长度为2048。相比之下,对于表4中的“MoE++ 7B/(16+4)E”模型,前900B令牌应用策略2,最后100B令牌应用策略3,总训练预算为1T令牌。
表B:训练超参数。表3中的所有MoE++(除了具有8路流水线并行的“MoE++ 7B/(16+4)E”)和MoE模型都使用策略1进行训练,预算为100B tokens。相比之下,对于表4中的“MoE++ 7B/(16+4)E”模型,前900B tokens应用策略2,最后100B tokens应用策略3,总训练预算为1T tokens。

表B
表B

C 定量评估细节

评估框架。我们对MoE++与香草MoE和密集模型进行了比较。评估在多个关键基准上进行,使用了Eleuther AI Language Model Evaluation Harness¶(【15, Leo Gao, et al. A framework for few-shot language model evaluation, 2024, URL https://zenodo.org/records/12608602】),这是一个用于在广泛任务上测试生成式语言模型的统一框架。用于评估的基准包括:

  • ARC(【8, Peter Clark, et al. Think you have solved question answering? try arc, the ai2 reasoning challenge. 2018, arXiv preprint arXiv:1803.05457】)是一个多项选择问答资源,包含从3到9年级科学考试中提取的问题。它分为两个部分:Easy和Challenge,后者包含需要推理的更难问题。我们报告了ARC Easy(ARC-E)的0-shot准确率和ARC Challenge(ARC-C(25))的25-shot准确率。
  • LAMBADA(【38, Denis Paperno, et al. The lambada dataset: Word prediction requiring a broad discourse context. 2016, ACL】)是一个开放式填空任务,要求预测每个段落最后一句话中缺失的目标词。我们报告了LAMBADA的0-shot准确率。
  • LogiQA(【34, Jian Liu, et al. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. 2020, IJCAI】)包含8,678个问答实例,涵盖各种类型的演绎推理。我们报告了LogiQA的0-shot准确率。
  • PIQA(【4, Yonatan Bisk, et al. Piqa: Reasoning about physical commonsense in natural language. 2020, AAAI】)是一个为常识推理设计的数据集,旨在评估当前模型的物理知识。我们报告了PIQA的0-shot准确率。
  • SciQ(【52, Johannes Welbl, et al. Crowdsourcing multiple choice science questions. 2017, arXiv preprint arXiv:1707.06209】)包括13,679个众包的科学考试问题,涵盖物理、化学和生物等科目。我们报告了SciQ的0-shot准确率。
  • WinoGrande(【44, Keisuke Sakaguchi, et al. Winogrande: An adversarial winograd schema challenge at scale. 2021, Communications of the ACM】)是一个大规模数据集,包含44,000个问题,旨在增加规模和难度。我们报告了WinoGrande的0-shot准确率。
  • HellaSwag(【58, Rowan Zellers, et al. Hellaswag: Can a machine really finish your sentence? 2019, ACL】)是一个具有挑战性的数据集,旨在评估常识性自然语言推断(NLI)。我们报告了10-shot HellaSwag(HellaSwag (10))的准确率。
  • MMLU(【19, Dan Hendrycks, et al. Measuring massive multitask language understanding. 2021, ICLR】)是一个旨在评估模型在预训练期间获得的知识的基准,涵盖57个学科。我们报告了5-shot MMLU(MMLU (5))的准确率。
  • Natural Questions (NQ)(【26, Tom Kwiatkowski, et al. Natural questions: a benchmark for question answering research. 2019, Transactions of the Association for Computational Linguistics】)是一个基于真实、匿名的谷歌查询的问答数据集。我们报告了32-shot Natural Questions(NQ (32))的精确匹配分数,以衡量模型中的事实知识。
  • BoolQ(【7, Christopher Clark, et al. Boolq: Exploring the surprising difficulty of natural yes/no questions. 2019, NAACL】)是一个问答数据集,包含15,942个是/否问题。我们报告了32-shot BoolQ(BoolQ (32))的准确率。

D 附加定性分析

跨层专家负载分布分析。为了探索MoE++模型中所有层在不同任务上的专家负载分布,我们在图A、图B、图C、图D和图E中提供了任务级别的专家负载分布可视化。这些可视化揭示了几个有趣的发现:

  • 我们观察到不同层之间,特别是相邻层之间的专家负载存在相关性。例如,当第j层激活了大部分FFN专家时,第j+1层也很有可能以类似的大比例激活FFN专家。
  • 我们发现,与中间层相比,浅层和最终层的专家分配模式在不同任务间变化更为显著。这表明模型主要通过其浅层和最终层来适应不同的任务。未来的工作可以集中于在这些层设计更复杂的结构,以增强模型对多样化任务的适应性。
  • 不同任务中每个令牌激活的FFN专家数量存在显著差异,但这并不一定是更简单的任务激活更少的FFN专家。例如,ARC Challenge任务通常比ARC Easy激活更多的FFN专家。这些结果表明,MoE++模型是根据知识内容和令牌级别的复杂性来分配专家,而不是根据整体任务难度。
  • 在所有专家类型中,零专家的平均激活次数最高,其中更简单的任务显示出更高的平均激活次数。例如,ARC Easy任务比ARC Challenge任务激活更多的零专家。这表明零专家的激活水平可以作为模型对任务难度的指示器。
  • 我们还观察到,MoE++模型中所有层的专家分配在不同任务主题之间差异显著,这表明MoE++模型通过采用不同的专家分配模式来处理不同主题的任务。

图A:任务级别的专家负载分布可视化。结果来自“MoE++ 7B/(16+4)E”模型的第1层到第5层,超参数τ设置为0.75。
图A:任务级别的专家负载分布可视化。结果来自“MoE++ 7B/(16+4)E”模型的第1层到第5层,超参数τ设置为0.75。

图B:任务级别的专家负载分布可视化。结果来自“MoE++ 7B/(16+4)E”模型的第6层到第10层,超参数τ设置为0.75。
图B:任务级别的专家负载分布可视化。结果来自“MoE++ 7B/(16+4)E”模型的第6层到第10层,超参数τ设置为0.75。

图C:任务级别的专家负载分布可视化。结果来自“MoE++ 7B/(16+4)E”模型的第11层到第15层,超参数τ设置为0.75。
图C:任务级别的专家负载分布可视化。结果来自“MoE++ 7B/(16+4)E”模型的第11层到第15层,超参数τ设置为0.75。

图D:任务级别的专家负载分布可视化。结果来自“MoE++ 7B/(16+4)E”模型的第16层到第20层,超参数τ设置为0.75。
图D:任务级别的专家负载分布可视化。结果来自“MoE++ 7B/(16+4)E”模型的第16层到第20层,超参数τ设置为0.75。

图E:任务级别的专家负载分布可视化。结果来自“MoE++ 7B/(16+4)E”模型的第21层到第24层,超参数τ设置为0.75。
图E:任务级别的专家负载分布可视化。结果来自“MoE++ 7B/(16+4)E”模型的第21层到第24层,超参数τ设置为0.75。