Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer/zh

Research Paper
Authors	Noam Shazeer; Azalia Mirhoseini; Krzysztof Maziarz; Andy Davis; Quoc Le; Geoffrey Hinton; Jeff Dean
Year	2017
Topic area	Machine Learning
Difficulty	Research
arXiv	1701.06538
PDF	Download PDF

This page is a translated version of the page Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer and the translation is 100% complete.

Other languages:

English
Español
中文

SummarySource

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer 是 Noam Shazeer、Azalia Mirhoseini、Krzysztof Maziarz、Andy Davis、Quoc Le、Geoffrey Hinton 和 Jeff Dean 于 2017 年发表的论文，提出了用于深度神经网络的 Sparsely-Gated Mixture-of-Experts（MoE，稀疏门控混合专家）层。该论文表明，由数千个前馈子网络组成、并由可训练的门控网络针对每个样本仅选择少量子网络的单一层，可以在保持每样本计算量大致不变的情况下，将模型容量提升 1000 倍以上。该方法在堆叠的 LSTM 层之间应用于语言建模和机器翻译任务，参数量高达 1370 亿的 MoE 模型在更低的计算成本下取得了优于稠密基线的最新最优结果。

概述

神经网络的容量——即其吸收信息的能力——受到参数数量的限制，但参数数量通常与每样本的计算量绑定。条件计算曾被提出作为打破这种耦合的方法：对每个输入只激活网络的一个稀疏子集，使容量在不按比例增加 FLOPs 的情况下扩展。早期工作已经识别出这一思路，但由于算法和系统层面的挑战（活跃子网络之间批量大小不均、网络带宽压力、专家坍缩）而未能在大规模上实现。

本文首次实现了大规模的条件计算。Sparsely-Gated Mixture-of-Experts 层由多达数千个专家前馈网络以及一个门控网络组成；门控网络对每个输入产生一个稀疏的专家权重向量。只有门值非零的专家会被计算，因此即使专家数量（以及参数量）扩展到数十亿级，每样本的成本仍然有界。该层以卷积方式插入——在每个时间步以相同方式应用——位于作者模型中两层堆叠 LSTM 之间。

主要贡献

Sparsely-Gated MoE 层。 单一神经网络层，最多包含 131,072 个专家，并配备 top-k 门控网络，每个样本仅激活少量恒定数量的专家，从而将容量与计算解耦。
Noisy Top-K Gating。 一种门控机制，在 top-k 选择之前加入可训练且可调的高斯噪声，同时提供稀疏性和随机负载均衡。
解决 shrinking-batch 问题。 结合数据并行与模型并行，并在时间维度上以卷积方式应用 MoE，使每个专家的有效批量足够大，能够在 GPU 上高效执行。
软负载均衡损失。 辅助性的 "importance" 和 "load" 损失，防止少数专家主导，并确保整个专家群体的均衡利用。
前所未有规模的实证表现。 参数量高达 1370 亿的 MoE 模型，在 1 Billion Word Benchmark 上将困惑度降低了 24%，并在 WMT'14 En→Fr、En→De 及多语言 En→{Fr, De, Es, It, Pt, ...} 翻译任务上超越了此前的最新最优结果——所有这些都以低于稠密基线的计算成本实现。

方法

MoE 层

一个 MoE 层包含 $$ n $$ 个专家网络 $E_1, \dots, E_n$ （通常是两层前馈 ReLU 网络），以及一个产生 $$ n $$ 维稀疏向量的门控网络 $$ G $$ 。对于输入 $$ x $$ ，该层的输出为

y = \sum_{i=1}^{n} G(x)_i \, E_i(x).

凡 $$ G(x)_i = 0 $$ 的位置，对应的专家便被跳过。在 top-k 门控下，每个样本只评估 k 个专家，因此无论 $$ n $$ 多大，每样本的计算量保持恒定。

Noisy Top-K Gating

门控网络在一个可学习的线性投影上施加 softmax，并加入可训练的加性高斯噪声，然后在归一化前只保留前 k 个最大的 logits：

H(x)_i = (x \cdot W_g)_i + \epsilon \cdot \mathrm{softplus}((x \cdot W_{noise})_i),\quad \epsilon \sim \mathcal{N}(0,1)

G(x) = \mathrm{Softmax}(\mathrm{KeepTopK}(H(x), k)).

噪声既能稀疏化激活，也有助于在训练期间平衡专家间的负载；将除前 k 个之外的所有项设为 $-\infty$ ，可在推理时保证严格的稀疏性。

性能工程

当从 $$ n $$ 个专家中为每个 $$ b $$ 个样本选择 $$ k $$ 个专家时，每个专家平均会看到 $$ kb/n $$ 个样本，这对于在 GPU 上高效执行矩阵乘法来说太小。作者通过以下方式解决这一问题：

混合数据并行与模型并行。 将 MoE 层在 $$ d $$ 个设备上分片，使每个专家恰好驻留在一个设备上，所有 $$ d $$ 个数据并行副本同步地将各自选中的样本路由到所需的专家。每个专家的有效批量变为 $k \cdot b \cdot d / n$ ，可随集群规模线性扩展。
卷积式应用。 MoE 在 LSTM 的每个时间步独立应用，将每个专家的批量再乘以展开后的序列长度。
分层 MoE。 对于非常大的 $$ n $$ ，由一个主门控选择专家组，再由次级门控在该组内选择具体专家，使路由保持可处理。
带宽优化。 将专家设计为计算量较大（隐藏层较宽），使其计算量相对于输入字节数的比值能够吸收设备间样本传输的开销。

负载均衡损失

在训练目标中加入两个辅助损失，防止门控网络坍缩到少数偏好的专家：

L_{importance}(X) = w_{importance} \cdot \mathrm{CV}\!\left(\sum_{x \in X} G(x)\right)^2,

其中 CV 是变异系数。第二个损失 $L_{load}$ 用基于带噪门控分布构建的平滑可微估计量，使每个专家所收到的样本数量均衡，而不仅仅是按门权重加权的 importance。两者共同保证每个专家的权重相等且样本数量相等。

结果

语言建模

在 1 Billion Word Language Modeling Benchmark 上，作者训练了具有 4、32、256、1024、4096、32768 和 65536 个专家的 MoE 模型，全部匹配在约 8M ops/timestep。具有 4096 个专家的模型相对于此前最佳已发表模型在测试集困惑度上降低了 24%，而参数量高达 40 亿的更高计算量变体则以远低于此前稠密模型的计算量刷新了困惑度最新最优记录。在 16–32 块 Tesla K40 GPU 上的计算效率达到 0.74–1.56 TFLOPS/GPU。

在 100 Billion Word Google News Corpus 上——其规模约为 1 Billion Word 的 100 倍——他们训练了具有 131,072 个专家、共计 1370 亿参数 的 MoE，并随着容量的扩大持续观察到困惑度的改进，表明在训练数据充足时，MoE 带来的收益会持续叠加。

机器翻译

在 WMT'14 En→Fr 和 En→De 上，加入 MoE 的堆叠 LSTM 编码器/解码器以更少的训练计算量超越了此前 Google Neural Machine Translation（GNMT）的最新最优 BLEU 分数。在覆盖十二个语言对的多语言生产翻译系统中，单一的 MoE 增强模型在每个语言对上的 BLEU 都优于针对每个语言对单独训练的专用 GNMT 模型。

影响

本文是现代 MoE 文献的奠基性工作之一。它确立了一套配方——稀疏 top-k 门控、基于噪声的负载均衡、专家分片、软均衡损失——后续工作在此基础上直接构建。它直接启发了 GShard（2020），将同样的思想扩展到用于多语言翻译的 Transformer 编码器；也启发了 Switch Transformer（2021），后者将门控简化为 top-1 路由并扩展到超过一万亿参数。它还为 GLaM、ST-MoE、Mixtral 8×7B 以及众多前沿大语言模型所采用的 MoE 主干提供了基础。论文标题中的 "outrageously large" 后来看反而显得保守：到 2020 年代初，参数量高达数千亿的 MoE 层已被常规部署到生产中的语言模型中，本文提出的条件计算原则也成为以参数量换取每 token 成本的核心工具。

参见

Multi-gate Mixture-of-Experts —— 用于多任务学习的任务条件 MoE。
Attention Is All You Need —— Transformer 架构，后由 GShard 和 Switch Transformer 在其中插入 MoE 层。
Dropout: A Simple Way to Prevent Neural Networks from Overfitting —— 一种相关的随机条件计算形式。
Language Models are Few-Shot Learners —— 大型稠密语言模型，MoE 扩展工作的目标即是以更低成本超越此类模型。

参考文献

Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., & Dean, J. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. International Conference on Learning Representations (ICLR).
Bengio, Y., Léonard, N., & Courville, A. (2013). Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation.
Eigen, D., Ranzato, M. A., & Sutskever, I. (2013). Learning Factored Representations in a Deep Mixture of Experts.
Jacobs, R. A., Jordan, M. I., Nowlan, S. J., & Hinton, G. E. (1991). Adaptive Mixtures of Local Experts. Neural Computation, 3(1), 79–87.
Jordan, M. I., & Jacobs, R. A. (1994). Hierarchical Mixtures of Experts and the EM Algorithm. Neural Computation, 6(2), 181–214.
Lepikhin, D. et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.
Fedus, W., Zoph, B., & Shazeer, N. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.