陈丹琦团队提出最新MoE架构Lory
陈丹琦团队提出了一种名为 Lory 的新型混合专家(Mixture-of-Experts,简称 MoE)模型,专门用于自回归语言模型的预训练。以下是对文章内容的详细解读。
论文: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 引入了两个关键技术:
因果分段路由策略:在保持语言模型自回归特性的同时,实现了高效的专家合并操作。
基于相似性的数据批处理方法:通过在训练实例中分组相似文档,鼓励专家专门化。
引言 (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 模型在不同领域上的专家路由权重。
文章通过实验验证了 Lory 模型在自回归语言模型预训练中的有效性,并展示了其在多个下游任务中的性能提升。此外,文章还探讨了专家的利用和专门化,以及如何通过不同的训练和推理策略进一步提高模型性能。最后,作者提出了未来工作的方向,并就使用 Lory 方法可能带来的伦理问题进行了讨论。
本文主要来自kimi解读,具体了解请阅读原论文。
微信扫码关注该文公众号作者