Lory: Fully Differentiable Mixture-of-Experts for Autoregressive Language Model Pre-training

混合专家（MoE）模型有助于实现高效的扩展；然而，训练路由器网络引入了优化非可微分、离散目标的挑战。最近，提出了一种完全可微分的MoE架构SMEAR（Muqeeth等人，2023年），它在参数空间中软合并专家；然而，它的有效性仅在分类任务的下游微调中得到了证明。在本文中，我们提出了Lory，这是第一种将这种架构扩展到自回归语言模型预训练的方法。Lory引入了两种关键技术：（1）一种因果分段路由策略，可以在保持语言模型自回归性质的同时实现高效的专家合并操作；（2）一种基于相似性的数据批处理方法，通过将相似的文档分组来鼓励专家专业化。我们从头开始预训练了一系列Lory模型，使用了150B个标记，最多有32个专家和30B（1.5B活动）参数。实验结果显示，在困惑度（+13.9%）和各种下游任务（+1.5%-11.1%）上，与参数匹配的密集模型相比，Lory模型取得了显著的性能提升。尽管是在段级路由的情况下，Lory模型的性能与令牌级路由的最先进的MoE模型相比具有竞争力。我们进一步证明，Lory中训练的专家在没有监督的情况下捕捉到了领域级别的专业化。我们的工作突显了完全可微分的MoE架构在语言模型预训练中的潜力，并提倡在这个领域进行未来的研究。

本论文旨在解决Mixture-of-experts（MoE）模型在训练路由网络时的优化问题，提出了一种可扩展的、全可微分的MoE架构Lory，用于自回归语言模型预训练。

Lory引入了两种关键技术：因果分段路由策略和基于相似性的数据分批方法，以提高专家合并操作的效率并鼓励专家专业化，同时保持语言模型的自回归性质。

Lory模型在150B令牌的预训练中获得了显著的性能提升，包括困惑度和下游任务的提高。Lory模型的专家捕捉到了领域级别的专业化，而无需监督。此外，本文提出的全可微分MoE架构为语言模型预训练提供了潜在的可能性。

最近的相关研究包括Muqeeth等人提出的可全可微分MoE架构SMEAR，以及其他基于MoE的语言模型预训练方法，如Gshard和Switch Transformers。