这篇AI论文提出了Soft MoE:一种全可微稀疏Transformer,旨在解决这些挑战,同时保持MoE的优势

This AI paper proposes Soft MoE, a fully differentiable sparse Transformer, to address these challenges while maintaining the advantages of MoE.

更大的Transformer需要更多的计算成本才能良好运行。最近的研究表明,模型大小和训练数据必须同时进行扩展,以充分利用训练计算资源。稀疏的专家混合是一种可能的替代方法,可以在不产生完全计算成本的情况下实现模型的可扩展性。最近的语言、视觉和多模型已经开发出了在整个网络中稀疏激活令牌路径的方法。选择应用于每个输入令牌的模块是稀疏MoE Transformer的核心离散优化挑战。

这些模块通常是多层感知机(MLP),被称为专家。线性规划、强化学习、确定性固定规则、优化传输、每个令牌的贪婪前k个专家和每个专家的贪婪前k个令牌等方法,都是用于确定适当的令牌到专家配对的方法。启发式辅助损失通常需要平衡专家利用率和减少未分配的令牌。小的推理批处理大小、唯一的输入或者迁移学习可能会在分布外设置中加剧这些问题。来自Google DeepMind的研究人员提供了一种名为Soft MoE的新策略,该策略解决了其中几个问题。

Soft MoE通过组合令牌而不是使用寻找令牌和专家之间良好硬分配的稀疏和离散路由器来执行软分配。它们特别构建了所有令牌的几个加权平均值,这些加权平均值的权重依赖于令牌和专家,然后通过相关的专家处理每个加权平均值。在稀疏MoE的核心离散过程引起的大部分问题在软MoE模型中是不存在的。施加一些期望行为并且依赖于路由分数的辅助损失是流行的稀疏MoE方法的梯度的常见来源,这些方法通过将专家输出与选择的路由分数进行后乘以学习路由器参数。

根据观察,这些算法通常与随机固定路由相似。Soft MoE通过根据每个输入令牌立即更新每个路由参数来避免这个问题。他们观察到,大量的输入令牌可以同时改变网络中的离散路径,从而在训练过程中造成训练问题。在训练路由器时,软路由可以提供稳定性。由于大量的专家,硬路由也可能很困难,因为大多数工作只使用少量的专家进行训练。他们证明了Soft MoE可以扩展到数千个专家,并且构建成平衡的。

最后但并非最不重要的是,在推理过程中没有批处理效应,一个输入可能会影响多个输入的路由和预测。虽然训练时间只需大约一半,但Soft MoE L/16在上游、少样本和微调方面优于ViT H/14,并且推理速度更快。此外,经过相同数量的训练后,Soft MoE B/16在上游评估指标上击败了ViT H/14,并且在少样本和微调方面与ViT H/14相匹配。即使Soft MoE B/16的参数比ViT H/14多5.5倍,它的推理速度也快5.7倍。