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

    From Marovi AI
    Other languages:
    SummarySource
    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

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

    Noam Shazeer Google Brain, {noam,azalia,andydavis,qvl,geoffhinton,jeff}@google.com Azalia Mirhoseini 同等主要贡献者作为 Google Brain Residency 计划成员完成的工作 (g.co/brainresidency) Google Brain, {noam,azalia,andydavis,qvl,geoffhinton,jeff}@google.com Krzysztof Maziarz 雅盖隆大学,克拉科夫,krzysztof.maziarz@student.uj.edu.pl Andy Davis Google Brain, {noam,azalia,andydavis,qvl,geoffhinton,jeff}@google.com Quoc Le Google Brain, {noam,azalia,andydavis,qvl,geoffhinton,jeff}@google.com Geoffrey Hinton Google Brain, {noam,azalia,andydavis,qvl,geoffhinton,jeff}@google.com Jeff Dean Google Brain, {noam,azalia,andydavis,qvl,geoffhinton,jeff}@google.com

    摘要

    神经网络吸收信息的能力受到其参数数量的限制。条件计算(即按样本激活网络的不同部分)在理论上被提出作为一种在不按比例增加计算量的情况下大幅提升模型容量的方法。然而在实践中,存在显著的算法和性能挑战。在这项工作中,我们解决了这些挑战,最终实现了条件计算的潜力,在现代 GPU 集群上以仅有的少许计算效率损失实现了模型容量超过 1000 倍的提升。我们引入了 Sparsely-Gated Mixture-of-Experts(MoE)层,由多达数千个前馈子网络组成。一个可训练的门控网络决定每个样本要使用的专家的稀疏组合。我们将 MoE 应用于语言建模和机器翻译任务,在这些任务中,模型容量对吸收训练语料库中海量知识至关重要。我们提出了一些模型架构,其中具有高达 1370 亿参数的 MoE 以卷积方式应用于堆叠的 LSTM 层之间。在大型语言建模和机器翻译基准上,这些模型以更低的计算成本取得了显著优于当时最新最优的结果。

    1 引言与相关工作

    1.1 条件计算

    在训练数据和模型规模上同时利用规模化一直是深度学习成功的核心。当数据集足够大时,增加神经网络的容量(参数数量)可以带来更好的预测准确率。这在文本 (Sutskever et al., 2014; Bahdanau et al., 2014; Jozefowicz et al., 2016; Wu et al., 2016)、图像 (Krizhevsky et al., 2012; Le et al., 2012) 和语音 (Hinton et al., 2012; Amodei et al., 2015) 等领域均已得到验证。对于典型的深度学习模型——整个模型针对每个样本都被激活——这会随模型规模和训练样本数的增加导致训练成本近似呈二次增长。不幸的是,计算能力和分布式计算的进展不足以满足这种需求。

    已有多种条件计算的形式被提出,以期在不按比例增加计算成本的前提下提升模型容量 (Davis & Arel, 2013; Bengio et al., 2013; Eigen et al., 2013; Ludovic Denoyer, 2014; Cho & Bengio, 2014; Bengio et al., 2015; Almahairi et al., 2015)。在这些方案中,网络的大部分按样本被激活或停用。门控决策可以是二值的,也可以是稀疏且连续的;可以是随机的,也可以是确定性的。为训练门控决策,研究者提出了多种强化学习和反向传播的形式。

    尽管这些思路在理论上很有前景,但迄今为止还没有任何工作在模型容量、训练时间或模型质量上取得大幅提升。我们将其归因于以下挑战的综合作用:

    • 现代计算设备,尤其是 GPU,对算术运算远比对分支运算快。上述大多数工作都意识到了这一点,并提出根据每次门控决策开启/关闭网络的大块。

    • 大批量对性能至关重要,因为它能摊销参数传输和更新的开销。条件计算会减小条件激活部分的批量大小。

    • 网络带宽可能成为瓶颈。一个 GPU 集群的计算能力可能比设备间总带宽高出数千倍。要在计算上保持高效,算法的计算需求与网络需求之比必须超过这一比例。embedding 层——可视为一种条件计算形式——正是受此问题制约。由于 embedding 通常需要通过网络发送,(样本, 参数) 交互的次数受网络带宽而非计算容量限制。

    • 视具体方案而定,可能需要损失项以达到所需的逐块和/或逐样本稀疏度。Bengio et al. (2015) 使用了三个此类损失项。这些问题可能同时影响模型质量与负载均衡。

    • 模型容量对超大数据集最为关键。现有关于条件计算的文献主要处理相对较小的图像识别数据集,最多包含约 60 万张图像。很难想象这些图像的标签能提供足够的信号来充分训练具有数百万、更不用说数十亿参数的模型。

    在这项工作中,我们首次同时解决了上述全部挑战,并最终实现了条件计算的承诺。我们以仅极小的计算效率损失获得了模型容量超过 1000 倍的提升,并在公开的语言建模和翻译数据集上显著刷新了当时的最新最优结果。

    1.2 我们的方法:稀疏门控的 Mixture-of-Experts 层

    我们对条件计算的方法是引入一种新的通用神经网络组件:Sparsely-Gated Mixture-of-Experts(MoE)层。MoE 由若干个专家组成,每个都是一个简单的前馈神经网络;同时还有一个可训练的门控网络,用于为每个输入选择一组稀疏的专家组合(见图 1)。网络的所有部分通过反向传播联合训练。

    Refer to caption

    尽管所提出的技术是通用的,本文聚焦于已知能从超大模型中显著获益的语言建模和机器翻译任务。具体而言,我们将 MoE 以卷积方式应用于堆叠的 LSTM 层之间 (Hochreiter & Schmidhuber, 1997),如图 1 所示。MoE 在文本的每个位置被调用一次,并在每个位置可能选择不同的专家组合。不同的专家往往会高度专精于句法和语义特征(参见附录 E9)。在语言建模和机器翻译两类基准上,我们以仅占已发表最佳结果一小部分的计算量便取得了更优的成绩。

    1.3 关于 Mixtures of Experts 的相关工作

    自二十多年前被提出以来 (Jacobs et al., 1991; Jordan & Jacobs, 1994),mixture-of-experts 方法已被广泛研究。研究者提出了多种类型的专家架构,例如 SVM (Collobert et al., 2002)、高斯过程 (Tresp, 2001; Theis & Bethge, 2015; Deisenroth & Ng, 2015)、狄利克雷过程 (Shahbaba & Neal, 2009) 以及深度网络。其他工作则关注不同的专家配置,例如层次化结构 (Yao et al., 2009)、无限数量的专家 (Rasmussen & Ghahramani, 2002) 以及按序添加专家 (Aljundi et al., 2016)。Garmash & Monz (2016) 为机器翻译提出了一种以 mixture of experts 形式呈现的集成模型,其门控网络在一个预训练的集成 NMT 模型之上进行训练。

    上述工作针对的是顶层的 mixture of experts,即整个模型本身就是 mixture of experts。Eigen et al. (2013) 引入了将多个带有自身门控网络的 MoE 作为深度模型组成部分的思路。这种做法在直觉上更强大,因为复杂问题可能包含许多子问题,每个子问题需要不同的专家。他们在结论中也提到了引入稀疏性的潜力,从而使 MoE 成为条件计算的载体。

    我们的工作建立在将 MoE 作为通用神经网络组件这一用法之上。Eigen et al. (2013) 使用两个堆叠的 MoE,从而允许两组门控决策;而我们将 MoE 以卷积方式应用,使得文本的每个位置都可以做出不同的门控决策。我们还实现了稀疏门控,并证明它是一种大幅扩展模型容量的实用手段。

    2 Mixture-of-Experts 层的结构

    Mixture-of-Experts(MoE)层由一组 $ {\textstyle n} $ 个 "expert networks" $ {\textstyle E_{1},\cdots,E_{n}} $ 以及一个 "gating network" $ {\textstyle G} $ 组成,门控网络的输出是一个稀疏的 $ {\textstyle n} $ 维向量。图 1 给出了 MoE 模块的总体视图。专家本身就是神经网络,每个都有自己的参数。原则上我们只要求专家接受相同尺寸的输入并产生相同尺寸的输出,但在本文的初步研究中,我们将其限制为架构相同、参数独立的前馈网络。

    $ {\textstyle G\hspace{0pt}{(x)}} $ 表示门控网络的输出,$ {\textstyle E_{i}\hspace{0pt}{(x)}} $ 表示第 $ {\textstyle i} $ 个专家网络对给定输入 $ {\textstyle x} $ 的输出。MoE 模块的输出 $ {\textstyle y} $ 可以写为:

    $ {\displaystyle y = {\sum\limits_{i = 1}^{n}{G\hspace{0pt}{(x)}_{i}\hspace{0pt}E_{i}\hspace{0pt}{(x)}}}} $ (1)

    我们基于 $ {\textstyle G\hspace{0pt}{(x)}} $ 输出的稀疏性来节省计算量。凡 $ {\textstyle {G\hspace{0pt}{(x)}_{i}} = 0} $ 之处,无需计算 $ {\textstyle E_{i}\hspace{0pt}{(x)}} $。在实验中,我们最多使用了数千个专家,但每个样本只需评估其中少数几个。如果专家数量非常大,可以使用两层分层 MoE 来降低分支因子:分层 MoE 中,主门控网络选择若干 "experts" 的稀疏加权组合,而每个所谓的 "expert" 本身又是带有自身门控网络的二级 mixture-of-experts。下文聚焦于普通 MoE,关于分层 MoE 的更多细节见附录 B

    我们的实现与其他条件计算模型有关。专家为简单权重矩阵的 MoE 与 (Cho & Bengio, 2014) 中提出的参数化权重矩阵类似。专家拥有一层隐藏层的 MoE 则类似于 (Bengio et al., 2015) 中描述的块级 dropout,其中被 dropout 的层夹在两个完全激活的层之间。

    2.1 门控网络

    Softmax 门控:

    一种简单的非稀疏门控函数选择 (Jordan & Jacobs, 1994) 是将输入乘以一个可训练的权重矩阵 $ {\textstyle W_{g}} $,然后应用 $ {\textstyle S\hspace{0pt}o\hspace{0pt}f\hspace{0pt}t\hspace{0pt}m\hspace{0pt}a\hspace{0pt}x} $ 函数。

    $ {\displaystyle {G_{\sigma}\hspace{0pt}{(x)}} = {S\hspace{0pt}o\hspace{0pt}f\hspace{0pt}t\hspace{0pt}m\hspace{0pt}a\hspace{0pt}x\hspace{0pt}{({x \cdot W_{g}})}}} $ (2)

    Noisy Top-K 门控:

    我们在 Softmax 门控网络中加入两个成分:稀疏性与噪声。在应用 softmax 之前,我们加入可调的高斯噪声,然后只保留 top k 个最大值,将其余设为 $ {\textstyle - \infty} $(从而使对应的门值等于 $ {\textstyle 0} $)。稀疏性如前所述用于节省计算。虽然这种稀疏方式在理论上会在门控函数的输出中造成一些令人担忧的不连续性,但我们在实践中尚未观察到这成为问题。噪声项有助于负载均衡,详见附录 A。每个分量的噪声幅度由第二个可训练权重矩阵 $ {\textstyle W_{n\hspace{0pt}o\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}} $ 控制。

    $ {\displaystyle {G\hspace{0pt}{(x)}} = {S\hspace{0pt}o\hspace{0pt}f\hspace{0pt}t\hspace{0pt}m\hspace{0pt}a\hspace{0pt}x\hspace{0pt}{({K\hspace{0pt}e\hspace{0pt}e\hspace{0pt}p\hspace{0pt}T\hspace{0pt}o\hspace{0pt}p\hspace{0pt}K\hspace{0pt}{({H\hspace{0pt}{(x)}},k)}})}}} $ (3)
    $ {\displaystyle {H\hspace{0pt}{(x)}_{i}} = {{({x \cdot W_{g}})}_{i} + {{{S\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}d\hspace{0pt}a\hspace{0pt}r\hspace{0pt}d\hspace{0pt}N\hspace{0pt}o\hspace{0pt}r\hspace{0pt}m\hspace{0pt}a\hspace{0pt}l\hspace{0pt}{()}} \cdot S}\hspace{0pt}o\hspace{0pt}f\hspace{0pt}t\hspace{0pt}p\hspace{0pt}l\hspace{0pt}u\hspace{0pt}s\hspace{0pt}{({({x \cdot W_{n\hspace{0pt}o\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}})}_{i})}}}} $ (4)
    $ {\displaystyle {K\hspace{0pt}e\hspace{0pt}e\hspace{0pt}p\hspace{0pt}T\hspace{0pt}o\hspace{0pt}p\hspace{0pt}K\hspace{0pt}{(v,k)}_{i}} = \begin{cases} v_{i} & {\text{if~}v_{i}\text{~is in the top~}k\text{~elements of~}v\text{.}} \\ {- \infty} & \text{otherwise.} \end{cases}} $ (5)

    训练门控网络

    我们通过简单的反向传播与模型其余部分一起训练门控网络。若选择 $ {\textstyle k > 1} $,则 top k 个专家的门值对门控网络权重具有非零梯度。这种偶发性敏感的行为在 (Bengio et al., 2013) 中关于 noisy rectifier 的部分有所描述。梯度也会通过门控网络反向传播到其输入。我们在此与 (Bengio et al., 2015) 不同——他们使用布尔门并采用类似 REINFORCE 的方法来训练门控网络。

    3 解决性能挑战

    3.1 批量缩小问题

    在现代 CPU 和 GPU 上,大批量对计算效率至关重要,以便摊销参数加载与更新的开销。如果门控网络为每个样本从 $ {\textstyle n} $ 个专家中挑选 $ {\textstyle k} $ 个,那么对于一个 $ {\textstyle b} $ 个样本的批次,每个专家平均只收到大约 $ {\textstyle \frac{k\hspace{0pt}b}{n} \ll b} $ 个样本,规模小得多。这使得朴素的 MoE 实现随着专家数量的增加变得非常低效。解决这一批量缩小问题的办法是把原始批量尽可能加大。然而批量大小往往受到前向与反向传播之间存储激活所需内存的限制。我们提出以下扩大批量大小的技术:

    数据并行与模型并行的混合:

    在传统的分布式训练设置中,多个模型副本在不同设备上异步处理不同的数据批次,参数通过一组 parameter server 进行同步。在我们的方法中,这些不同的批次同步运行,以便能在 MoE 层合并。我们按常规的数据并行方案分布模型的标准层和门控网络,但每个专家只保留一份共享副本。MoE 层中的每个专家接收一个合并的批次,由所有数据并行输入批次中的相关样本组成。同一组设备同时充当数据并行副本(用于标准层和门控网络)和模型并行分片(每个分片承载一部分专家)。如果模型分布在 $ {\textstyle d} $ 个设备上,每个设备处理大小为 $ {\textstyle b} $ 的批次,那么每个专家收到的批次大约为 $ {\textstyle \frac{k\hspace{0pt}b\hspace{0pt}d}{n}} $ 个样本。因此每个专家批量大小提升了 $ {\textstyle d} $ 倍。

    对于分层 MoE(见 第 B 节),主门控网络采用数据并行,次级 MoE 采用模型并行。每个次级 MoE 驻留在单个设备上。

    这一技术使我们能通过按比例增加训练集群中的设备数来增加专家数量(从而增加参数数量)。总批量大小随之增加,每个专家的批量大小保持恒定;每设备的内存和带宽需求也保持不变,每步耗时不变,处理与模型参数总数相同数量级训练样本所需的时间同样不变。我们的目标是用一万亿词的语料训练一个一万亿参数的模型。截至本文写作时,我们尚未将系统扩展到如此规模,但通过添加更多硬件应当可以实现。

    利用卷积特性:

    在我们的语言模型中,我们将同一 MoE 应用到上一层的每个时间步上。如果等待上一层完成,我们就可以把所有时间步一起作为一个大批次应用 MoE。这样做会将 MoE 层的输入批量按展开时间步数倍数地放大。

    为递归 MoE 增大批量大小:

    我们怀疑功能更强的模型可能会以递归方式应用 MoE。例如,LSTM 或其他 RNN 的权重矩阵可以被 MoE 替代。遗憾的是,这类模型破坏了上一段中的卷积技巧,因为某个时间步对 MoE 的输入依赖于上一时间步 MoE 的输出。Gruslys et al. (2016) 描述了一种以重计算前向激活为代价、大幅减少展开 RNN 中存储激活数量的技术,这将允许批量大小大幅增加。

    3.2 网络带宽

    分布式计算中另一个重要的性能问题是网络带宽。由于专家是固定的(见上文)且门控参数数量很少,大部分通信涉及通过网络发送专家的输入和输出。为了保持计算效率,专家计算量与其输入/输出大小之比必须超过计算设备的计算能力与网络能力之比。对于 GPU 而言,这一比值可能高达数千比一。在我们的实验中,我们使用具有一层 ReLU 激活、含数千单元隐藏层的专家。由于专家的权重矩阵大小为 $ {\textstyle i\hspace{0pt}n\hspace{0pt}p\hspace{0pt}u\hspace{0pt}t} $_$ {\textstyle {{s\hspace{0pt}i\hspace{0pt}z\hspace{0pt}e} \times h}\hspace{0pt}i\hspace{0pt}d\hspace{0pt}d\hspace{0pt}e\hspace{0pt}n} $_$ {\textstyle s\hspace{0pt}i\hspace{0pt}z\hspace{0pt}e} $$ {\textstyle h\hspace{0pt}i\hspace{0pt}d\hspace{0pt}d\hspace{0pt}e\hspace{0pt}n} $_$ {\textstyle {{s\hspace{0pt}i\hspace{0pt}z\hspace{0pt}e} \times o}\hspace{0pt}u\hspace{0pt}t\hspace{0pt}p\hspace{0pt}u\hspace{0pt}t} $_$ {\textstyle s\hspace{0pt}i\hspace{0pt}z\hspace{0pt}e} $,计算量与输入/输出之比等于隐藏层大小。便利的是,我们可以简单地通过使用更大的隐藏层或更多的隐藏层来提高计算效率。

    4 专家利用的均衡

    我们观察到,门控网络往往收敛到这样一种状态:始终对同样少数几个专家产生较大的权重。这种不平衡是自我强化的,因为受偏好的专家训练得更快,从而进一步被门控网络选中。Eigen et al. (2013) 描述了相同的现象,并在训练初期使用硬约束来避免这一局部最小值。Bengio et al. (2015) 在每个门的批次平均值上加入了软约束。111Bengio et al. (2015) 还包含两个额外的损失项。其一控制每个样本的稀疏性,由于固定的 $ {\textstyle k} $ 已经强制了这一点,我们并不需要它。第三个损失鼓励门值的多样性。在我们的实验中,我们发现随着专家专精化,门值会自然分化(形成良性循环),无需强制门值多样性。

    我们采用软约束的方法。我们将一个专家相对于一个训练样本批次的 importance 定义为该专家门值在批次内的总和。我们定义一个额外损失 $ {\textstyle L_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e}} $,将其加到模型的总损失函数中。该损失等于 importance 值集合的变异系数的平方乘以一个手动调整的缩放因子 $ {\textstyle w_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e}} $,从而促使所有专家具有同等的 importance。

    $ {\displaystyle {I\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e\hspace{0pt}{(X)}} = {\sum\limits_{x \in X}{G\hspace{0pt}{(x)}}}} $ (6)
    $ {\displaystyle {L_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e}\hspace{0pt}{(X)}} = {{w_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e} \cdot C}\hspace{0pt}V\hspace{0pt}{({I\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e\hspace{0pt}{(X)}})}^{2}}} $ (7)

    虽然该损失函数可以确保 importance 相等,但专家仍可能收到数量差异很大的样本。例如,一个专家可能收到少量带大权重的样本,而另一个专家可能收到大量带小权重的样本。这在分布式硬件上会导致内存和性能问题。为解决该问题,我们引入第二个损失函数 $ {\textstyle L_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d}} $,确保负载均衡。附录 A 给出了该函数的定义以及实验结果。

    5 实验

    5.1 1 Billion Word Language Modeling Benchmark

    数据集:

    该数据集由 (Chelba et al., 2013) 引入,由打乱的新闻文章中的唯一句子组成,总计约 8.29 亿词,词表为 793,471 个词。

    此前的最新最优:

    此前已发表的最佳结果 (Jozefowicz et al., 2016) 使用由一层或多层堆叠的 Long Short-Term Memory (LSTM) 层 (Hochreiter & Schmidhuber, 1997; Gers et al., 2000) 组成的模型。这些模型 LSTM 层的参数量从 2M 到 151M 不等。质量随参数量大幅提升,但计算成本也随之增加。这些模型的结果构成了图 2-右的顶线。

    MoE 模型:

    我们的模型由两层堆叠的 LSTM 以及位于其间的一层 MoE 组成(见图 1)。我们改变各层的大小以及专家数量。关于模型架构、训练方案、额外基线和结果的完整细节见附录 C

    低计算量、不同容量:

    为研究增加容量的效果,我们训练了一系列 MoE 模型,其计算量大致相同:在前向传播中,每个训练样本每时间步约有 800 万次乘加运算(不计 softmax 层)。我们将这一指标称为 (ops/timestep)。我们训练了包含 4、32 和 256 个专家的扁平 MoE 模型,以及包含 256、1024 和 4096 个专家的分层 MoE 模型。每个专家约有 100 万个参数。对所有 MoE 层,每个输入激活 4 个专家。

    这些模型的结果如图 2-左所示。具有 4 个始终激活专家的模型表现(不出意外地)与计算量匹配的基线模型类似,而其中最大的模型(4096 个专家)在测试集上取得了令人瞩目的、降低 24% 的困惑度。

    Refer to caption Refer to caption

    Test Test #Parameters ops/timestep Training TFLOPS
    Perplexity Perplexity excluding embedding Time /GPU
    10 epochs 100 epochs and softmax layers 10 epochs
    Best Published Results 34.7 30.6 151 million 151 million 59 hours, 32 k40s 1.09
    Low-Budget MoE Model 34.1 4303 million 8.9 million 15 hours, 16 k40s 0.74
    Medium-Budget MoE Model 31.3 4313 million 33.8 million 17 hours, 32 k40s 1.22
    High-Budget MoE Model 28.0 4371 million 142.7 million 47 hours, 32 k40s 1.56

    不同计算量、高容量:

    除了上一节中最大的模型外,我们还训练了另外两个具有类似高容量(40 亿参数)但计算预算更大的 MoE 模型。这些模型具有更大的 LSTM 以及更少但更大的专家。详细信息见附录 C.2。这三个模型的结果构成了图 2-右的底线。表 1 将这些模型的结果与该数据集上此前最佳已发表结果进行了对比。即使这其中最快的模型,在训练 epoch 数相同的情况下也超过了已发表的最佳结果,尽管只需要 6% 的计算量。

    计算效率:

    我们使用 TensorFlow (Abadi et al., 2016) 在配备 16–32 块 Tesla K40 GPU 的集群上训练模型。对于每个模型,我们将处理一个训练批次所需的浮点运算次数除以观察到的步耗时和集群中的 GPU 数,得到以 TFLOPS/GPU 计的计算效率。这里使用的运算计数高于我们在 ops/timestep 中报告的值,因为我们包含了反向传播、softmax 层基于 importance sampling 的训练,并将一次 multiply-and-add 计为两次独立运算。对于我们所有的 MoE 模型,专家中的浮点运算占总运算的 37% 至 46%。

    对于不含 MoE 的基线模型,观察到的计算效率在 1.07–1.29 TFLOPS/GPU 之间。对于低计算量 MoE 模型,效率在 0.74–0.90 TFLOPS/GPU 之间,除了 4 专家模型没有充分利用可用并行性。我们计算量最高的 MoE 模型效率更高,达 1.56 TFLOPS/GPU,可能得益于更大的矩阵。这些数字相比 NVIDIA 宣称的理论峰值 4.29 TFLOPS/GPU 已是相当可观。详细结果见附录 C7

    5.2 100 Billion Word Google News Corpus

    Refer to caption

    在 1 Billion Word 语料上,当 MoE 层中的参数量超过 10 亿时,进一步增加容量看起来收益递减,如图 2-左所示。我们假设在更大的训练集上,更高的容量将带来显著的质量提升。

    我们构造了一个类似的训练集,由 Google 内部新闻语料中打乱后的唯一句子组成,总计约 1000 亿词。与上一节类似,我们测试了一系列计算成本约为 8M ops/timestep 的模型。除一个基线 LSTM 模型外,我们还训练了在 MoE 层中包含 32、256、1024、4096、16384、65536 和 131072 个专家的扩展模型。这对应于 MoE 层中最多 1370 亿参数。架构、训练和结果细节见附录 D

    结果:

    3 展示了在 100 亿词(顶线)与 1000 亿词(底线)训练后,测试困惑度随容量的变化。在完整 1000 亿词上训练时,测试困惑度在专家数达到 65536(680 亿参数)之前持续显著改善,比计算匹配基线低 39%,但在 131072 个专家时反而劣化,这可能是过度稀疏的结果。两条线之间不断扩大的差距(不出意外地)表明,模型容量的增加在更大的训练集上获益更多。

    即便是 65536 个专家(层稀疏度 99.994%),模型的计算效率仍维持在尚可的 0.72 TFLOPS/GPU。

    5.3 机器翻译(单一语言对)

    模型架构:

    我们的模型是 (Wu et al., 2016) 中描述的 GNMT 模型的修改版本。为了减少计算,我们将编码器和解码器中的 LSTM 层数分别从 9 和 8 减少到 3 和 2。我们在编码器(第 2 与第 3 层之间)和解码器(第 1 与第 2 层之间)都插入了 MoE 层。每个 MoE 层最多包含 2048 个专家,每个专家约有两百万参数,整体为模型增加约 80 亿参数。关于模型架构、测试过程和结果的更多细节见附录 E

    数据集:

    我们在 WMT'14 En$ {\textstyle \rightarrow} $Fr 和 En$ {\textstyle \rightarrow} $De 语料上进行评测,二者的训练集分别包含 3600 万和 500 万对句子。实验协议也与 (Wu et al., 2016) 类似:newstest2014 作为与此前工作 (Luong et al., 2015a; Zhou et al., 2016; Wu et al., 2016) 比较的测试集,newstest2012 与 newstest2013 的组合作为开发集。我们还在 Google 生产环境的英语到法语数据上测试了同一模型。

    Model Test Test ops/timenstep Total Training
    Perplexity BLEU #Parameters Time
    MoE with 2048 Experts 2.69 40.35 85M 8.7B 3 days/64 k40s
    MoE with 2048 Experts (longer training) 2.63 40.56 85M 8.7B 6 days/64 k40s
    GNMT (Wu et al., 2016) 2.79 39.22 214M 278M 6 days/96 k80s
    GNMT+RL (Wu et al., 2016) 2.96 39.92 214M 278M 6 days/96 k80s
    PBMT (Durrani et al., 2014) 37.0
    LSTM (6-layer) (Luong et al., 2015b) 31.5
    LSTM (6-layer+PosUnk) (Luong et al., 2015b) 33.1
    DeepAtt (Zhou et al., 2016) 37.7
    DeepAtt+PosUnk (Zhou et al., 2016) 39.2
    Model Test Test ops/timestep Total Training
    Perplexity BLEU #Parameters Time
    MoE with 2048 Experts 4.64 26.03 85M 8.7B 1 day/64 k40s
    GNMT (Wu et al., 2016) 5.25 24.91 214M 278M 1 day/96 k80s
    GNMT +RL (Wu et al., 2016) 8.08 24.66 214M 278M 1 day/96 k80s
    PBMT (Durrani et al., 2014) 20.7
    DeepAtt (Zhou et al., 2016) 20.6
    Model Eval Eval Test Test ops/timestep Total Training
    Perplexity BLEU Perplexity BLEU #Parameters Time
    MoE with 2048 Experts 2.60 37.27 2.69 36.57 85M 8.7B 1 day/64 k40s
    GNMT (Wu et al., 2016) 2.78 35.80 2.87 35.56 214M 278M 6 days/96 k80s

    结果:

    234 展示了我们最大模型与已发表结果的对比。我们的方法在 WMT'14 En$ {\textstyle \rightarrow} $Fr 和 En$ {\textstyle \rightarrow} $De 基准上分别取得了 40.56 和 26.03 的 BLEU 分数。由于我们的模型没有使用 RL 细调,这些结果相对于 (Wu et al., 2016) 中的强基线分别带来了 1.34 和 1.12 BLEU 的显著提升,困惑度分数也更优。222所报告的困惑度针对我们模型与 GNMT 共同使用的 tokenization。在 Google 生产环境数据集上,我们的模型即便仅训练了六分之一的时间,仍取得了高出 1.01 的测试 BLEU。

    5.4 多语言机器翻译

    数据集:

    (Johnson et al., 2016) 在涵盖 12 个语言对的非常大的合并数据集上训练单一 GNMT (Wu et al., 2016) 模型。结果略逊于分别为每个语言对单独训练的 12 个单语言对 GNMT 模型。这并不令人意外,因为这 12 个模型加起来的容量与训练量是单一模型的 12 倍。我们以单一 MoE 增强模型重复该实验。模型架构细节见附录 E。我们使用与 (Johnson et al., 2016) 相同的数据集进行训练,并处理相同数量的训练样本(约 30 亿对句子)。由于我们的模型计算预算更低,训练时间更短。

    结果:

    单语言对 GNMT 模型、多语言 GNMT 模型与多语言 MoE 模型的结果见表 5。MoE 模型在开发集上的困惑度比多语言 GNMT 模型低 19%。在 BLEU 分数上,MoE 模型在 12 个语言对中的 11 个上显著优于多语言 GNMT 模型(最多高出 5.84 分),并在 12 个语言对中的 8 个上甚至超过了单语种 GNMT 模型。英语 $ {\textstyle \rightarrow} $ 韩语上的较差表现似乎是严重过训练的结果,因为对较稀有的语言对,训练语料中少量真实样本被大量上采样。

    GNMT-Mono GNMT-Multi MoE-Multi MoE-Multi vs.
    GNMT-Multi
    Parameters 278M / model 278M 8.7B
    ops/timestep 212M 212M 102M
    training time, hardware various 21 days, 96 k20s 12 days, 64 k40s
    Perplexity (dev) 4.14 3.35 -19%
    French $ {\textstyle \rightarrow} $ English Test BLEU 36.47 34.40 37.46 +3.06
    German $ {\textstyle \rightarrow} $ English Test BLEU 31.77 31.17 34.80 +3.63
    Japanese $ {\textstyle \rightarrow} $ English Test BLEU 23.41 21.62 25.91 +4.29
    Korean $ {\textstyle \rightarrow} $ English Test BLEU 25.42 22.87 28.71 +5.84
    Portuguese $ {\textstyle \rightarrow} $ English Test BLEU 44.40 42.53 46.13 +3.60
    Spanish $ {\textstyle \rightarrow} $ English Test BLEU 38.00 36.04 39.39 +3.35
    English $ {\textstyle \rightarrow} $ French Test BLEU 35.37 34.00 36.59 +2.59
    English $ {\textstyle \rightarrow} $ German Test BLEU 26.43 23.15 24.53 +1.38
    English $ {\textstyle \rightarrow} $ Japanese Test BLEU 23.66 21.10 22.78 +1.68
    English $ {\textstyle \rightarrow} $ Korean Test BLEU 19.75 18.41 16.62 -1.79
    English $ {\textstyle \rightarrow} $ Portuguese Test BLEU 38.40 37.35 37.90 +0.55
    English $ {\textstyle \rightarrow} $ Spanish Test BLEU 34.50 34.25 36.21 +1.96

    6 结论

    这项工作首次证明了条件计算可在深度网络中带来巨大收益。我们仔细识别了条件计算的设计要点与挑战,并通过算法与工程相结合的方案加以解决。尽管我们聚焦于文本,条件计算只要有足够大的训练集,也可能在其他领域带来帮助。我们期待在未来几年看到条件计算的更多新颖实现与应用。

    致谢

    感谢 Google Brain 与 Google Translate 团队所有为本项目提供帮助的成员,特别是 Zhifeng Chen、Yonghui Wu 和 Melvin Johnson。也感谢匿名 ICLR 审稿人对改进本文提出的有益建议。

    参考文献

    • Abadi et al. (2016) Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Gregory S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian J. Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Józefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mané, Rajat Monga, Sherry Moore, Derek Gordon Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul A. Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda B. Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. Tensorflow: Large-scale machine learning on heterogeneous distributed systems. CoRR, abs/1603.04467, 2016. URL http://arxiv.org/abs/1603.04467.
    • Aljundi et al. (2016) Rahaf Aljundi, Punarjay Chakravarty, and Tinne Tuytelaars. Expert gate: Lifelong learning with a network of experts. CoRR, abs/1611.06194, 2016. URL http://arxiv.org/abs/1611.06194.
    • Almahairi et al. (2015) A. Almahairi, N. Ballas, T. Cooijmans, Y. Zheng, H. Larochelle, and A. Courville. Dynamic Capacity Networks. ArXiv e-prints, November 2015.
    • Amodei et al. (2015) Dario Amodei, Rishita Anubhai, Eric Battenberg, Carl Case, Jared Casper, Bryan Catanzaro, Jingdong Chen, Mike Chrzanowski, Adam Coates, Greg Diamos, Erich Elsen, Jesse Engel, Linxi Fan, Christopher Fougner, Tony Han, Awni Y. Hannun, Billy Jun, Patrick LeGresley, Libby Lin, Sharan Narang, Andrew Y. Ng, Sherjil Ozair, Ryan Prenger, Jonathan Raiman, Sanjeev Satheesh, David Seetapun, Shubho Sengupta, Yi Wang, Zhiqian Wang, Chong Wang, Bo Xiao, Dani Yogatama, Jun Zhan, and Zhenyao Zhu. Deep speech 2: End-to-end speech recognition in english and mandarin. arXiv preprint arXiv:1512.02595, 2015.
    • Bahdanau et al. (2014) Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473, 2014.
    • Bengio et al. (2015) Emmanuel Bengio, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models. arXiv preprint arXiv:1511.06297, 2015.
    • Bengio et al. (2013) Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.
    • Chelba et al. (2013) Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv:1312.3005, 2013.
    • Cho & Bengio (2014) K. Cho and Y. Bengio. Exponentially Increasing the Capacity-to-Computation Ratio for Conditional Computation in Deep Learning. ArXiv e-prints, June 2014.
    • Collobert et al. (2002) Ronan Collobert, Samy Bengio, and Yoshua Bengio. A parallel mixture of SVMs for very large scale problems. Neural Computing, 2002.
    • Davis & Arel (2013) Andrew Davis and Itamar Arel. Low-rank approximations for conditional feedforward computation in deep neural networks. arXiv preprint arXiv:1312.4461, 2013.
    • Deisenroth & Ng (2015) Marc Peter Deisenroth and Jun Wei Ng. Distributed Gaussian processes. In ICML, 2015.
    • Duchi et al. (2010) John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization, 2010.
    • Durrani et al. (2014) Nadir Durrani, Barry Haddow, Philipp Koehn, and Kenneth Heafield. Edinburgh’s phrase-based machine translation systems for wmt-14. In Proceedings of the Ninth Workshop on Statistical Machine Translation, 2014.
    • Eigen et al. (2013) David Eigen, Marc’Aurelio Ranzato, and Ilya Sutskever. Learning factored representations in a deep mixture of experts. arXiv preprint arXiv:1312.4314, 2013.
    • Garmash & Monz (2016) Ekaterina Garmash and Christof Monz. Ensemble learning for multi-source neural machine translation. In staff.science.uva.nl/c.monz, 2016.
    • Gers et al. (2000) Felix A. Gers, Jürgen A. Schmidhuber, and Fred A. Cummins. Learning to forget: Continual prediction with lstm. Neural Computation, 2000.
    • Gruslys et al. (2016) Audrunas Gruslys, Rémi Munos, Ivo Danihelka, Marc Lanctot, and Alex Graves. Memory-efficient backpropagation through time. CoRR, abs/1606.03401, 2016. URL http://arxiv.org/abs/1606.03401.
    • He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. IEEE Conference on Computer Vision and Pattern Recognition, 2015.
    • Hinton et al. (2012) Geoffrey Hinton, Li Deng, Dong Yu, George E. Dahl, Abdel-rahman Mohamed, Navdeep Jaitly, Andrew Senior, Vincent Vanhoucke, Patrick Nguyen, Tara N. Sainath, et al. Deep neural networks for acoustic modeling in speech recognition: The shared views of four research groups. IEEE Signal Processing Magazine, 2012.
    • Hochreiter & Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 1997.
    • Ioffe & Szegedy (2015) Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03167, 2015.
    • Jacobs et al. (1991) Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive mixtures of local experts. Neural Computing, 1991.
    • Johnson et al. (2016) Melvin Johnson, Mike Schuster, Quoc V. Le, Maxim Krikun, Yonghui Wu, Zhifeng Chen, Nikhil Thorat, Fernanda B. Viégas, Martin Wattenberg, Greg Corrado, Macduff Hughes, and Jeffrey Dean. Google’s multilingual neural machine translation system: Enabling zero-shot translation. CoRR, abs/1611.04558, 2016. URL http://arxiv.org/abs/1611.04558.
    • Jordan & Jacobs (1994) Michael I. Jordan and Robert A. Jacobs. Hierarchical mixtures of experts and the EM algorithm. Neural Computing, 1994.
    • Jozefowicz et al. (2016) Rafal Jozefowicz, Oriol Vinyals, Mike Schuster, Noam Shazeer, and Yonghui Wu. Exploring the limits of language modeling. arXiv preprint arXiv:1602.02410, 2016.
    • Kingma & Ba (2015) Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
    • Kneser & Ney (1995) Reinhard Kneser and Hermann. Ney. Improved backingoff for m-gram language modeling., 1995.
    • Krizhevsky et al. (2012) Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep convolutional neural networks. In NIPS, 2012.
    • Le et al. (2012) Quoc V. Le, Marc’Aurelio Ranzato, Rajat Monga, Matthieu Devin, Kai Chen, Greg S. Corrado, Jeffrey Dean, and Andrew Y. Ng. Building high-level features using large scale unsupervised learning. In ICML, 2012.
    • Ludovic Denoyer (2014) Patrick Gallinari Ludovic Denoyer. Deep sequential neural network. arXiv preprint arXiv:1410.0510, 2014.
    • Luong et al. (2015a) Minh-Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention-based neural machine translation. EMNLP, 2015a.
    • Luong et al. (2015b) Minh-Thang Luong, Ilya Sutskever, Quoc V. Le, Oriol Vinyals, and Wojciech Zaremba. Addressing the rare word problem in neural machine translation. ACL, 2015b.
    • Rasmussen & Ghahramani (2002) Carl Edward Rasmussen and Zoubin Ghahramani. Infinite mixtures of Gaussian process experts. NIPS, 2002.
    • Sak et al. (2014) Hasim Sak, Andrew W Senior, and Françoise Beaufays. Long short-term memory recurrent neural network architectures for large scale acoustic modeling. In INTERSPEECH, pp.  338–342, 2014.
    • Schuster & Nakajima (2012) Mike Schuster and Kaisuke Nakajima. Japanese and Korean voice search. ICASSP, 2012.
    • Shahbaba & Neal (2009) Babak Shahbaba and Radford Neal. Nonlinear models using dirichlet process mixtures. JMLR, 2009.
    • Sutskever et al. (2014) Ilya Sutskever, Oriol Vinyals, and Quoc V. Le. Sequence to sequence learning with neural networks. In NIPS, 2014.
    • Theis & Bethge (2015) Lucas Theis and Matthias Bethge. Generative image modeling using spatial LSTMs. In NIPS, 2015.
    • Tresp (2001) Volker Tresp. Mixtures of Gaussian Processes. In NIPS, 2001.
    • Wu et al. (2016) Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V. Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, Jeff Klingner, Apurva Shah, Melvin Johnson, Xiaobing Liu, Łukasz Kaiser, Stephan Gouws, Yoshikiyo Kato, Taku Kudo, Hideto Kazawa, Keith Stevens, George Kurian, Nishant Patil, Wei Wang, Cliff Young, Jason Smith, Jason Riesa, Alex Rudnick, Oriol Vinyals, Greg Corrado, Macduff Hughes, and Jeffrey Dean. Google’s neural machine translation system: Bridging the gap between human and machine translation. arXiv preprint arXiv:1609.08144, 2016.
    • Yao et al. (2009) Bangpeng Yao, Dirk Walther, Diane Beck, and Li Fei-fei. Hierarchical mixture of classification experts uncovers interactions between brain regions. In NIPS. 2009.
    • Zaremba et al. (2014) Wojciech Zaremba, Ilya Sutskever, and Oriol Vinyals. Recurrent neural network regularization. arXiv preprint arXiv:1409.2329, 2014.
    • Zhou et al. (2016) Jie Zhou, Ying Cao, Xuguang Wang, Peng Li, and Wei Xu. Deep recurrent models with fast-forward connections for neural machine translation. arXiv preprint arXiv:1606.04199, 2016.

    附录

    A 负载均衡损失

    正如第 4 节所讨论的,出于负载均衡的目的,我们希望定义一个额外的损失函数,鼓励各专家接收数量大致相等的训练样本。不幸的是,专家所接收样本的数量是离散量,无法用于反向传播。我们因此定义一个对一个输入批次 $ {\textstyle X} $ 中分配给每个专家样本数的平滑估计量 $ {\textstyle L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d\hspace{0pt}{(X)}} $。其平滑性允许我们通过该估计量进行反向传播——这正是门控函数中噪声项的用途。我们将 $ {\textstyle P\hspace{0pt}{(x,i)}} $ 定义为:在元素 $ {\textstyle i} $ 上重新随机采样噪声、其它元素的噪声采样保持不变的情况下,$ {\textstyle G\hspace{0pt}{(x)}_{i}} $ 不为零的概率。要计算 $ {\textstyle P\hspace{0pt}{(x,i)}} $,注意 $ {\textstyle G\hspace{0pt}{(x)}_{i}} $ 不为零当且仅当 $ {\textstyle H\hspace{0pt}{(x)}_{i}} $ 大于 $ {\textstyle H\hspace{0pt}{(x)}} $ 中除自身之外的第 $ {\textstyle k^{t\hspace{0pt}h}} $ 大元素。该概率推导得:

    $ {\textstyle P{(x,i)} = Pr\left( {(x \cdot W_{g})}_{i} + StandardNormal{()} \cdot Softplus{({(x \cdot W_{n\hspace{0pt}o\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e})}_{i})} \right.} $ (8)
    $ {\textstyle \left. > kth\_ excluding{(H{(x)},k,i)} \right)} $

    其中 $ {\textstyle k\hspace{0pt}t\hspace{0pt}h\hspace{0pt}\_\hspace{0pt}e\hspace{0pt}x\hspace{0pt}c\hspace{0pt}l\hspace{0pt}u\hspace{0pt}d\hspace{0pt}i\hspace{0pt}n\hspace{0pt}g\hspace{0pt}{(v,k,i)}} $ 表示 $ {\textstyle v} $ 中除分量 $ {\textstyle i} $ 之外的第 k 大分量。化简后得到:

    $ {\displaystyle {P\hspace{0pt}{(x,i)}} = {\Phi\hspace{0pt}\left( \frac{{({x \cdot W_{g}})}_{i} - {k\hspace{0pt}t\hspace{0pt}h\hspace{0pt}\_\hspace{0pt}e\hspace{0pt}x\hspace{0pt}c\hspace{0pt}l\hspace{0pt}u\hspace{0pt}d\hspace{0pt}i\hspace{0pt}n\hspace{0pt}g\hspace{0pt}{({H\hspace{0pt}{(x)}},k,i)}}}{S\hspace{0pt}o\hspace{0pt}f\hspace{0pt}t\hspace{0pt}p\hspace{0pt}l\hspace{0pt}u\hspace{0pt}s\hspace{0pt}{({({x \cdot W_{n\hspace{0pt}o\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}})}_{i})}} \right)}} $ (9)

    其中 $ {\textstyle \Phi} $ 是标准正态分布的 CDF。

    $ {\displaystyle {L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d\hspace{0pt}{(X)}_{i}} = {\sum\limits_{x \in X}{P\hspace{0pt}{(x,i)}}}} $ (10)

    于是我们可以将 load loss 定义为负载向量的变异系数的平方乘以一个手动调整的缩放因子 $ {\textstyle w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d}} $

    $ {\displaystyle {L_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d}\hspace{0pt}{(X)}} = {{w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d} \cdot C}\hspace{0pt}V\hspace{0pt}{({L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d\hspace{0pt}{(X)}})}^{2}}} $ (11)

    初始负载不均衡:

    为避免内存溢出错误,我们需要让网络初始化在专家负载大致相等的状态(因为软约束需要一些时间才会生效)。为此,我们将矩阵 $ {\textstyle W_{g}} $$ {\textstyle W_{n\hspace{0pt}o\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}} $ 初始化为全零,这样就只有噪声而没有信号。

    实验:

    我们训练了一组架构相同的模型(附录 C 所描述的 MoE-256 模型),使用不同的 $ {\textstyle w_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e}} $$ {\textstyle w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d}} $ 取值。每个模型训练 10 个 epoch 后测量测试集困惑度。我们还测量了 $ {\textstyle I\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e} $$ {\textstyle L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d} $ 的变异系数,以及最过载专家的负载与平均负载之比。最后这一指标对在分布式硬件上的负载均衡尤为重要。所有这些指标均在多个训练批次上进行平均。

    $ {\textstyle w_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e}} $ $ {\textstyle w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d}} $ Test Perplexity $ {\textstyle C\hspace{0pt}V\hspace{0pt}{({I\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e\hspace{0pt}{(X)}})}} $ $ {\textstyle C\hspace{0pt}V\hspace{0pt}{({L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d\hspace{0pt}{(X)}})}} $ $ {\textstyle \frac{m\hspace{0pt}a\hspace{0pt}x\hspace{0pt}{({L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d\hspace{0pt}{(X)}})}}{m\hspace{0pt}e\hspace{0pt}a\hspace{0pt}n\hspace{0pt}{({L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d\hspace{0pt}{(X)}})}}} $
    0.0 0.0 39.8 3.04 3.01 17.80
    0.2 0.0 35.6 0.06 0.17 1.47
    0.0 0.2 35.7 0.22 0.04 1.15
    0.1 0.1 35.6 0.06 0.05 1.14
    0.01 0.01 35.7 0.48 0.11 1.37
    1.0 1.0 35.7 0.03 0.02 1.07

    结果:

    结果见表 6。任何至少包含两种损失中之一的组合,模型质量都非常接近;而完全不使用任何损失则差得多。$ {\textstyle w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d}} $ 取值较高的模型在最过载专家上的负载更低。

    B 分层 Mixture of Experts

    如果专家数量非常大,可以使用两层分层 MoE 来降低分支因子。在分层 MoE 中,主门控网络选择若干 "experts" 的稀疏加权组合,而每个所谓的 "expert" 本身又是带有自身门控网络的二级 mixture-of-experts。333 我们没有发现需要更深层级的层次结构。若分层 MoE 由 $ {\textstyle a} $ 组、每组 $ {\textstyle b} $ 个专家构成,我们将主门控网络记为 $ {\textstyle G_{p\hspace{0pt}r\hspace{0pt}i\hspace{0pt}m\hspace{0pt}a\hspace{0pt}r\hspace{0pt}y}} $,次级门控网络记为 $ {\textstyle (G_{1},G_{2}..G_{a})} $,专家网络记为 $ {\textstyle (E_{0,0},E_{0,1}..E_{a,b})} $。MoE 的输出为:

    $ {\displaystyle y_{H} = {\sum\limits_{i = 1}^{a}{\sum\limits_{j = 1}^{b}{{{{{G_{p\hspace{0pt}r\hspace{0pt}i\hspace{0pt}m\hspace{0pt}a\hspace{0pt}r\hspace{0pt}y}\hspace{0pt}{(x)}_{i}} \cdot G_{i}}\hspace{0pt}{(x)}_{j}} \cdot E_{i,j}}\hspace{0pt}{(x)}}}}} $ (12)

    我们的专家利用率指标变为如下形式:

    $ {\displaystyle {I\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e_{H}\hspace{0pt}{(X)}_{i,j}} = {\sum\limits_{x \in X}{{{G_{p\hspace{0pt}r\hspace{0pt}i\hspace{0pt}m\hspace{0pt}a\hspace{0pt}r\hspace{0pt}y}\hspace{0pt}{(x)}_{i}} \cdot G_{i}}\hspace{0pt}{(x)}_{j}}}} $ (13)
    $ {\displaystyle {L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{H}\hspace{0pt}{(X)}_{i,j}} = \frac{{{L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{p\hspace{0pt}r\hspace{0pt}i\hspace{0pt}m\hspace{0pt}a\hspace{0pt}r\hspace{0pt}y}\hspace{0pt}{(X)}_{i}} \cdot L}\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{i}\hspace{0pt}{(X^{(i)})}_{j}}{|X^{(i)}|}} $ (14)

    $ {\textstyle L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{p\hspace{0pt}r\hspace{0pt}i\hspace{0pt}m\hspace{0pt}a\hspace{0pt}r\hspace{0pt}y}} $$ {\textstyle L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{i}} $ 分别表示主门控网络和第 $ {\textstyle i^{t\hspace{0pt}h}} $ 个次级门控网络的 $ {\textstyle L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d} $ 函数。$ {\textstyle X^{(i)}} $ 表示 $ {\textstyle X} $ 中满足 $ {\textstyle {G_{p\hspace{0pt}r\hspace{0pt}i\hspace{0pt}m\hspace{0pt}a\hspace{0pt}r\hspace{0pt}y}\hspace{0pt}{(x)}_{i}} > 0} $ 的子集。

    看起来直接令 $ {\textstyle {L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{H}\hspace{0pt}{(X)}_{i,j}} = {L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{i}\hspace{0pt}{(X_{i})}_{j}}} $ 会更简单,但这样关于主门控网络就没有梯度,因此我们采用上述形式。

    C 1 Billion Word Language Modeling Benchmark —— 实验细节

    C.1 每时间步 8M 运算的模型

    模型架构:

    我们的模型由五层组成:一层 word embedding 层、一层循环 Long Short-Term Memory (LSTM) 层 (Hochreiter & Schmidhuber, 1997; Gers et al., 2000)、一层 MoE 层、第二层 LSTM 以及 softmax 层。embedding 层的维度、各 LSTM 层的单元数以及 MoE 层的输入输出维度都等于 512。对除 softmax 之外的每一层,我们对其输出施加 dropout (Zaremba et al., 2014),以概率 $ {\textstyle D\hspace{0pt}r\hspace{0pt}o\hspace{0pt}p\hspace{0pt}P\hspace{0pt}r\hspace{0pt}o\hspace{0pt}b} $ 丢弃每个激活值,否则除以 $ {\textstyle ({1 - {D\hspace{0pt}r\hspace{0pt}o\hspace{0pt}p\hspace{0pt}P\hspace{0pt}r\hspace{0pt}o\hspace{0pt}b}})} $。dropout 之后,将上一层的输出加到该层输出上。这条残差连接有助于梯度流动 (He et al., 2015)。

    MoE 层架构:

    MoE 层中的每个专家是一个前馈网络,具有一层 1024 维的 ReLU 激活隐藏层和一层 512 维的输出层。因此每个专家包含 $ {\textstyle {{\lbrack{512 \ast 1024}\rbrack} + {\lbrack{1024 \ast 512}\rbrack}} = {1\hspace{0pt}M}} $ 个参数。MoE 层的输出在 dropout 之前经过 sigmoid。我们在不同模型间改变专家数量:使用 4、32、256 个专家的普通 MoE 层,以及 256、1024、4096 个专家的分层 MoE 层。我们将相应模型称为 MoE-4、MoE-32、MoE-256、MoE-256-h、MoE-1024-h 和 MoE-4096-h。对于分层 MoE 层,第一层的分支因子为 16,对应集群中的 GPU 数。我们对普通 MoE 层使用 Noisy-Top-K Gating(见第 2.1 节)并取 $ {\textstyle k = 4} $,对分层 MoE 在每一层取 $ {\textstyle k = 2} $。因此每个样本恰好被 4 个专家处理,总计 4M ops/timestep。两层 LSTM 各贡献 2M ops/timestep,使总量达到目标的 8M。

    计算量匹配的基线:

    MoE-4 模型不采用稀疏性,因为 4 个专家始终都被激活。此外,我们还训练了四个计算量匹配且不含稀疏性的额外基线模型:

    • MoE-1-Wide:MoE 层由单一 "expert" 组成,包含一层 4096 维的 ReLU 激活隐藏层。

    • MoE-1-Deep:MoE 层由单一 "expert" 组成,包含四层 ReLU 激活隐藏层,每层维度为 $ {\textstyle 1024} $

    • 4xLSTM-512:我们用两层额外的 512 单元 LSTM 替换 MoE 层。

    • LSTM-2048-512:模型包含一层 2048 单元的 LSTM(不含 MoE)。LSTM 输出被投影到 512 维 (Sak et al., 2014)。LSTM 的下一时间步接收投影后的输出。该模型与 (Jozefowicz et al., 2016) 中发表的模型之一相同。我们重新运行以适应训练方案的差异,并得到了与原文非常接近的结果。

    训练:

    模型在由 16 块 K40 GPU 组成的集群上训练,采用第 3 节描述的同步方法。每个批次由若干句子组成,总计约 30 万词。出于时间原因,我们将训练限制为 10 个 epoch(27,000 步)。除 MoE-4 之外的所有模型的训练耗时为 12–16 小时;MoE-4 耗时 18 小时(因为所有专家计算仅在 16 块 GPU 中的 4 块上完成)。我们使用 Adam 优化器 (Kingma & Ba, 2015)。基础学习率在前 1000 个训练步内线性增加,之后按步数倒数平方根的比例衰减。Softmax 输出层使用 importance sampling 高效训练,方式类似于 (Jozefowicz et al., 2016) 中的模型。我们对每个模型以 0.1 为步长搜索最佳 dropout 概率。

    为确保专家利用的均衡,我们设 $ {\textstyle w_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e} = 0.1} $$ {\textstyle w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d} = 0.1} $,如第 4 节与附录 A 所述。

    结果:

    我们在 (Chelba et al., 2013; Jozefowicz et al., 2016) 所用的留出数据集上以困惑度评估模型。我们遵循标准做法,对包含句末标记在内的所有词求和。结果见表 7。对每个模型,我们报告测试困惑度、计算预算、参数量、$ {\textstyle D\hspace{0pt}r\hspace{0pt}o\hspace{0pt}p\hspace{0pt}P\hspace{0pt}r\hspace{0pt}o\hspace{0pt}b} $ 取值以及计算效率。

    Model Test Test ops/timestep #Params excluding Total $ {\textstyle D\hspace{0pt}r\hspace{0pt}o\hspace{0pt}p} $- TFLOPS
    Perplexity Perplexity (millions) embed. & softmax #Params $ {\textstyle P\hspace{0pt}r\hspace{0pt}o\hspace{0pt}b} $ per GPU
    10 epochs (final) (millions) (billions) (observed)
    Kneser-Ney 5-gram* 67.6 0.00001 1.8
    LSTM-512-512* 54.1 2.4 2.4 0.8 0.1
    LSTM-1024-512* 48.2 4.7 4.7 0.8 0.1
    LSTM-2048-512* 45.0 43.7 9.4 9.4 0.8 0.1 0.61
    LSTM-2048-512 44.7 9.4 9.4 0.8 0.1 1.21
    4xLSTM-512 46.0 8.4 8.4 0.8 0.1 1.07
    MoE-1-Wide 46.1 8.4 8.4 0.8 0.1 1.29
    MoE-1-Deep 45.7 8.4 8.4 0.8 0.1 1.29
    MoE-4 45.0 8.4 8.4 0.8 0.1 0.52
    MoE-32 39.7 8.4 37.8 0.9 0.1 0.87
    MoE-256 35.7 8.6 272.9 1.1 0.1 0.81
    MoE-256-h 36.0 8.4 272.9 1.1 0.1 0.89
    MoE-1024-h 34.6 8.5 1079.0 1.9 0.2 0.90
    MoE-4096-h 34.1 8.9 4303.4 5.1 0.2 0.74
    2xLSTM-8192-1024* 34.7 30.6 151.0 151.0 1.8 0.25 1.09
    MoE-34M 31.3 33.8 4313.9 6.0 0.3 1.22
    MoE-143M 28.0 142.7 4371.1 6.0 0.4 1.56

    C.2 更昂贵的模型

    我们运行了另外两个模型 (MoE-34M 与 MoE-143M),以研究在大型 MoE 层存在时增加计算量的效果。它们的计算预算分别为 34M 与 143M ops/timestep。与上述模型类似,它们也是在两层 LSTM 之间使用一层 MoE。embedding 层维度以及 MoE 层的输入输出维度均设为 1024(而不是 512)。MoE-34M 的 LSTM 层为 1024 单元;MoE-143M 的 LSTM 层为 4096 单元并带一个 1024 维的输出投影 (Sak et al., 2014)。MoE-34M 使用 1024 个专家的分层 MoE 层,每个专家的隐藏层大小为 2048。MoE-143M 使用 256 个专家的分层 MoE 层,每个专家的隐藏层大小为 8192。两个模型在 MoE 层上都有 4B 参数。我们为每个模型搜索最佳的 $ {\textstyle D\hspace{0pt}r\hspace{0pt}o\hspace{0pt}p\hspace{0pt}P\hspace{0pt}r\hspace{0pt}o\hspace{0pt}b} $,并为每个模型训练 10 个 epoch。

    两个模型分别取得了 $ {\textstyle 31.3} $$ {\textstyle 28.0} $ 的测试困惑度,表明即便在存在大型 MoE 的情况下,继续增加计算量仍然有用。结果见表 7 末尾。两者中较大的模型在计算预算上与文献中已发表的最佳模型相近,训练时间也相近。在训练 10 个 epoch 后比较,我们的模型测试困惑度低 $ {\textstyle 18\%} $

    D 100 Billion Word Google News Corpus —— 实验细节

    模型架构:

    这些模型在结构上与上一节描述的每时间步 8M 运算的模型类似。我们改变模型中的专家数量:使用 32 个专家的普通 MoE 层,以及 256、1024、4096、16384、65536 和 131072 个专家的分层 MoE 层。对于分层 MoE 层,第一层的分支因子分别为 32、32、64、128、256 和 256。

    训练:

    模型在 32 块 Tesla K40 GPU 集群上训练,最后两个模型则分别在 64 和 128 块 GPU 的集群上训练,以容纳所有参数。所有模型的训练批量大小约为 250 万词。模型在约 1000 亿词的语料上做一次性遍历训练。

    我们实施了若干内存优化,使每 GPU 能容纳多达 10 亿参数。首先,我们不保存专家隐藏层的激活,而是在反向传播时重新计算。其次,我们修改了专家参数上的优化器,使其所需的辅助存储更少:

    Adam 优化器 (Kingma & Ba, 2015) 维护逐参数梯度的一阶与二阶矩估计,这会使所需内存增至三倍。为避免保留一阶矩估计,我们设 $ {\textstyle \beta_{1} = 0} $。为缩减二阶矩估计的大小,我们用一个分解近似来替代它。对一个参数矩阵,我们不再维护完整的二阶矩矩阵,而是只维护该矩阵按行与按列的均值向量。每一步,估计矩阵都取这两个向量的外积除以其中任一向量的均值。该技术也可类似地应用于 Adagrad (Duchi et al., 2010)。

    Model Test Test ops/timestep #Params excluding Total TFLOPS
    Perplexity Perplexity (millions) embed. & softmax #Params per GPU
    .1 epochs 1 epoch (millions) (billions) (observed)
    Kneser-Ney 5-gram 67.1 45.3 0.00001 76.0
    4xLSTM-512 54.5 47.0 8.4 8.4 0.1 1.23
    MoE-32 48.5 40.4 8.4 37.8 0.1 0.83
    MoE-256-h 42.8 35.3 8.4 272.9 0.4 1.11
    MoE-1024-h 40.3 32.7 8.5 1079.0 1.2 1.14
    MoE-4096-h 38.9 30.9 8.6 4303.4 4.4 1.07
    MoE-16384-h 38.2 29.7 8.8 17201.0 17.3 0.96
    MoE-65536-h 38.2 28.9 9.2 68791.0 68.9 0.72
    MoE-131072-h 39.8 29.2 9.7 137577.6 137.7 0.30

    结果:

    我们在留出数据集上以困惑度评估模型,结果见表 8。在 1000 亿词训练后,680 亿参数的 MoE 模型的困惑度比基线模型低 39%。值得注意的是,最大模型测得的计算效率 (0.30 TFLOPS/GPU) 与其它模型相比非常低。这很可能是因为,为了与其它模型可比,我们没有按 GPU 数按比例增大训练批量大小。为便于比较,我们还给出了由 4 个 LSTM 组成的计算量匹配基线的结果,以及一个未剪枝的、使用 Kneser-Ney 平滑的 5-gram 模型 (Kneser & Ney, 1995) 的结果。444虽然语料原始规模为 1300 亿词,但神经模型最多只训练了 1000 亿词。所报告的 Kneser-Ney 5-gram 模型分别在 130 亿和 1300 亿词上训练,因此相对其它报告结果略占优势。

    E 机器翻译 —— 实验细节

    单一语言对 MoE 模型的架构:

    我们的模型是 (Wu et al., 2016) 中所述 GNMT 模型的修改版本。为减少计算,我们将编码器与解码器中的 LSTM 层数从 9 与 8 分别降至 3 与 2。我们在编码器(第 2 与 3 层之间)以及解码器(第 1 与 2 层之间)都插入了 MoE 层。我们在编码器与解码器之间使用注意力机制,第一层解码器 LSTM 同时接收并向注意力提供输入555出于性能考虑,我们使用了与 (Wu et al., 2016) 略有不同的注意力函数——见附录 G。模型中所有层的输入输出维度均为 512。LSTM 层有 2048 个隐藏单元,并带 512 维的输出投影。我们在所有 LSTM 与 MoE 层周围添加残差连接以促进梯度流动 (He et al., 2015)。与 GNMT 类似,为有效处理稀有词,我们在系统的输入与输出中使用了子词单元(也称 "wordpieces")(Schuster & Nakajima, 2012)。

    我们使用了一个 32K wordpieces 的源端与目标端共享词表,并采用了 (Wu et al., 2016) 提出的相同 beam search 技术。

    我们训练了 MoE 层中专家数量不同的多个模型。除一个不含 MoE 层的基线模型外,我们还训练了含 32 个专家的扁平 MoE 层模型,以及含 512 与 2048 个专家的分层 MoE 层模型。扁平 MoE 层使用 $ {\textstyle k = 4} $,分层 MoE 模型在门控网络的每一层使用 $ {\textstyle k = 2} $。因此每个输入在每个 MoE 层都恰好被 4 个专家处理。MoE 层中的每个专家是一个前馈网络,含一层 2048 维的 ReLU 激活隐藏层,因此每个专家包含 $ {\textstyle {{\lbrack{512 \ast 2048}\rbrack} + {\lbrack{2048 \ast 512}\rbrack}} = {2\hspace{0pt}M}} $ 个参数。MoE 层的输出经过一个 sigmoid 函数。我们使用附录 F 中描述的严格平衡门控函数。

    多语言 MoE 模型的架构:

    我们使用了与单语言对模型相同的架构,但有以下不同:使用第 2.1 节描述的 noisy-top-k gating,而非附录 F 的方案。编码器与解码器中的 MoE 层为非分层 MoE,$ {\textstyle n = 512} $ 个专家,$ {\textstyle k = 2} $。每个专家具有更大的隐藏层,大小为 $ {\textstyle 8192} $。这使 MoE 层的计算量翻倍,整个模型的计算预算从 85M ops/timestep 提升到 102M ops/timestep。

    训练:

    我们使用 Adam 优化器训练网络 (Kingma & Ba, 2015)。基础学习率在前 2000 个训练步内线性增加,在接下来的 8000 步内保持常数,之后按步数倒数平方根的比例衰减。对单语言对模型,类似 (Wu et al., 2016),我们对所有 embedding、LSTM 和 MoE 层的输出施加 dropout (Zaremba et al., 2014),$ {\textstyle {D\hspace{0pt}r\hspace{0pt}o\hspace{0pt}p\hspace{0pt}P\hspace{0pt}r\hspace{0pt}o\hspace{0pt}b} = 0.4} $。训练按照第 3 节描述在最多 64 块 GPU 的集群上同步进行。每个训练批次由若干句子对组成,每 GPU 约 16000 词。

    为确保专家利用的均衡,我们设 $ {\textstyle w_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e} = 0.01} $$ {\textstyle w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d} = 0.01} $,如第 4 节和附录 A 所述。

    评测指标:

    我们使用困惑度与标准 BLEU 分数评估模型。报告的是由 multi-bleu.pl 脚本(取自 Moses 在 Github 上的公开实现)计算的 tokenized BLEU 分数;(Luong et al., 2015a) 也使用了同一脚本。

    结果:

    5.3 节中的表 234 给出了我们的结果与其它已发表方法的对比。图 4 展示了不同专家数量模型的测试困惑度随处理的(训练数据中)源句词数变化的曲线。从图中可以看出,随着专家数接近 2048,我们模型的测试困惑度持续改善。

    Refer to caption Refer to caption

    我们发现专家确实会按句法和/或语义高度专精化,如表 9 所示。例如,当不定冠词 "a" 在表示重要性或领导关系的动词短语中引入直接宾语时,会激活某个特定的专家。

    Expert 381 Expert 752 Expert 2004
    … with researchers , … … plays a core … … with rapidly growing …
    … to innovation . … plays a critical … … under static conditions …
    … tics researchers . … provides a legislative … … to swift ly …
    … the generation of … … play a leading … … to dras tically …
    … technology innovations is … … assume a leadership … … the rapid and …
    … technological innovations , … … plays a central … … the fast est …
    … support innovation throughout … … taken a leading … … the Quick Method …
    … role innovation will … … established a reconciliation … … rec urrent ) …
    … research scienti st … … played a vital … … provides quick access …
    … promoting innovation where … … have a central … … of volatile organic …

    F 严格平衡的门控

    由于我们基础设施中一些此后已被修复的特殊性,当我们运行某些机器翻译实验时,模型在每个专家恰好收到相同批量大小时运行得更快。为适应这一点,我们使用了一种不同的门控函数,描述如下。

    回顾一下,我们将 softmax 门控函数定义为:

    $ {\displaystyle {G_{\sigma}\hspace{0pt}{(x)}} = {S\hspace{0pt}o\hspace{0pt}f\hspace{0pt}t\hspace{0pt}m\hspace{0pt}a\hspace{0pt}x\hspace{0pt}{({x \cdot W_{g}})}}} $ (15)

    Sparse Gating(替代表述):

    为得到稀疏门控向量,我们将 $ {\textstyle G_{\sigma}\hspace{0pt}{(x)}} $ 与一个稀疏掩码 $ {\textstyle M\hspace{0pt}{({G_{\sigma}\hspace{0pt}{(x)}})}} $ 逐分量相乘,并对输出归一化。掩码本身是 $ {\textstyle G_{\sigma}\hspace{0pt}{(x)}} $ 的函数,指定每个输入样本被分配到哪些专家:

    $ {\displaystyle {G\hspace{0pt}{(x)}_{i}} = \frac{G_{\sigma}\hspace{0pt}{(x)}_{i}\hspace{0pt}M\hspace{0pt}{({G_{\sigma}\hspace{0pt}{(x)}})}_{i}}{\sum_{j = 1}^{n}{G_{\sigma}\hspace{0pt}{(x)}_{j}\hspace{0pt}M\hspace{0pt}{({G_{\sigma}\hspace{0pt}{(x)}})}_{j}}}} $ (16)

    Top-K 掩码:

    在这一表述下,要实现 top-k 门控,我们令 $ {\textstyle {M\hspace{0pt}{(v)}} = {T\hspace{0pt}o\hspace{0pt}p\hspace{0pt}K\hspace{0pt}{(v,k)}}} $,其中:

    $ {\displaystyle {T\hspace{0pt}o\hspace{0pt}p\hspace{0pt}K\hspace{0pt}{(v,k)}_{i}} = \begin{cases} 1 & {\text{if~}v_{i}\text{~is in the top~}k\text{~elements of~}v\text{.}} \\ 0 & \text{otherwise.} \end{cases}} $ (17)

    批量级掩码:

    为强制每个专家收到完全相同数量的样本,我们引入一个替代的掩码函数 $ {\textstyle M_{b\hspace{0pt}a\hspace{0pt}t\hspace{0pt}c\hspace{0pt}h\hspace{0pt}w\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}\hspace{0pt}{(X,m)}} $,它对一批输入向量整体操作。我们不再按样本保留 top $ {\textstyle k} $ 个值,而是在整个训练批次上按专家保留 top $ {\textstyle m} $ 个值,其中 $ {\textstyle m = \frac{k\hspace{0pt}{|X|}}{n}} $,从而每个样本平均被发送到 $ {\textstyle k} $ 个专家。

    $ {\displaystyle {M_{b\hspace{0pt}a\hspace{0pt}t\hspace{0pt}c\hspace{0pt}h\hspace{0pt}w\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}\hspace{0pt}{(X,m)}_{j,i}} = \begin{cases} 1 & {\text{if~}X_{j,i}\text{~is in the top~}m\text{~values for to expert~}i} \\ 0 & \text{otherwise} \end{cases}} $ (18)

    正如我们的实验所表明、并在 (Ioffe & Szegedy, 2015) 中也观察到的那样,在训练阶段使用按批次的函数(如 $ {\textstyle M_{b\hspace{0pt}a\hspace{0pt}t\hspace{0pt}c\hspace{0pt}h\hspace{0pt}w\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}} $)需要在推理时进行修改,因为推理时未必有大批量样本。我们的方案是训练一个逐专家的阈值向量 $ {\textstyle T} $ 以近似批次级掩码的效果。在推理时使用如下掩码:

    $ {\displaystyle {M_{t\hspace{0pt}h\hspace{0pt}r\hspace{0pt}e\hspace{0pt}s\hspace{0pt}h\hspace{0pt}o\hspace{0pt}l\hspace{0pt}d}\hspace{0pt}{(x,T)}_{i}} = \begin{cases} 1 & {\text{if~}{x_{i} > T_{i}}} \\ 0 & \text{otherwise} \end{cases}} $ (19)

    为学习这些阈值,我们在训练时施加一个额外的损失项,当批次级掩码与阈值掩码完全相同时该损失最小。

    $ {\displaystyle {L_{b\hspace{0pt}a\hspace{0pt}t\hspace{0pt}c\hspace{0pt}h\hspace{0pt}w\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}\hspace{0pt}{(X,T,m)}} = {\sum\limits_{j = 1}^{|X|}{\sum\limits_{i = 1}^{n}{{({{M_{t\hspace{0pt}h\hspace{0pt}r\hspace{0pt}e\hspace{0pt}s\hspace{0pt}h\hspace{0pt}o\hspace{0pt}l\hspace{0pt}d}\hspace{0pt}{(x,T)}_{i}} - {M_{b\hspace{0pt}a\hspace{0pt}t\hspace{0pt}c\hspace{0pt}h\hspace{0pt}w\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}\hspace{0pt}{(X,m)}_{j,i}}})}\hspace{0pt}{({X_{j,i} - T_{i}})}}}}} $ (20)

    G 注意力函数

    GNMT (Wu et al., 2016) 中描述的注意力机制涉及一个学习得到的 "Attention Function" $ {\textstyle A\hspace{0pt}{(x_{i},y_{j})}} $,它接受 "source vector" $ {\textstyle x_{i}} $ 与 "target vector" $ {\textstyle y_{j}} $,并需对每个源端时间步 $ {\textstyle i} $ 与目标端时间步 $ {\textstyle j} $ 都计算。在 GNMT 中,该注意力函数实现为一个具有 $ {\textstyle n} $ 维隐藏层的前馈神经网络,可写为:

    $ {\displaystyle {A_{G\hspace{0pt}N\hspace{0pt}M\hspace{0pt}T}\hspace{0pt}{(x_{i},y_{j})}} = {\sum\limits_{d = 1}^{n}{V_{d}\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}h\hspace{0pt}{({{({x_{i}\hspace{0pt}U})}_{d} + {({y_{j}\hspace{0pt}W})}_{d}})}}}} $ (21)

    其中 $ {\textstyle U} $$ {\textstyle W} $ 是可训练的权重矩阵,$ {\textstyle V} $ 是一个可训练的权重向量。

    出于性能考虑,我们的模型使用了略有不同的注意力函数:

    $ {\displaystyle {A\hspace{0pt}{(x_{i},y_{j})}} = {\sum\limits_{d = 1}^{n}{V_{d}\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}h\hspace{0pt}{({({x_{i}\hspace{0pt}U})}_{d})}\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}h\hspace{0pt}{({({y_{j}\hspace{0pt}W})}_{d})}}}} $ (22)

    使用我们的注意力函数,可以借助优化后的矩阵乘法同时对多个源端时间步与多个目标端时间步计算注意力。我们发现两种函数在质量上差别不大。