随着大模型涌现出令人惊艳的性能，模型大小已经成为影响模型性能的关键因素之一。通常，对 Transformer 模型来说，模型越大，性能就会越好，但计算成本也会增加。近期有研究表明，模型大小和训练数据必须一起扩展，才能最佳地使用给定的训练计算预算。
稀数混合专家模型（MoE）是一种很有前途的替代方案，可以在计算成本较少的情况下，扩展模型的大小。稀疏 MoE Transformer 有一个关键的离散优化问题：决定每个输入 token 应该使用哪些模块。这些模块通常是称为专家的 MLP。
为了让 token 与专家良好匹配，人们设计了许多方法，包括线性规划、强化学习、最优传输（optimal transport）等。在许多情况下，需要启发式辅助损失（auxiliary loss）来平衡专家的利用率并最大限度地减少未分配的 token。这些挑战在分布外场景中可能会加剧。
现在，来自 Google DeepMind 的研究团队提出了一种称为「Soft MoE」的新方法 ，解决了许多问题。
论文地址：https://arxiv.org/pdf/2308.00951.pdf
Soft MoE 不采用稀疏且离散的路由器在 token 和专家之间进行硬分配（hard assignment），而是通过混合 token 来执行软分配（soft assignment）。值得注意的是，这种方法会计算所有 token 的多个加权平均值（weighted average），其中权重取决于 token 和专家，然后由相应的专家处理每个加权平均值。
常见的稀疏 MoE 算法通常会学习一些路由器参数，但这些算法的效果有时甚至不如随机固定路由。在 Soft MoE 中，由于每个路由（或混合）参数都是根据单个输入 token 直接更新的，因此可以在训练路由器期间提供稳定性。研究团队还观察到，在训练期间，大部分输入 token 可以同时改变网络中的离散路由。
此外，硬路由（hard routing）在专家模块数量较多时可能具有挑战性，因此大多数研究的训练只有几十个专家模块。相比之下，Soft MoE 可扩展至数千个专家模块，并且可以通过构建实现平衡。最后，Soft MoE 在推理时不存在批次效应（batch-effect）。
该研究进行了一系列实验来探究 Soft MoE 方法的实际效果。实验结果表明，Soft MoE L/16 在上游任务、少样本任务和微调方面击败了 ViT H/14，并且 Soft MoE L/16 仅需要一半的训练时间，推理速度还是 ViT H/14 的 2 倍。值得注意的是，尽管 Soft MoE B/16 的参数量是 ViT H/14 的 5.5 倍，但 Soft MoE B/16 的推理速度却是 ViT H/14 的 5.7 倍。
