探索OLMoE：首个100%开源的混合专家模型，拥有7B参数与1B稀疏激活的高效推理成本

训练代码、中间checkpoint、训练日志和训练数据都已经开源。尽管大语言模型（LM）在各类任务中取得了显著进展，但在训练和推理方面，性能与成本之间的权衡仍然是一个问题。对于很多学者和开发者而言，高性能的语言模型由于高昂的成本而难以触及。一种改善成本-性能的方法是采用稀疏激活混合专家（MoE）。MoE在每一层都有多个专家，每次只激活其中的一部分，从而使得MoE模型比密集模型更高效。因此，许多前沿模型如Gemini-1.5和GPT-4等都采用了MoE。

然而，大多数MoE模型都是闭源的，即便有些模型公开了权重，但对于训练数据和代码等信息却很少公开。这种缺乏开放资源的情况阻碍了低成本开源模型的构建，使其难以达到闭源模型的水平。为解决这一问题，来自艾伦人工智能研究院和Contextual AI等机构的研究者推出了OLMoE，这是一种完全开源的混合专家语言模型，具有卓越的性能。

研究团队使用5.1万亿个token预训练了OLMoE-1B-7B模型，该模型共有69亿参数，每个输入token只激活13亿参数。这使得OLMoE-1B-7B与具有大约1B参数的密集模型（如OLMo 1B或TinyLlama 1B）具有相似的推理成本，但需要更多的GPU内存来存储7B的总参数。实验表明，MoE的训练速度比具有相同激活参数数量的密集LM快约2倍。OLMoE-1B-7B在多个基准测试中表现出色，超越了其他开源1B模型。

通过指令和偏好调优，研究团队还创建了OLMoE-1B-7B-INSTRUCT，该模型在MMLU、GSM8k、HumanEval等多个基准上超越了包括Llama2-13B-Chat在内的各种更大指令模型。受控实验显示，MoE的关键设计决策之一是使用细粒度路由和小专家：每层使用64个小专家，每次激活8个。此外，路由算法的选择也很重要，研究发现基于token的无丢弃路由优于基于专家的路由。

最后，研究团队分析了OLMoE-1B-7B中的路由行为，发现路由在预训练早期就饱和了，专家很少被共同激活，并且表现出领域和词汇的专业化。通过这些努力，研究者们希望这个完全开源的MoE能够推动更多研究和分析，提高对这些模型的理解。

训练代码、中间检查点（每5000步）、训练日志和训练数据都已经开源。研究者表示，OLMoE是第一个100%开源的混合专家LLM。预训练与自适应预训练架构OLMoE由N_L个transformer层组成，仅包含解码器。对于OLMo这样的密集模型，原本模型中的单一前馈网络被N_E个小前馈网络（专家）组成的混合专家网络替代。对于每个输入token x，只有k个专家被选中并激活，负责处理这个输入。路由器r是一个经过训练的线性层，将输入映射到k个专家上。对路由器的输出应用softmax函数，计算N_E个专家的路由概率。每个选定的专家E_i处理输入x，其输出乘以其各自的路由概率。将所有选定专家的结果相加，构成模型单个层的MoE模块输出。MoE模型的训练通常涉及将已有的大型密集模型转换成稀疏模型，即“稀疏升级”。在这个过程中，需要改变模型的训练目标，例如调整辅助负载平衡和路由器的损失函数。

论文作者使用了总计69亿参数中的13亿活跃参数，每层有64个专家，每次激活8个。他们使用了一种名为“无丢弃token”的路由方法：对于每个输入token，路由器网络将分配8个专家来处理它。论文作者引入了两个辅助损失函数：负载平衡损失和路由器z损失，来训练OLMoE-1B-7B。他们给这两个损失函数分别设定了权重α和β，然后将其与主要学习目标（交叉熵损失）结合，最终计算的损失函数为：

训练数据方面，论文作者使用了来自两个不同来源的数据集：DCLM和Dolma 1.7。这些数据集包括多种类型的数据，如网络爬取的数据、编程问题解答、数学问题解答和学术论文等。他们将这些数据混合起来，创建了一个名为OLMoE-MIX的新数据集。对于数据的处理，论文作者使用了过滤器去除了包含太多重复token的、GitHub上星标少于2的项目以及某些词出现频率过高的文档。他们将在每轮训练开始前随机混洗数据，总计超过5万亿个token。在“退火”阶段（最后100B个token），他们首先重新混洗整个数据集，然后按照此前OLMo论文中的方法，将学习率线性衰减到0。

自适应论文作者从指令调优和偏好调优两方面，基于之前的开放模型，构造了OLMoE-1B-7B-INSTRUCT。在指令调优集中，他们增加了更多的代码和数学数据，以提高模型在这些领域的性能。GPT-4和Llama 3在预训练阶段使用了像GSM8k或MATH这样的数学数据集的样本。按照这个思路，论文作者还添加了“No Robots”和“Daring Anteater”的一个子集。这些数据集不仅质量高还更多样，这是拓展模型适应性的两个关键因素。

实验该研究的评估程序由三部分组成：预训练期间、预训练之后和自适应之后。预训练期间：如图3所示，该研究在预训练期间使用当前最佳的OLMo模型在常用下游任务上对OLMoE-1B-7B的性能进行了基准测试。研究团队发现，在所有任务中，OLMoE-1B-7B比密集OLMo模型以更少的计算量获得了更好的性能。尽管OLMoE-1B-7B使用了不到一半的FLOP进行训练并且仅使用1B个激活参数，但OLMoE-1B-7B在训练结束时可与OLMo-7B媲美，甚至优于OLMo-7B。预训练之后：在表4中，该研究在常见的下游任务上对OLMoE-1B-7B进行了基准测试。研究发现OLMoE-1B-7B在使用少于2B个激活参数的模型中表现最好，使其成为许多LM用例中最经济的选择。如果预算较大，Qwen1.5-3B-14B具有更强的性能，但其激活参数和总参数比OLMoE-1B-7B多一倍以上。研究发现，尽管每条前向传播所需的计算量减少了约6-7倍，但OLMoE-1B-7B的性能优于一些具有7B参数的密集LM，例如Llama2-7B，但不如其他LM，例如Llama3.1-8B。上图1比较了OLMoE-1B-7B和其他LM的MMLU性能和激活参数，表明OLMoE-1B-7B是其成本范围内最先进的。自适应之后：在表5中，该研究对OLMoE-1B-7B的指令（SFT）和偏好（DPO）调优进行了基准测试。SFT在所有测量任务上都改进了本文的模型。DPO在大多数任务上都有帮助，尤其是AlpacaEval，这与先前研究的结果一致。DPO模型（称之为OLMoE-1B-7B-INSTRUCT）在所有基准测试模型中具有最高平均值。

本文来源：