Dense Backpropagation Improves Training for Sparse Mixture-of-Experts

专家混合模型（MoE）的预训练相比密集型 Transformer 的预训练更具可扩展性，因为 MoE 学会将输入路由到其前馈参数的一个稀疏子集。然而，这也意味着 MoE 只能接收到稀疏的反向更新，从而导致训练不稳定和性能次优。我们提出了一种轻量级的近似方法，在继续稀疏激活其参数的同时，为 MoE 路由器提供一个稠密的梯度更新。我们将这种方法称为默认 MoE（Default MoE），它通过使用指数移动平均值来替代缺失的专家激活，该平均值基于训练过程中之前见过的专家输出计算得出。这使得路由器能够从每个 token 的所有专家中接收信号，从而显著提升训练性能。我们的默认 MoE 在多种设置下超越了标准的 TopK 路由方法，并且不需要显著增加计算开销。代码：https://github.com/vatsal0/default-moe。

论文试图解决Mixture of Experts (MoE) 模型在预训练过程中由于稀疏反向传播更新导致的训练不稳定性和性能次优的问题。这是一个现有问题，但该研究提出了新的解决方案。

关键思路是提出了一种称为Default MoE的方法，通过用专家输出的历史指数移动平均值（EMA）作为默认输出来填补未激活专家的缺失激活。这种方法使路由器能够接收来自每个专家的密集梯度信号，同时保持参数的稀疏激活，从而提高训练性能。相比现有的TopK路由方法，这种思路能够在不显著增加计算开销的情况下改善模型性能。

论文展示了Default MoE方法在多种设置下优于标准TopK路由方法。实验设计包括对比不同路由策略下的模型性能，并验证了Default MoE在减少训练波动和提升收敛速度方面的优势。研究使用了常见的NLP数据集进行验证，并且代码已开源至GitHub (https://github.com/vatsal0/default-moe)，便于后续研究者复现和改进。值得进一步研究的方向包括将Default MoE扩展到更大的模型规模以及探索其在多模态任务中的应用。

最近的相关研究包括Google提出的Switch Transformer，它通过稀疏激活机制实现了大规模语言模型的高效训练；以及GShard和V-MoE等方法，它们也在尝试优化MoE模型的可扩展性和效率。此外，还有一些工作专注于改进MoE的路由算法，例如Load Balancing for Mixture-of-Experts and Beyond (NeurIPS 2021)。这些研究共同推动了MoE架构在大规模机器学习中的应用。