[论文解读] EvoMoE: An Evolutional Mixture-of-Experts Training Framework via Dense-To-Sparse Gate

Xiaonan Nie, Xupeng Miao+8 more|arXiv (Cornell University)|Dec 29, 2021

Expert finding and Q&A systems被引用 7

一句话总结

EvoMoE 提出了一种两阶段训练框架，通过专家多样化和门控稀疏化阶段，将单个专家逐步演化为稀疏的专家混合模型（MoE）。它引入了一种从密集到稀疏的门控机制（DTS-Gate），初始时为密集路由，随后通过 Gumbel-Softmax 中的温度衰减逐步变为稀疏，从而提升训练稳定性和性能，在 GLUE 上实现了 SOTA 结果，相比基线模型最高提升 0.562 和 0.403。

ABSTRACT

Mixture-of-experts (MoE) is becoming popular due to its success in improving the model quality, especially in Transformers. By routing tokens with a sparse gate to a few experts (i.e., a small pieces of the full model), MoE can easily increase the model parameters to a very large scale while keeping the computation cost in a constant level. Most existing works just initialize some random experts, set a fixed gating strategy (e.g., Top-k), and train the model from scratch in an ad-hoc way. We identify that these MoE models are suffering from the immature experts and unstable sparse gate, which are harmful to the convergence performance. In this paper, we propose an efficient end-to-end MoE training framework called EvoMoE. EvoMoE starts from training one single expert and gradually evolves into a large and sparse MoE structure. EvoMoE mainly contains two phases: the expert-diversify phase to train the base expert for a while and spawn multiple diverse experts from it, and the gate-sparsify phase to learn an adaptive sparse gate and activate a dynamic number of experts. EvoMoE naturally decouples the joint learning of both the experts and the sparse gate and focuses on learning the basic knowledge with a single expert at the early training stage. Then it diversifies the experts and continues to train the MoE with a novel Dense-to-Sparse gate (DTS-Gate). Specifically, instead of using a permanent sparse gate, DTS-Gate begins as a dense gate that routes tokens to all experts, then gradually and adaptively becomes sparser while routes to fewer experts. Evaluations are conducted on three popular models and tasks, including RoBERTa for masked language modeling task, GPT for language modeling task and Transformer for machine translation task. The results show that EvoMoE outperforms existing baselines, including Switch, BASE Layer, Hash Layer and StableMoE.

研究动机与目标

• 解决现有 MoE 训练中因随机初始化专家和稀疏门控导致的不稳定与收敛困难问题。
• 通过解耦专家与门控学习，提升训练效率与模型质量。
• 通过自适应稀疏化实现从密集路由到稀疏路由的平稳过渡。
• 相比现有 MoE 方法，实现更优的性能与 FLOPs 效率。
• 探索适用于大规模 MoE 模型的更稳定、可扩展的训练范式。

提出的方法

• 在专家多样化阶段，先训练一个基础专家，随后通过随机掩码将其分化为多个具有结构多样性的专家。
• 在门控稀疏化阶段，引入一种从密集到稀疏的门控机制（DTS-Gate），初始时对所有专家进行软性密集路由，随后通过 Gumbel-Softmax 中的温度衰减逐步实现稀疏化。
• DTS-Gate 采用基于内容的路由机制，仅激活门控权重超过学习阈值的专家，实现自适应与动态的专家选择。
• 该框架解耦了专家与门控的训练：专家在引入稀疏路由前先进行共享预训练。
• 在 Gumbel-Softmax 中应用温度衰减调度，逐步增强门控输出的稀疏性。
• 该方法端到端地应用于多种模型，包括 RoBERTa、GPT 和 Transformer，在多样化的 NLP 任务上均取得效果。

实验结果

研究问题

• RQ1 从单个专家逐步演化为稀疏 MoE 架构，是否能提升训练稳定性和模型性能？
• RQ2 从密集到稀疏的门控机制是否相比固定稀疏门控（如 Top-k）具有更好的收敛性与路由准确性？
• RQ3 从共享基础专家出发的专家多样化是否能增强知识泛化能力并降低训练不稳定性？
• RQ4 与 Switch、BASELayer 和 StableMoE 等现有 MoE 方法相比，EvoMoE 在性能与 FLOPs 效率方面表现如何？
• RQ5 解耦专家与门控训练在多大程度上改善了 MoE 模型的学习动态？

主要发现

• EvoMoE 在 GLUE 基准上相比现有 MoE 基线模型，平均得分最高提升 0.562 和 0.403，展现出更优的模型质量。
• DTS-Gate 机制实现了更平滑的路由适应，相比固定 Top-k 门控，对随机初始化的敏感性更低。
• 专家多样化阶段促进了更稳定且高效的专家专业化，提升了整体模型鲁棒性。
• EvoMoE 在 FLOPs 效率上优于密集 MoE，且与稀疏 MoE 基线相比性能具有竞争力。
• 该框架在多个任务上均表现出一致的性能提升，包括掩码语言建模、语言建模和机器翻译。
• 消融实验确认，专家多样化与门控稀疏化两个阶段对实现最优性能均至关重要。