Dense Backpropagation Improves Training for Sparse Mixture-of-Experts

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

论文试图解决Mixture of Experts (MoE) 模型在预训练过程中由于稀疏反向传播更新导致的训练不稳定和性能次优的问题。这是一个现有问题，但通过改进路由机制来提升MoE模型性能是一个值得探索的方向。

论文提出了一种名为Default MoE的方法，其核心思路是在前向传播中保持参数稀疏激活的同时，在反向传播中为路由器提供密集梯度更新。具体而言，通过用专家输出的历史指数移动平均值（EMA）作为默认输出来填补未被选中的专家的激活，从而使路由器能够从所有专家中接收信号。相比传统的TopK路由方法，这种方法能够在不显著增加计算开销的情况下改善训练性能。

实验表明Default MoE在多种设置下优于标准的TopK路由方法，并且代码已开源（https://github.com/vatsal0/default-moe），便于复现和进一步研究。此外，该方法设计巧妙，仅通过轻量级的近似手段就解决了稀疏更新带来的问题，为进一步优化MoE模型提供了新思路。未来可以探索不同EMA策略或更复杂的默认输出生成方法。

近期关于MoE的研究包括Google提出的Switch Transformer，它通过稀疏激活实现了超大规模模型的高效训练；以及GShard和V-MoE等方法，它们分别关注于分布式训练和变分路由策略。其他相关工作还包括《Sparse Expert Models for Efficient Inference》和《Rethinking the Value of Experts in Mixture-of-Experts》等论文，这些研究均致力于提高MoE模型的效率和效果。