[论文解读] From Sparse to Soft Mixtures of Experts

Joan Puigcerver, Carlos Riquelme+2 more|arXiv (Cornell University)|Aug 2, 2023

Domain Adaptation and Few-Shot Learning被引用 23

一句话总结

Soft MoE 引入了一个完全可微的稀疏 Transformer，可以将输入在所有 token 间以软方式混合到每个专家，从而在降低推理成本和稳定性的前提下实现可扩展容量，并在视觉任务中超越 ViTs 和现有 MoEs。

ABSTRACT

Sparse mixture of expert architectures (MoEs) scale model capacity without significant 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 propose Soft 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 MoEs, experts in Soft MoE only process a subset of the (combined) tokens, enabling larger model capacity (and performance) at lower inference cost. In the context of visual recognition, Soft MoE greatly outperforms dense Transformers (ViTs) and popular MoEs (Tokens Choice and Experts Choice). Furthermore, Soft MoE scales well: Soft MoE Huge/14 with 128 experts in 16 MoE layers has over 40x more parameters than ViT Huge/14, with only 2% increased inference time, and substantially better quality.

研究动机与目标

• 在不产生高昂计算或内存成本的前提下，推动 Transformer 模型的扩展。
• 在保持专家特化优势的同时，解决传统稀疏 MoE 的不稳定性和 token 丢弃问题。
• 提出一种完全可微的软路由机制，使成千上万的专家成为可能。
• 在图像分类任务上展示该方法，并与 ViTs 及现有 MoEs 进行对比。

提出的方法

• 将 Soft MoE 定义为一个完全可微的层，通过软分配来实现路由。
• 使用 token-slot 与 slot-token 轴上的 softmax 计算派发权重和合并权重（方程式与论文中的（1）–（3）类似）。
• 用与其对应的专家处理每个输入 slot（通常是一个 MLP）。
• 对输入和路由参数进行每 token 和每 slot 的 L2 归一化以稳定训练。
• 用 Soft MoE 块替换 Transformer 中的密集 MLP 块，通过总 slot 数来控制计算。
• 提供一个简单的 JAX 实现，并在 Google Research GitHub（vmoe）引用完整代码。

Figure 1 : Main differences between Sparse and Soft MoE layers. While the router in Sparse MoE layers (left) learns to assign individual input tokens to each of the available slots, in Soft MoE layers (right) each slot is the result of a (different) weighted average of all the input tokens. Learning

实验结果

研究问题

• RQ1 Soft MoE 是否能够在训练和推理预算范围内实现与密集 ViTs 及现有稀疏 MoEs 相当或更高的准确性？
• RQ2 软路由是否在规模化时缓解令牌丢弃和专家失衡等经典 MoE 问题？
• RQ3 Soft MoE 如何随专家数量和每个专家的 slot 数扩展，以及在视觉任务中的最佳配置是什么？
• RQ4 Soft MoE 能否将优势扩展到下游任务，如图像文本对比学习？
• RQ5 与密集与稀疏基线相比，在训练时间、FLOPs 和实测用时上的权衡是什么？

主要发现

• Soft MoE 在多种模型尺寸上，在训练成本与性能的帕累托前沿上同时优于密集 ViTs 和流行的稀疏 MoEs。
• Soft MoE Base/16 在经过相似训练后，与 ViT-Huge/14 相比推理成本降低 10.5x、实测时间快 5.7x，同时达到相当的性能。
• Soft MoE 128 专家和 16 层的参数数量可以比 ViT Huge/14 多出超过 40 倍，而推理成本仅增加约 2%，并且性能显著更好。
• 长期规模实验表明，在可比计算预算下，Soft MoE 模型优于 Vision Transformers，在较小骨干网络上收益明显，在更大规模上结果具有竞争力甚至优于。
• Soft MoE B/16 与 L/16 变体在上游和微调上取得强劲结果，与 ViT 基线相比在推理时间上获得较大加速（例如 Soft MoE L/16 的表现优于密集 H/14 且更快）。
• Soft MoE 能随着专家数量扩展良好，若每个专家只有一个 slot，拥有数百到数千个专家也能在成本不成问题的情况下提升性能，与某些稀疏 MoE 变体不同。

Figure 2 : The Soft MoE routing algorithm. Soft MoE first computes scores or logits for every pair of input token and slot, based on some learnable per-slot parameters. These logits are then normalized per slot (columns) and every slot computes a linear combination of all the input tokens based on t