语言模型 (LM) 是自然语言处理领域近期许多突破的驱动力。T5 、LaMDA、GPT-3和PaLM等模型在各种语言任务上都表现出色。虽然有多种因素可以提高 LM 的性能，但最近的一些研究表明，扩大模型规模对于揭示新兴能力至关重要。换句话说，有些情况可以通过小模型来解决，而其他情况似乎可以从扩大规模中受益。
尽管最近的努力使得 LM 能够高效地在大量数据上进行训练，但经过训练的模型在实际使用中仍然速度慢且成本高昂。在推理时生成文本时，大多数自回归 LM 输出的内容类似于我们说话和写作的方式（逐字逐句），根据前面的单词预测每个新单词。这个过程无法并行化，因为 LM 需要先完成一个单词的预测，然后才能开始计算下一个单词。此外，考虑到模型的数十亿个参数，预测每个单词都需要大量计算。
在NeurIPS 2022上展示的 “自信自适应语言建模”中，我们介绍了一种通过提高推理效率来加速 LM 文本生成的新方法。我们的方法名为 CALM，其动机是直觉认为某些下一个单词的预测比其他单词更容易。在写句子时，某些延续很简单，而其他延续可能需要更多努力。当前的 LM 为所有预测投入相同的计算能力。相反，CALM 会在生成时间步骤中动态分配计算工作量。通过有选择地将更多计算资源仅分配给更难的预测，CALM 可以更快地生成文本，同时保持输出质量。
自信的自适应语言建模
如果可能，CALM 会跳过某些预测的一些计算工作。为了证明这一点，我们使用了流行的编码器-解码器T5架构。编码器读取输入文本（例如，要总结的新闻文章）并将文本转换为密集表示。然后，解码器通过逐字预测来输出摘要。编码器和解码器都包含一长串 Transformer层。每层都包含注意力和前馈模块，并进行多次矩阵乘法。这些层逐渐修改最终用于预测下一个单词的隐藏表示。
CALM 不会等待所有解码器层完成，而是尝试在某个中间层之后更早地预测下一个单词。为了决定是坚持某个预测还是将预测推迟到后面的层，我们会测量模型对中间预测的置信度。只有当模型足够自信地认为预测不会改变时，才会跳过其余的计算。为了量化什么是“足够自信”，我们校准了一个阈值，该阈值在统计上满足整个输出序列的任意质量保证。
使用常规语言模型（顶部）和 CALM（底部）生成文本。 CALM 尝试做出早期预测。 一旦信心足够（深蓝色调），它就会跳过并节省时间。
具有早期退出的语言模型
为 LM 启用这种早期退出策略需要对训练和推理过程进行最少的修改。在训练期间，我们鼓励模型在中间层产生有意义的表示。我们的学习损失函数不是仅使用顶层进行预测，而是对所有层的预测取加权平均值，为顶层分配更高的权重。我们的实验表明，这显著改善了中间层预测，同时保留了整个模型的性能。在一个模型变体中，我们还包括一个小型早期退出分类器，该分类器经过训练可对局部中间层预测是否与顶层一致进行分类。我们在第二个快速步骤中训练此分类器，其中我们冻结模型的其余部分。
模型训练完成后，我们需要一种允许提前退出的方法。首先，我们定义一个局部置信度度量，用于捕获模型对其中期预测的置信度。我们探索了三种置信度度量（在下面的结果部分中描述）：（1）softmax 响应，从 softmax 分布中取出最大预测概率；（2）状态传播，当前隐藏表示与前一层隐藏表示之间的余弦距离；（3）提前退出分类器，专门为预测局部一致性而训练的分类器的输出。我们发现 softmax 响应在统计上很强，同时计算简单快捷。另外两种替代方案在浮点运算（FLOPS）方面更轻松。
另一个挑战是，每一层的自注意力都依赖于前几个单词的隐藏状态。如果我们提前退出某些单词预测，这些隐藏状态可能会丢失。相反，我们会回头关注最后一个计算层的隐藏状态。
最后，我们设置了提前退出的局部置信度阈值。在下一节中，我们将描述我们用于寻找良好阈值的控制过程。作为第一步，我们通过基于一个有用的观察来简化这个无限的搜索空间：在生成过程开始时犯的错误更为有害，因为它们会影响所有后续输出。因此，我们从较高（更保守）的阈值开始，并随着时间的推移逐渐降低它。我们使用具有用户定义温度的负指数来控制此衰减率。我们发现这可以更好地控制性能效率权衡（每个质量级别获得的加速）。
可靠地控制加速模型的质量
提前退出决策必须是局部的；它们需要在预测每个单词时发生。然而，在实践中，最终输出应该是全局一致的或与原始模型相当的。例如，如果原始完整模型生成“音乐会很精彩，时间很长”，人们会接受 CALM 切换形容词的顺序并输出“音乐会很长，很精彩”。然而，在局部层面，“精彩”一词被替换为“长”。因此，这两个输出是全局一致的，但包含一些局部不一致。我们以学习然后测试(LTT) 框架为基础，将基于本地置信度的决策与全局一致的输出联系起来。
在 CALM 中，通过 LTT 校准，从用户定义的对完整输出文本的一致性约束中得出早期退出决策的局部每时间步置信度阈值。红色框表示 CALM 使用了解码器的大部分层来进行该特定预测。绿色框表示 CALM 仅使用几个 Transformer 层节省了时间。完整句子显示在本文的最后一个示例中。
首先，我们定义并制定两种类型的一致性约束供选择：
文本一致性：我们限制了 CALM 输出和完整模型输出之间的预期文本距离。这不需要任何标记数据。
风险一致性：我们限制了 CALM 与完整模型相比的预期损失增长。这需要参考输出进行比较。
对于每一个约束，我们都可以设置允许的容差，并校准置信度阈值以允许提前退出，同时以任意高的概率可靠地满足我们定义的约束。
CALM 节省推理时间
我们对三个流行的生成数据集进行了实验：用于摘要的CNN/DM 、用于机器翻译的WMT和用于问答的SQuAD 。我们使用 8 层编码器-解码器模型评估三个置信度度量（softmax 响应、状态传播和早期退出分类器）中的每一个。为了评估全局序列级性能，我们使用标准Rouge-L、BLEU和Token-F1分数来衡量与人工编写的参考之间的距离。我们表明，平均仅使用三分之一或一半的层就可以保持完整的模型性能。CALM 通过在预测时间步长上动态分配计算工作来实现这一点。
作为近似上限，我们还使用局部 oracle 置信度度量来计算预测，这使得从第一层退出时可以得出与顶层相同的预测。在这三个任务中，当平均仅使用 1.5 个解码器层时，oracle 度量可以保持完整的模型性能。与 CALM 相比，静态基线对所有预测使用相同数量的层，需要 3 到 7 层（取决于数据集）来保持其性能。这说明了为什么动态分配计算工作很重要。只有一小部分预测需要大部分模型的复杂性，而对于其他预测来说，更少的预测就足够了。
每个任务的性能与所使用的解码器层的平均数量相对应。
最后，我们还发现 CALM 可以实现实际的加速。在 TPU 上进行基准测试时，我们节省了近一半的计算时间，同时保持了输出质量。
生成的新闻摘要示例。顶部单元格显示参考人工编写的摘要。下面是完整模型（8 层）的预测，后面是两个不同的 CALM 输出示例。第一个 CALM 输出比完整模型快 2.9 倍，第二个输出比完整模型快 3.6 倍，以 TPU 为基准。
结论
CALM 允许使用 LM 更快地生成文本，而不会降低输出文本的质量。这是通过动态修改每个生成时间步长的计算量来实现的，允许模型在足够自信时提前退出计算序列。
随着语言模型的规模不断扩大，研究如何有效地使用它们变得至关重要。CALM 是正交的，可以与许多与效率相关的努力相结合，包括模型量化、蒸馏、稀疏性、有效分区和分布式控制流。
致谢
我很荣幸能与 Adam Fisch、Ionel Gog、Seungyeon Kim、Jai Gupta、Mostafa Dehghani、Dara Bahri、Vinh Q. Tran、Yi Tay 和 Donald Metzler 合作。我们还要感谢 Anselm Levskaya、Hyung Won Chung、Tao Wang、Paul Barham、Michael Isard、Orhan Firat、Carlos Riquelme、Aditya Menon、Zhifeng Chen、Sanjiv Kumar 和 Jeff Dean 提供的有益讨论和反馈。最后，我们感谢 Tom Small 为本博文制作动画。
评论