谷歌研究人员开发出新型人工智能技术，不再将脑力浪费在无用的单词上

链接表

1. 介绍
2. 背景
3. 实现深度混合 Transformer

   • 3.1. 定义计算预算

   • 3.2. 绕过变压器块

   • 3.3. 路由方案

   • 3.4. 路由实现

   • 3.5. 抽样和 3.6. 训练方法

4. 结果
   • 4.1. 训练、isoFLOP 比较
   • 4.2. 自回归评估和 4.3. 深度与专家混合 (MoDE)
5. 讨论和参考文献

基于 Transformer 的语言模型将 FLOP 均匀地分布在输入序列中。在这项工作中，我们证明 Transformer 可以学习动态地将 FLOP（或计算）分配到序列中的特定位置，从而优化模型深度中不同层的序列分配。我们的方法通过限制可以参与给定层上的自注意力和 MLP 计算的 token（𝑘）数量来强制执行总计算预算。要处理的 token 由网络使用 top-𝑘 路由机制确定。由于 𝑘 是先验定义的，因此与其他条件计算技术不同，此简单过程使用具有已知张量大小的静态计算图。然而，由于 𝑘 token 的身份是不固定的，因此该方法可以在时间和模型深度维度上非均匀地消耗 FLOP。因此，计算支出在总和上是完全可预测的，但在 token 级别是动态的和上下文相关的。以这种方式训练的模型不仅学会了动态分配计算，而且做得很高效。这些模型在等效的 FLOPS 和训练挂钟时间方面与基线性能相匹配，但每次前向传递只需要一小部分 FLOP，并且在训练后采样期间的速度可以提高 50% 以上。

1. 简介

并非所有问题都需要相同的时间或精力来解决。类似地，在语言建模中，并非所有标记和序列都需要相同的时间或精力来准确做出预测。然而，Transformer 模型在前向传递中每个标记消耗的计算量相同。理想情况下，Transformer 会通过不花费不必要的计算来使用较小的总计算预算。

条件计算是一种尝试通过仅在需要时消耗计算量来减少总计算量的技术（Bengio 等人，2016 年；Bengio，2013 年；Bengio 等人，2013 年）。各种算法都提供了何时以及使用多少计算量的解决方案（Ainslie 等人，2023 年；Bapna 等人，2020 年；Fedus 等人，2022 年）。然而，这个具有挑战性的问题的一般公式可能不适用于现有的硬件约束，因为它们倾向于引入动态计算图（Dehghani 等人，2018 年；Graves，2016 年）。最有前途的条件计算方法可能是与我们当前的硬件堆栈相协调的方法，它优先考虑静态计算图和已知的张量大小，这些张量大小被选择来最大化硬件利用率。

在这里，我们考虑使用静态计算预算进行语言建模的问题，该预算可以比 vanilla Transformer 使用的预算少。网络必须学习如何通过在每个层中针对每个 token 做出决策来动态分配可用计算，决定从可用预算中花费计算资源。在我们的实现中，总计算量是用户定义的，并且在训练之前是不变的，而不是网络动态决策的函数。因此，可以提前预期和利用硬件效率提升（例如减少内存占用或减少每次前向传递的 FLOP）。正如我们将展示的那样，这些收益可以在不牺牲整体性能的情况下获得。

我们利用一种类似于混合专家 (MoE) 变压器的方法，其中动态令牌级路由决策是在整个网络深度上做出的。与 MoE 不同，我们选择将计算应用于令牌（就像标准变压器的情况一样），或者将其通过残差连接（保持不变并节省计算）。与 MoE 相比，我们还将这种路由应用于前向 MLP 和多头注意力。由于这也会影响我们处理的键和查询，因此路由不仅会决定要更新哪些令牌，还会决定哪些令牌可供处理。我们将此策略称为深度混合 (MoD)，以强调各个令牌如何通过变压器的深度穿过不同数量的层或块（见图 1）。

MoD 技术还允许人们在性能和速度之间进行权衡。一方面，可以训练一个 MoD 变压器，在最终对数概率训练目标上，在同等训练 FLOP（isoFLOP）的情况下，比 vanilla 变压器提高多达 1.5%，同时花费同等数量的时钟时间进行训练。另一方面，可以训练一个 MoD 变压器，使训练损失与 isoFLOP 最优 vanilla 变压器相当，但每次前向传递只使用一小部分 FLOP（高达 50%），因此速度更快。总之，这些结果表明，MoD 变压器可以学习智能路由（即跳过不必要的计算），因为它们可以在每次前向传递占用较少的 FLOP 的情况下，实现每个序列相等或更好的对数概率。