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

论文 · arXiv 1701.06538

超大规模神经网络：稀疏门控混合专家层 (Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer)

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-01-23cs.LGcs.CLcs.NEstat.ML arXiv PDF

用 XYZNOT 阅读（全文翻译与 AI 讲解）

三句话总结

🎯 稠密神经网络在扩展模型容量以吸收海量数据集时，会遭受计算成本的二次方级激增。
🔧 提出了稀疏门控混合专家（MoE）层，利用 Noisy Top-K 门控和跨分布式 GPU 分片的辅助负载均衡损失。
📈 将语言和翻译模型扩展至高达 1370 亿参数，以更低的计算成本将已发表最佳 LM 困惑度降低了 24% 至 39%。

摘要

神经网络吸收信息的能力受限于其参数量。条件计算（Conditional computation）在理论上被提出作为一种大幅增加模型容量而无需成比例增加计算量的方法，即根据每个样本动态激活网络的部分组件。然而在实践中，存在显著的算法和性能挑战。在这项工作中，我们解决了这些挑战，终于兑现了条件计算的承诺，在现代 GPU 集群上仅损失极少计算效率的情况下，实现了超过 1000 倍的模型容量提升。我们提出了稀疏门控混合专家层（Sparsely-Gated Mixture-of-Experts layer，简称 MoE），它由多达数千个前馈子网络组成。一个可训练的门控网络决定为每个样本采用这些专家的稀疏组合。我们将 MoE 应用于语言建模和机器翻译任务中，在这些任务中，模型容量对于吸收训练语料库中蕴含的大量知识至关重要。我们提出的模型架构中，参数量高达 1370 亿的 MoE 以卷积方式应用在堆叠的 LSTM 层之间。在大型语言建模和机器翻译基准测试中，这些模型以更低的计算成本取得了显著优于最先进水平（state-of-the-art）的结果。

英文原文摘要

The capacity of a neural network to absorb information is limited by its number of parameters. Conditional computation, where parts of the network are active on a per-example basis, has been proposed in theory as a way of dramatically increasing model capacity without a proportional increase in computation. In practice, however, there are significant algorithmic and performance challenges. In this work, we address these challenges and finally realize the promise of conditional computation, achieving greater than 1000x improvements in model capacity with only minor losses in computational efficiency on modern GPU clusters. We introduce a Sparsely-Gated Mixture-of-Experts layer (MoE), consisting of up to thousands of feed-forward sub-networks. A trainable gating network determines a sparse combination of these experts to use for each example. We apply the MoE to the tasks of language modeling and machine translation, where model capacity is critical for absorbing the vast quantities of knowledge available in the training corpora. We present model architectures in which a MoE with up to 137 billion parameters is applied convolutionally between stacked LSTM layers. On large language modeling and machine translation benchmarks, these models achieve significantly better results than state-of-the-art at lower computational cost.

论文解读

问题与动机

深度学习模型在海量数据集上训练时，极大地受益于参数容量的提升。然而，标准架构对每个输入 token 或样本都会激活整个网络。随着网络规模和数据集体积的双双扩展，这会导致计算成本出现不可承受的激增。

尽管条件计算（按样本仅激活模型的子集）在理论上已被提出，但早期实现在现代硬件加速器上步履维艰，原因包括：

GPU 偏好统一规整的算术运算，而非分支操作。
“批次萎缩问题”（shrinking batch problem）：将样本子集路由后，会导致各个子网络分配到的批次极小且效率低下。
设备间网络带宽瓶颈。
路由不均衡，导致少数子网络过载而其他子网络处于饥饿状态。

核心思想

稀疏门控 MoE 层 (Sparsely-Gated MoE Layer)： 前馈子网络（“专家”）被置于循环层（例如堆叠的 LSTM）之间。每个位置仅评估稀疏的 top- $k$ 选定专家。
Noisy Top-K 门控 (Noisy Top-K Gating)： 在 top- $k$ 选择和 softmax 之前向门控 logits 添加可调的高斯噪声，以提供连续且可反向传播的梯度，并促进跨专家的均衡探索。
数据并行与模型并行相结合： 标准层和门控在各设备间进行数据并行，而专家网络则在 GPU 间进行模型并行切分。所有设备上的批次被组合起来，使各个专家能够接收到足够大的聚合微批次（mini-batches）。
均衡约束： 针对专家重要性（importance）和专家负载（load）施加显式的辅助损失，防止门控网络退化为过度利用少数专家的非均衡解。

方法（包含核心公式）

MoE 层由 $n$ 个前馈专家网络 $E_{1}, \dots, E_{n}$ 以及一个门控网络 $G (x)$ 组成。该层的输出为： $y = i = 1 \sum n G (x)_{i} E_{i} (x)$ 只要 $G (x)_{i} = 0$ ，即可避免相应的计算。

Noisy Top-K 门控

为了将稀疏性与可微门控相结合，在投影线性变换中引入了噪声： $H (x)_{i} = (x \cdot W_{g})_{i} + StandardNormal () \cdot Softplus ((x \cdot W_{noise})_{i})$ $KeepTopK (v, k)_{i} = {v_{i} - \infty 如果 v_{i} 属于 v 中前 k 个元素其他情况$ $G (x) = Softmax (KeepTopK (H (x), k))$

负载均衡损失

为了促使每个专家的总门控权重均等，通过变异系数（ $C V$ ）在批次 $X$ 上定义了重要性损失： $Importance (X) = x \in X \sum G (x)$ $L_{importance} (X) = w_{importance} \cdot C V (Importance (X))^{2}$ 为了确保分配给每个专家的样本数量均衡，平滑负载估计器利用了正态分布累积分布函数 $Φ$ ： $P (x, i) = Φ (\frac{( x \cdot W _{g} ) _{i} - kth_excluding ( H ( x ) , k , i )}{Softplus (( x \cdot W _{noise} ) _{i} )})$ $Load (X)_{i} = x \in X \sum P (x, i)$ $L_{load} (X) = w_{load} \cdot C V (Load (X))^{2}$

结果（包含核心指标数据）

10 亿词基准测试 (1 Billion Word Benchmark)： 计算量匹配的 MoE-4096-h 模型取得了 34.1 的测试困惑度（相比计算量相当的基线 LSTM 降低了 24%）。高预算 MoE 模型在 32 块 K40 GPU 上用时 47 小时达到了 28.0 的测试困惑度，超越了先前 34.7 的最先进水平（后者在 32 块 K40 GPU 上耗时 59 小时）。
1000 亿词语料库： 扩展至 65,536 个专家（680 亿参数）时，测试困惑度相比稠密基线降低了 39%（28.9 对比 47.0），同时保持了 0.72 TFLOPS/GPU 的运算吞吐。
机器翻译： MoE 增强版 GNMT 在 WMT'14 英 $\to$ 法任务上取得了 40.56 BLEU（相比无强化学习的 GNMT 提升 +1.34 BLEU），在 WMT'14 英 $\to$ 德任务上达到了 26.03 BLEU（提升 +1.12 BLEU），且每个时间步所需的计算操作量减少了一半以上。
多语言翻译： 单个多语言 MoE 模型在 12 个语言对中的 11 个上均超越了多语言 GNMT 基线（在韩语 $\to$ 英语上最高提升 +5.84 BLEU），并将验证集困惑度降低了 19%。

局限性与未来方向

极致稀疏性可能导致收益递减甚至性能倒退；在 1000 亿词数据集上，从 65,536 个专家扩展到 131,072 个专家时，困惑度出现了劣化。
分布式通信开销仍然依赖于隐藏层大小与输入/输出尺寸的比例，这使得较小的专家容量在带宽利用效率上偏低。
在 RNN 递归步内应用 MoE 会破坏跨序列位置的批展开（batch-unrolling）优化，需要采用激活值重计算方案。

关键术语

Sparsely-Gated Mixture-of-Experts Layer (MoE)（稀疏门控混合专家层）: 一种包含多个并行专家子网络的神经网络层，由可训练的门控网络为每个输入样本仅选择专家的一个稀疏子集。
Conditional Computation（条件计算）: 一种架构范式，其中神经网络的各个部分按样本动态激活或休眠，以在不按比例增加计算量的情况下提升参数容量。
Noisy Top-K Gating（含噪 Top-K 门控）: 一种门控机制，在仅保留 top-k 个值并应用 softmax 之前，向 logits 添加可调高斯噪声，从而实现反向传播并防止路由退化。
Shrinking Batch Problem（批次萎缩问题）: 当输入被稀疏划分到众多专家时，各个子网络接收到的有效批大小降低，从而导致硬件算术效率下降的现象。
Importance Loss（重要性损失）: 一种辅助损失，通过惩罚批次内各专家门控值之和的变异系数平方，以促进专家分配的均等。
Load Loss（负载损失）: 一种辅助目标函数，通过最小化跨专家的 token 估计概率加权分配数量的方差，从而平衡硬件负载。
Hierarchical Mixture of Experts（分层混合专家）: 一种两级 MoE 结构，主门控网络将输入路由至次级 MoE 集群，以降低分支因子和分布式通信开销。

常见问题

门控网络如何实现离散 top-k 路由的反向传播？

由于 top-k 筛选是在 softmax 之前应用的，选定专家的非零门控值对输入权重依然保留连续导数，使得梯度能够正常回传流经门控网络。

跨分布式设备时的批次萎缩问题是如何解决的？

该模型混合使用了数据并行和模型并行：标准层在各设备间处理独立的微批次，但所有设备上的 token 会在路由到分布在不同 GPU 上的驻留专家之前进行聚合。

为什么门控网络中需要引入噪声？

注入的高斯噪声在训练初期促进了对各专家的探索，并为负载均衡损失提供了平滑的估计器，防止了自我强化的专家利用退化。

当模型扩展到 131,072 个专家时会发生什么？

困惑度从 28.9（65,536 个专家时）轻微退化至 29.2（131,072 个专家时），且计算效率下降至 0.30 TFLOPS/GPU，这很可能是由于极致稀疏性以及批大小固定未同步扩展导致的。

本页解读由 AI 根据论文生成，可能有错误；细节请以原文为准。

用 XYZNOT 阅读（全文翻译与 AI 讲解）