陈丹琦团队提出最新MoE架构Lory

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

链接：https://arxiv.org/pdf/2405.03133

摘要 (Abstract)

• 背景: 传统的 MoE 模型通过稀疏激活来扩大模型规模，但训练过程中路由网络的优化是一个挑战，因为它涉及到不可微分的离散目标。
• SMEAR: 之前的研究提出了一种完全可微分的 MoE 架构 SMEAR，它在参数空间中软合并专家，但只在分类任务的微调中展示了有效性。
• Lory: 本文提出的 Lory 是首个将此类完全可微分的 MoE 架构扩展到自回归语言模型预训练的方法。Lory 引入了两个关键技术：
1. 因果分段路由策略：在保持语言模型自回归特性的同时，实现了高效的专家合并操作。
2. 基于相似性的数据批处理方法：通过在训练实例中分组相似文档，鼓励专家专门化。

引言 (Introduction)

• MoE 架构: 介绍了 MoE 架构及其在模型大小扩展中的作用，同时指出了训练路由网络的挑战。
• 现有解决方案: 讨论了现有的一些解决方案，如开关路由、top-k 专家选择路由和线性规划等。

预备知识 (Preliminaries)

• 稀疏激活 MoE: 解释了 Transformer 基础的 MoE 语言模型是如何工作的，以及如何通过路由网络计算路由权重。
• 通过专家合并的完全可微分 MoE 架构: 讨论了如何通过在参数空间中计算所有专家 FFN 的加权平均值来创建“合并 FFN”，从而实现 MoE 架构的完全可微分性。

Lory 方法

• 核心技术: Lory 的核心是专家合并，提出了因果分段路由方法和数据批处理策略，以实现高效计算和专家专门化。
• 符号表示: 文章定义了符号，如输入序列、路由网络、专家 FFN 参数等。
• 高效专家合并: 通过因果分段路由策略，减少了合并操作的数量，并通过停止梯度操作防止信息泄露。
• 数据批处理: 通过相似性基础的数据批处理技术，鼓励了专家对不同领域或主题的专门化。

实验 (Experiments)

• 模型设置: 描述了实验中使用的模型配置，包括不同数量的活跃参数和专家。
• 训练细节: 包括优化器选择、学习率调度、训练数据集和评估数据集。
• 主要结果: 展示了 Lory 模型在训练效率、语言建模和下游任务中的性能，并与密集模型进行了比较。

分析和消融研究 (Analysis and Ablation Studies)

• 因果分段路由的重要性: 通过比较因果分段路由和前缀路由，强调了每个分段提供强训练信号的重要性。
• 相似性基础数据批处理的重要性: 展示了使用相似性基础批处理方法对 MoE 模型性能提升的影响。
• 与现有 MoE 模型的比较: 将 Lory 与现有的 Expert Choice (EC) MoE 方法进行了比较。
• 专家利用和专门化: 分析了专家的利用情况和专门化程度，展示了在不同领域专家的路由权重。
  

相关工作 (Related Work)

• MoE: 讨论了 MoE 模型的发展历程和挑战。
• 基于相似性的数据批处理: 提到了其他研究中使用类似数据批处理方法的工作。

结论 (Conclusion)

• Lory 的贡献: 强调了 Lory 在自回归语言模型预训练中的潜力和优势。
• 未来工作: 提出了进一步扩展 Lory、结合不同级别的路由策略和开发高效解码方法的未来研究方向。

伦理声明 (Ethics Statement)

• 社会影响: 讨论了使用 Lory 方法训练的语言模型可能带来的潜在社会影响，包括错误信息的传播风险和隐私问题。

附录 (Appendix)

• 伪代码: 提供了因果分段路由策略的伪代码。
• 计算开销: 分析了 MoE 层与密集层相比的计算开销。
• 数据批处理细节: 描述了相似性基础数据批处理的具体实现方法。
• 模型配置: 列出了实验中使用的模型架构和大小。
• 7B 模型实验: 提供了 7B 模型的实验设置和结果。
• 专家专门化: 展示了 0.3B/8E 模型在不同领域上的专家路由权重。

本文主要来自kimi解读，具体了解请阅读原论文。