OLMoE : 开放式专家混合语言模型

摘要

我们介绍了 OLMoE，一个完全开放的、最先进的语言模型，它利用稀疏专家混合 (MoE)。 OLMoE-1B-7B 有 70 亿 (B) 个参数，但每个输入符元只使用 1B 个参数。 我们在 5 万亿个符元上对其进行了预训练，并进一步调整它以创建 OLMoE-1B-7B-Instruct。 我们的模型在所有具有类似活跃参数的可用模型中表现最佳，甚至超过了 Llama2-13B-Chat 和 DeepSeekMoE-16B 等更大的模型。 我们展示了关于 MoE 训练的各种实验，分析了模型中的路由，显示了高度专业化，并将我们工作的所有方面开源：模型权重、训练数据、代码和日志。

	Weights	https://hf.co/allenai/OLMoE-1B-7B-0924
	Data	{https://hf.co/datasets/allenai/OLMoE-mix-0924}
	Code	https://github.com/allenai/OLMoE
	Logs	https://wandb.ai/ai2-llm/olmoe/reports/
		OLMoE-1B-7B-0924--Vmlldzo4OTcyMjU3

图 1: 开放式 MoE 和密集型 LMs 的性能、成本和开放程度。 模型名称包含四舍五入的参数计数：MoE 的 model-active-total 和密集型 LMs 的 model-total。 #ckpts 是可用中间检查点的数量。 我们突出显示 MMLU 作为整体性能的总结；有关更多结果，请参见 §3。 OLMoE-1B-7B 在具有类似活跃参数计数的模型中表现最佳，并且是最开放的 MoE。

1 引言

尽管大型语言模型 (LM) 在各种任务上取得了重大进展，但在训练和推理中，性能和成本之间仍然存在明显的权衡。 由于构建和部署高性能 LM 成本高昂，因此许多学者和开源开发人员无法获得这些模型。 改善成本性能权衡的一种方法是使用稀疏激活的专家混合 (MoE) [152]。 MoE 在每一层都有多个专家，每次只激活其中一个子集（参见 Figure 2）。 这使得 MoE 比具有类似参数总数的密集模型更有效率，后者对每个输入都激活所有参数 [204]。 因此，工业界领先的模型使用 MoE，包括 Gemini-1.5 [173] 和据称的 GPT-4 [29]。

但是，大多数 MoE 模型是闭源的：虽然一些模型公开了模型权重 [43, 78, 156, 176, 178]，但它们对训练数据、代码或配方提供的信息有限甚至没有（参见图 1）。 虽然此前已有一些努力使语言建模研究完全公开 [18, 64, 88, 102, 192, 208]，但它们在很大程度上局限于密集 LM。 尽管 MoE 需要 更多 的开放性，因为它们为 LM 添加了复杂的新设计问题，例如使用多少总参数和活动参数，是否使用许多小的专家或少数大的专家，专家是否应该共享，以及使用什么路由算法。 缺乏关于这些细节的公开资源和发现，阻碍了该领域构建接近闭源领先模型能力的成本效益高的开放 MoE。

为了解决这些问题，我们引入了OLMoE，这是一种完全开放的专家混合语言模型，在同等规模的模型中表现出最先进的性能。 具体来说，我们对OLMoE-1B-7B进行了预训练，使用 5.1 万亿个符元，总参数量为 69 亿，其中每个输入符元仅激活 13 亿个参数。 这导致了与使用具有约 10 亿个参数的密集模型（例如 OLMo 1B[64] 或 TinyLlama 1B[209]）类似的推理成本，但需要更多的 GPU 内存来存储其 70 亿个总参数。 我们的实验表明，MoE 的训练速度比具有相同激活参数的密集 LLM$\sim$2$\times$快。 在Figure 1中，我们展示了OLMoE-1B-7B显着优于所有开放的 10 亿模型，并与具有显着更高推理成本和内存存储的密集模型表现出竞争力（例如，与 Llama2-13B 相似的 MMLU 分数，后者是$\sim$10$\times$倍的成本）。 通过指令和偏好调优，我们创建了OLMoE-1B-7B-Instruct，我们发现它在常见的基准测试（MMLU、GSM8k、HumanEval 等）上超过了各种较大的指令模型，包括 Llama2-13B-Chat[181]、OLMo-7B-Instruct（0724）和 DeepSeekMoE-16B[42]。

我们全面的控制实验集突出了 MoE（见Table 1）和 LLM 的关键设计选择。 使 MoE 具有高性能的一个关键设计决策是使用细粒度路由和粒度专家[42]：我们在每一层中使用 64 个小型专家，其中 8 个被激活。 路由算法的选择也很重要：我们发现无丢弃[58] 基于符元的路由[152] 优于基于专家的路由[218]。 我们的发现还包括挑战先前工作的内容，例如共享专家的无效性[42] 以及将预训练的密集 LLM 稀疏升级到 MoE[84] 的有限益处，除非在小计算预算下。 最后，我们分析了OLMoE-1B-7B 中的路由行为，发现路由在预训练的早期阶段就饱和了，专家很少被共同激活，并且专家表现出领域和词汇专业化。

我们希望我们完全开放的 MoE 有助于更多研究和分析，以提高我们对这些模型的理解。 我们发布了训练代码、中间检查点（每 5000 步）、训练日志和训练数据，并使用开源许可证（Apache 2.0http://www.apache.org/licenses/LICENSE-2.0 或 ODC-By 1.0https://opendatacommons.org/licenses/by/1-0/）。

2 预训练和自适应

图 2: 密集 LLM 和 MoE 模型（如OLMoE）的架构比较。 此图省略了一些细节，例如，OLMoE-1B-7B 也使用 QK-Norm (§4.2.5)。

Design choice	Description	Exper-iment	OLMoE-1B-7B
Active params	# active parameters per input token	§4.1.1	1.3B active
Total params	Total # of parameters in the model	§4.1.1	6.9B total
Expert granularity	Using fine-grained small experts vs. a few large experts [39]	§4.1.2	64 small experts with 8 activated
Expert sharing	Whether or not to include a shared expert [39]	§4.1.3	No shared expert
Routing algorithm	How inputs are assigned to experts, e.g., assignment on a per token basis (e.g., 2 experts per token) or per expert basis (e.g., 2 tokens per expert), and whether or not all tokens get assigned or some get dropped [58, 218]	§4.1.4	Dropless [58] MoE with token choice
Sparse upcycling	Whether to start from a dense model [84, 210]	§4.1.5	Not used
Load balancing loss	Auxiliary loss to penalize unequal assignment to experts that may harm performance [152]	§4.1.6	Used with weight 0.01
Router z-loss	Auxiliary loss to penalize large logits in the router that may cause instabilities [220]	§4.1.7	Used with weight 0.001

表 1： 基于我们的实验，OLMoE-1B-7B 的主要 MoE 设计选择和我们的设置。 OLMoE-1B-7B 的完整配置见 Appendix B。

Source	Doc Type	GPT-NeoX	Words (billions)	UTF-8	Documents     (millions)
		tokens		bytes
		(billions)		(GB)
DCLM-Baseline [89]	web pages	3,860	3,380	16,700	2,950
StarCoder [91, 83]	code	101	63.9	325	78.7
peS2o [162, 161]	STEM papers	57.2	51.3	268	38.8
arXiv [36]	STEM papers	21.1	23.5	88.8	1.55
OpenWebMath [129]	math web pages	12.7	10.2	42.4	2.91
Algebraic Stack [11]	math proofs code	12.6	9.6	39.3	2.83
English Wikipedia & Wikibooks [161]	encyclopedic	3.69	3.16	16.2	6.17
Total		4,060	3,530	17,400	3,080

表 2： 用于 OLMoE-1B-7B 的预训练数据组成。 StarCoder、peS2o 和 Wikipedia 部分来自 Dolma 1.7 [161]。 我们数据的链接见 Appendix A。

预训练架构

OLMoE 是一种仅解码器的 LM，由 $N_L$ Transformer [183] 层组成。 像 OLMo [64] 这样的密集模型中的前馈网络 (FFN) 被一个 MoE 模块取代，该模块包含 $N_E$ 个称为专家的较小 FFN 模块，其中 $k$ 个专家的子集被激活，用于每个处理的输入符元 $x$（另请参见 Figure 2）：

$$\text{MoE module}(x)=\sum _{i=1}^{\text{Top}-k(r(x))}\mathrm{softmax}{\left(r(x)\right)}_i{E}_i(x)$$ (1)

其中 $r$，称为路由器，是一个从输入 logits 映射到所选 $k$ 专家的学习线性层。 对路由器输出应用 softmax 来计算所有 $N_E$ 专家的路由概率。 每个选定的专家 $E_i$ 处理输入 $x$，其输出然后乘以其相应的路由概率。 然后将结果在所有选定的 Top-$k$ 专家中进行汇总，以构成 MoE 模块的输出，该模块是模型的 $N_L$ 个总层中的一个单层。 在设计 MoE 模型时，关键决策包括确定激活和总参数的数量、专家的设计（例如，粒度，是否包含共享专家）以及路由算法的选择。 此外，训练 MoE 模型可能涉及从密集模型初始化（稀疏升级）并改变训练目标，例如包括辅助负载均衡和路由器 z 损失。 与这些设计选择相关的实验在 §4.1 中；Table 1 显示了我们最终的决定。

总之，我们从总共 69 亿个参数中使用 13 亿个活动参数，每层激活 8 个专家，共 64 个专家。 我们使用无滴标记选择路由 [58]：对于每个输入标记，学习的路由器网络确定 8 个专家来处理它。 我们从头开始训练 OLMoE-1B-7B，使用两个辅助损失：负载均衡损失 (${ℒ}_{LB}$) [152] 和路由器 z 损失 (${ℒ}_{RZ}$) [220]，我们在 §4.1.6 和 §4.1.7 中分别定义和实验了它们。 我们将它们与各自的损失权重相乘，$\alpha$ 和 $\beta$，并将它们与交叉熵损失 (${ℒ}_{\text{CE}}$) 线性相加，得到我们的最终训练损失：

$$ℒ={ℒ}_{\text{CE}}+\alpha {ℒ}_{\text{LB}}+\beta {ℒ}_{\text{RZ}}$$ (2)

我们 OLMoE-1B-7B 的完整预训练配置在 Appendix B 中。

Source	Domain	Samples
Instruction Tuning
Tulu 2 SFT Mix [75]	Various	326,154
No Robots [138]	Various	9,500
CodeFeedback-Filtered-Instruction [213]	Coding	156,526
MetaMathQA [203]	Math	98,750
Advanced (non-chat) subset of Daring Anteater [187]	Various	17,082
Preference Tuning (DPO [136])
UltraFeedback [38] binarized and filtered for TruthfulQA [98] contamination	Various	60,800

表 3： OLMoE-1B-7B 的适应训练数据。 我们数据的链接在 Appendix A 中。

预训练数据

我们使用了来自 DCLM [89] 和 Dolma 1.7 [161] 的混合数据，其中包括：（1）经过质量过滤的 Common Crawl 子集，称为 DCLM-Baseline，（2）StarCoder、Algebraic Stack 和 arXiv，在 DCLM 和 Dolma 1.7 中都有使用，以及（3）来自 Dolma 1.7 的 peS2o 和 Wikipedia。 我们将我们的预训练数据集称为 OLMoE-Mix。

对上述所有来源，我们应用了一个过滤器，该过滤器删除所有包含 32 个或更多重复 n 元组的文档，其中 n 元组是 1 到 13 个标记的任何跨度。 对于 StarCoder 子集，我们还删除任何来自 GitHub 上星级少于 2 个的存储库的文档，或者其最常出现的词占文档的 30% 以上，或者其前 2 个最常出现的词占文档的 50% 以上的文档。

我们在每个 epoch 开始时随机对所有样本进行洗牌，并总共训练 5.133T 个标记（遵循 Muennighoff 等人 [120] 的 1.3 个 epoch）。 在我们的退火阶段（最终 100B Token ），我们首先重新整理整个数据集，然后将学习率线性衰减到 0，遵循之前的工作 [64, 89]。 我们的预训练数据统计信息在 Table 2 中。

适应

我们通过遵循标准适应方案创建 OLMoE-1B-7B-Instruct，该方案分为 指令调优 [117, 189, 147, 154, 205]，然后是 偏好调优 [31, 15, 136, 54]，建立在先前开放模型 [182, 75, 186] 的基础上。 在我们的指令调优数据集中，我们添加了更多代码和数学数据以提高下游编码和数学应用程序的性能。 其他模型，例如 GPT-4 [126] 和 Llama 3 [50]，在预训练期间包括来自 GSM8k [35] 或 MATH [70] 等数学数据集的样本。 我们还包括 No Robots 和 Daring Anteater 的一个子集，因为它们质量很高并且增加了多样性，这是成功适应的两个关键因素 [186, 215, 103, 119]。 我们在 Table 3 中描述了我们的适应数据集，并在 Appendix B 中描述了超参数。

3 结果

我们的评估程序包括三个部分：预训练期间、预训练后 和 适应后。 我们在 Appendix C 中详细说明了每个部分的设置。

图 3： 评估 OLMoE-1B-7B 和当前最好的 OLMo 模型在预训练期间的表现。 OLMoE-1B-7B 与 OLMo 模型不同，它的 MoE 架构、几个训练超参数和训练数据集不同，参见 §2。 Appendix E 中提供了该图的一个版本，该图的 x 轴为标记，退火开始处为标记。更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-OLMoE-1B-7B-vs-OLMo-7B-vs-OLMo-1B--Vmlldzo4OTcyMjEz

预训练期间

在 Figure 3 中，我们对 OLMoE-1B-7B 在预训练期间的性能进行了基准测试，并将其与当前最佳的 OLMo 模型 [64] 在常用的下游任务上进行了比较。 我们发现，在所有任务中，OLMoE-1B-7B 在计算量 (FLOPs) 较少的情况下，其性能优于密集的 OLMo 模型。 尽管 OLMoE-1B-7B 在训练中使用的 FLOP 数不到一半且仅使用 1B，但在训练结束时，OLMoE-1B-7B 与 OLMo-7B 匹配或优于 OLMo-7B活动参数。 这可能是由于我们在 OLMo 设置中对数据集和建模所做的更改，包括与 MoE 相关的更改、稳定性和性能改进，这些更改在 Appendix B 中进行了概述。 Appendix E 包含训练和验证损失图，这些图显示了非常平滑的损失曲线，在我们的预训练 5T 个符元中没有出现主要的损失峰值。

	Active	Open	MMLU	Hella-	ARC-	ARC-	PIQA	Wino-
	params	Data		Swag	Chall.	Easy		Grande
LMs with $\sim$7-9B active parameters
Llama2-7B [181]	6.7B		46.2	78.9	54.2	84.0	77.5	71.7
OLMo-7B (0724) [64]	6.9B		54.9	80.5	68.0	85.7	79.3	73.2
Mistral-7B [77]	7.3B		64.0	83.0	78.6	90.8	82.8	77.9
DCLM-7B [89]	6.9B		64.4	82.3	79.8	92.3	80.1	77.3
Llama3.1-8B [50]	8.0B		66.9	81.6	79.5	91.7	81.1	76.6
Gemma2-9B [175]	9.2B		70.6	87.3	89.5	95.5	86.1	78.8
LMs with $\sim$2-3B active parameters
OpenMoE-3B-9B [198]	2.9B		27.4	44.4	29.3	50.6	63.3	51.9
StableLM-2B [16]	1.6B		40.4	70.3	50.6	75.3	75.6	65.8
DeepSeek-3B-16B [39]	2.9B		45.5	80.4	53.4	82.7	80.1	73.2
JetMoE-2B-9B [156]	2.2B		49.1	81.7	61.4	81.9	80.3	70.7
Gemma2-3B [175]	2.6B		53.3	74.6	67.5	84.3	78.5	71.8
Qwen1.5-3B-14B [178]	2.7B		62.4	80.0	77.4	91.6	81.0	72.3
LMs with $\sim$1B active parameters
Pythia-1B [18]	1.1B		31.1	48.0	31.4	63.4	68.9	52.7
OLMo-1B (0724) [64]	1.3B		32.1	67.5	36.4	53.5	74.0	62.9
TinyLlama-1B [209]	1.1B		33.6	60.8	38.1	69.5	71.7	60.1
DCLM-1B [89]	1.4B		48.5	75.1	57.6	79.5	76.6	68.1
OLMoE-1B-7B	1.3B		54.1	80.0	62.1	84.2	79.8	70.2

表 4： OLMoE-1B-7B 预训练后与较大的 MoE 和密集的 LMs 相比。 我们将其与密集的 LMs 进行比较，这些 LMs 在活动参数（1B，近似速度和成本）或总参数（7B，近似内存需求）方面与 OLMoE-1B-7B 相似。 模型名称包含四舍五入的参数计数：模型-活动-总计 用于 MoE，模型-总计 用于密集的 LMs（这导致与官方名称存在一些差异，例如，虽然被称为“Gemma2-2B”，但它实际上有 2.6B 个活动和总参数 [175]）。 Chall.   = Challenge. 我们自己使用 5 个少样本进行所有评估，有关详细信息，请参阅 Appendix C。

				Human-	Alpaca-
Task ($\to$)	MMLU	GSM8k	BBH	Eval	Eval 1.0	XSTest	IFEval	Avg
Setup ($\to$)	0-shot	8-shot CoT	3-shot	0-shot	0-shot	0-shot	0-shot
Metric ($\to$)	EM	EM	EM	Pass@10	%win	F1	Loose Acc
OLMo-1B (0724)	25.0	7.0	22.5	16.0	-	67.6	20.5	-
+SFT	36.0	12.5	27.2	21.2	41.5	81.9	26.1	35.9
+DPO	36.7	12.5	30.6	22.0	50.9	79.8	24.2	37.4
OLMo-7B (0724)	50.8	32.5	36.9	32.3	-	80.8	19.6	-
+SFT	54.2	25.0	35.7	38.5	70.9	86.1	39.7	49.3
+DPO	52.8	9.0	16.6	35.0	83.5	87.5	37.9	49.1
JetMoE-2B-9B	45.6	43.0	37.2	54.6	-	68.2	20.0	-
+SFT	46.1	53.5	35.6	64.8	69.3	55.6	30.5	50.4
DeepSeek-3B-16B	37.7	18.5	39.4	48.3	-	65.9	13.5	-
+Chat	48.5	46.5	40.8	70.1	74.8	85.6	32.3	57.0
Qwen1.5-3B-14B	60.4	13.5	27.2	60.2	-	73.4	20.9	-
+Chat	58.9	55.5	21.3	59.7	83.9	85.6	36.2	57.3
OLMoE-1B-7B	49.8	3.0	33.6	22.4	-	59.7	16.6	-
+SFT	51.4	40.5	38.0	51.6	69.2	84.1	43.3	54.0
+DPO	51.9	45.5	37.0	54.8	84.0	82.6	48.1	57.7

表 5： OLMoE-1B-7B 适应后与其他模型相比。 我们发现 JetMoE 聊天模型 (https://hf.co/jetmoe/jetmoe-8b-chat) 存在问题，导致得分随机，因此我们将其排除在外。 名称包含四舍五入后的参数数量：model-active-total 用于 MoE，model-total 用于稠密语言模型。 我们自己运行所有评估 (Appendix C)。 模型使用不同的混合方式进行适应，例如，OLMoE 在用于 OLMo 模型的管道改进版本上进行训练。

预训练后

在 Table 4 中，我们在常见的下游任务上对 OLMoE-1B-7B 进行了基准测试。 我们发现 OLMoE-1B-7B 在使用小于 20 亿个活动参数的模型中表现最佳，使其成为许多语言模型用例中最经济的选择。 对于更大的预算，Qwen1.5-3B-14B 具有更强的性能，但其活动参数和总参数是 OLMoE-1B-7B 的两倍多。 我们发现，尽管 OLMoE-1B-7B 每个前向传递所需的计算量减少了 $\sim$6-7$\times$ 倍，但它仍然优于一些具有 70 亿个参数的稠密语言模型，例如 Llama2-7B [181]，但不及其他模型，例如 Llama3.1-8B [50]。 Figure 1 比较了 OLMoE-1B-7B 和其他语言模型的 MMLU 性能与活动参数，活动参数是模型价值与其成本的代理指标。 OLMoE-1B-7B 是其成本范围内的最先进技术。

适应后

在 Table 5 中，我们对 OLMoE-1B-7B 的指令（SFT）和偏好（DPO）微调进行了基准测试。 SFT 在所有衡量任务中都改善了我们的模型。 我们在 GSM8k 上观察到 $>$10$\times$ 的提升，这可能是由于我们包含了额外的数学数据来弥补预训练期间数学数据的相对较少 (§2)。 DPO 在大多数任务上都有帮助，特别是 AlpacaEval，这与之前工作的发现一致 [186, 75, 121]。 我们将 DPO 模型称为 OLMoE-1B-7B-Instruct，它在所有基准测试模型中具有最高的平均值。 我们发现它在性能上超过了 Qwen1.5-3B-14B 的聊天版本，尽管 Qwen 具有 $>$2$\times$ 个更多的参数，并且它的预训练模型在 Table 4 中的表现优于 OLMoE-1B-7B。 AlpacaEval 上 84% 的得分也优于排行榜上更大的密集模型， 例如 Llama2-13B-Chat [181]。

4 尝试使用替代设计选择

在本节中，我们介绍了导致 OLMoE-1B-7B 的预训练和适应实验。 我们将它们分组到针对专家混合模型的特定设置的实验中 (§4.1)，适用于密集 LM 和 MoE 的设置的实验中 (§4.2)，以及适应实验中 (§4.3)。 在预训练实验中，我们经常使用 MMLU Var，这是 MMLU [69] 的一个版本，它具有不同的少样本和不同的格式，在训练期间可以更早地提供信号。 我们在 Appendix C 中描述了我们的完整评估设置，并在 Appendix F 中提供了额外的实验。 每个实验都链接到一个 Weights & Biases 报告，其中包含更多验证和下游结果，以及运行的完整配置。 为了隔离更改的影响并最大程度地减少混杂因素，我们每个实验只改变一个超参数。 然而，由于超参数数量众多，一些结果可能会在不同的配置下发生变化，我们无法保证每个超参数选择的正确性。 由于我们改变了基本模型以纳入成功的发现，因此模型在不同的实验中不可比较。

4.1 MoE 特定的预训练设置

4.1.1 专家混合模型与   密集

图 4： MoE 与   稠密。 我们训练了一个具有 13 亿参数的稠密模型和一个具有 13 亿个活动参数、69 亿个总参数的 MoE 模型，每个模型都在 128 个 H100 GPU 上训练。 除了与 MoE 相关的更改之外，我们使用相同的配置对这两个模型进行了 1300 亿个符元的训练。 MoE 包含 64 个专家，其中 8 个被激活，FFN 维度为 1024，而稠密模型的 FFN 维度为 8192。 因此，两者具有相同数量的活动参数。 顶部： MoE 在 $\sim$3$\times$ 个较少的符元（或 FLOP，因为两者具有相同的活动参数，忽略了微不足道的路由器参数）的情况下达到了最终的稠密性能。 底部： 由于一些内存开销，这相当于 $\sim$2$\times$ 倍的更快训练速度。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-MoE-vs-Dense--Vmlldzo4OTM0Mjkx

先前的工作报告了 MoE 相对于密集模型的各种加速：Artetxe 等人 [10] 报告 MoE 需要 2–4$\times$ 更少的计算来匹配密集模型，MoMa [99] 在语言任务中表现出 2.6$\times$ FLOP 节省，Arctic [159] 产生 4$\times$ FLOP 节省，但对于非常不同的密集和 MoE 配置，以及开关 Transformer [56] 使用 MoE 训练速度要快 2-7$\times$，但对于编码器-解码器模型，而其他作品研究仅解码器 LM [ 135]。

在 Figure 4 中，我们在受控设置中比较了 MoE 和稠密模型。 我们发现，我们的 MoE 在 $\sim$3$\times$ 个较少的符元情况下达到了稠密模型的性能，相当于 $\sim$3$\times$ 倍的更少计算量，以 FLOP 衡量。 但是，由于使用其 70 亿个总参数训练 MoE 所带来的额外内存开销，它每秒处理的符元数量少于稠密模型（MoE 每个 GPU 每秒 23600 个符元，而稠密模型为 37500 个）。 因此，在训练时间方面，它仅 $\sim$2$\times$ 倍地快于达到稠密模型的性能。 可能存在一些优化，可以使加速更接近 3$\times$ Token 加速，我们将其留给未来的工作。 基于这些结果，我们选择了一个具有 69 亿个总参数和 13 亿个活动参数的 MoE 配置，分别与 OLMo-7B 的总参数数量和 OLMo-1B 的活动参数数量相匹配。

4.1.2 专家粒度

图 5: 专家粒度。 我们同时调整专家数量和 FFN 维度，以确保活动参数和总参数以及计算成本保持不变。 例如，对于 64 个专家，FFN 维度为 1,024，激活 8 个专家，而对于 32 个专家，FFN 维度为 2,048，激活 4 个专家。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-Granularity--Vmlldzo4OTIxOTE4

Dai 等人 [39] 提议使用小型细粒度专家，以允许更多专家组合，从而使模型更灵活。 例如，Mixtral 模型 [78] 使用每层 8 个专家、激活 2 个专家的常见配置。 这允许每层 $\left(\genfrac{}{}{0pt}{}{8}{2}\right)=28$ 种组合。 通过将每个专家的大小减半，从而将专家数量增加一倍以保持相同的计算和参数预算，我们可以将可能的组合增加到 $\left(\genfrac{}{}{0pt}{}{16}{4}\right)=1,820$。 Krajewski 等人 [85] 研究了计算最佳粒度配置，发现更高的计算预算需要更细粒度的专家。

在 Figure 5 中，我们观察到更细粒度的专家可以改善训练损失、验证损失和下游性能。 8 个专家配置使用 1 个活动专家，这产生了 $\left(\genfrac{}{}{0pt}{}{8}{1}\right)=8$ 种组合。 通过将每个专家的规模缩小四倍，但将数量增加到 32 个，并激活 4 个 ($\left(\genfrac{}{}{0pt}{}{32}{4}\right)=35,960$ 种组合)，我们观察到在约 1300 亿个符元时，HellaSwag 和 MMLU 上的性能提高了约 10%。 然而，我们发现粒度收益递减。 将专家数量进一步增加到 64 个，并激活 8 个 ($\left(\genfrac{}{}{0pt}{}{64}{8}\right)=4,426,165,368$ 种组合)，下游指标仅提高了 1-2%。 对于我们的 OLMoE-1B-7B 计算预算 的 $3\times {10}^{22}$，Krajewski 等人 [85] 预计最佳专家数量为 256 ($G=32$ 在他们的论文中)。 然而，他们的预测是针对计算最优模型 [71, 32]，而我们针对 5T 个符元进行训练，这远远超出了通常认为对我们模型规模而言最佳的范围。 因此，他们的预测可能不适用于我们的设置，而 我们坚持使用 64 个专家来进行 OLMoE-1B-7B，这也归因于 Figure 5 中的收益递减。

4.1.3 共享专家

图 6： 共享专家。 两种设置都具有相同数量的活动参数和总参数，并使用相同数量的 FLOP。 在 32 个路由专家中，有 4 个被激活，而在另一个模型的 31 个路由专家中，有 3 个被激活，因为它有一个始终处于活动状态的共享专家。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-Expert-sharing--Vmlldzo4OTIyMjQz

Dai 等人 [39] 提出了使用一个共享/固定专家进行训练，该专家始终与路由专家一起使用。 其直觉是鼓励共享专家学习共同信息，并允许其他路由专家学习更专业的知识。 这将减少专家之间的冗余，从而导致更好的模型，因为它可以存储更多总信息。

在 Figure 6 中，我们对具有单个共享专家和单个路由专家的基准测试，与具有两个路由专家进行比较。 虽然两种设置都导致相似的性能，但共享专家表现略差。 共享专家会消除模型的灵活性，因此与 §4.1.2 中的发现相矛盾，该发现表明允许更多专家组合会提高性能。 具体来说，Figure 6 中的两个模型在每层有 $\left(\genfrac{}{}{0pt}{}{32}{4}\right)=35,960$ 和 $\left(\genfrac{}{}{0pt}{}{31}{3}\right)=4,495$ 种可能的组合。 因此，删除其中一个路由专家并将其变成一个共享专家，消除了将近 90% 的可能组合。 这可能起到了与在共享专家中隔离通用知识的潜在好处相反的作用。 基于这些结果，我们没有在 OLMoE-1B-7B 中使用共享专家， 但我们确实认为专家被更频繁地激活，甚至始终被激活的想法是有价值的。 然而，我们认为，与其通过共享专家来强制执行这种行为，不如让模型学习这种行为。 由于需要负载平衡损失 (§4.1.6) 来惩罚模型如果符元没有在专家之间平均分配，这在当前设置下很困难。 未来潜在的工作可以探索删除负载平衡损失，以允许更灵活地使用专家。

4.1.4 专家选择与   符元选择

图 7： 专家选择 (EC) 与符元选择 (TC)。 两个模型在每隔一层都有一个 8 个专家的 MoE。 对于 TC，每个符元激活 2 个专家，而对于 EC，容量因子为 2。 因此，两个模型都使用相同数量的活动参数。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-EC-vs-TC--Vmlldzo4MzkzMDM3

MoE 路由器决定哪个专家处理每个输入符元 (§2)。 存在两种常见的类型 [101]: 专家选择 (EC) [218] 和 符元选择 (TC) [152]。 对于 EC，每个专家从传入的序列中选择固定数量的符元。 按设计，这会导致每个专家处理相同数量的符元。 这是 EC 的主要优势，因为它确保了完美的负载均衡，这提高了训练吞吐量并消除了对负载均衡损失的需求。 EC 的主要缺点是它不容易用于自回归生成，其中在每个步骤中处理单个符元，而不是在一个 [141] 中处理整个序列。 另一个潜在的缺点是 EC 会导致符元丢失，其中某些符元没有被任何专家选中，这会损害性能 [58]。 同时，它会导致某些符元被多个专家处理，这也可能是有益的，因为它允许模型为某些符元分配更多计算 [218]。 对于 TC，每个符元选择固定数量的专家。 这会导致许多符元选择同一个专家，从而损害训练效率。 因此，通常使用 TC 和负载均衡损失 [152] 来鼓励均匀分布。

在 Figure 7 中，我们对 EC 和 TC 进行了基准测试。 我们发现，对于相同的符元预算，TC 在所有显示的任务以及其他任务（如 PIQA、SciQ 等）上都优于 EC，这些任务我们在 https://wandb.ai/ai2-llm/olmoe/reports/Plot-EC-vs-TC--Vmlldzo4MzkzMDM3 中进行了报告。 虽然 Zhou 等人 [218] 发现 EC 更好，但我们的配置略有不同，因为我们使用的是无丢失 MoEs [58] 和负载均衡损失。 因此，我们的 TC 变体预计将比 Zhou 等人 [218] 中的 TC 变体表现更好。 我们证实了 EC 在 29,400 个符元/秒/设备的速度下比 TC 的 24,400 个符元/秒/设备的速度快约 20% [218]。 EC 在多模态设置中可能更有益 [99]，因为丢弃噪声图像符元可能比丢弃文本符元危害更小。 因此，尽管 对于本版本的 OLMoE 我们坚持使用 TC，但我们可能会在未来的多模态模型中重新使用 EC。

4.1.5 稀疏升级

图 8： 稀疏升级。 我们将 OLMo-1B (0724) 在 2T 个符元上升级到一个包含 8 个专家的 MoE，其中激活了 2 个，并对其进行额外的 6100 亿个符元的训练。 我们将其与一个从头开始训练 6100 亿个符元的模型进行比较。 除了这个区别之外，这两个模型都使用相同的配置，其中包含一些导致不稳定的次优设置，例如没有 QK-Norm (§4.2.5) 和没有截断正态初始化 (§4.2.2)。 更多结果、日志和配置： https://wandb.ai/ai2-llm/olmoe/reports/Plot-Scratch-vs-Upcycle--Vmlldzo4NDIyOTc4

Komatsuzaki 等人 [84] 提出通过稀疏升级将密集模型转换为专家混合模型：（1）将密集 MLP 克隆到每个所需的专家以构成 MoE 层。 (2) 在每个 MoE 层之前添加一个新初始化的路由器。 (3) 预训练继续使用新模型，以便克隆的 MLP 可以逐渐专门化于不同的内容，并且可以学习路由器。 他们发现，升级方法在语言模型性能方面优于从头开始训练的语言模型，最高可达稀疏模型升级自的原始密集检查点计算预算的 120%。 例如，如果在 2 万亿个符元上稀疏升级 13 亿个参数的模型，那么只有在 2.4 万亿个符元时，从头开始训练的 MoE 才能赶上升级的模型。 也就是说，稀疏升级的模型将被训练另外 4000 亿个符元，从而节省相当于高达 2T 个符元的计算量。 其他作品，如 MiniCPM [73]、Qwen2 [200] 以及据报道的 Mixtral [25, 78] 已经采用稀疏升级，但只分享了关于其配置的有限信息。

在 Figure 8 中，我们将稀疏升级 OLMo-1B (0724) [64] 与从头开始训练 MoE 进行比较。 我们发现，在 5000 亿个符元之后，从头开始训练的 MoE 在 Figure 8 中的指标和我们位于 https://wandb.ai/ai2-llm/olmoe/reports/Plot-Scratch-vs-Upcycle--Vmlldzo4NDIyOTc4 的额外指标上已经赶上了升级后的模型。 在大约 6000 亿个符元时，从头开始训练的 MoE 开始超越升级后的 MoE。 因此，它只需要原始密集模型计算预算的 25% 就能赶上，而不是 Komatsuzaki 等人 [84] 中报告的 120%。 但是，他们使用了专家选择路由并研究了编码器-解码器模型 [137]。 同时，我们使用符元选择路由 (§4.1.4) 和仅解码器模型 (§2)。 此外，我们升级了一个已经过渡训练 [57] 的模型，即一个为 2T 个符元训练的 1B 模型。 它的参数可能已经在密集模型的最佳范围内，这可能会限制升级后可能的额外探索。 这促使我们尝试在 Appendix F 中概述的升级权重中添加噪声，但我们发现它并没有带来更好的性能。 升级的一个主要缺点是，升级后的 MoE 受密集模型的某些超参数约束。 具体来说，OLMo-1B (0724) 在没有 QK-Norm 和正常初始化的情况下进行了训练，这两种方法在我们的实验中都损害了稳定性 (§4.2.5，§4.2.2)。 虽然可以简单地添加新的 QK-Norm 并从头开始训练它们，就像从头开始训练的新路由层一样，但在升级时不可能更改原始密集模型的初始化。 因此，由于我们想要更改这些超参数，并且还希望将 OLMoE-1B-7B 训练到密集模型计算预算的 250% 左右（5T vs 2T 个符元）， 我们不使用升级。

4.1.6 负载平衡损失

图 9： 应用负载平衡损失 (LBL) 的影响。 两个模型的训练损失图都排除了负载均衡损失。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-LBL-vs-No-LBL--Vmlldzo4OTkyNDg4

图 10： 在使用或不使用第一个 MoE 层的负载均衡损失的情况下，训练期间的专家分配。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-LBL-vs-No-LBL--Vmlldzo4OTkyNDg4

Shazeer 等人 [152] 提出了负载均衡损失，如果模型不平衡，即如果模型将所有符元路由到少数专家，则会对模型进行惩罚。 这是基于这样的观察：没有这种惩罚，模型往往只更新每一层中选定的少数专家 [52, 17]。 为了计算负载均衡损失 (${ℒ}_{\text{LB}}$)，我们将路由到一个专家 $E_i$ 的符元分数 $f_i$ 乘以分配给 $E_i$ 的总路由概率 $P_i$，并将其对一批中的专家数量 $N_E$ 求和：

$${ℒ}_{\text{LB}}=N_E\cdot \sum _{i=1}^{N_E}{f}_i\cdot P_i$$ (3)

损失进一步按 $N_E$ 和损失权重 $\alpha$（参见 Equation 2）进行缩放，这是一个可选的权重，用于确定损失的大小，通常设置为 0.01 [220, 198]。 我们没有尝试改变 0.01 的权重。

在 Figure 9 中，我们研究了使用辅助负载均衡损失对性能的影响。 我们发现，在训练损失和验证损失方面，使用负载均衡损失即使在只有几十亿个符元之后也能带来更好的性能。 我们仍然测量负载均衡损失，即使它没有被使用（“没有 LBL”），我们发现，虽然它最初会激增，但它在接下来的几十亿个符元中会慢慢下降。 这种行为在 Figure 10（左侧）中也可见，其中最初第一层中的所有符元都分配给第 6 个专家（粉色）。 最终，模型也开始将一些符元分配给第一个专家（黄色）。 但是，所有其他专家基本上保持不变，因此是“死重”，占用 GPU 内存但不使用。 鉴于这些结果，我们使用辅助负载平衡损失，权重为 0.01，遵循先前的工作 [152, 156]。 然而，消除负载平衡损失是未来研究的一个重要方向，因为它通过强制模型近似地平等地使用所有专家来限制模型的灵活性。 这可能会阻止专家专门研究某些数据域，并且可能是先前工作未能找到专家专门化强有力证据的原因 [78, 220]。

4.1.7 路由器 Z 损失

图 11： 路由器 Z 损失。 我们比较添加权重为 0.001 的路由器 Z 损失与不添加额外 Z 损失的情况。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-Zloss-vs-none--Vmlldzo4NDM4NjUz

Zoph 等人 [220] 提出了路由器 Z 损失，以提高 MoE 模型的稳定性和质量。 这种辅助损失会对进入门控网络的大 logits 进行惩罚。 这种大型 logits 会导致 MoE 层中发生的矩阵乘法出现数值溢出。 它通过对路由层之前的 logits 进行指数运算 $x_j$，然后在专家数量上求和 $N_E$，并在批次上求平均 $B$ 来计算，从而使较大的 logits 导致更大的损失：

$${ℒ}_{\text{RZ}}(x)=\frac{1}{B}\cdot \sum _{i=1}^B{\left(\log \sum _{j=1}^{N_E}\exp (x_j^{(i)})\right)}^2$$ (4)

损失进一步乘以一个可选的损失权重 $\beta$（参见 Equation 2），以确定损失的大小，通常设置为 0.001 [220, 156]。 我们没有尝试改变 0.001 的权重。

在 Figure 11 中，我们确认，在训练损失、验证损失和下游性能方面，添加路由器 Z 损失可以提高稳定性（减少峰值）和质量（降低损失和提高下游性能）。 因此，尽管它将吞吐量降低了 $\sim$2% ，我们还是在 OLMoE-1B-7B 中使用权重为 0.001 的路由器 Z 损失，如 Zoph 等人 [220] 所述。

4.2 通用预训练设置

4.2.1 数据集实验

图 12： OLMoE-Mix 与 Dolma 1.7 相比。 我们将我们在 §2 中描述的数据混合与用于训练先前 OLMo 模型的 Dolma 1.7 进行比较。 较低的训练损失并不意味着一个数据集比另一个更好，而是表明哪个数据集更容易让模型学习。 更多结果、日志和配置： https://wandb.ai/ai2-llm/olmoe/reports/Plot-Dolma-1-7-vs-Dolma-OLMoE--Vmlldzo4OTIxNTg5

Li 等人 [89] 发布了 DCLM-Baseline 数据集，并确定它在 MMLU [69] 等常见基准测试中比 Dolma 1.7 和其他数据集产生更好的语言模型。 这促使我们将他们的 DCLM 数据集与我们认为高质量的 Dolma 1.7 中的一些组件混合；参见 §2。 在 Figure 12 中，我们在受控环境中将我们的混合，OLMoE-Mix，与 Dolma 1.7 进行比较。 我们发现 OLMoE-Mix 在所有三个下游指标上都取得了明显提升，尤其是 MMLU。 DCLM-Baseline 是通过针对 MMLU 和其他下游指标的一系列数据集消融而创建的，这解释了这些结果。 我们还比较了根据 Appendix F 中的详细说明将 Reddit 和 FLAN 添加到我们的混合中，但没有发现一致的性能提升。 我们没有强烈的直觉来解释为什么添加这些数据集没有帮助，并且在未来的迭代中可能需要一种更自动化的数据集混合方法 [100, 4]。 我们使用 DCLM-Baseline 和 Dolma 1.7 的混合进行预训练，称为 OLMoE-Mix。

4.2.2 初始化

图 13： 初始化。 我们将标准差 (std) 为 0.02 的普通初始化与最大 (最小) 截断值为 0.06 (–0.06) 的截断正态初始化进行比较，对应于三个标准差 (3$\times$0.02)。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-Init--Vmlldzo4NDIzMzM5

之前关于专家混合的工作很少分享它们的初始化策略。 即使是这项工作之前最开放的 MoE，JetMoE [156] 和 OpenMoE [198]，也没有提到它们的初始化方案。 对于 DeepSeekMoE [39] 和 DeepSeekV2 [43]，作者分享他们使用标准差 (std) 为 0.006 的正态初始化。 对于密集的语言模型，标准差为 0.02 的正态初始化已被广泛使用，正如 Shoeybi 等人 [157] 所推广的那样。

在 Figure 13 中，我们发现截断正态初始化比常规正态初始化导致更稳定的训练和更好的性能。 两种初始化之间的差异直到大约 4500 亿个符元才变得明显，其中使用正态初始化的模型开始发散。 尽管这两个模型都使用了相同的配置，除了权重初始化的差异。 在进行数百亿个符元的训练之前，实验才能提供明确的信号，这是预训练消融的关键挑战之一。 我们对 OLMoE-1B-7B 使用截断正态初始化。

4.2.3 RMSNorm

图 14： 非参数层归一化与 RMSNorm 的比较。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-LN--Vmlldzo4NDQyMTAz

图 15： 衰减 RMSNorm 参数。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-Decay-LN--Vmlldzo4NDQ1NDYy

图 16： 使用 RMS 或非参数归一化训练时的梯度总范数。 我们在 750 亿个符元处增加了 RMS 运行的日志间隔，因此它的厚度发生了变化。

OLMo [64] 使用非参数层归一化 [12]，主要是因为它比常用的 RMSNorm [207, 112] 快得多。 这是一种不寻常的选择，因为大多数 LMs 使用 RMSNorm，例如 Llama [180, 181, 50]、Gemma [174, 175] 和 Qwen [13, 200] 模型系列。

在Figure 14 中，我们观察到用参数化 RMSNorm 替换 OLMo 中的非参数层归一化会导致更好的性能。 这可能是因为非参数层归一化会导致大量梯度峰值，如Figure 16 所示。 我们将梯度裁剪到 1.0，这可以防止这些峰值导致非常大且可能具有破坏性的参数更新。 但是，裁剪后的梯度仍然可能损害模型的性能，因为它们不再是真正的梯度。 因此，尽管 RMSNorm 使我们的训练吞吐量降低了 15%， 我们使用 RMSNorm 训练了我们的最终模型。 我们将 RMSNorm 参数纳入权重衰减，因为我们发现它表现略好 (Figure 15)，即使将其排除是常见做法。

4.2.4 衰减嵌入参数

图 17： 衰减嵌入参数。 更多结果、日志和配置：https://api.wandb.ai/links/ai2-llm/3h22onp5

与 RMSNorm 参数类似 (§4.2.3)，嵌入参数通常从权重衰减中排除。 在 Figure 17 中，我们发现它们是否被衰减对性能只有轻微影响，衰减略好。 因此，为了简单起见， 我们对 OLMoE-1B-7B 中的所有参数进行权重衰减，包括嵌入和 RMSNorm。

4.2.5 QK-Norm

图 18： 查询-键层归一化 (QK-Norm)。 两种模型都使用非参数层归一化。 QK-Norm 对应于对查询和键投影的额外层归一化。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-QKNorm-vs-none--Vmlldzo4NDIzMzE2

一些工作报告称，在查询和键投影后添加层归一化 (“QK-Norm”) 会提高稳定性 [171, 112, 44]。 QK-Norm 可以防止后续注意力操作导致非常大的 logits，这些 logits 可能导致数值溢出并使网络不稳定，尤其是在低精度训练时。 类似于模型中其他位置的层归一化，QK-Norm 可以是非参数化的，或者使用参数化的 RMSNorm (§4.2.3)。

在 Figure 18 中，我们比较了使用 QK-Norm 和在查询和键投影之后不使用归一化。 我们发现 QK-Norm 带来了一些稳定性和性能提升。 我们使用 OLMo [64] 中使用的非参数化层归一化执行此实验，而我们使用参数化 RMS 层归一化 [207] 用于 OLMoE-1B-7B (§4.2.3)。 为了确保 QK-Norm 的益处不是与非参数化层归一化比较的结果，我们使用 RMS 层归一化进行了另一个实验，仍然发现 QK-Norm 会导致训练损失略微下降，并防止梯度范数出现大幅飙升。 因此，我们为 OLMoE-1B-7B 使用 QK-Norm，尽管它将吞吐量降低了近 10%。

4.2.6 AdamW Epsilon

图 19: AdamW epsilon。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-AdamW-eps--Vmlldzo4NDc5MDg0

Groeneveld 等人 [64] 在 AdamW 优化器中使用 1E-05 的 epsilon（“eps”）值来训练 OLMo。 更大的 eps 值会导致优化器步长更小，但可能更稳定 [82]。

在 Figure 19 中，我们发现将 eps 降低到推荐的默认值 1E-08 [82] 会显著提高性能，同时运行保持稳定。 因此，我们最终将 eps 设置为 1E-08。

4.3 适应设置

Data ($\downarrow$)	OLMoE-1B-7B
	After pretraining	After SFT
SFT data	12.22	12.16
Github	13.85	14.85
Wikipedia	14.48	14.24
C4	9.09	9.13

表 6: 在使用负载均衡损失权重 $\alpha$ 进行缩放之前，在各自语料库的子集上进行负载均衡损失 (Equation 3)。 尽管我们在预训练期间使用了负载均衡损失，但我们在 SFT 期间没有使用它。

我们使用在 Appendix C 中描述的评估设置来试验适应性的微小设计选择。 (1) 辅助损失： Zoph 等人 [220] 发现，在常规微调期间使用辅助负载均衡损失 (§4.1.6) 会带来微小的性能提升。 但是，对于指令调优，Shen 等人 [154] 没有找到支持使用负载均衡或路由器 z 损失的明确证据，只有微小的性能差异，既支持也反对辅助损失。 在 Table 7 中，我们展示了在适应期间使用负载均衡损失的实验，发现不使用它会带来更好的性能（指令调优 (SFT) 后的性能为 54.0 对 52.8，偏好调优 (DPO) 后的性能为 57.7 对 57.1）。 停用负载均衡损失的一个潜在问题是它可能会损害专家之间的平衡，并使某些专家变成沉重的负担，正如我们在预训练期间在 §4.1.6 中观察到的那样。 然而，当我们在 Table 6 中测量 SFT 数据 (§2) 上的负载均衡损失时，我们发现损失在 SFT 期间实际上略有下降 (12.16 对 12.22)。 这可能是因为某些符元被路由到哪些专家是在预训练的早期阶段决定的，正如我们后来在分析部分 (§5.1) 中发现的那样。 我们还可视化了预训练后模型专家的激活模式，以及在 Appendix G (Figure 33) 中没有负载均衡训练的 SFT 和 DPO 后模型的激活模式，发现分布保持在同一水平。 因此，由于我们没有负载均衡适应的模型表现更好，并且我们发现它不会对路由产生重大影响，我们在适应期间不使用负载均衡。 (2) 退火检查点： 我们还尝试使用预退火检查点 (§2) 进行适应，发现退火后检查点会导致更好的性能（SFT 后的性能为 53.8 对 54.0，DPO 后的性能为 56.3 对 57.7），因此 我们使用退火后的检查点。 (3) 偏好算法： 自 DPO (直接偏好优化) 发布以来 [136]，各种偏好算法已被提出 [54, 72, 113]。 我们使用 KTO [54] 进行实验，发现它在我们的设置中与 DPO 相匹配（Table 7）（Appendix B）。 虽然我们发布了这两个模型，但我们最终使用 DPO 作为我们的 OLMoE-1B-7B-Instruct 模型，因为它在 AlpacaEval 上得分更高，与我们其他基准相比，AlpacaEval 的数据污染可能性更小 [197]。

				Human-	Alpaca-
Task ($\to$)	MMLU	GSM8k	BBH	Eval	Eval 1.0	XSTest	IFEval	Avg
Setup ($\to$)	0-shot	8-shot CoT	0-shot	0-shot	0-shot	0-shot	0-shot	0-shot
Metric ($\to$)	EM	EM	EM	Pass@10	%win	F1	Loose Acc
OLMoE-1B-7B w/o annealing	49.0	2.0	31.5	18.9	-	62.1	18.5	-
+SFT	50.2	43.0	35.6	55.5	68.9	83.8	39.7	53.8
+DPO	50.9	36.0	35.8	58.8	81.7	83.2	47.9	56.3
OLMoE-1B-7B	49.8	3.0	33.6	22.4	-	59.7	16.6	-
+SFT	51.4	40.5	38.0	51.6	69.2	84.1	43.3	54.0
+DPO	51.9	45.5	37.0	54.8	84.0	82.6	48.1	57.7
+KTO	51.2	45.5	34.1	57.1	81.6	86.6	47.5	57.7
+SFT (load balancing)	50.9	36.5	35.7	52.4	66.9	84.8	42.3	52.8
+DPO (load balancing)	51.1	42.5	39.3	55.6	82.9	82.1	46.0	57.1

表 7: OLMoE-1B-7B 的自适应实验。 我们比较了在自适应之前使用预训练检查点（在退火之前）、在额外 100B 个符元退火之后使用检查点，以及在额外 100B 个符元退火之后并在自适应期间使用负载均衡损失（§4.1.6）使用检查点。 我们将 DPO/KTO 应用于各自的 SFT 模型。

5 MoE 分析

通过推进开放且成本效益高的模型（§1），OLMoE-1B-7B 为 LMs 和 MoEs 的新研究提供了可能。 利用我们发布的中间检查点、数据和代码，我们定义并分析了 MoE 特有的四个属性：路由器饱和（§5.1）、专家共同激活（§5.2）、领域专业化（§5.3）和 词汇专业化（§5.4）。

5.1 路由器饱和

图 20: 预训练期间的路由器饱和度，在 C4 验证数据的随机 0.5% 上进行测量。 我们通过将路由与四个中间检查点（预训练的 1%、10%、20% 和 40%）处的 top-$k$ 专家进行比较来计算饱和度，并将这些与最终预训练检查点进行比较（Equation 5）。

我们将路由器饱和度定义为在某个时间点 $t$ 的某个中间检查点处激活的专家数量，该数量与在相同数据集上激活的某个最终检查点的专家 ID 匹配的比例：

$$\text{Router Saturation}(t)=\frac{1}{N}\sum _{i=1}^N\frac{|{ℰ}_i^{(t)}\cap {ℰ}_i^{(T)}|}{k},$$ (5)

其中:

• $N$: 数据集中的总符元数。

• $k$：每个输入 Token 激活的顶级$k$专家数量。 虽然我们在 $k=8$ (§2) 上训练，但我们也通过仅查看具有最高路由概率的专家来分析 $k=1$。

• ${ℰ}_i^{(t)}$: 在第 $t$ 个检查点激活的第 $i$ 个符元的 $k$ 专家集合。

• ${ℰ}_i^{(T)}$: 在最终检查点 $T$ 激活的第 $i$ 个符元的 $k$ 专家集合。

• $|{ℰ}_i^{(t)}\cap {ℰ}_i^{(T)}|$: 在第 $t$ 个和最终检查点之间，第 $i$ 个符元激活的共同专家数量。

因此，路由器饱和对应于路由器权重是否仍在学习哪个专家将处理某些数据。 值为 100% 表示中间检查点的路由器将路由到与最终检查点路由器相同的专家。 但是，即使在 100% 饱和的情况下，路由器权重仍然可以改变并调整每个专家的确切路由概率。 这些概率用于缩放模型中相应专家的输出。 对于具有 64 个专家的 OLMoE-1B-7B，随机路由等于 $1/64=1.6\%$ 的饱和度，对于 $k=1$ 和 $8/64=12.5\%$ 则为 $k=8$。

在 Figure 20 中，我们发现经过 1% 的预训练（5000 步或 200 亿个符元），高达 $\sim$60% 的路由到前 8 个激活的专家已经饱和（右）。 因此，模型已经对给定输入数据使用相同的 8 个专家，就像它在预训练结束时一样。 这种早期饱和与先前的工作 [198] 一致。 在 40% 的预训练中，饱和度达到 $\sim$80%。 但是，哪个顶级-1 专家具有最高的路由概率饱和度较慢（左）。 我们发现，在预训练期间，较后层的路由饱和度较早。 第 0 层是一个异常值，其饱和速度明显慢于其他层。 Dai 等人[39] 在第一层中没有使用 MoE，因为他们发现负载均衡在第一层的收敛速度较慢。 这可能与我们关于饱和度的发现有关。 由于第一层的路由饱和速度较慢，某些输入数据被路由到的专家会经常发生变化。 这些变化可能会导致一个专家突然比其他专家获得更多的数据，从而影响负载均衡。 我们很高兴未来能够通过基于我们的公开发布来进一步研究第一层中发生的事情。

5.2 专家协同激活

图 21： OLMoE-1B-7B 在 C4 验证数据的随机 0.5% 上的专家协同激活。 我们在 x 轴和 y 轴上通过专家 ID 显示了 32 个具有最高最大协同激活分数的专家。

我们将专家协同激活定义为两个特定专家 $E_i$ 和 $E_j$ 在其中一个专家的总激活次数中同时被激活的比例：

$$\text{Expert co-activation}(E_i,E_j)=\frac{N_{E_i,E_j}}{N_{E_i}},$$ (6)

其中：

• $E_i$：第一个专家。

• $E_j$：第二个专家。

• $N_{E_i,E_j}$：专家 $E_i$ 和 $E_j$ 一起被激活的次数。

• $N_{E_i}$：专家 $E_i$ 被激活的总次数。

100% 的共激活表示如果 $E_i$ 被激活，$E_j$ 也总是被激活。 0% 的值表示专家从不共现。 如果多个专家对具有高共激活，则可能表明这些专家可以合并，从将它们分开中获益较少。 在分布式设置中，我们可以将高度共激活的专家放置在同一个设备上，以减少模型推理过程中的通信成本。

在 Figure 21 中，我们发现一层中专家之间没有强烈的共激活，只有少数例外。 这可能表明不同专家之间几乎没有冗余。 总体而言，第 7 层和第 15 层显示出相似的共激活模式，其中有几个由 3 或 2 个专家组成的群体往往会一起被激活。 我们在 §5.4 中调查了激活这些专家的符元。 此外，在 Appendix G (Figure 35) 中，我们调查了不同层中的专家而不是同一层中的专家是否倾向于一起处理符元。

5.3 领域专业化

图 22： OLMoE-1B-7B（顶部）与 Mixtral-8x7B（底部）的领域专业化。 我们可视化了来自不同领域的符元在预训练结束时被路由到 64 个 (OLMoE) 或 8 个 (Mixtral) 专家中的频率。 我们考虑被路由到任何 $k=8$ (OLMoE) 或 $k=2$ (Mixtral) 活动专家中的符元 (Equation 7)。 水平灰色线对应于随机机会或均匀路由（对于 OLMoE-1B-7B，每层有 8 个活动专家，总共 64 个专家，每个专家为 8/64=12.5%，对于 Mixtral，每层有 2 个活动专家，总共 8 个专家，每个专家为 2/8=25%）。 有关 $k=1$ 结果，请参见 Figure 34。

我们将领域专业化定义为来自特定领域$D$的符元被路由到特定专家$E_i$的比例：

$$\text{Domain specialization}(E_i,D)=\frac{N_{E_i,D}^{(k)}}{N_D},$$ (7)

其中：

• $E_i$：模型中第$i$个专家。

• $D$：数据源自的领域。

• $k$：考虑的专家数量（例如，$k=8$表示考虑具有最高路由概率的前8个专家）。

• $N_{E_i,D}^{(k)}$：来自领域$D$的符元数量，其中$E_i$位于所选前-$k$个专家中。

• $N_D$：由 MoE 处理的来自领域$D$的符元总数。

因此，领域专业化指的是专家$E_i$对领域$D$的专业化。100% 的值表示来自该领域的所有数据都被路由到$E_i$，而 0% 表示专家从未用于该领域，可以从模型中删除而不会影响该领域的性能。

在Figure 22（顶部）中，我们找到了许多专家示例，它们在特定领域中被激活的程度远高于或低于随机机会。 例如，对于 arXiv，它具有非常特殊的分布，包含大量的科学文本，第 0 层中的第一个专家几乎 100% 专业化。 这表明OLMoE-1B-7B中专家的知识几乎没有冗余，因为它们专门针对不同类型的数据。 GitHub 和 arXiv 通常在第 7 层中一起被激活，我们将在§5.4中进一步探讨这一点。 对于通用领域，例如 C4 [137]，它是一个包含各种类型数据的网络爬取，OLMoE-1B-7B中的专家激活更加平衡。 这突出表明负载均衡（§4.1.6）按预期工作，并且模型对通用数据进行了所有专家的适当使用。 但是，Figure 22（底部）中的 Mixtral-8x7B [78] 在独特和通用领域之间表现出很少的领域专业化。 专家在所有层和领域中被激活的程度接近统一路由基线。 因此，Mixtral 中的专家之间可能存在更多冗余，因为它们可能包含类似的知识。 我们假设这是由于 Mixtral 是从 Mistral [25] 中升级而来的。 从密集模型初始化可能会限制专家在专业化方面的可能性，因为它们都从同一个局部最优值开始。 这可能是为什么从头开始训练最终在我们预训练实验中胜过升级的原因 (§4.1.5)。

5.4 词汇专业化

图 23: OLMoE-1B-7B 在各层和专家之间词汇专业化。 为了计算每层的词汇专业化（左图），我们对该层中每个专家的专业化取平均值。 虚线（右图）对应于如图所示的第 7 层的平均值。 我们展示了 64 个专家中的前 32 个。 该图用于 $k=1$ (Equation 8)，我们提供 $k=8$ 以及与 Mixtral-8x7B 的比较 Appendix G。

Expert ID	Input token IDs	Predicted output token IDs
27	(100%) (100%) 3 (100%) (100%) (100%) (100%) (100%) (100%) (100%) (100%) (100%)	(100%) § (100%) (100%) (100%) (100%) (100%) (100%) (100%) (100%) (100%) (100%)
58	(“ (100%) (” (100%) ‘ (94%) ’ (92%) “ (92%) ( (92%) ” (90%) ’ (89%) “ (88%) $ (87%) [ (87%) £ (86%)	such (100%) 486 (100%) see (95%) which (91%) driving (91%) UK (90%) who (88%) including (88%) normal (88%)
7	Him (100%) inde (100%) Jesus (98%) God (90%) pray (81%) Holy (80%) Quran (80%) God (77%) Lord (76%) glory (75%) Spirit (66%) Christ (65%)	rella (100%) Him (94%) sin (90%) prince (80%) glory (72%) Jesus (69%) Lord (68%) Christ (65%) Spirit (55%) Holy (53%) God (50%) Prayer (50%)
37	Sunday (100%) Tuesday (100%) Thursday (100%) Olympic (100%) Christmas (100%) rugby (100%) Championship (100%) weekends (100%)	days (91%) anniversary (90%) month (88%) week (84%) mpi (83%) semester (81%) mand (80%) Olympics (78%) cent (76%) season (76%) perm (75%)
43	Armenian (100%) ijan (100%) enia (96%) Iraq (95%) Iranian (92%) Iran (92%) Saudi (90%) northern (90%) Lebanon (90%) Singapore (88%) Turkey (88%) Asia (87%) Egypt (86%) western (86%)	enia (90%) invasion (80%) Arabia (76%) irregular (66%) regions (64%) border (63%) Kong (61%) ians (61%) bases (60%) Republic (59%) Ireland (58%) Korea (58%) War (55%) Carolina (52%)
4	sq (89%) Main (70%) reversal (69%) YR (63%) GC (56%) Overall (50%) 79 (50%) main (50%) RE (46%) PCR (46%) tomb (45%) normal (43%) intensity (41%) Overall (41%) median (41%)	YR (90%) Character (88%) sq (77%) Os (76%) GHz (71%) fluence (60%) amycin (60%) pixels (56%) = (53%) arc (52%) Story (52%) = (51%) anth (50%) GHz (50%) cm (46%)
0	ESM (100%) icillin (100%) agra (98%) aust (96%) asa (93%) pills (92%) mg (85%) uk (82%) login (82%) doc (81%) generic (81%) cd (81%) Essay (81%) password (81%) Content (80%)	*, (100%) sil (96%) pills (91%) vi (90%) xen (87%) pharmacy (87%) gener (85%) aust (82%) mg (75%) Content (75%) uk (73%) THAT (73%) dispens (68%) icillin (68%) generic (66%)
3	grandmother (92%) brother (91%) Daisy (83%) daughter (78%) mum (75%) father (72%) wife (70%) husband (70%) lady (63%) dad (62%) boy (61%)	hood (36%) mother (35%) inde (31%) boy (29%) girl (28%) married (27%) tri (21%) Gab (20%) died (18%) taught (14%) lived (13%) knew (10%)
48	compared (42%) !) (41%) Then (41%) ’, (40%) ), (35%) ”, (35%) instead (33%)	except (60%) tennis (41%) Marks (40%) Dunn (33%) tears (30%) Arizona (30%)
23	…. (58%) Therefore (55%) So (46%) !!! (46%) And (44%) According (41%) .” (41%) !! (40%) ?” (38%) But (38%)	(53%) Republican (50%) Jack (47%) THIS (40%) Democratic (40%) according (39%) So (38%) Step (33%)

表 8: OLMoE-1B-7B 第 7 层的词汇专业化。 我们使用 $k=1$ (Equation 8) 以及 C4 验证数据的随机 0.5%，不包括出现次数 10 次的词元 ID。

我们将词汇专业化定义为词元 ID $x$（也称为词汇元素）的比例，这些词元 ID 被路由到该层中所有专家中的特定专家 $E_i$:

$$\text{Vocabulary specialization}(E_i,x)=\frac{N_{x,E_i}^{(k)}}{N_x},$$ (8)

其中:

• $E_i$: 模型中的第 $i$ 个专家。

• $x$: 正在分析的符元 ID。

• $k$: 考虑的专家数量（例如，$k=8$ 表示考虑具有最高路由概率的前 8 个专家）。

• $N_{x,E_i}$: 输入数据被路由到 $E_i$ 用于 $x$ 的次数。

• $N_x$: 输入数据在所有专家中被路由到 $x$ 的总次数。

因此，词汇专业化指的是特定专家在一个词汇项上的专业程度。 我们区分这种专业化的输入和输出变体，其中 $x$ 是输入符元 ID 或下一个输出符元 ID（要么是真实下一个符元 ID，要么是模型预测的符元 ID）。 值为 100% 表示对于该词汇元素的所有出现，输入数据都被路由到 $E_i$，而 0% 表示对于该词汇元素完全不相关的专家，并且可以在该符元 ID 出现时有效地从模型中移除，而不会影响性能。

在 Figure 23 中，我们发现词汇专业化在后层更高，类似于后层如何更早地饱和（§5.1）。 后层也更多地专门针对预测的输出符元 ID 而不是输入符元 ID，即路由更多地由模型即将预测的符元决定，而不是原始输入符元。 这是直观的，因为在早期层，关于模型将预测哪个符元存在更大的不确定性。 在 $\sim$90% 时，专家 27 的专业化程度最高，我们在 Table 8 中发现，它激活了许多非字母符元，例如西里尔字母和梵文字母。 专家 43 显示出在输入和输出符元上对地理术语的专业化。 专家 48 和 23 都专注于连接词，例如 Then 和 Therefore。 这可能是因为它们通常与 Figure 21（中间）中 60% 的高共激活一起处理符元。 基于我们在 §5.3 中的发现，即对于 GitHub 和 arXiv，通常第 7 层中的相同专家被激活，我们在 Table 8 中显示了一个这样的专家（专家 ID 4）。 它似乎专门用于测量，例如 sq、YR（年）和 GHz。 这些是与 arXiv 领域相对应的科学论文中的常用术语，也可能出现在与测量相关的 GitHub 代码中。 它们不太可能出现在书籍中，这解释了为什么在 Figure 22 中书籍数据的第 7 层中专家 ID 4 的激活率很低。 专家 3 是Figure 22 中图书数据第 7 层最活跃的三个专家之一（第 7 层的第四个黄色条）。 当查看其在 Table 8 中对家庭术语的专门化时，这一点得到了共鸣，家庭术语在书籍中比在科学论文或代码中更为常见。 总体而言，领域专门化和词汇专门化紧密相连，因为领域通常以其独特的词语分布为特征。 在 Appendix G（Figure 32）中，我们通过比较不同领域和专家 ID 之间的词汇专门化程度，将它们更紧密地联系起来。 在 Appendix G（Figure 30、Figure 31）中，我们还发现 OLMoE-1B-7B 表现出比 Mixtral-8x7B 更强的词汇专门化。

6 相关工作

MoE 的进展

当前的 LMs 仍然在很大程度上遵循 Transformer 架构 [183]，只有一些架构上的改变被广泛采用，例如仅解码器训练 [135]、SwiGLU 激活 [151, 41]、RoPE [164]、MQA/GQA [150, 3] 和 RMSNorm [207]。 通过专家混合实现模型稀疏性是目前正在积极探索的一种修改方法，虽然早期已经有一些应用，但大多数 LMs，包括 Llama 3 [50]，仍然依赖于密集架构。 自从专家混合层被引入以来 [152]，在改进稀疏门控专家混合层方面取得了很大进展：新的路由技术 [87, 144, 221, 65, 76, 49, 214, 194, 123]、细粒度专家分割 [39, 68]、稳定性 [220] 和效率 [86, 139, 48, 217, 90, 166, 127, 143] 的改进。 在这项工作中，我们进行了许多实验，以提供对训练专家混合 LMs 的见解。 随后，我们训练了 OLMoE-1B-7B，共计 5T 个符元。 据我们所知，之前没有任何 MoE 被过度训练到这种程度 [57]，这使得 OLMoE-1B-7B 成为研究 MoE 与密集模型的性能饱和度的最佳测试平台。 通过 OLMoE，我们希望促进这类和其他研究，以帮助该领域揭示 MoE 是否应该被纳入所有未来的 LLM，以及以何种精确配置。

开放 LLM

在不同程度的开放性下，已经提出了各种模型系列，通常根据模型权重是否可用进行分类。 封闭权重 模型包括 GPT [24, 126]、Gemini [172, 173]、PaLM [30, 9]、Reka [179]，开放权重 模型包括 Llama [180, 181, 50]、Mistral [77, 78]、Gemma [174, 175]、Falcon [8, 130]、MPT [177]、Qwen [13, 200]、GLM [61]、Yi [2]、DeepSeek [42, 43, 39]、Nemotron [128, 125, 188]、InternLM [26]、Baichuan [199]、Phi [67, 93, 1]、StableLM [16]、OPT [211]。 然而，除了模型权重之外，训练数据和代码是使这些模型的科学研究成为可能 [104, 105] 并广泛传播其益处的关键 [23]。 除了模型权重之外，还有一些版本包括数据和代码，我们将其称为 “完全开源”：BLOOM [192, 149, 122, 202]、GPT-NeoX [21, 22, 184]、StarCoder [91, 108, 5, 119, 219]、Pythia [18]、OLMo [64]、LLM360 [102]、Cerebras-GPT [46]、DCLM [89]、MAP-Neo [208]、RWKV [131, 132] 和 SmolLM [6]。 对于混合专家模型，只有 OpenMoE [198] 旨在完全开源，但其糟糕的表现限制了其实用性。 我们发布 OLMoE-1B-7B 作为第一个完全开源的最先进混合专家 LLM：模型权重、数据、代码和日志。

7 结论

我们开源 OLMoE-1B-7B 和 OLMoE-1B-7B-Instruct，包括模型、数据、代码和日志。 我们的模型拥有 1B 个活动参数和 7B 个总参数，在具有类似活动参数数量的模型中，我们的模型产生了最先进的性能，甚至超过了更大的模型，包括 DeepSeekMoE-16B 和 Llama2-13B-Chat。 我们分享了各种训练实验，并定义和分析了我们的模型的路由器饱和、专家共同激活、领域和词汇专业化。 通过我们的完全开放发布，我们力求帮助该领域构建更好的 MoE。 我们对 OLMoE 的新迭代感到兴奋，以缩小前沿模型和完全开放模型之间的差距。

作者贡献

Niklas Muennighoff 提出了该项目并领导了该项目。 他负责预训练实验，对模型进行预训练，协助进行适应和分析，并撰写了论文的大部分内容。

Luca Soldaini 创建了预训练数据集，并就预训练问题提供了建议。

Dirk Groeneveld 就预训练提供了建议，尤其是稳定性和吞吐量改进方面。

Kyle Lo 协助创建预训练数据集，分析数据实验，就数据和框架方面提供建议，并帮助编辑论文。

Jacob Morrison 共同创建了适应数据集，执行了大多数适应实验，并帮助编辑论文。

Sewon Min 分析了路由器饱和度、专家关联度和词汇专业化，并帮助框架和编辑论文。

Weijia Shi 分析了领域和词汇专业化，在项目的各个阶段提供建议，并帮助编辑论文。

Pete Walsh 就预训练提供了建议，尤其是稳定性和吞吐量改进方面。

Oyvind Tafjord 执行了 OLMES 评估。

Nathan Lambert 共同创建了适应数据集，就适应问题提供建议，并帮助编辑论文。

Yuling Gu 执行了 OLMES 评估，并帮助编辑论文。

Shane Arora 上传了模型，并帮助进行代码审查。

Akshita Bhagia 支持稳定性调查，并帮助进行 DCLM 评估。

Dustin Schwenk 支持稳定性调查。

David Wadden 执行了 DCLM 评估，并帮助制作 Weights & Biases 报告。

Alexander Wettig 分析了负载均衡、路由和领域专业化，并帮助编辑论文。

Binyuan Hui 就预训练方面提供了建议。

Tim Dettmers 就分析和推理实验方面提供了建议。

Douwe Kiela 就框架方面提供了建议。

Ali Farhadi 就预训练和框架方面提供了建议。

Noah A. Smith 就预训练方面提供了建议，并帮助构建和编辑了论文。

Pang Wei Koh 就分析方面提供了建议，并帮助构建和编辑了论文。

Amanpreet Singh 就预训练和框架方面提供了建议，并帮助编辑了论文。

Hannaneh Hajishirzi 负责指导和建议整体工作，并帮助构建和编辑了论文。

致谢

OLMoE 的完成离不开许多个人和机构的支持。 我们感谢艾伦人工智能研究所、Contextual AI 和华盛顿大学的团队成员的支持，尤其是 Aditya Kusupati、Ananya Harsh Jha、Caitlin Wittlif、Carissa Schoenick、Costa Huang、Crystal Nam、David Atkinson、Emma Strubell、Faeze Brahman、Hamish Ivison、Karel D’Oosterlinck、Matt Latzke、Ian Magnusson、Jack Merullo、Jay Chen、Jennifer Dumas、Jiacheng Liu、Johann Dahm、Luke Zettlemoyer、Michael Schmitz、Michael Wilson、Pradeep Dasigi、Sahil Verma、Sam Skjonsberg、Sophie Lebrecht、Stas Bekman、Taira Anderson、Valentina Pyatkin、Yanai Elazar、Yizhong Wang 和 Yoganand Chandrasekhar。 我们还要感谢 Armen Aghajanyan、Akshat Shrivastava、Colin Raffel、Haokun Liu、Ludwig Schmidt 和 Shayne Longpre。 PWK 受新加坡国家研究基金会和新加坡数字发展与创新部国家人工智能小组的资助，该小组参与了人工智能访问教授计划（项目编号 AIVP-2024-001）。

附录 A 人工制品

Artifact	Public link
OLMoE-1B-7B	https://hf.co/allenai/OLMoE-1B-7B-0924
OLMoE-1B-7B-Instruct	https://hf.co/allenai/OLMoE-1B-7B-0924-Instruct
OLMoE-1B-7B-SFT	https://hf.co/allenai/OLMoE-1B-7B-0924-SFT
OLMoE-Mix	https://hf.co/datasets/allenai/OLMoE-mix-0924
SFT data	https://hf.co/datasets/allenai/
	tulu-v3.1-mix-preview-4096-OLMoE
KTO/DPO data	https://hf.co/datasets/allenai/
	ultrafeedback_binarized_cleaned
Code	https://github.com/allenai/OLMoE
Logs	https://wandb.ai/ai2-llm/olmoe/reports/
	OLMoE-1B-7B-0924--Vmlldzo4OTcyMjU3
BLOOM-7B	https://hf.co/bigscience/bloom-7b1
DeepSeekMoE-3B-16B	https://hf.co/deepseek-ai/deepseek-moe-16b-base
DeepSeekMoE-3B-16B+chat	https://hf.co/deepseek-ai/deepseek-moe-16b-chat
DCLM-1B	https://hf.co/TRI-ML/DCLM-1B
DCLM-7B	https://hf.co/TRI-ML/DCLM-7B
Falcon-7B	https://hf.co/tiiuae/falcon-7b
Gemma2-3B	https://hf.co/google/gemma-2-2b
Gemma2-9B	https://hf.co/google/gemma-2-9b
JetMoE-2B-9B	https://hf.co/jetmoe/jetmoe-8b
JetMoE-2B-9B+SFT	https://hf.co/jetmoe/jetmoe-8b-sft
JetMoE-2B-9B+Chat	https://hf.co/jetmoe/jetmoe-8b-chat
Llama-7B	https://hf.co/huggyllama/llama-7b
Llama2-7B	https://hf.co/meta-llama/Llama-2-7b-hf
Llama3.1-8B	https://hf.co/meta-llama/Meta-Llama-3.1-8B
MPT-7B	https://hf.co/mosaicml/mpt-7b
Mistral-7B	https://hf.co/mistralai/Mistral-7B-v0.1
Mixtral-8x7B	https://hf.co/mistralai/Mixtral-8x7B-v0.1
OLMo-1B (0724)	https://hf.co/allenai/OLMo-1B-0724-hf
OLMo-7B (0724)	https://hf.co/allenai/OLMo-7B-0724-hf
OpenMoE-3B-9B	https://hf.co/OrionZheng/openmoe-8b
Pythia-7B	https://hf.co/EleutherAI/pythia-6.9b
Qwen1.5-3B-14B	https://hf.co/Qwen/Qwen1.5-MoE-A2.7B
Qwen1.5-3B-14B+Chat	https://hf.co/Qwen/Qwen1.5-MoE-A2.7B-Chat
StableLM2-2B	https://hf.co/stabilityai/stablelm-2-1_6b
TinyLlama-1B	https://hf.co/TinyLlama/TinyLlama_v1.1

表格 9： 本文发布和使用的所有工件。 我们从本文中使用给定工件的名称（例如 Figure 1）指向其可获取的 URL。

附录 B 训练配置

预训练

我们在 Appendix B 中展示了 OLMoE-1B-7B 的预训练超参数配置，并与其他相关模型进行了比较。 我们遵循 Groeneveld 等人 [64] 使用 AdamW 优化器 [106]，通过 PyTorch FSDP [212] 和混合精度训练 [115] 使用 ZeRO [140]。 我们与 Groeneveld 等人 [64] 的主要模型设置不同之处在于： (1) 与 MoE 相关的更改： OLMoE-1B-7B 是一个使用无 dropout 的专家混合模型 [58] 的稀疏激活解码器 Transformer [183]。 与大多数之前的 MoE 不同，我们使用了一种高粒度 [39, 85]，其中有 64 个小型专家，FFN 维度仅为 1,024，而不是几个大型专家。 我们还使用两个辅助损失：路由器 z 损失 [220] 和负载平衡损失 [152]。 (2) 稳定性改进： (a) 我们使用截断正态分布初始化，标准差为 0.02，最小（最大）截止值为 -0.06（0.06），对应于三个标准差。 (b) 我们使用 QK 规范化 [171, 112, 44]。 (c) 我们使用 RMSNorm [207] 而不是 Groeneveld 等人 [64] 中使用的非参数层归一化。 (3) 性能改进： 除了影响性能的一些稳定性改进外，我们还将 AdamW epsilon 从 Groeneveld 等人 [64] 中使用的 1.0E-05 降低到 1.0E-08，以加快收敛速度。 最后，我们训练 OLMoE-1B-7B 比所有以前的 OLMo 模型都要长得多，总共 5T 个符元，因此超过一个 epoch（1.2），遵循 Muennighoff 等人 [120]。 我们在开始第二个 epoch 之前对预训练数据集进行洗牌。 对于最后的 100B 个符元，我们将学习率从 5.0E-04 线性衰减到 0（“退火”）。 我们在 §4 中对许多这些设置进行了实验。

适应

为了微调，我们使用 Open Instruct [186, 75]。 我们将所有 SFT 样本过滤到少于 4096 个符元的长度，以匹配模型的序列长度。 遵循 Muennighoff 等人 [121]，我们在 SFT 期间在符元级别聚合损失，以提高对长生成任务（如 AlpacaEval）的性能。 我们在 BF16 中进行微调，全局批次大小为 128（4 个 H100 节点，每个节点有 8 个 GPU，每个设备的批次大小为 2，以及 2 个梯度累积步骤）。 我们训练 2 个 epoch，学习率恒定为 2.0E-5。 对于 DPO [136]，我们将全局批次大小减少到 32（4 个 H100 节点，每个节点有 8 个 GPU，每个设备的批次大小为 1）。 我们训练 3 个 epoch，学习率为 5.0E-7，DPO beta 为 0.1。 我们适应的模型建立在我们退火的检查点之上，并且我们在 SFT 和 DPO 期间都包含负载平衡损失，这是基于我们在 §4.3 中的实验。 我们的偏好调整配方经过大量优化，适用于 DPO，这是基于 Ivison 等人 [75] 的大量实验，因此对于 KTO [54]，我们在 Appendix F 中对一些设置进行了实验。 我们最终的 KTO 适应使用了与 DPO 相同的超参数，除了我们使用 RMSProp 优化器而不是 Adam，我们使用 Adam 用于 SFT 和 DPO，以及我们将 KTO 的训练持续时间减少到 1.3 个 epoch（5,000 步），而不是用于 DPO 的 3 个 epoch。

硬件

我们在 256 个 H100 GPU 上对 OLMoE-1B-7B 进行预训练，大约持续 10 天，使用 NV-link 互连跨 GPU 和 InfiniBand 互连跨节点。 我们也使用 H100 GPU 进行所有实验，但有些实验使用一个集群，该集群在节点之间使用 GCP TCPx 互连。 为了进行适应，我们使用 32 个 H100 GPU 进行 33 个小时的指令调优，以及 14 个小时的偏好调优，通过 DPO 实现。 对于 KTO 适应，我们使用 8 个 H100 GPU 进行 30 个小时的训练。

	OLMoE-1B-7B	JetMoE	OpenMoE	OLMo-1B (0724)
Dimension	2,048	2,048	2,048	2,048
Activation	SwiGLU	SwiGLU	SwiGLU	SwiGLU
FFN dimension	1,024	5,632	8,192	8,192
Vocab size	50,304	32,000	256,384	50,304
Attn heads	16	16	24	16
Num layers	16	24	32	16
Layer norm type	RMSNorm	RMSNorm	RMSNorm	non-parametric
Layer norm eps	1.0E-05	1.0E-05	1.0E-06	1.0E-05
QK-Norm	yes	no	no	no
Pos emb.	RoPE	RoPE	RoPE	RoPE
RoPE $\theta$	10,000	10,000	10,000	10,000
Attention variant	full	MoA	full	full
Biases	-	MLP & Attn	-	-
Weight tying	no	yes	no	no
Init dist	trunc normal	?	?	normal
Init std	0.02	0.02	varies	varies
Init trunc	3$\times$std	-	-	-
MoE layers	Every	Every	Every 6th	-
MoE layer type	dMoE	dMoE	ST-MoE	-
# Experts	64	8	32	1
# Activated	8	2	2	1
# Vocab params	103M	66M	525M	103M
# Active params	1.3B	2.2B	2.6B	1.3B
# Total params	6.9B	8.5B	8.7B	1.3B
Sequence length	4,096	4,096	2,048	4,096
Batch size (samples)	1,024	1,024	2,048	512
Batch size (tokens)	$\sim$4M	$\sim$4M	$\sim$4M	$\sim$2M
warmup steps	2,000	2,500	10,000	2,000
peak LR	4.0E-04	5.0E-04	0.01	4.0E-04
minimum LR	5.0E-04	5.0E-05	-	5.0E-05
optimizer	AdamW	AdamW	Adafactor	AdamW
weight decay	0.1	0.1	0.0	0.1
beta1	0.9	?	0.9	0.9
beta2	0.95	?	-	0.95
AdamW epsilon	1.0E-08	?	-	1.0E-05
LR schedule	cosine	WSD	Inv Sq Root	cosine
gradient clipping	global 1.0	global 1.0	global 1.0	global 1.0
gradient reduce dtype	FP32	?	?	FP32
optimizer state dtype	FP32	?	?	FP32
LBL weight	0.01	0.01	0.01	-
Router z-loss weight	0.001	0.001	0.0001	-
Pretraining tokens	5,033B	1,000B	1,100B	2,000B
Annealing tokens	100B	250B	-	50B
Annealing schedule	linear	-	-	linear
Annealing min LR	0	-	-	0

表 10： OLMoE-1B-7B 和从头开始训练的类似模型的预训练超参数。 我们重点介绍了 OLMoE-1B-7B 与 OLMo-1B 不同的行。 活动参数包括词汇参数。 “？” = 未公开设置，FFN = 前馈网络，Attn = 注意力，LR = 学习率，WSD = 权重稳定衰减 [73]，LBL = 负载均衡损失，Inv Sq Root = 逆平方根衰减 [153]，trunc = 截断，std = 标准差，“varies” = 与层或权重相关的 std。

附录 C 评估设置

Dataset ($\downarrow$)	During pretraining				After pretraining (OLMES [66])
	Format	Shot	Norm	Split	Format	Shot	CF Norm	Split
ARC-C [34]	CF	0	token	val	max(MCF,CF)	5	pmi	test
ARC-E [34]	CF	0	none	val	max(MCF,CF)	5	character	test
BoolQ [33]	CF	0	none	val	max(MCF,CF)	5	none	val
COPA [62]	CF	0	none	val	-	-	-	-
CSQA [168]	CF	0	token	val	max(MCF,CF)	5	pmi	val
HellaSwag [206]	CF	0	token	val	max(MCF,CF)	5	character	val
MMLU [69]	MCF	5	none	val	max(MCF,CF)	5	character	test
MMLU Var	CF	0-5	token	val	-	-	-	-
OBQA [116]	CF	0	token	val	max(MCF,CF)	5	pmi	test
PIQA [20]	CF	0	token	val	max(MCF,CF)	5	character	val
SciQ [191]	CF	0	none	val	-	-	-	-
SocialIQA [148]	CF	0	token	val	max(MCF,CF)	5	character	val
Winogrande [146]	CF	0	none	val	max(MCF,CF)	5	none	val

表 11： 预训练期间和之后的下游评估摘要（OLMES）。 ARC-C 和 ARC-E 指的是 ARC 挑战和简单，CSQA=常识问答，OBQA=开放式问答，CF=完成/封闭公式，MCF=多项选择公式，pmi=点互信息，Var=变体，指从 0-5 变化的少样本使用。

在预训练期间

我们使用类似于 Groeneveld 等人 [64] 的循环内评估设置进行评估，并增加了更多任务，例如常识问答、PIQA 和 MMLU 的不同实现。 遵循 Groeneveld 等人 [64]，对于大多数任务，我们使用完成/封闭公式 (CF) 进行 0-shot 评估，使用语言模型概率对每个答案字符串进行排名。 在概率归一化方面，当仅基于概率进行排序时，要么没有归一化 (none)，要么按答案中符元数量进行归一化 (token)，因为这可能会严重偏向较短的答案 [24]. 对于 MMLU，循环内评估还包括一种设置，其中我们通过将 0-样本到 5-样本设置的范围包含在一起来增加实例总数，因为我们发现这在训练过程中提供了更平滑的趋势 (“MMLU Var”)。 我们还包括 MMLU 的多项选择公式 (MCF) 版本，对像 A/B/C/D 这样的答案标签进行预测，这通常直到训练后期才会开始上升，因为模型直到后期才获得多项选择能力 (对于 OLMoE-1B-7B 来说，在 Figure 25 中大约在 1 万亿个符元时)。 我们还在 Paloma 的选定验证集上评估困惑度 [110, 142, 59, 161, 95, 114]. 用于预训练期间评估的所有代码都位于 https://github.com/allenai/OLMo/tree/61ac104d616ec5435db225796e5c7532c9abd95a/olmo/eval.

预训练后 - OLMES

我们按照 OLMES 评估标准 [66] 进行评估，使用原始论文中的任务套件。 OLMES (开放式语言模型评估标准) 是一种用于可重复的语言模型评估的标准，它开放、实用且有文档，提供基于实验和文献结果的建议 [19, 60, 63]. 它旨在支持较小的基础模型之间的比较，这些模型需要多项选择题的完形填空公式，而较大的模型则可以利用多项选择公式。 为了使我们的评估具有可重复性，我们在提示格式、上下文中示例的选择、概率归一化、任务公式以及所有其他细节方面都遵循 OLMES。 我们在 Table 4 中总结了此设置，并参考 Gu 等人 [66] 以获取更多详细信息。

预训练后 - DCLM

对于 Table 13 中关于 DCLM 任务 [89] 的结果，我们严格遵循他们的设置，使用作者在 https://github.com/mlfoundations/dclm 上发布的评估代码。 “核心”结果是其评估代码中的 低方差 任务，而 “扩展” 对应于 繁重 的任务。

适应后

在监督微调和直接偏好优化之后，我们使用评估的一个子集以及与 Ivison 等人 [75] 和 Wang 等人 [186] 中使用的相同总体设置来评估模型。 在我们的评估套件中，我们涵盖了广泛的模型能力，包括编码 (HumanEval [28] )、一般和数学推理 (Big Bench Hard [167] 、GSM8k [35] )、世界知识 (MMLU)、一般指令遵循 (AlpacaEval 1.0 [92] ，不包括长度控制变体 [51] )、精确指令遵循 (IFEval [216] ) 和安全性 (XSTest [145] )。 关于每个基准的更多细节，请参考 Wang 等人 [186] 。

附录 D 模型的开放性

我们列出了 Figure 1 中总结的各种模型的开放性。 我们排除了 Switch Transformers [56] ，因为它是在三年前发布的，与最近的 MoE 模型（MLM 目标、编码器-解码器等）非常不同。

Grok-86B-314B [195]

• 

  模型： 他们的模型在开源 Apache 2.0 许可证下授权。

• 

  数据： 不可用。

• 

  代码： 不可用。

• 

  日志： 不可用。

Mixtral-39B-141B 和 Mixtral-13B-42B [78]

• 

  模型： 他们的模型在开源 Apache 2.0 许可证下授权。

• 

  数据： 不可用。

• 

  代码： 不可用。

• 

  日志： 不可用。

DBRX-36B-132B [40]

• 

  模型： 该模型根据自定义的非开源许可证许可，并附带其他用例限制。

• 

  数据： 不可用。

• 

  代码： 他们使用其公开库 LLM-foundry、composer 和 megablocks 的闭源自定义改编版本。

• 

  日志： 不可用。

Skywork-MoE-22B-146B [190]

• 

  模型： 该模型根据自定义的非开源许可证许可。

• 

  数据： 不可用。

• 

  代码： 不可用。

• 

  日志： 不可用。

DeepSeekV2-21B-236B [43] 和 DeepSeekMoE-3B-14B [39]

• 

  模型： 这些模型根据自定义的非开源许可证许可。

• 

  数据： 不可用。

• 

  代码： 不可用。

• 

  日志： 不可用。

Arctic-17B-480B [160]

• 

  模型： 该模型根据开源的 Apache 2.0 许可证许可。

• 

  数据： 他们描述了他们的混合物，但没有发布它。

• 

  代码： 不可用。

• 

  日志： 不可用。

Qwen2-14B-57B [178]

• 

  模型： 该模型在开源 Apache 2.0 许可证下授权。

• 

  数据： 不可用。

• 

  代码： 不可用。

• 

  日志： 不可用。

Jamba-12B-52B [96]

• 

  模型： 该模型在开源 Apache 2.0 许可证下授权。

• 

  数据： 不可用。

• 

  代码： 不可用。

• 

  日志： 不可用。

Qwen1.5-3B-14B [178]

• 

  模型： 该模型在自定义非开源许可证下授权。

• 

  数据： 不可用。

• 

  代码： 不可用。

• 

  日志： 不可用。

JetMoE-2B-9B [156]

• 

  模型： 该模型采用开源 Apache 2.0 许可。

• 

  数据： 他们描述了他们的混合模型，但没有发布它。

• 

  代码： 他们公开发布了他们对 megablocks 的修改版本， 但是他们的 Megatron-LM 训练代码不可用。

• 

  日志： 不可用。

OpenMoE-2B-9B [198]

• 

  模型： 该模型采用开源 Apache 2.0 许可。

• 

  数据： 他们提供了用于重建其数据的脚本。

• 

  代码： 他们公开了他们的代码。

• 

  日志： 不可用。

OLMoE-1B-7B

• 

  模型： 该模型采用开源 Apache 2.0 许可。

• 

  数据： 该数据采用开源 ODC-By 1.0 许可。

• 

  代码： 该代码采用开源 Apache 2.0 许可。

• 

  日志： 日志与代码具有相同的开源许可证（Apache 2.0）。

附录 E 附加评估

图 24： OLMoE-1B-7B 在训练过程中的损失。 Books、Reddit 和 Stack [83] 数据集来自 Dolma 1.7 [161]，通过 Paloma [110] 获取。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-OLMoE-1B-7B--Vmlldzo4OTcyMjU3

图 25： 预训练期间 OLMoE-1B-7B 和当前最佳 OLMo 模型的评估。 灰色垂直线对应于各个运行进入退火的位置，第一条线对应 OLMo-7B，第二条线对应 OLMo-1B，第三条线对应 OLMoE-1B-7B。 Figure 3 是该图的版本，其中训练 FLOPs 作为 x 轴。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-OLMoE-1B-7B-vs-OLMo-7B-vs-OLMo-1B--Vmlldzo4OTcyMjEz

Model	ARC_C	ARC_E	BoolQ	CSQA	HSwag	MMLU	OBQA	PIQA	SIQA	WinoG	Avg
LMs with $\sim$7-9B active parameters
Mistral-7B	${78.6}^{†}$	${90.8}^{†}$	89.3	${72.4}^{†}$	83.0	${64.0}^{†}$	${80.6}^{†}$	82.8	${71.3}^{†}$	77.9	79.1
OLMo-7B (0724)	${68.0}^{†}$	${85.7}^{†}$	85.3	${85.4}^{†}$	80.5	${54.9}^{†}$	${67.6}^{†}$	79.3	${76.1}^{†}$	73.2	75.6
DCLM-7B	${79.8}^{†}$	${92.3}^{†}$	87.0	77.0	82.3	${64.4}^{†}$	${79.6}^{†}$	80.1	${71.2}^{†}$	77.3	79.1
Llama2-7B	54.2	84.0	86.1	74.2	78.9	${46.2}^{†}$	57.8	77.5	59.6	71.7	69.0
Llama3.1-8B	${79.5}^{†}$	${91.7}^{†}$	88.5	${74.3}^{†}$	81.6	${66.9}^{†}$	${78.6}^{†}$	81.1	${71.4}^{†}$	76.6	79.0
Gemma2-9B	${89.5}^{†}$	${95.5}^{†}$	89.4	${78.8}^{†}$	${87.3}^{†}$	${70.6}^{†}$	${88.4}^{†}$	${86.1}^{†}$	${76.0}^{†}$	78.8	84.0
LMs with $\sim$2-3B active parameters
StableLM-2B	${50.6}^{†}$	75.3	82.3	${70.4}^{†}$	70.3	${40.4}^{†}$	${56.6}^{†}$	75.6	${64.3}^{†}$	65.8	65.1
Gemma2-3B	${67.5}^{†}$	${84.3}^{†}$	83.6	${66.4}^{†}$	74.6	${53.3}^{†}$	${68.8}^{†}$	78.5	${64.7}^{†}$	71.8	71.4
JetMoE-2B-9B	${61.4}^{†}$	${81.9}^{†}$	85.7	${75.3}^{†}$	81.7	${49.1}^{†}$	${68.0}^{†}$	80.3	${71.3}^{†}$	70.7	72.5
OpenMoE-3B-9B	29.3	50.6	63.2	21.5	44.4	27.4	34.6	63.3	42.9	${51.9}^{†}$	42.9
DeepSeek-3B-16B	53.4	82.7	81.9	72.7	80.4	${45.5}^{†}$	58.4	80.1	59.9	73.2	68.8
Qwen1.5-3B-14B	${77.4}^{†}$	${91.6}^{†}$	85.0	${81.4}^{†}$	80.0	${62.4}^{†}$	${80.6}^{†}$	81.0	${74.1}^{†}$	72.3	78.6
LMs with $\sim$1B active parameters
OLMo-1B (0724)	36.4	53.5	66.8	42.4	67.5	32.1	44.2	74.0	45.2	62.9	52.5
TinyLlama-1B	38.1	69.5	63.6	61.1	60.8	33.6	45.0	71.7	50.4	60.1	55.4
Pythia-1B	31.4	63.4	${56.8}^{†}$	50.9	48.0	31.1	40.4	68.9	46.4	52.7	49.0
DCLM-1B	${57.6}^{†}$	79.5	80.9	71.3	75.1	${48.5}^{†}$	${60.0}^{†}$	76.6	${60.5}^{†}$	68.1	67.8
OLMoE-1B-7B	${62.1}^{†}$	84.2	79.2	72.9	80.0	${54.1}^{†}$	${65.4}^{†}$	79.8	${63.0}^{†}$	70.2	71.1

表 12： OLMES 上的更多结果。 ^† 表示使用 MCF 分数，见 Appendix C。

OLMoE-1B-7B checkpoint ($\to$)	step 1,200,000	step 1,220,000	annealed	OLMo-1B	OLMo-7B
AGI Eval LSAT-AR∗	24.3	26.5	28.7	28.3	28.3
AGI Eval LSAT-LR	40.2	38.6	37.3	30.2	42.9
AGI Eval LSAT-RC	47.4	43.7	46.6	23.5	61.6
AGI Eval SAT-En	55.3	54.9	52.9	28.2	73.8
AGI Eval SAT-Math CoT	5.5	4.1	6.4	1.8	6.8
AQuA CoT	2.4	2.9	2.0	2.9	6.1
ARC Challenge∗	53.3	53.4	53.8	34.6	48.1
ARC Easy∗	77.1	78.5	77.7	64.4	75.9
BBQ	49.8	48.3	50.6	45.8	67.2
BigBench CS Algorithms∗	47.1	50.2	47.2	47.5	53.6
BigBench Conceptual Combinations	51.5	50.5	56.3	31.1	68.0
BigBench Conlang Translation	3.7	6.1	7.3	4.3	7.3
BigBench Dyck Languages∗	19.3	15.9	21.5	26.6	22.2
BigBench Elementary Math QA	26.2	27.0	26.9	26.2	30.4
BigBench Language Identification∗	31.9	34.0	31.0	27.0	39.1
BigBench Logical Deduction	26.6	25.3	24.6	23.6	27.3
BigBench Misconceptions	59.8	55.3	62.6	55.7	58.0
BigBench Novel Concepts	62.5	62.5	65.6	43.8	53.1
BigBench Operators∗	36.2	34.3	33.8	23.8	45.2
BigBench QA Wikidata∗	68.2	68.8	69.2	67.0	69.9
BigBench Repeat Copy Logic∗	15.6	15.6	18.8	3.1	9.4
BigBench Strange Stories	66.7	68.4	69.5	53.4	66.1
BigBench Strategy QA	56.2	58.1	57.0	51.5	68.6
BigBench Understanding Fables	47.1	44.4	47.6	28.0	61.4
BoolQ∗	73.3	72.8	73.2	63.7	83.9
COPA∗	81.0	80.0	78.0	75.0	77.0
CoQA∗	43.7	44.4	43.7	3.4	45.4
CommonsenseQA∗	67.2	67.0	69.3	19.6	86.0
Enterprise PII Classification	52.3	53.7	52.2	57.3	50.6
GPQA Diamond	22.2	21.2	19.7	19.7	20.2
GPQA Main	24.8	22.3	22.5	20.3	23.0
GSM8K CoT	6.4	7.4	7.4	4.9	30.6
HellaSwag 0-shot∗	76.0	76.0	77.0	65.8	76.7
HellaSwag 10-shot∗	77.6	77.5	78.6	66.3	78.9
Jeopardy∗	48.8	48.7	50.3	22.6	46.5
LAMBADA∗	72.7	72.2	73.3	61.1	71.8
LogiQA	34.9	34.3	34.6	28.7	31.0
MMLU Few-shot	52.2	51.9	53.3	28.4	55.1
MMLU Zero-shot	41.6	42.7	43.3	26.2	50.0
Math QA	26.4	27.1	27.5	24.1	29.8
OpenBookQA∗	41.4	44.0	44.8	36.6	43.4
PIQA∗	81.3	81.2	82.0	76.4	81.7
PubMedQA	56.1	46.6	57.9	0.2	57.9
SQuAD∗	52.9	52.4	52.4	0.0	65.5
SVAMP CoT	30.0	28.0	33.0	14.3	44.7
Simple Arithmetic, no spaces	17.6	18.1	20.1	1.2	15.3
Simple Arithmetic, with spaces	19.5	20.6	22.1	1.8	16.0
Social IQA	71.5	70.7	69.3	69.5	84.4
Trivia QA	54.2	53.0	55.9	25.1	51.8
Winogender Female	50.0	46.7	50.0	41.7	58.3
Winogender Male	55.0	58.3	60.0	63.3	58.3
Winograd∗	82.8	83.2	84.6	79.9	83.2
Winogrande∗	68.0	68.5	69.0	61.8	67.6
Core	46.3	46.5	47.2	30.2	49.8
Extended	31.3	30.9	32.5	16.9	37.0

表 13： DCLM 在核心和扩展任务子集 [89] 上的评估指标。 ^∗=核心任务。 “退火”是我们用于 OLMoE-1B-7B 的最终预训练检查点，它是从第 1,200,000 步的检查点退火而来的。 我们让非退火预训练运行训练时间更长，从而产生了 1,220,000 检查点。

附录 F 额外实验

图 26： 将 Reddit 或 FLAN 添加到 OLMoE-Mix 中。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-Adding-Reddit-FLAN--Vmlldzo4OTg1NTg4

将 Reddit 或 FLAN 添加到 OLMoE-Mix 中

在 Figure 26 中，我们对将 Reddit 或 FLAN [189] 的 Dolma 1.7 [161] 子集添加到我们的预训练数据混合中进行了基准测试（§2）。 总体而言，我们没有发现任何一种能够带来一致的收益，因此我们在最终数据混合中不使用它们。

图 27： 负载均衡精度。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-FP32-LBL--Vmlldzo4NDMxNDA4

负载均衡精度

Fedus 等人 [56] 选择性地执行与路由相关的操作以全精度 (FP32) 进行，以提高稳定性。 在 Figure 27 中，我们测试了以全精度计算负载均衡损失是否可以提高稳定性，但没有发现它能减少峰值。 因此，我们坚持使用 bfloat16 (BF16)。

图 28： 在升级后的检查点中添加噪声。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-Noise-upcycle---Vmlldzo4NDA3MzI2

噪声升级

为了创建 Qwen2-MoE [200, 178, 13]，作者在继续训练升级后的设置 [84]之前，在前馈网络中添加了 50% 的高斯噪声。 Komatsuzaki 等人 [84] 还报告说，他们尝试添加噪声，但没有发现它有益。 在 Figure 28 中，我们通过随机用从标准差为 0.02 的正态分布中抽取的数字替换每个 MLP 的 50% 来试验常规升级与添加噪声。 我们发现，在 7000 亿个符元之后，无噪声变体仍然表现略好，但两者似乎都收敛到相同的性能。 如果进一步训练，噪声变体最终可能优于无噪声变体，但在那一点上，从头开始训练 MoE 可能会更有意义 (§4.1.5)。

图 29： 在各层之间共享同一个 MoE 与常规密集 LM。 MoE 中的专家数量与其层数相同。 因此，由于 MoE 在各层之间共享，它与密集模型具有相同数量的总参数和活动参数。 更多结果、日志和配置：https://wandb.ai/ai2-llm/olmoe/reports/Plot-Shared-vs-Dense--Vmlldzo4NDI0MTc5

共享层

一些工作在 Universal Transformers 的背景下研究了权重在层之间共享的混合专家 [169, 37, 45]。 我们在 Figure 29 中测试了层共享混合专家是否能够胜过非共享密集模型。 层共享 MoE 使用了一种负载均衡损失，它是在模型级别而不是层级别应用的。 这通过允许模型完全停用某些层的特定专家，甚至通过始终为每一层激活一个独立的专家来模拟密集模型，从而为模型提供了更大的灵活性。 这使得它成为密集模型的推广，我们的假设是它可能比密集模型表现得更好。 然而，在实践中，我们发现两者都与常规的密集模型表现相似，甚至在验证损失和 HellaSwag 上保持着微小的优势。 层共享 MoE 的一个可能的优势是它们可以实现更好的推理负载均衡。 如果提示不断出现，那么新传入的提示可以与已经通过几个层的先前提示一起批处理，并一起发送到 MoE 模块，因为 MoE 模块无论是在第一层还是最后一层都是一样的。 共享还将训练期间的吞吐量降低了约 20%，这进一步促使我们决定不将其用于 OLMoE-1B-7B。

KTO 实验

在 Table 14 中，我们对用于 KTO [54] 的步骤数量（5,000 vs. 10,000）和优化器（Adam [82] vs. RMS）进行了实验。 基于这些实验，我们在 §4.3 中使用 RMS 优化器和 5,000 步的检查点。

				Human-	Alpaca-
Task ($\to$)	MMLU	GSM8k	BBH	Eval	Eval 1.0	XSTest	IFEval	Avg
Setup ($\to$)	0-shot	8-shot CoT	0-shot	0-shot	0-shot	0-shot	0-shot	0-shot
Metric ($\to$)	EM	EM	EM	Pass@10	%win	F1	Loose Acc
KTO, 5,000 steps, RMS	51.2	45.5	34.1	57.1	81.6	86.6	47.5	57.7
KTO, 10,000 steps, RMS	51.0	41.0	34.7	53.8	81.0	62.3	47.5	54.2
KTO, 5,000 steps, Adam	51.2	42.0	35.3	55.6	81.0	84.5	46.6	56.0
KTO, 10,000 steps, Adam	51.0	43.0	34.1	54.9	79.7	62.7	47.5	53.3

表 14: KTO 自适应实验。 5,000 和 10,000 步分别对应于我们自适应数据集（§2）的 1.3 和 2.6 个 epoch。

附录 G 额外分析

图 30: 考虑所有 8 个被激活的专家时，OLMoE-1B-7B 的词汇专业化。 等效于 Equation 8 中的 $k=8$。

图 31: 考虑所有 2 个被激活的专家时，Mixtral-8x7B 的词汇专业化。 等效于 Equation 8 中的 $k=2$。

图 32: OLMoE-1B-7B (顶部) 和 Mixtral-8x7B (底部) 跨领域的词汇专业化。 我们可视化了符元 ID 被路由到特定专家的频率。 我们只包含在各个语料库中至少出现 8 次的 ID。 垂直灰色线对应于均匀路由 (OLMoE-1B-7B 为 8/64=12.5%，因为它有 64 个专家，其中 8 个被激活；Mixtral 为 2/8=25%，因为它有 8 个专家，其中 2 个被激活)。 例如，在所有被路由到 OLMoE-1B-7B 专家 0 且至少出现 8 次的 GitHub 符元 ID 中，$\sim$40% 的它们以 $\sim$100% 的概率被路由到专家 0 (左上角)，表明专家 0 专注于这些符元 ID。 对于 OLMoE-1B-7B，路由概率极值 (0% 或 100%) 的频率很高，表明这些专家专门关注某些符元 ID，尤其是 特定领域 (§5.3)，如 GitHub 和 arXiv。

图 33: 适应后选择性层中的负载不平衡。 我们可视化了来自我们的指令调优数据集 (§2) 的符元被路由到 64 个专家中的 8 个活动专家 ($k=1$ 在 Equation 7 中) 的频率。 水平灰色线对应均匀路由（8/64=每个专家12.5%）。 虽然我们在没有损失平衡损失的情况下运行 SFT 和 DPO (§4.3)，但我们观察到负载分布没有发生实质性变化。

图 34： OLMoE-1B-7B（顶部）与 Mixtral-8x7B（底部）的领域专业化，对应于顶级-$1$ 路由的专家。 我们可视化了预训练结束时来自不同领域的符元被路由到 64 (OLMoE) 或 8（Mixtral）个专家的频率。 与 Figure 22 不同，这里我们只考虑路由到顶级 1 个专家的符元 ($k=1$ in Equation 7)。 水平灰色线对应均匀路由（每个专家 1/64=1.56% 用于 OLMoE-1B-7B 和 1/8=12.5% 用于 Mixtral）。

图 35： OLMoE-1B-7B 跨层的符元路由。 我们可视化了来自不同领域的符元在顶级 1 路由下跨层被路由到一对专家的频率，对应于 Figure 34。 每个矩形的尺寸与专家接收的符元总数成正比，而两个专家之间的流动显示了路由到两个专家的符元比例。 我们只显示接收符元超过随机机会 50% 的专家，并对更大的流动使用更强的颜色。 我们观察到专家对之间跨层协调的一些实例，例如，第 7 层的专家 27 和第 15 层的专家 57 共同处理了大量的维基百科符元。 在此可视化中，层 0 $\to$ 7 和 7 $\to$ 15 之间的流是独立的。

图 36： Mixtral-8x7B 符元在层之间的路由。 我们可视化了来自不同领域的符元在顶级路由下被路由到一对专家中的频率，对应于 Figure 34。 每个矩形的尺寸与专家接收的符元总数成正比，而两个专家之间的流动显示了路由到两个专家的符元比例。 在此可视化中，第 0 $\to$ 7 层和第 7 $\to$ 15 层之间的流动是独立的。

附录 H 局限性和未来工作

我们重点介绍了此版本 OLMoE-1B-7B 的四个主要局限性。 我们期待在未来的 OLMoE 版本中解决这些问题。

更多参数

OLMoE-1B-7B 拥有 7B 个参数，其中 1B 个参数为每个输入符元激活。 这个小尺寸使得 OLMoE-1B-7B 使用起来非常便宜，然而我们在本工作中证明了它优于更昂贵的模型 (Figure 1)。 但是，对每个输入符元仅使用 1B 个参数也限制了 OLMoE-1B-7B 的能力，如其与使用 $>$ 7$\times$ 个更多参数的模型相比的表现所见，例如 §3 中的 Llama3.1-8B。 虽然可能不需要更多参数来匹配 8B 模型及以上 [80]，但从短期来看，添加参数是提高 OLMoE 性能的一种简单方法，至少允许模型利用每个输入超过 1B 个参数，这可以通过递归 [45] 或代理工作流程 [185, 201] 实现。 同样，更改参数分配到例如词汇参数与非词汇参数是另一个改进方法 [170]。

更多数据

我们训练 OLMoE-1B-7B 5 万亿符元，然而，一些最近的稠密模型训练时间明显更长，例如 Llama 3 训练 15 万亿符元 [50]。 据我们所知，还没有像 OLMoE-1B-7B 那样经过过度训练 [57] 的大型 MoE。 特别是，采用 OLMoE-1B-7B 的活动参数，我们的符元乘数 [57] 约为 5,000（5T / 1B）。 训练更长时间可能会有益处，但 MoE 的过度训练在多大程度上有效，以及它与密集模型有何不同，仍需要更多研究 [7]。

多模态

OLMoE-1B-7B 是一种纯文本的大语言模型，因此它无法接收或生成图像或音频等其他模态的输入或输出。 这限制了它在这些模型的各种多模态用例中的实用性 [74, 165, 27, 81, 118, 134, 14, 47, 50]。 早期已经有一些关于开放式多模态 MoE 的研究 [124, 97, 94, 155, 111, 193]，我们期待着使未来的 OLMoE 版本成为其中的一部分。

多语言

我们在主要以英语为主的语料库上预训练了 OLMoE-1B-7B，并在英语任务上进行评估。 这可能会严重限制我们模型在非英语语言模型研究中的实用性 [107, 158, 222, 53, 163, 196]。 虽然已经有一些关于训练特定语言的 LM 的研究 [109, 55]，但更有可能的是，由于数据限制 [120]，当我们添加更多数据来构建未来的 OLMoE 版本时，我们将混合更多非英语数据。 这可能使未来的 OLMoE 模型在非英语语言中表现得更好。