摘要
大型预训练语言模型在那些可以"一次性完成"的任务上表现出色，例如生成逼真的文本（Brown et al., 2020）或合成计算机程序（Chen et al., 2021；Austin et al., 2021）。然而，它们在需要无限多步计算的任务上表现不佳，比如整数加法（Brown et al., 2020）或程序执行（Austin et al., 2021）。令人惊讶的是，我们发现这些模型在被要求"逐步"执行操作、展示中间计算结果时，即使在 few-shot 场景下，也能够完成复杂的多步计算。具体来说，我们通过要求 transformer 将中间计算步骤输出到"草稿纸"（scratchpad）中，来训练其执行多步计算。在一系列从长整数加法到任意程序执行的复杂任务上，我们展示了草稿纸机制极大地提升了语言模型进行多步计算的能力。
1 引言
基于 transformer 的大型语言模型表现出令人惊讶的能力（Devlin et al., 2019；Brown et al., 2020），包括生成可解决简单编程问题的代码的能力（Chen et al., 2021；Austin et al., 2021）。然而，这些模型在执行多步算法计算时表现不佳，特别是当任务需要精确推理和无限计算时。例如，GPT-3 无法在 few-shot 场景下对超过三位数的加法问题做出正确回答（Brown et al., 2020）。同样，大规模语言模型难以预测 Python 代码的执行结果，即便这些代码本身是模型能够写出的解决方案（Austin et al., 2021）。此外，标准的循环神经网络和图神经网络在预测带有循环的简单程序输出时，也无法系统性泛化（Bieber et al., 2020）。也就是说，语言模型在一定程度上可以写出代码，但似乎并不能准确地表达它们所写代码的语义，因为它们无法预测代码的执行结果。这促使研究者探索可以执行算法推理的网络（Graves et al., 2014；Zaremba & Sutskever, 2014；Bieber et al., 2020）。能够准确表达程序语义的神经网络有望支持多种下游任务，包括程序生成（Devlin et al., 2017）、程序分析（Allamanis et al., 2018）以及其他算法推理任务（Velickovic & Blundell, 2021）。
为什么大型语言模型在算法推理任务上表现不佳？我们认为这至少部分源于 transformer 架构在这些任务中的应用方式：模型被要求在一次前向传播中完成整个任务。在给定固定层数和固定计算时间的前提下，模型在输出结果前无法根据问题的难度动态调整计算资源。已有研究（Graves, 2016；Banino et al., 2021）探索了允许根据子任务动态分配计算时间的神经网络架构。而在本研究中，我们提出了一种不同的方法------一种可以利用现有 transformer 架构和支持 few-shot 能力的大型语言模型的方法------我们修改的是任务设计，而不是模型本身或训练方式。
我们的方法很简单：允许模型在生成最终答案之前，输出任意长度的中间 token 序列，我们称之为草稿纸（scratchpad）。例如，在加法问题中，草稿纸包含了标准长加法算法中的中间结果（见图 2）。为了训练模型，我们将算法的中间步骤编码为文本，并使用标准的监督学习方法进行训练。
本文的贡献如下：
- 我们在第 2 节提出了 transformer 的"草稿纸"概念，以便在不修改底层架构的前提下，提升其处理复杂离散计算任务的能力。
- 我们在第 3 节展示了草稿纸机制如何帮助 transformer 在 fine-tuning 场景下学习执行长加法，尤其是在泛化到更大问题实例时效果更佳。
- 我们在第 4 节进一步发现，草稿纸机制也帮助 transformer 执行更高层次的任务：多项式求值。这一结论在 few-shot 和 fine-tuning 场景下都成立。
- 最后，在第 5 节中，我们将视角拓展到更一般的场景，展示了让 transformer 按行输出带有局部变量注释的完整程序执行轨迹，能够显著提升其预测程序在给定输入下执行结果的能力。某种程度上，这一应用涵盖了前述所有任务。
2 方法
在本研究中，我们考虑两个相关的问题：算法归纳（Graves 等，2014；2016；Kurach 等，2016；Kaiser & Sutskever，2016）和学习执行（Zaremba & Sutskever，2014；Bieber 等，2020）。这两个问题的目标都是让神经网络学习模拟一个函数 f f f，这个函数在"算法意义"上是可表示为一个简短程序的，例如加法或多项式求值，基于其输入输出行为进行学习。
在神经算法归纳中，目标是学习一个单一算法，每个训练样本提供一个输入和期望输出，二者均表示为字符串。因此，训练数据为 D = ( x i , f ( x i ) ) i = 1 N D = {(x_i, f(x_i))}_{i=1}^N D=(xi,f(xi))i=1N。
在学习执行任务中，我们希望模型能够输出一个程序（以源码形式表示）在某个输入上的执行结果。如果每个 π i \pi_i πi 是程序 f i f_i fi 的源码，那么训练数据为 D = ( π i , x i , f i ( x i ) ) i = 1 N D = {(\pi_i, x_i, f_i(x_i))}_{i=1}^N D=(πi,xi,fi(xi))i=1N（通常每个 f i f_i fi 会有多个输入输出示例，但为了简化记号，我们在此省略）。
本论文的主要思想是：为了解决某个算法任务，我们只需将该算法的中间步骤编码为文本，并训练模型将这些中间步骤输出到一个我们称之为"scratchpad"的缓冲区中。例如，考虑学习长加法的算法归纳任务。为了教会模型如何将 29 与 57 相加，一个训练样本可能如图 2 所示，其中文本明确写出了小学长加法算法的各个步骤。
学习执行任务也可以用类似的方式进行编码，不同之处在于我们现在在输入、scratchpad 和期望输出之前添加了源码 π i \pi_i πi。图 1 展示了一个用于学习执行任务的训练样本示例。
在训练时，模型将接收输入和目标，用于标准的基于似然的训练。在测试时，模型只会接收到输入，并被要求预测目标，例如通过 beam search 或 temperature sampling 来完成。在原理上，任何序列模型都可以用于该任务。在本研究中，我们选择使用仅包含解码器的 Transformer 语言模型，但其他序列模型同样可能有效，例如 encoder-decoder 模型（Raffel 等，2019）或循环神经网络。
添加 scratchpad 有几个潜在优势：首先，模型具有自适应计算时间，也就是说，模型现在可以根据输入任务的复杂度，处理信息所需的时间可以自行调整。其次，模型可以将中间计算状态存储在 scratch 缓冲区中，并通过注意其上下文来回溯这些状态。这消除了需要将所有中间状态存储在激活中的需求。第三，通过强制模型从生成模型中采样并输出具体的中间状态，我们的目标是减少小错误的传播与累积，因为状态被量化为 token 的嵌入表示。在使用递归机制以支持扩展计算的方法中（例如 Neural Turing Machines（Graves 等，2014）），常会出现错误积累的问题。最后，检查模型生成的 scratchpad 输出可以帮助我们识别常见错误，并通过修改 scratchpad 格式来加以纠正。在本研究中，我们发现这种对错误的可解释能力非常有用。
在所有实验中，我们使用了预训练的稠密 decoder-only Transformer 语言模型，模型规模从 2 百万个参数到 1370 亿个参数不等。这些模型是在网页文档和对话数据上预训练的，与 Austin 等人（2021）中使用的模型一致。
3 加法
作为第一个任务，我们考虑整数加法。基础加法任务的输入是两个数字，目标是它们的和。例如：
输入： 2 9 + 5 7
目标： 8 6
我们通过在目标中加入长加法算法的中间步骤来实现 scratchpad，如图 2 所示。我们训练了多个模型来解决输入长度为 1 到 8 位数的整数加法问题。然后在以下两类任务中测试模型性能：一是 分布内测试任务 （输入长度最多为 8 位）；二是 分布外测试任务（输入为 9 位或 10 位数）。
这些模型在 100,000 个样本上进行了微调，训练了 5000 步，每个 batch 大小为 32。我们构建了 10,000 个分布内测试样本，以及每种分布外任务（9 位和 10 位数加法）各 1,000 个测试样本。我们考察了模型规模（从 2M 到 1B 参数）对性能的影响，并将其与一个基线进行比较：该基线只包含输入和目标数字，但不包含中间的 scratchpad 步骤。
结果
图 3 展示了 scratchpad 算法与基线方法的性能对比。可以看到，在超过某个临界模型规模之后，模型能够使用 scratchpad 成功解决加法任务，而未使用 scratchpad 训练的模型即便在最大规模下也无法完成任务。在分布外的任务（9 位和 10 位数加法）中，未使用 scratchpad 的模型完全失败，而使用了 scratchpad 的模型则随着模型规模的增加表现出持续的改进。
4 多项式求值
接下来我们关注一个稍微更高级的任务：多项式求值。受到 Saxton 等人（2019）中"多项式求值"子问题的启发，我们生成了一个多项式数据集，次数小于或等于三，系数和输入均为整数，输入范围限制在 [−10,10]，输出范围限制在 [−1000,1000]。我们生成了一个包含 10,000 个多项式的训练集和一个包含 2,000 个多项式的测试集。图 4 展示了该任务中一个带有 scratchpad 的目标示例，每一项多项式都被单独求值。和上一节类似，我们比较了直接执行和使用 scratchpad 的结果。在本实验中，我们在 few-shot 设定下使用一个 137B 参数的预训练 decoder-only 模型进行评估，因为先前工作表明非常大的模型可能能够在 few-shot 条件下完成三位数以内的加法和乘法（Brown et al., 2020），few-shot 提示中包含了 n = 4 n=4 n=4 个示例问题。我们也在微调设定下使用一个 8B 参数的模型，在训练集上进行了 2000 步微调。两种评估结果如表 1 所示。我们发现，在 few-shot 和微调两种设定下，使用 scratchpad 的执行方式都显著优于直接执行。
5 执行 Python 程序
我们已经展示了 scratchpad 可以帮助算法归纳，即帮助模型通过直接的算法特定监督学习实现某个算法。但对于每个新任务都需要手动设计中间状态，这是次优的做法。本节中，我们评估模型是否能够通过执行任意代码来学习实现新的算法。为测试这一能力，我们采用 Austin 等人（2021）的任务设置，即让语言模型预测在特定输入下执行给定 Python 程序的结果。此前研究表明，语言模型在该任务上表现较差，即使是对模型能解决的编程任务的程序也如此。在这里，我们展示了 scratchpad 技术能显著提升语言模型执行程序的能力。
直接执行预测 我们的主要基线是 Austin 等人（2021）探讨的直接执行预测方法。模型给出一个函数的源代码，要求预测该函数在特定输入上的输出。例如，图 1 中的函数输入为字符串 s 和字符 ch，功能是移除字符串 s 中字符 ch 的第一个和最后一个实例。该任务的直接执行提示和目标展示在图 1 的"直接执行预测"框中。如果模型正确输出目标字符串，则该任务被视为已解决。
通过 scratchpad 跟踪执行预测 如上所述，直接执行预测要求模型在一次推理中正确输出执行整个函数的结果。Austin 等人（2021）证明直接执行预测在 Python 程序上表现较差。因此，我们设计了一个 scratchpad 格式的执行任务，模型先预测程序执行过程中计算出的中间状态序列，再基于这些状态预测最终输出。具体来说，我们训练模型预测交替出现的序列，依次为 1）执行的源代码行序列，2）每行执行后局部变量的状态。我们称此对象为程序的"执行跟踪"，它能追踪控制流（执行的操作序列）及状态变化。执行跟踪以字符串形式表示，代码行直接复现，状态信息用 JSON 字典表示。例如，图 1 中"scratchpad 跟踪"框包含了上述函数的跟踪提示和目标。
具体而言，对于每个待跟踪的函数，提示包括函数定义，随后是一行调用该函数并传入具体输入的代码：output = 函数名(输入值)，其中函数名和输入值被相应替换。在图 1 中，可以看到 removeOcc("PHP","P") 的正确输出被写在跟踪的最后一行，赋值给变量 "output"。如果跟踪中最后一行赋给 "output" 变量的值与目标输出语义一致（这里是 "output": "P"），则认为该跟踪示例得到了正确的执行输出。若所有给定的输入-输出示例均正确执行，则认为该任务执行正确。我们还可以测试模型预测的跟踪与真实跟踪是否"完全匹配"，方法包括 a) 语义比较跟踪中每个状态与对应真实状态，b) 比较预测的源代码行序列与真实序列。
温馨提示：
阅读全文请访问"AI深语解构 " Show your work：展示你的推理过程-语言模型中的中间计算草稿机制