一段话总结
论文《FROM SPARSE TO SOFT MIXTURES OF EXPERTS》提出Soft MoE，一种全可微的稀疏Transformer，解决了稀疏混合专家（MoE）架构训练不稳定、令牌丢弃、专家数量难以扩展和微调效果不佳等问题。Soft MoE通过对所有输入令牌进行加权平均实现软分配，让每个专家处理不同加权组合的令牌。实验表明，在视觉识别任务中，Soft MoE在训练成本、推理速度和性能上优于密集Transformer和其他流行的MoE方法，且能有效扩展到其他任务，但在自回归解码和内存消耗方面存在局限性。
详细总结
-
研究背景：大Transformer模型性能提升的同时计算成本增加，稀疏混合专家（MoE）架构可在不显著增加成本的情况下扩展模型容量，但存在训练不稳定、令牌丢弃、专家数量难以扩展和微调效果不佳等问题。
-
Soft MoE算法
- 算法描述：输入令牌$X\in \mathbb{R}^{m×d}$，每个MoE层有$n$个专家函数$f_{i}$，通过计算$D_{ij}=\frac{exp((X\Phi){ij})}{\sum^{m}exp((X\Phi){i'j})}$得到输入插槽$\tilde{X}=D^{\top}X$，再经专家函数处理后，用$C=\frac{exp((X\Phi){ij})}{\sum^{n\cdot p}exp((X\Phi)_{ij'})}$计算输出令牌$Y = C\tilde{Y}$。通常用Soft MoE块替换Transformer中后半部分的MLP块。
- 算法特性：全可微，避免了离散分配的优化和实现问题；不存在令牌丢弃和专家不平衡现象；速度快，计算成本取决于插槽数量而非专家数量，且避免了慢且不适合硬件加速的排序或top - k操作；兼具稀疏和密集模型的特点；具有序列级确定性。
- 实现细节：时间复杂度为$O(mnpd + npk)$，选择$p = O(m/n)$时可降为$O(m^{2}d+mk)$；采用L2归一化解决模型维度$d$增大时的稳定性问题；专家数量增加时采用标准技术分布式训练模型。
-
实验
- 训练和评估数据：在JFT - 4B数据集上预训练，用上游验证精度、ImageNet 10-shot准确率评估，在ImageNet - 1k数据集上微调后用验证集准确率评估。
- 稀疏路由算法对比：Tokens Choice每个令牌选择得分最高的$K$个专家；Experts Choice专家选择得分最高的$C$个令牌。对比发现Soft MoE在训练成本/性能帕累托前沿上表现更优。
- 模型训练实验：训练多种模型，结果表明Soft MoE在固定训练预算下性能优于ViT，如Soft MoE S/16在JFT和10-shot ImageNet上表现优于ViT B/16，且训练和推理成本更低。
- 消融实验：研究插槽和专家数量变化对模型的影响，发现Soft MoE每个令牌插槽数相同时，专家越多性能越好；对比不同路由方式，证明Soft MoE的有效性。
-
对比学习：在图像文本对比学习任务中，Soft MoE在零样本学习的下游数据集上表现优于ViT，且从数据增强中受益更多。
-
相关工作：与减少序列长度的方法、多头注意力机制、其他MoE方法存在差异，Soft MoE目的不是减少序列长度，专家为非线性且组合$d$维向量，计算成本相对较低。
-
局限性：难以应用于自回归解码，因需保留因果关系且避免令牌和插槽索引相关；一个专家对应一个插槽更优，导致模型内存需求可能较大。
| 模型 | 参数 | 训练步数 | 训练TPUv3天数 | 训练ExaFLOP | 推理ms/img | 推理GFLOP/img | JFT @1 P@1 | INet 10shot | INet finetune |
|---|---|---|---|---|---|---|---|---|---|
| ViT S/16 | 33M | 4M (50k) | 153.5 | 227.1 | 0.5 | 9.2 | 51.3 | 67.6 | 84.0 |
| ViT B/16 | 108M | 4M (50k) | 410.1 | 864.1 | 1.3 | 35.1 | 56.2 | 76.8 | 86.6 |
| Soft MoE S/16 128E | 933M | 4M (50k) | 175.1 | 211.9 | 0.7 | 8.6 | 58.1 | 78.8 | 86.8 |
| Soft MoE B/16 128E | 3.7B | 4M (50k) | 449.5 | 786.4 | 1.5 | 32.0 | 60.0 | 82.0 | 88.0 |
| --- |
关键问题
-
Soft MoE与传统稀疏MoE相比，在训练和推理方面有哪些优势？
- 答案：在训练方面，Soft MoE全可微，避免了离散分配带来的优化和实现问题，不存在令牌丢弃和专家不平衡现象；在推理方面，速度更快，计算成本取决于插槽数量而非专家数量，避免了排序或top - k等慢操作，且在相同性能下推理速度比ViT快很多，如Soft MoE B/16推理速度比ViT H/14快5.7倍。
-
从实验结果看，Soft MoE在不同规模模型中的表现如何？
- 答案：在小模型如Soft MoE S/16中，训练时间更短，在JFT和10-shot ImageNet上性能优于ViT B/16；中等模型Soft MoE B/16在少样本学习和微调任务中表现出色，匹配或超越ViT H/14，且推理速度更快、成本更低；大模型Soft MoE L/16性能超过所有ViT模型，推理速度比ViT H/14快近2倍。
-
Soft MoE在应用中存在哪些局限性？
- 答案：一是在自回归解码方面存在困难，由于需要学习合并所有输入令牌，难以在训练中保留过去和未来令牌之间的因果关系，使用因果掩码时需避免引入令牌和插槽索引的相关性；二是内存消耗问题，一个专家对应一个插槽往往是最优选择，这使得模型在利用大量专家时，内存需求可能会显著增加。