当代深度学习模型在从自然语言到计算机视觉等许多领域都取得了显著的成功。Transformer神经网络(Transformer) 是一种流行的深度学习架构，如今已成为自然语言处理中大多数任务的基础，并且也开始扩展到其他领域的应用，例如计算机视觉、机器人技术和自动驾驶。此外，它们构成了所有当前最先进语言模型的支柱。
Transformer 网络规模的扩大带来了性能的提升，并出现了小型网络中不存在的行为。然而，规模的扩大往往伴随着计算成本和推理延迟的大幅增加。一个自然而然的问题是，我们能否在不产生计算负担的情况下获得更大模型的好处。
在NeurIPS 2023被选为焦点的 “高效 Transformer 的交替更新”中，我们介绍了 AltUp，这是一种在不增加计算成本的情况下利用增加的 token 表示的方法。AltUp 易于实现，广泛适用于任何 Transformer 架构，并且只需要极少的超参数调整。例如，在 770M 参数的 T5-Large 模型上使用 AltUp 的变体，添加约 100 个参数即可产生质量明显更好的模型。
背景
为了理解如何实现这一点，我们深入研究了 Transformer 的工作原理。首先，它们将输入划分为一系列标记。然后将每个标记映射到一个嵌入向量（通过嵌入表），称为标记嵌入。我们将这个向量的维度称为标记表示维度。然后，Transformer 通过使用其网络参数应用一系列计算模块（称为层）来操作这个标记嵌入序列。每个 Transformer 层中的参数数量是层宽度的函数，而层宽度由标记表示维度决定。
为了在不增加计算负担的情况下实现规模效益，先前的研究（例如稀疏混合专家 (Sparse MoE) 模型（例如Switch Transformer、Expert Choice、V-MoE））主要侧重于通过有条件地激活基于输入的子集来高效地扩大网络参数（在自注意力和前馈层中）。这使我们能够扩大网络规模，而无需显著增加每个输入的计算量。但是，通过有条件地激活 token 表示向量的部分来扩大 token 表示维度本身的研究还存在空白。
最近的研究（例如，缩放定律和无限宽度网络）已经从经验和理论上证实，更宽的标记表示有助于学习更复杂的函数。这种现象在功能不断增强的现代架构中也很明显。例如，在 T5 模型中，表示维度从 512（小）增加到 768（基础）和 1024（分别对应具有 770M、3B 和 11B 参数的模型），在PaLM 模型中从 4096（8B）增加到 8192（64B）和 18432（540B）。加宽的表示维度也显著提高了双编码器检索模型的性能。然而，简单地加宽表示向量需要相应地增加模型维度，这会使前馈计算的计算量 二次方增加。
方法
AltUp 的工作原理是将扩展的表示向量划分为大小相等的块，在每一层仅处理一个块，并使用高效的预测校正机制来推断其他块的输出（如下图右侧所示）。这使得 AltUp 能够同时保持模型维度（因此计算成本）大致不变，并利用增加的标记维度。增加的标记维度允许模型将更多信息打包到每个标记的嵌入中。通过保持每个转换器层的宽度不变，AltUp 避免了计算成本的二次增加，否则会出现表示的简单扩展。
不使用 AltUp（左）和使用 AltUp（右）扩展 token 表示的图示。由于层宽增加，这种扩展导致 vanilla Transformer 中的计算量几乎呈二次方增长。相比之下，交替更新保持层宽不变，并通过在每一层对表示的子块进行操作来高效计算输出。
更具体地说，每个层的输入是两个或更多个块，其中一个块被传递到 1x 宽度的 Transformer 层（见下图）。我们将这个块称为“激活”块。此计算将产生激活块的精确输出。同时，我们调用一个轻量级预测器来计算所有输入块的加权组合。预测值以及激活块的计算值被传递给轻量级校正器，该校正器根据观察值更新预测。这种校正机制使未激活的块能够根据激活的块进行更新。预测和校正步骤都只涉及有限数量的向量加法和乘法，因此比常规 Transformer 层快得多。我们注意到，这个过程可以推广到任意数量的块。
预测器和校正器计算：预测器将子块与可训练标量系数混合；校正器返回预测器输出和变压器输出的加权平均值。预测器和校正器执行标量向量乘法，与变压器相比，计算成本可以忽略不计。预测器输出具有标量混合系数 p i, j 的块的线性混合，校正器将预测器输出和变压器输出与权重 g i相结合。
从更高层次来看，AltUp 与稀疏 MoE 类似，因为它是一种以条件访问（外部）参数的形式向模型添加容量的方法。在稀疏 MoE 中，附加参数采用前馈网络 (FFN) 专家的形式，条件性与输入有关。在 AltUp 中，外部参数来自扩展的嵌入表，条件性采用交替逐块激活表示向量的形式，如上图所示。因此，AltUp 具有与稀疏 MoE 模型相同的基础。
AltUp 相对于稀疏 MoE 的一个优势是它不需要分片，因为引入的额外参数数量是嵌入表大小的2倍，而嵌入表通常只占整个模型大小的一小部分。此外，由于 AltUp 专注于有条件地激活更广泛 token 表示的部分，因此它可以与 MoE 等正交技术协同应用，以获得互补的性能提升。
评估
AltUp 在各种基准语言任务上针对 T5 模型进行了评估。在相同准确度下，使用 AltUp 增强的模型均比外推密集模型更快。例如，我们观察到使用 AltUp 增强的 T5 Large 模型在GLUE、SuperGLUE、SQuAD和Trivia-QA基准上分别可使速度提高 27%、39%、87% 和 29%。
对不同大小和流行基准的 T5 模型上的 AltUp 进行评估。在相同准确度下，AltUp 始终能够带来相对于基线的显著加速。延迟是在具有 8 个核心的TPUv3上测量的。加速定义为延迟变化除以 AltUp 延迟（B = T5 Base，L = T5 Large，XL = T5 XL 模型）。
当我们将 AltUp 应用于更大的模型时，其相对性能会得到改善——比较 T5 Base + AltUp 与 T5 Large + AltUp 的相对加速。这证明了 AltUp 的可扩展性及其在更大模型上的性能提升。总体而言，在所有评估的模型大小和基准上，AltUp 始终能够使模型的预测性能优于速度相同的相应基线模型。
扩展：回收的 AltUp
AltUp 公式增加了少量的每层计算，但它确实需要使用更宽的嵌入表。在词汇量（即标记器可以生成的不同标记的数量）非常大的情况下，这可能会导致初始嵌入查找和最终线性 + softmax 操作增加大量计算。非常大的词汇量也可能导致增加大量不受欢迎的嵌入参数。为了解决这个问题，Recycled-AltUp 是 AltUp 的扩展，它通过保持嵌入表的宽度相同来避免这些计算和参数成本。
K = 2的 Recycled-AltUp 架构图。
在 Recycled-AltUp 中，我们不会加宽初始 token 嵌入，而是将嵌入复制K次以形成更宽的 token 表示。因此，Recycled-AltUp 相对于基线转换器几乎没有添加任何额外参数，同时受益于更宽的 token 表示。
T5-B/L/XL 上的 Recycled-AltUp 与基线相比。Recycled-AltUp 可显著提高预训练性能，且不会造成任何明显的减速。
我们还在 T5 基础模型、大型模型和 XL 模型上评估了 AltUp 的轻量级扩展 Recycled-AltUp（K = 2），并将其预训练准确度和速度与基线进行比较。由于 Recycled-AltUp 不需要扩展嵌入表维度，因此使用它进行增强的模型几乎具有与基线模型相同数量的可训练参数。与密集基线相比，我们再次观察到持续的改进。
AltUp 为何有效？
AltUp 通过向嵌入表添加和有效利用辅助参数，并在各层之间保持更高维度的表示 ，增加了模型的容量。我们认为，此计算的关键因素在于 AltUp 的预测机制，该机制执行不同块的集成。这种加权组合使得消息能够连续传递到整个向量，尽管在每个层中仅激活它的子块。另一方面，Recycled-AltUp 不会向 token 嵌入添加任何其他参数。但是，它仍然具有在更高维度的表示空间中模拟计算的好处，因为在从一个转换器层移动到另一个转换器层时会保持更高维度的表示向量。我们推测，这通过增强网络中的信息流来帮助训练。一个有趣的研究方向是探索 Recycled-AltUp 的好处是否可以完全通过更有利的训练动态来解释。
致谢
我们感谢我们的合作者谷歌研究中心的 Cenk Baykal、Dylan Cutler 和 Rina Panigrahy，以及加州大学伯克利分校的 Nikhil Ghosh（在谷歌研究实习期间完成的工作）。
评论