Sparse mixture of expert architectures (MoEs) scale model capacity without large increases in training or inference costs. Despite their success, MoEs suffer from a number of issues: training instability, token dropping, inability to scale the number of experts, or ineffective finetuning. In this work, we proposeSoft MoE, a fully-differentiable sparse Transformer that addresses these challenges, while maintaining the benefits of MoEs. Soft MoE performs an implicit soft assignment by passing different weighted combinations of all input tokens to each expert. As in other MoE works, experts in Soft MoE only process a subset of the (combined) tokens, enabling larger model capacity at lower inference cost. In the context of visual recognition, Soft MoE greatly outperforms standard Transformers (ViTs) and popular MoE variants (Tokens Choice and Experts Choice). For example, Soft MoE-Base/16 requires 10.5x lower inference cost (5.7x lower wall-clock time) than ViT-Huge/14 while matching its performance after similar training. Soft MoE also scales well: Soft MoE Huge/14 with 128 experts in 16 MoE layers has over 40x more parameters than ViT Huge/14, while inference time cost grows by only 2%, and it performs substantially better.
翻译:稀疏混合专家架构(MoE)在不大幅增加训练或推理成本的情况下扩展模型容量。尽管取得了成功,MoE仍面临诸多问题:训练不稳定、令牌丢弃、无法扩展专家数量以及微调效果不佳。在本文中,我们提出Soft MoE,一种全可微分的稀疏Transformer,它在保持MoE优点的同时解决了这些挑战。Soft MoE通过将输入令牌的不同加权组合传递给每个专家来执行隐式软分配。与其他MoE工作类似,Soft MoE中的专家仅处理(组合后)令牌的子集,从而以更低的推理成本实现更大的模型容量。在视觉识别场景中,Soft MoE大幅优于标准Transformer(ViT)和流行的MoE变体(令牌选择与专家选择)。例如,Soft MoE-Base/16的推理成本比ViT-Huge/14低10.5倍(墙钟时间低5.7倍),同时在类似训练后匹配其性能。Soft MoE还具有良好的可扩展性:在16个MoE层中使用128个专家的Soft MoE Huge/14,其参数量比ViT Huge/14高出40倍以上,而推理时间成本仅增加2%,且性能显著更优。