Nexus: Specialization meets Adaptability for Efficiently Training Mixture of Experts

📄 arXiv: 2408.15901v1 📥 PDF

作者: Nikolas Gritsch, Qizhen Zhang, Acyr Locatelli, Sara Hooker, Ahmet Üstün

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

发布日期: 2024-08-28

💡 一句话要点

Nexus：通过自适应路由实现高效的混合专家模型训练

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

关键词: 混合专家模型 自适应路由 领域自适应 模型升级改造 高效训练

📋 核心要点

  1. 现有大型语言模型难以同时兼顾效率、专业化和对新数据分布的适应性，MoE架构是潜在的解决方案。
  2. Nexus通过自适应路由机制，学习从领域表示中投影专家嵌入，从而灵活地添加新专家。
  3. 实验表明，Nexus在初始升级改造和扩展新专家方面均优于基线模型，提升显著。

📝 摘要（中文）

大型语言模型难以兼顾效率、专业化和对新数据分布的适应性。混合专家（MoE）架构因其固有的条件计算特性而备受关注，使其能够实现这些理想的属性。本文重点在于将密集专家模型“升级改造”为MoE，旨在提高专业化水平，同时增加适应新任务的能力。我们提出了Nexus，一种增强的MoE架构，具有自适应路由，模型学习从领域表示中投影专家嵌入。这种方法允许Nexus在初始升级改造后灵活地添加通过单独训练的密集模型获得的新专家，而无需对未见过的数据领域进行大规模MoE训练。实验表明，Nexus在初始升级改造方面比基线模型相对提高了2.1%，在使用有限的微调数据扩展MoE与新专家方面相对提高了18.8%。Nexus的这种灵活性对于实现一个开放源代码生态系统至关重要，在该生态系统中，每个用户都可以根据自己的需求不断组装自己的MoE组合。

🔬 方法详解

问题定义：现有MoE模型在添加新专家时，通常需要大规模的重新训练，这限制了其灵活性和可扩展性。特别是当面对新的数据领域时，如何快速有效地集成新的专家知识是一个挑战。

核心思路：Nexus的核心思路是通过学习一个自适应的路由机制，将领域表示投影到专家嵌入空间。这样，新加入的专家可以通过单独训练的密集模型获得，而无需对整个MoE模型进行大规模的重新训练。这种方法提高了模型的适应性和灵活性。

技术框架：Nexus架构包含以下主要模块：1) 领域表示模块，用于提取输入数据的领域特征；2) 自适应路由模块，学习将领域表示投影到专家嵌入空间；3) 专家模块，由多个预训练的密集专家模型组成；4) 组合模块，根据路由结果将各个专家的输出进行加权组合。整体流程是：输入数据首先通过领域表示模块提取特征，然后通过自适应路由模块选择合适的专家，最后将各个专家的输出进行组合得到最终结果。

关键创新：Nexus的关键创新在于其自适应路由机制，它允许模型根据输入数据的领域特征动态地选择合适的专家。与传统的MoE模型相比，Nexus不需要对整个模型进行大规模的重新训练即可添加新的专家，从而提高了模型的适应性和灵活性。

关键设计：自适应路由模块使用一个神经网络来学习从领域表示到专家嵌入空间的映射。损失函数包括路由损失和专家损失，其中路由损失用于鼓励模型选择合适的专家，专家损失用于优化各个专家的性能。具体的网络结构和参数设置根据具体的任务和数据集进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Nexus在初始升级改造方面比基线模型相对提高了2.1%，在使用有限的微调数据扩展MoE与新专家方面相对提高了18.8%。这些结果表明，Nexus在提高模型效率和适应性方面具有显著优势。尤其是在添加新专家时，Nexus能够通过有限的微调数据快速适应新的数据领域，这对于实际应用具有重要意义。

🎯 应用场景

Nexus架构可应用于各种需要高效和灵活的专家组合的场景，例如：个性化推荐系统、多语言机器翻译、特定领域的知识问答等。通过Nexus，用户可以根据自己的需求不断组装自己的MoE组合，从而构建一个开放源代码的专家生态系统。这种方法可以显著降低模型训练的成本，并提高模型的适应性和可扩展性。

📄 摘要（原文）

Efficiency, specialization, and adaptability to new data distributions are qualities that are hard to combine in current Large Language Models. The Mixture of Experts (MoE) architecture has been the focus of significant research because its inherent conditional computation enables such desirable properties. In this work, we focus on "upcycling" dense expert models into an MoE, aiming to improve specialization while also adding the ability to adapt to new tasks easily. We introduce Nexus, an enhanced MoE architecture with adaptive routing where the model learns to project expert embeddings from domain representations. This approach allows Nexus to flexibly add new experts after the initial upcycling through separately trained dense models, without requiring large-scale MoE training for unseen data domains. Our experiments show that Nexus achieves a relative gain of up to 2.1% over the baseline for initial upcycling, and a 18.8% relative gain for extending the MoE with a new expert by using limited finetuning data. This flexibility of Nexus is crucial to enable an open-source ecosystem where every user continuously assembles their own MoE-mix according to their needs.