OLMoE: 开源的MoE语言模型(预训练&效果)

一、结论写在前面

论文标题：OLMoE: Open Mixture-of-Experts Language Models

论文链接：​​https://arxiv.org/pdf/2409.02060​​

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​​

论文开源了OLMoE-1B-7B和OLMoE-1B-7B-INSTRUCT，包括模型、数据、代码和日志。OLMOE-1B-7B拥有7B参数，但每个输入token仅使用1B参数。论文在5T token上对其进行预训练，并进一步Adaptation以创建OLMoE-1B-7B-INSTRUCT。

论文的模型在所有具有相似活跃参数的可用模型中表现最佳，甚至超越了Llama2-13B-Chat和DeepSeekMoE-16B等更大模型。

论文分享了各种训练实验，并定义和分析了路由器饱和度、专家协同激活、领域和词汇专业化。通过论文的完全开源发布，论文希望帮助该领域构建更好的专家混合模型。论文对OLMoE的新迭代感到兴奋，以缩小前沿模型与完全开源模型之间的差距。

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

二、论文的简单介绍

2.1 论文的背景

高性能的语言模型对许多学者和开源开发者来说是很难接触，因为它们的构建和部署成本过高。改善成本性能权衡的一种方法是使用稀疏激活的专家混合(Mixture-of-Experts，MoEs)模型。MoEs在每一层都有多个专家，每次只激活其中的一个子集(见图2)。这使得MoEs比具有相似总参数数量的密集模型更加高效，因为密集模型需要为每个输入激活所有参数。出于这个原因，行业前沿模型使用MoEs，包括Gemini-1.5和据报道的GPT-4。

然而，大多数MoE模型都是闭源的:虽然有些模型公开发布了模型权重，但它们提供的关于训练数据、代码或配方的信息有限或没有(见图1)。虽然之前有一些努力使语言建模研究完全可访问，但它们主要局限于密集LMs。这种情况出现的原因是，MoEs需要更多的开放性，因为它们为LMs增加了复杂的新设计问题，例如使用多少总参数与活跃参数，是使用多个小专家还是少数大专家，专家是否应该共享，以及使用什么路由算法。缺乏关于这些细节的开放资源和发现，阻碍了该领域构建接近闭源前沿模型能力的高效开放MoEs。

为了解决这些问题，论文引入了OLMOE，一个完全开放的专家混合语言模型，在同等规模的模型中具有最先进的性能。具体来说，论文预训练了OLMOE-1B-7B，总共训练了5.1T token，总参数为6.9B，其中只有1.3B参数为每个输入token激活。这导致推理成本与使用约1B参数的密集模型(如OLMo 1B或TinyLlama 1B)相似，但需要更多GPU内存来存储其7B总参数。

论文的实验表明，MoEs的训练速度比具有相同活跃参数数量的密集LMs快约2倍。在图1中，论文展示了OLMOE-1B-7B显著优于所有开放的10亿参数模型，并展现出与推理成本和内存存储显著更高的密集模型相当的性能(例如，在MMLU评分上与Llama2-13B相似，后者的成本约为前者的10倍)。通过指令和偏好微调，论文创建了OLMOE-1B-7B-INSTRUCT，论文发现它在常见基准测试(MMLU、GSM8k、HumanEval等)上超过了各种更大的指令模型，包括Llama2-13B-Chat、OLMo-7B-Instruct (0724)和DeepSeekMoE-16B。

论文全面的对照实验突出了MoEs(见表1)和LMs的关键设计选择。使MoEs性能出色的一个关键设计决策是使用细粒度路由和细粒度专家:论文在每层使用64个小专家，其中8个被激活。路由算法的选择也很重要:论文发现无丢弃的基于token的路由优于基于专家的路由。论文的发现还包括一些挑战先前工作的结果，例如共享专家的无效性，以及将预训练的密集LM稀疏地升级为MoE的有限益处，除非在小规模计算预算下。最后，论文分析了OLMOE-1B-7B的路由行为，发现路由在预训练早期就饱和了，专家很少被共同激活，并且专家表现出领域和词汇专门化。    

图2：密集型语言模型（LMs）与OLMoE等MoE模型的架构对比。图中省略了一些细节，例如，OLMOE-1B-7B还使用了QK-Norm。

表1: MoE的关键设计选择及基于论文实验的OLMOE-1B-7B设置    

表2：OLMOE-1B-7B预训练数据的组成。StarCoder、peS2o和Wikipedia部分来自Dolma 1.7 。

表3：OLMOE-13.7的Adaptation训练数据。

2.2 预训练和Adaptation(Pretraining and Adaptation)

预训练架构 OLMOE是一个仅有解码器的语言模型，由NL个transformer层组成。在像OLMo这样的密集模型中的前馈网络(FFN)被替换为一个MoE模块，该模块由NE个较小的FFN模块组成，这些模块被称为专家。对于每个处理的输入token x，会激活其中的k个专家子集(另见图2):

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

总结来说，论文使用了总共6.9B参数中的1.3B活跃参数，每层激活了8个专家，总共64个专家。论文采用了无丢弃的token选择路由[58]：对于每个输入token，学习到的路由网络决定由8个专家来处理。论文从零开始训练OLMOE-1B-7B，使用了两个辅助损失：负载均衡损失( L_LB) 和路由z-损失( L_RZ) 。论文将它们分别乘以各自的损失权重alpha和beta，并与交叉熵损失(L_CE )线性相加，得到最终的训练损失：

OLMOE-1B-7B的完整预训练配置在附录B中。

预训练数据 论文使用了来自DCLM [89]和Dolma 1.7 [161]的混合数据，包括以下内容：

(1) 经过质量过滤的Common Crawl子集，称为DCLM-Baseline;

(2) 在DCLM和Dolma 1. 7中都使用的StarCoder、Algebraic Stack和arXiv;

(3) 来自Dolma 1.7的peS2o和Wikipedia。论文将预训练数据集称为OLMOE-MIX。

对于上述所有来源，论文应用了一个过滤器，移除所有包含32个或更多重复n-gram的文档，其中n-gram是任意1到13个token的跨度。对于StarCoder子集，论文还移除了任何来自GitHub上少于2颗星的仓库的文档，或者其最频繁词占据文档超过30%的内容，或者其前两个最频繁词占据文档超过50%的内容。

论文在每个epoch开始时随机打乱所有样本，并总共训练5.133T个token（根据Muennighoff等人[120]的方法，训练1.3个epoch）。在论文的退火阶段（最后100B个token），论文首先重新打乱整个数据集，然后按照先前的工作64， 89，将学习率线性衰减到0。论文的预训练数据统计如表2所示。

Adaptation（Adaptation） 论文通过遵循标准的Adaptation配方创建了OLMoE-1B-7B-INSTRUCT，该配方分为指令微调和偏好微调，基于先前的开源模型。在论文的指令微调数据集中，论文增加了更多的代码和数学数据，以提升下游编码和数学应用的性能。其他模型，如GPT-4和Llama 3，在预训练期间包含了来自GSM8k或MATH等数学数据集的样本。论文还包含了No Robots和Daring Anteater的一个子集，因为它们质量高且增加了多样性，这是成功Adaptation的两个关键因素。论文在表3中描述了论文的Adaptation数据集。

2.3 结果

论文的评估程序包括三个部分：预训练期间、预训练后和Adaptation后(After adaptation)。    

图3: OLMOE-1B-7B和当前最佳OLMo模型在预训练期间的评估。OLMOE-1B-7B与OLMo模型在MoE架构、几个训练超参数和训练数据集方面有所不同。更多结果、日志和配置:​​https://wandb.ai/ai2-llm/olmoe/reports/Plot-OLMoE-1B-7B-vs-OLMo-7B-vs-OLMo-1B--Vmlldzo4OTcyMjEz​​

预训练期间(During pretraining) 在图3中，论文对比了OLMOE-1B-7B在预训练期间与当前最佳OLMo模型在常用下游任务上的性能表现。论文发现，在所有任务中，OLMOE-1B-7B都以更少的计算量(FLOPs)达到了比密集OLMo模型更好的性能。

尽管OLMOE-1B-7B在训练中使用的FLOPs不到一半，且只使用了1B的活跃参数，但在训练结束时，它的表现与OLMo-7B相当或更优。这可能是由于论文对OLMo设置进行的数据集和建模变更，包括MoE相关的改变、稳定性和性能改进，这些在附录B中有详细说明。    

表4: 预训练后的OLMOE-1B-7B与更大的MoEs和密集LMs的比较。论文将其与在活跃参数(1B，近似速度和成本)或总参数(7B，近似内存需求)上与OLMOE-1B-7B相似的密集LMs进行比较。模型名称包含四舍五入的参数计数:MoEs的格式为model-active-total，密集LMs的格式为model-total(这导致与官方名称有些差异，例如，虽然被称为"Gemma2-2B"，但它实际上有26亿活跃和总参数[175])。Chall. = Challenge(挑战)。论文自己使用5个少样本示例进行所有评估

预训练后(After pretraining) 在表4中，论文对OLMoE-1B-7B在常见的下游任务上进行了基准测试。论文发现，在使用的活跃参数少于2B的模型中，OLMoE-1B-7B表现最佳，使其成为许多语言模型使用场景中最经济的选择。对预算更大的情况，Qwenl.5 3B-14B的性能更强，但其活跃参数和总参数数量是OLMoE-1B-7B的两倍以上。论文发现，尽管每次前向传递所需的计算量减少了6-7倍，OLMoE-1B-7B仍优于一些具有7B参数的密集模型，如Llama2-7B ，但在其他模型如Llama3.1-8B面前稍显不足。图1比较了OLMoE-1B-7B与其他语言模型在MMLU性能与活跃参数（作为模型成本的代理）上的表现。OLMoE-1B-7B在其成本范围内处于领先地位。

Adaptation后(After adaptation) 在表5中，论文基准测试了OLMoE-1B-7B的指令（SFT）和偏好（DPO）调优。SFT在所有测量的任务上提升了论文的模型性能。论文在GSM8k上观察到>10×的增益，这可能归因于论文在预训练期间加入了额外的数学数据，以弥补相对较少的数学数据（S2）。DPO在大多数任务上都有帮助，尤其是在AlpacaEval上，这与先前的工作[186，75，121]中的发现一致。    

论文的DPO模型，论文称之为OLMoE-1B-7B-INSTRUCT在所有基准模型中具有最高的平均分。论文发现它优于Qwen1.5-3B-14B的聊天版本，尽管Qwen的参数量超过2倍，并且在表4中其预训练模型优于OLMoE-1B-7B。在AlpacaEval上的84%得分也优于排行榜上更大的密集模型，例如Llama2-13B-Chat。

表5：OLMOE-1B-7BAdaptation后与其他模型的比较。论文发现JetMoE聊天模型（​​​https://hf.co/jetmoe/jetmoe-8b-chat​​）存在问题，导致随机评分，因此论文将其排除。名称包含四舍五入的参数计数：模型-活跃-总参数（对于MoEs）和模型-总参数（对于密集LMs）。论文自行运行所有评估。模型使用不同的Adaptation混合，例如，OLMoE是在用于OLMo模型的改进管道上训练的。

本文转载自​​AI帝国​​，作者： 无影寺 ​​​​