Mixture-of-Experts with Expert Choice Routing：专家混合模型与专家选择路由

摘要

稀疏激活的专家混合模型（MoE）允许在保持每个token或每个样本计算量不变的情况下，大幅增加参数数量。然而，糟糕的专家路由策略可能导致某些专家未被充分训练，从而使得专家在特定任务上过度或不足专业化。先前的研究通过使用top-k函数为每个token分配固定数量的专家，无视不同token的重要性差异。为了解决这个问题，我们提出了一种异构专家混合模型，采用专家选择方法。我们并非让token选择top-k个专家，而是让专家选择top-k个token。结果是，每个token可以路由到不同数量的专家，并且每个专家可以有固定的桶大小。我们通过使用与Switch Transformer的top-1和GShard的top-2门控相同的计算资源，系统地研究了预训练加速效果，发现我们的方法使得训练收敛时间提高了2倍以上。在相同计算成本下，我们的方法在GLUE和SuperGLUE基准测试中精调了11个任务，表现更佳。对于较小的激活成本，我们的方法在11个任务中的7个任务上优于T5密集模型。

1 引言

扩大模型容量、数据集大小和训练时间已被证明能够大幅提升计算机视觉架构[4, 11, 13, 14]和神经语言模型[2, 20, 26, 27]的性能。最终模型的质量已被发现与数据量、模型大小和计算时间呈幂次关系[16, 20]。然而，随着我们对绿色AI[29]的关注加深，训练效率（即使用总计算量比现有最先进系