论文笔记｜OUTRAGEOUSLY LARGE NEURAL NETWORKS- THE SPARSELY-GATED MIXTURE-OF-EXPERTS LAYER

ICLR 2017

摘要

神经网络吸收信息的能力受到其参数数量的限制。条件计算，即网络的某些部分在每个示例的基础上处于活动状态，在理论上已被提出作为一种在不按比例增加计算量的情况下大幅增加模型容量的方法。然而，在实践中，存在重大的算法和性能挑战。在这项工作中，我们解决了这些挑战，最终实现了条件计算的承诺，模型容量提高了 1000 倍以上，而现代 GPU 集群的计算效率仅略有损失。我们引入了稀疏门控专家混合层（MoE），由多达数千个前馈子网络组成。可训练的门控网络确定用于每个示例的这些专家的稀疏组合。我们将 MoE 应用于语言建模和机器翻译的任务，其中模型容量对于吸收训练语料库中的大量可用知识至关重要。我们提出了模型架构，其中具有多达 1370 亿个参数的 MoE 在堆叠的 LSTM 层之间以卷积方式应用。在大型语言建模和机器翻译基准上，这些模型以较低的计算成本取得了比最先进模型更好的结果。

动机

使用条件计算方法实现在增加模型容量的同时，减少计算的开销。实现了稀疏门控专家混合层。

模型结构

经典的moe结构表示：$y={\sum }_{i=1}^{n}G(x{)}_i{E}_i(x)$

如果专家组数量很大，使用分层moe。

路由网络

• Softmax Gating: 最简单和原始的方法 $G_{\sigma }(x)=Softmax(x\cdot W_g)$

• Noisy Top-K Gating: 添加噪声，平衡训练 $$G(x)=Softmax(KeepTopK(H(x),k))$$
  $$H(x{)}_i=(x\cdot W_g{)}_i+StandardNormal()\cdot Softplus((x\cdot W_{noise}{)}_i)$$
  $$KeepTopK(v,k{)}_i=\{\begin{array}{ll}{v}_i,& \text{if}{v}_i\text{is in the top k elements of v}\\ -\infty ,& \text{otherwise.}\end{array}$$

平衡专家利用

作者观察到，门控网络倾向于收敛到一种状态，即它总是为相同的少数专家产生较大的权重。这种不平衡是自我强化的，因为受青睐的专家训练得更快，因此更容易被门控网络选择。
bengio也提到了这个问题，解决方法是使用了两个损失。其中一个控制每个example的稀疏性，我们不需要它，因为它是由 k 的固定值强制执行的。第三个损失鼓励了门值的多样性。在我们的实验中，我们发现随着专家的专业化，门值自然会多样化（良性循环），并且我们不需要强制门值的多样性。

We take a soft constraint approach：定义了重要性损失，一个专家的重要性定义为batch数据中该专家的gate值的加和。这个额外的损失鼓励全部的专家有相同的重要性。
$$Importance(X)=\sum _{x\in X}G(x)$$
$$L_{importance}(X)=w_{importance}\cdot CV(Importance(X){)}^2$$
CV：离散系数

虽然这种方法可以使得样本均匀的分给不同的专家，但是，专家的gate weight也是不均匀的，会造成在分布式硬件上内存和性能问题，因此引入了另外的损失：load loss

两个损失的图解参考：【手撕LLM-sMoE】离GPT4又近了一步

总结

• sMoE 的经典工作之一
• 辅助loss是有效的设计

参考：
https://zhuanlan.zhihu.com/p/335024684
【手撕LLM-sMoE】离GPT4又近了一步
源码