Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer【小白读论文】

摘要：
神经网络接收信息的能力受到其参数数量的限制。 条件计算（激活网络的部分子网络进行计算）是一种在不按比例增加计算量的情况下大幅增加模型容量的方法。 然而，在实践中，存在重大的算法和性能挑战。 在这篇文章，我们解决了这些挑战，模型容量提高了 1000 倍以上。 我们引入了稀疏门控专家混合层（MoE），由多达数千个前馈子网络组成。 可训练的门控网络确定这些专家的稀疏组合以用于每个example。 我们将 MoE 应用于语言建模和机器翻译的任务，其中模型容量对于吸收训练语料库中的大量可用知识至关重要。 我们提出了模型架构，MoE有1370 亿个参数，在堆叠的 LSTM 层之间以卷积方式应用。 在大型语言建模和机器翻译基准上，这些模型以较低的计算成本取得了比最先进模型更好的结果。
 
1、Introduction：
1.1、条件计算
利用训练数据和模型规模对于深度学习的成功至关重要。 当数据集足够大时，增加神经网络的容量（参数数量）可以提供更好的预测精度。 对于典型的深度学习模型，每个示例都会激活整个模型，随着模型大小和训练示例数量的增加，这会导致训练成本大致呈二次方激增。 但计算能力和分布式计算的进步不足以满足这种需求。 人们提出了各种形式的条件计算，使得模型容量增加而计算成本不按比例增加。 在这些方案中，网络的大部分在每个示例上处于活动或不活动状态。 门控决策可以是二元的、稀疏的、连续的、随机的或确定性的，人们提出了各种形式的强化学习和反向传播来训练门控决策。
虽然这些想法在理论上很有前景，但迄今为止还没有任何工作证明模型容量、训练时间或模型质量有巨大的改进。 我们将此归咎于以下挑战：

• 现代计算设备，尤其是 GPU，在算术方面比在分支（branching）方面快得多。 上面的大多数工作都认识到这一点，并建议通过每个门控决策来打开/关闭网络的chunk。
• 大批量大小对于性能至关重要，因为它们可以分摊参数传输和更新的成本。 条件计算减少了网络中条件活动块的批量大小。
• 网络带宽可能成为瓶颈。GPU 集群的计算能力可能比设备间网络带宽的总和大数千倍。 为了提高计算效率，算法的相对计算与网络需求必须超过该比率。 由于embeddings 通常需要通过网络发送，因此（例如参数）交互的数量受到网络带宽而不是计算能力的限制。
• 根据方案的不同，可能需要损失项才能实现每个块或每个示例所需的稀疏度水平。 这些问题可能会影响模型质量和负载平衡。
• 模型容量对于大数据集最为关键。 现有的条件计算涉及相对较小的图像识别数据集，最多包含 600,000 张图像。 很难想象这些图像的标签提供了足够的信号来充分训练具有数百万个参数的模型，更不用说数十亿个参数了。

在我们的工作中，首次解决了上述所有挑战，并最终实现了条件计算的预期结果。 我们的模型容量提高了 1000 倍以上，而计算效率仅略有损失，并显着提高了公共语言建模和翻译数据集的最新结果。

1.2、所提方法：稀疏门控混合专家层
我们的条件计算方法是引入一种新型通用神经网络组件：稀疏门控混合专家层 (MoE)。 MoE 由多个专家组成，每个专家都是一个简单的前馈神经网络，以及一个可训练的门控网络，该网络使用专家的稀疏组合来处理每个输入（见图 1）。 网络的所有部分都通过反向传播联合训练。虽然引入的技术是通用的，但在本文中，我们重点关注语言建模和机器翻译任务，众所周知，这些任务受益于非常大的模型。 特别是，我们在堆叠的 LSTM 层之间以卷积方式应用 MoE，如图 1 所示。MoE 对文本中的每个位置调用一次，在每个位置选择可能不同的专家组合。 不同的专家往往会根据语法和语义变得高度专业化。

1.3、专家混合的相关工作
自二十多年前提出以来，专家混合方法一直是许多研究的主题。 人们提出了不同类型的专家架构，例如 SVM、高斯过程、狄利克雷过程、 和深层网络。 其他工作侧重于不同的专家配置，例如层次结构、无限数量的专家以及按顺序添加专家。 Garmash提出了一种机器翻译专家混合形式的集成模型。 门控网络在预先训练的集成 NMT 模型上进行训练。 上述作品涉及 top-level 专家的组合。 专家的混合就是整个模型。 Eigen 等人介绍了使用多个 MoE 及其自己的门控网络作为深度模型的一部分。 后一种方法（Eigen的MoE作为网络的一部分）更强大，因为复杂的问题可能包含许多子问题，每个子问题都需要不同的专家。 他们还引入稀疏性，将 MoE 转变为计算工具。 我们的工作建立在使用 MoE 作为通用神经网络组件的基础上。 Eigen 使用两个堆叠的 MoE 允许两组门控决策，我们的 卷积MoE 在文本中的每个位置有不同的门控决策。 我们还实现了稀疏门控，并展示了它可以大幅增加模型容量。
 
2、专家混合层的结构
专家混合层（MoE）由一组 n 个“专家网络”E1、····、En 和一个“门网络”G 组成，其输出为 稀疏 n 维向量。 图 1 显示了 MoE 模块的概述。 专家本身就是神经网络，每个都有自己的参数。 虽然原则上我们只要求专家接受相同大小的输入并产生相同大小的输出，但在本文的初步研究中，我们限制模型是具有相同架构的前馈网络的情况，但具有 单独的参数。 对于给定的输入 x，我们用 G(x) 和 Ei(x) 表示门控网络的输出和第 i 个专家网络的输出。 MoE模块的输出y可以写成如下：

我们根据 G(x) 输出的稀疏性来节省计算。 当 G(x)i = 0 时，我们不需要计算 Ei(x)。 在我们的实验中，我们有多达数千名专家，但只需要为每个示例评估其中的一小部分。如果专家数量非常多，我们可以通过使用两级分层 MoE 来减少分支因子。 在分层 MoE 中，主门控网络选择“专家”的稀疏加权组合，每个“专家”本身就是具有自己的门控网络的辅助专家混合。下面我们重点关注普通 MoE。
我们实现的moe与其他条件计算模型相关，其专家是简单的权重矩阵，类似于（Cho & Bengio，2014）中提出的参数化权重矩阵。 层类似于 (Bengio et al., 2015) 中描述的逐块 dropout，其中 drop-out 层夹在完全激活的层之间。
 
2.1、门控网络
Softmax 门控：非稀疏门控函数的一个简单选择是将输入乘以可训练的权重矩阵 Wg，然后应用 Softmax 函数。

噪声 Top-K 门控：我们向 Softmax 门控网络添加两个组件，稀疏性和噪声。 在采用softmax函数之前，我们添加可调高斯噪声，然后仅保留前k个值，将其余值设置为−∞（相应的门值等于0）。 如上所述，稀疏性用于节省计算量。 噪声项有助于负载平衡，每个组件的噪声量由第二个可训练权重矩阵 Wnoise 控制。

训练门控网络：我们通过简单的反向传播来训练门控网络以及模型的其余部分。 如果我们选择 k > 1，则前 k 个专家的门值相对于门控网络的权重具有非零导数。梯度还通过门控网络反向传播到其输入。 我们的方法与（Bengio et al., 2015）不同，后者使用布尔门和 REINFORCE 式方法来训练门控网络。
 
3、解决性能挑战
3.1、批量缩小问题
在现代 CPU 和 GPU 上，为了提高计算效率，需要大batch size，以便分摊参数加载和更新的开销。如果门控网络为每个示例从 n 个专家中选择 k 个，那么对于一批 b 个示例，每个专家都会收到一小批大约 kb/n << b 个示例。 随着专家数量的增加，这会导致简单的 MoE 实施变得非常低效。 解决批量缩小问题的方法是使原始批量大小尽可能大。 然而，批量大小往往受到存储前向和后向传递之间的激活所需的内存的限制。 我们提出了以下增加批量大小的技术：
混合数据并行和模型并行：在传统的分布式训练设置中，不同设备上的模型的多个副本异步处理不同批次的数据，并且通过一组参数服务器同步参数 。 在我们的技术中，这些不同的批次同步运行，以便可以将它们组合到 MoE 层。 我们根据传统的数据并行方案分配模型和门控网络的标准层，但只保留每个专家的一个共享副本。 MoE 层中的每个专家都会收到一个组合批次，其中包含来自所有数据并行输入批次的相关示例。 同一组设备充当数据并行副本（用于标准层和门控网络）和模型并行分片（每个管理一部分专家）。如果模型分布在 d 个设备上，并且每个设备处理一批大小为 b 的批次，则每个专家都会收到一批大约 kbd/n 个示例。 因此，我们将专家批量大小提高了 d 倍。
在分层 MoE 的情况下，主门控网络采用数据并行性，次要 MoE 采用模型并行性。 每个辅助 MoE 驻留在一台设备上。
这项技术允许我们通过按比例增加训练集群中的设备数量来增加专家的数量（从而增加参数的数量）。 总批量大小增加，保持每个专家的批量大小不变。 每个设备的内存和带宽要求也保持不变，step time不变，处理多个训练示例所需的时间也等于模型中参数的数量。 我们的目标是在万亿单词语料库上训练万亿参数模型。 截至撰写本文时，我们尚未对系统进行扩展，但通过添加更多硬件应该可以实现。
利用卷积性：在我们的语言模型中，我们将相同的 MoE 应用于前一层的每个时间步。 如果我们等待上一层完成，我们可以将 MoE 作为一个大批次一起应用于所有时间步。 这样做会将 MoE 层的输入批次的大小增加到展开时间步数的一个因子。
增加循环 MoE 的批量大小：我们怀疑更强大的模型可能涉及循环应用 MoE。 例如，LSTM 或其他 RNN 的权重矩阵可以用 MoE 替换。 遗憾的是，这样的模型打破了上一段中的卷积技巧，因为 MoE 在一个时间步的输入取决于 MoE 在前一时间步的输出。 格鲁斯利斯等人描述了一种以重新计算前向激活为代价，大幅减少展开的 RNN 中存储的激活数量的技术。 这将允许批量大小的大幅增加。
3.2、网络带宽
分布式计算中另一个主要的性能问题是网络带宽。 由于专家是固定的并且门控参数的数量很少，因此大多数通信涉及通过网络发送专家的输入和输出。 为了保持计算效率，专家的计算与其输入和输出大小的比率必须超过计算设备的计算能力与网络容量的比率。 对于 GPU 来说，这可能是千比一。 在我们的实验中，我们使用的专家具有一个包含数千个 RELU 激活单元的隐藏层。 由于 Expert 中的权重矩阵的大小为 input_size×hidden_​​size 和hidden_​​size × output_size，因此输入和输出的计算比率等于隐藏层的大小。 我们可以简单地通过使用更大的隐藏层或更多的隐藏层来提高计算效率。
 
4、平衡专家利用率
我们观察到，门控网络倾向于收敛到一种状态，即它总是为相同的少数专家产生较大的权重。 这种不平衡是自我强化的，因为受青睐的专家训练得更快，因此更容易被门控网络选择。 Eigen(2013)描述了相同的现象，并在训练开始时使用硬约束来避免这种局部最小值。 本吉奥(2015) 对每个门的批量平均值添加软约束。
我们采用软约束方法。 我们将专家相对于一批训练样本的重要性定义为该专家的门值的批量总和。 我们定义一个额外的损失 $L_{importance}$，将其添加到模型的整体损失函数中。 该损失等于重要性值集合的变异系数的平方乘以手动调整的缩放因子 $w_{importance}$。 这种额外的损失鼓励所有专家都同等重视

虽然这种损失函数可以确保同等重要性，但专家可能仍然会收到数量截然不同的示例。 例如，一位专家可能会收到一些权重较大的示例，而另一位专家可能会收到许多权重较小的示例。 这可能会导致分布式硬件上的内存和性能问题。 为了解决这个问题，我们引入了第二个损失函数 Lload ，它确保平衡负载。
 
5、实验
5.1、10亿字语言建模基准数据集
5.2、1000 亿词汇的谷歌新闻语料库
5.3、机器翻译
5.4、多语言机器翻译
 
6、结论
这项工作首次展示了深度网络条件计算的重大成果。 我们仔细确定了条件计算的设计考虑因素和挑战，并通过算法和工程解决方案的结合来解决它们。 虽然我们专注于文本，但如果提供足够大的训练集，条件计算也可能在其他领域有所帮助。 我们期待在未来几年看到条件计算的许多新颖的实现和应用。