150B token从头训练，普林斯顿Meta发布完全可微MoE架构Lory

新智元报道

编辑：乔杨 好困

【新智元导读】前几天，普林斯顿大学联合Meta在arXiv上发表了他们最新的研究成果——Lory模型，论文提出构建完全可微的MoE模型，是一种预训练自回归语言模型的新方法。

不同于大多数模型使用字母缩略起名，论文作者在脚注中解释道，Lory是一种羽毛有彩虹颜色的鹦鹉，和「软MoE」的精神非常相似。

论文的作者团队也可以称之为「明星阵容」。

论文地址：https://arxiv.org/abs/2405.03133

主要作者之一陈丹琦是普林斯顿大学计算机科学系的助理教授，也是普林斯顿NLP小组共同领导人之一。她本科毕业于清华大学姚班，2018年在斯坦福大学获得博士学位，导师是大名鼎鼎的Christopher Manning。

斯坦福教授、NLP领域泰斗Dan Jurafsky曾这样评价她：「她在发现重要的研究问题上很有品位。她已经对该领域产生了非凡的影响，并且她的影响只会越来越大。」

Mike Lewis是Meta AI的一名研究科学家，他领导了Meta刚发布的大语言模型Llama 3的预训练工作。

他此前曾发表过多项有影响力的研究成果，包括Bart、Roberta、top-k采样等。

本文的第一作者是普林斯顿大学五年级博士生钟泽轩，导师是陈丹琪教授。

钟泽轩硕士毕业于伊利诺伊大学香槟分校，本科毕业于北京大学计算机系，曾在Meta AI和微软亚洲研究院实习，这项研究就是他在Meta实习期间完成的。

发布后，论文作者也在推特上提供了全文解读。

引入的关键技术包含两个方面，一是用因果分段路由策略取代token级别的路由，可以在保持语言模型自回归属性的同时实现高效的专家合并。

二是提出了基于相似性的数据批处理方法，如果仅仅是把随机选择的文本拼接在一起训练会导致低水平的专家模型，而将相似的文本进行分组可以使模型更加专业化。

基于这些方法，作者使用150B token的数据从头训练了一系列的Lory模型，活跃参数有0.3B和1.5B两个级别，含有最多32个专家。

与稠密模型相比，Lory的训练过程更为高效，可以用少2.5倍的步数实现相同的损失值。

研究团队使用上下文学习的方法评估Lory的能力，发现模型在常识推理、阅读理解、闭卷问答、文本分类等下游任务上都取得了很好的效果。

可以观察到，使用更多专家可以改进模型的表现。

相比目前MoE领域的SOTA模型Expert Choice（EC），Lory模型也表现出了有竞争力的性能。

2023年12月，一家名为Mistral AI的法国创业公司发布了一款性能媲美甚至优于GPT-3.5和Llama 2 70B的模型Mixtral 8x7B。

Mixtral使用了一种稀疏的MoE网络，不仅表现出了强大的性能，而且十分高效，推理速度相比Llama 2 70B提高了6倍，于是让MoE得到了开源社区的广泛关注。

甚至有人猜测，GPT-4可能也使用了MoE技术实现了超过一万亿参数的超大模型。

对于Transformer架构的语言模型，MoE主要有两个元素：

一是使用参数更为稀疏的MoE层代替密集的前馈网络层（FFN），其中每个专家都是一个独立的神经网络，甚至可以是MoE本身，从而形成层级式的MoE结构。

二是使用门控网络或路由机制决定token被发送到哪个专家，其中token的路由机制是决定MoE模型表现的关键点。

因果分段路由

虽然MoE的这种机制有助于高效扩展模型规模，但训练路由网络的过程会引入离散化、不可微的学习目标。2023年发布的SMEAR模型就已经开始探索解决方案，使用专家合并方法构建完全可微的MoE模型。

论文地址：https://arxiv.org/abs/2306.03745

然而，SMEAR使用的方法是将所有专家进行软合并，取其加权平均值，这适用于文本分类任务，但很难应用到自回归语言模型上。

于是，作者提出了使用分段路由的方法，对每一段语句而非每个token进行专家合并，有效减少了合并操作的数量。

如果仅仅使用当前语段进行路由，很可能导致语言模型遗漏跨语段的信息，所以论文提出采用类似于自回归的因果分段路由。

在为当前语段合并专家时，需要考虑前一个语段的信息，从而决定每个专家的路由权重。

消融实验的结果也证明，与因果分段路由的策略相比，单纯使用前缀进行路由会导致语言模型性能降低。

基于相似性的数据批处理

预训练语言模型的标准做法是将数据集中的文档随机拼接在一起，构造出固定长度的训练样本。

对于MoE模型而言，这种方法存在问题，相邻段的token可能来自非常不同且毫不相关的文档，可能会损害专家模型的专业化程度。

因此，受到ICLR 2024中一篇论文的启发，作者在Lory中采用了类似的技术，依次连接相似的文档来构造训练样本，使专家模型更「专注」地研究不同的领域或主题。

论文地址：https://arxiv.org/abs/2310.10638

实验表明，无论是随机批处理还是基于相似度批处理，Lory模型的效果都优于稠密模型，但使用基于相似度的方法可以得到更大的loss提升。

参考资料：

https://huggingface.co/blog/moe

https://mistral.ai/news/mixtral-of-experts/

https://medium.com/generative-ai-insights-for-business-leaders-and/is-gpt-4-a-mixture-of-experts-model-exploring-moe-architectures-for-language-models-9fcb50b8d122

https://twitter.com/ZexuanZhong/status/1787868307066126689