Upcycling Large Language Models into Mixture of Experts

📄 arXiv: 2410.07524v2 📥 PDF

作者: Ethan He, Abhinav Khattar, Ryan Prenger, Vijay Korthikanti, Zijie Yan, Tong Liu, Shiqing Fan, Ashwath Aithal, Mohammad Shoeybi, Bryan Catanzaro

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-10-10 (更新: 2025-06-16)

💡 一句话要点

提出虚拟组初始化和权重缩放方法，高效地将大型语言模型升级为混合专家模型。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 模型升级 语言模型 稀疏模型 虚拟组初始化

📋 核心要点

  1. 现有方法难以有效扩展预训练稠密语言模型的容量，且大规模升级为混合专家模型的技术尚不明确。
  2. 提出虚拟组初始化和权重缩放方法，支持细粒度MoE架构升级，并探索了更优的专家路由策略。
  3. 实验表明，升级方法优于持续稠密模型训练，且升级后的Nemotron-4 15B模型在MMLU上性能提升。

📝 摘要（中文）

本文研究了将预训练的稠密语言模型升级为稀疏混合专家(MoE)模型的方法，这是一种提高已训练模型容量的有效途径。然而，大规模升级的最佳技术仍不明确。本文对十亿参数规模语言模型的升级方法和超参数进行了广泛研究。我们提出了一种新颖的“虚拟组”初始化方案和权重缩放方法，以实现对细粒度MoE架构的升级。通过消融实验，我们发现升级优于持续的稠密模型训练。此外，我们表明softmax-then-topK专家路由优于topK-then-softmax方法，并且更高粒度的MoE可以帮助提高准确性。最后，我们将Nemotron-4 15B模型在1T tokens上进行了升级，并将其与在相同1T tokens上持续训练的相同模型版本进行了比较：持续训练的模型达到了65.3%的MMLU，而升级后的模型达到了67.6%。我们的结果为有效利用升级来构建MoE语言模型提供了见解和最佳实践。代码已开源。

🔬 方法详解

问题定义：论文旨在解决如何高效地将预训练的稠密语言模型转化为混合专家（MoE）模型的问题。现有方法在扩展模型容量方面存在效率瓶颈，并且缺乏针对大规模MoE模型升级的优化策略，导致性能提升受限。

核心思路：论文的核心思路是通过一种新颖的初始化和权重缩放方法，使得预训练的稠密模型能够平滑过渡到稀疏的MoE架构。这种方法旨在克服直接训练MoE模型的困难，并充分利用已有的预训练知识。

技术框架：整体流程包括以下几个关键步骤：1) 选择一个预训练的稠密语言模型；2) 使用提出的虚拟组初始化方案和权重缩放方法，将稠密模型转换为MoE模型；3) 使用softmax-then-topK路由策略选择专家；4) 在目标数据集上对升级后的MoE模型进行微调。

关键创新：论文的关键创新在于提出了“虚拟组”初始化方案和权重缩放方法。虚拟组初始化旨在模拟MoE结构，使得初始化后的模型更接近MoE模型的分布。权重缩放方法则用于调整专家网络的权重，以避免训练初期出现梯度消失或爆炸的问题。

关键设计：虚拟组初始化将稠密层的权重矩阵划分为多个“虚拟组”，每个组对应一个专家。权重缩放方法通过调整专家网络的输出权重，平衡不同专家之间的贡献。此外，论文还探索了softmax-then-topK和topK-then-softmax两种路由策略，并发现前者在性能上更优。具体而言，softmax-then-topK先对所有专家的输出进行softmax归一化，然后选择topK个专家进行加权平均。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，通过升级方法得到的Nemotron-4 15B模型在MMLU基准测试中达到了67.6%的准确率，而持续训练的相同模型仅达到65.3%。这表明升级方法能够有效提升模型性能，并且优于传统的持续训练方法。

🎯 应用场景

该研究成果可应用于构建更大规模、更高性能的语言模型，尤其是在计算资源有限的情况下，通过升级现有模型可以快速获得性能提升。潜在应用领域包括自然语言处理、机器翻译、文本生成、对话系统等。

📄 摘要（原文）

Upcycling pre-trained dense language models into sparse mixture-of-experts (MoE) models is an efficient approach to increase the model capacity of already trained models. However, optimal techniques for upcycling at scale remain unclear. In this work, we conduct an extensive study of upcycling methods and hyperparameters for billion-parameter scale language models. We propose a novel "virtual group" initialization scheme and weight scaling approach to enable upcycling into fine-grained MoE architectures. Through ablations, we find that upcycling outperforms continued dense model training. In addition, we show that softmax-then-topK expert routing improves over topK-then-softmax approach and higher granularity MoEs can help improve accuracy. Finally, we upcycled Nemotron-4 15B on 1T tokens and compared it to a continuously trained version of the same model on the same 1T tokens: the continuous trained model achieved 65.3% MMLU, whereas the upcycled model achieved 67.6%. Our results offer insights and best practices to effectively leverage upcycling for building MoE language models. Code is available.