Original: S HOW Y OUR W ORK : S CRATCHPADS FOR I NTERMEDIATE C OMPUTATION WITH L ANGUAGE M ODELS
Revised: S HOW Y OUR W ORK : S CRATCHPADS FOR I NTERMEDIATE C OMPUTATION WITH L ANGUAGE M ODELS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Large pre-trained language models perform remarkably well on tasks that can be done “in one pass”, such as generating realistic text (Brown et al., 2020) or synthesizing computer programs (Chen et al., 2021;
Revised: Large pre-trained language models perform remarkably well on tasks that can be done “in one pass”, such as generating realistic text (Brown et al., 2020) or synthesizing computer programs (Chen et al., 2021;

Original: Austin et al., 2021).
Revised: Austin et al., 2021).

Original: However, they struggle with tasks that require unbounded multi-step computation, such as adding integers (Brown et al., 2020) or executing programs (Austin et al., 2021).
Revised: However, they struggle with tasks that require unbounded multi-step computation, such as adding integers (Brown et al., 2020) or executing programs (Austin et al., 2021).

Original: Surprisingly, we ﬁnd that these same models are able to perform complex multistep computations—even in the few-shot regime—when asked to perform the operation “step by step”, showing the results of intermediate computations.
Revised: Surprisingly, we ﬁnd that these same models are able to perform complex multistep computations—even in the few-shot regime—when asked to perform the operation “step by step”, showing the results of intermediate computations.

Original: In particular, we train Transformers to perform multi-step computations by asking them to emit intermediate computation steps into a “scratchpad”.
Revised: In particular, we train Transformers to perform multi-step computations by asking them to emit intermediate computation steps into a “scratchpad”.

Original: On a series of increasingly complex tasks ranging from long addition to the execution of arbitrary programs, we show that scratchpads dramatically improve the ability of language models to perform multi-step computations.
Revised: On a series of increasingly complex tasks ranging from long addition to the execution of arbitrary programs, we show that scratchpads dramatically improve the ability of language models to perform multi-step computations.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Large Transformer-based language models exhibit surprisingly impressive capabilities (Devlin et al., 2019;
Revised: Large Transformer-based language models exhibit surprisingly impressive capabilities (Devlin et al., 2019;

Original: Brown et al., 2020), including the ability to generate code that solves simple programming problems (Chen et al., 2021; Austin et al., 2021).
Revised: Brown et al., 2020), including the ability to generate code that solves simple programming problems (Chen et al., 2021; Austin et al., 2021).

Original: However, these models struggle to perform multistep algorithmic calculations, especially those that require precise reasoning and unbounded computation.
Revised: However, these models struggle to perform multistep algorithmic calculations, especially those that require precise reasoning and unbounded computation.

Original: For example, GPT-3 struggles to perform few-shot addition on numbers with greater than three digits (Brown et al., 2020).
Revised: For example, GPT-3 struggles to perform few-shot addition on numbers with greater than three digits (Brown et al., 2020).

Original: Similarly, large-scale language models struggle to predict the result of executing Python code, even code which is a solution to a programming task the model is able to solve (Austin et al., 2021).
Revised: Similarly, large-scale language models struggle to predict the result of executing Python code, even code which is a solution to a programming task the model is able to solve (Austin et al., 2021).

Original: Likewise, standard recurrent and graph neural networks fail to systematically generalize when predicting the output of simple programs with loops (Bieber et al., 2020).
Revised: Likewise, standard recurrent and graph neural networks fail to systematically generalize when predicting the output of simple programs with loops (Bieber et al., 2020).

Original: So language models can to some extent write code, but do not seem to understand the code they write, because they cannot predict its execution.
Revised: So language models can to some extent write code, but do not seem to accurately represent the semantics of the code they write, because they cannot predict its execution.

Original: This has motivated research on networks that can perform algorithmic reasoning (Graves et al., 2014; Zaremba & Sutskever, 2014; Bieber et al., 2020).
Revised: This has motivated research on networks that can perform algorithmic reasoning (Graves et al., 2014; Zaremba & Sutskever, 2014; Bieber et al., 2020).

Original: Neural networks that understand programs could enable a variety of downstream tasks, including program synthesis (Devlin et al., 2017), program analysis (Allamanis et al., 2018), and other algorithmic reasoning tasks (Velickovic & Blundell, 2021).
Revised: Neural networks that accurately represent the semantics of programs could enable a variety of downstream tasks, including program synthesis (Devlin et al., 2017), program analysis (Allamanis et al., 2018), and other algorithmic reasoning tasks (Velickovic & Blundell, 2021).

Original: Why do large language models struggle with algorithmic reasoning tasks?
Revised: Why do large language models struggle with algorithmic reasoning tasks?

Original: We suggest that this is at least partly due to a limitation of the  Transformer architecture and a limitation of task design: the model is asked to perform these tasks in one forward pass.
Revised: We suggest that this is at least partly due to a limitation of the way the Transformer architecture is applied to these tasks: the model is asked to perform these tasks in one forward pass.

Original: Given a ﬁxed number of layers and a ﬁxed amount of computation time, the model cannot adapt the amount of compute spent on a problem to its difﬁculty before producing an output.
Revised: Given a ﬁxed number of layers and a ﬁxed amount of computation time, the model cannot adapt the amount of compute spent on a problem to its difﬁculty before producing an output.

Original: 1 Prior work (Graves, 2016; Banino et al., 2021) has explored neural architectures that explicitly allow for dynamically chosen amounts of computation time to be dedicated to different sub-tasks.
Revised: 1 Prior work (Graves, 2016; Banino et al., 2021) has explored neural architectures that explicitly allow for dynamically chosen amounts of computation time to be dedicated to different sub-tasks.

Original: In this work, we propose a different approach—one thatcan exploit existing Transformer architectures and large few-shot-capable language models—we modify the task design rather than the model or training procedure.
Revised: In this work, we propose a different approach—one that

Original: [Table]
Revised: [Table]

Original: 1,-1,-1): if (s[i] == ch): s =
Revised: 1,-1,-1): if (s[i] == ch): s =

Original: s[0 : i] + s[i + 1:] break return s output = remove_Occ("PHP","P")
Revised: s[0 : i] + s[i + 1:] break return s output = remove_Occ("PHP","P")

Original: What is the execution trace?
Revised: What is the execution trace?

Original: Consider the following Python function: def remove_Occ(s,ch): for i in range(len(s)): if (s[i] == ch): s =
Revised: Consider the following Python function: def remove_Occ(s,ch): for i in range(len(s)): if (s[i] == ch): s =

Original: s[0 : i] + s[i + 1:] break for i in range(len(s) -
Revised: s[0 : i] + s[i + 1:] break for i in range(len(s) -

Original: 1,-1,-1): if (s[i] == ch): s =
Revised: 1,-1,-1): if (s[i] == ch): s =

Original: s[0 : i] + s[i + 1:] break return s
Revised: s[0 : i] + s[i + 1:] break return s

Original: Fill in the ???
Revised: Fill in the ???

Original: below: assert remove_Occ("PHP","P")
Revised: below: assert remove_Occ("PHP","P")

Original: =
Revised: =

Original: = ???
Revised: = ???

Original: assert remove_Oct("PHP", "P") =
Revised: assert remove_Oct("PHP", "P") =

Original: = "H"
Revised: = "H"

Original: [Figure]
Revised: [Figure]

Original: PREDICTION

Original: SCRATCHPAD TRACING
Revised: SCRATCHPAD TRACING

Original: Large Language
Revised: Large [Figure]

Original: Model 
Revised: Model can exploit existing Transformer architectures and large few-shot-capable language models—we modify the task design rather than the model or training procedure.

Original: Our proposal is simple: Allow the model to produce an arbitrary sequence of intermediate tokens, which we call a scratchpad , before producing the ﬁnal answer.
Revised: Our proposal is simple: Allow the model to produce an arbitrary sequence of intermediate tokens, which we call a scratchpad , before producing the ﬁnal answer.

Original: For example, on addition problems, the scratchpad contains the intermediate results from a standard long addition algorithm
Revised: For example, on addition problems, the scratchpad contains the intermediate results from a standard long addition algorithm

Original: (see Figure 2).
Revised: (see Figure 2).

Original: To train the model, we encode the intermediate steps of the algorithm as text and use standard supervised training.
Revised: To train the model, we encode the intermediate steps of the algorithm as text and use standard supervised training.

Original: This paper makes the following contributions:
Revised: This paper makes the following contributions:

Original: • We introduce (Section 2) the notion of a “scratchpad” for Transformers, in order to make them better at performing complex discrete computations without modifying the underlying architecture.
Revised: • We introduce (Section 2) the notion of a “scratchpad” for Transformers, in order to make them better at performing complex discrete computations without modifying the underlying architecture.

Original: • We show (Section 3) that scratchpads help Transformers learn to perform long addition in the ﬁne-tuning regime, and in particular that they improve out-of-distribution generalization to larger problem instances.
Revised: • We show (Section 3) that scratchpads help Transformers learn to perform long addition in the ﬁne-tuning regime, and in particular that they improve out-of-distribution generalization to larger problem instances.

Original: • We also ﬁnd (Section 4) that scratchpads help Transformers perform a somewhat higher level task: polynomial evaluation.
Revised: • We also ﬁnd (Section 4) that scratchpads help Transformers perform a somewhat higher level task: polynomial evaluation.

Original: This is true in both the few-shot and ﬁne-tuning regimes.
Revised: This is true in both the few-shot and ﬁne-tuning regimes.

Original: Finally, we move to a much more general context and show (Section 5) that training Transformers to emit full program traces line by line annotated with local variables dramatically improves their ability to predict the result of executing a given computer program on a particular input.
Revised: Finally, we move to a much more general context and show (Section 5) that training Transformers to emit full program traces line by line annotated with local variables dramatically improves their ability to predict the result of executing a given computer program on a particular input.

Original: This application in some sense subsumes the others.
Revised: This application in some sense subsumes the others.

Original:  M ETHOD
Revised:  M ETHOD

Original: In this work we consider two related problems: algorithm induction (Graves et al., 2014; 2016;
Revised: In this work we consider two related problems: algorithm induction (Graves et al., 2014; 2016;

Original: Kurach et al., 2016; Kaiser & Sutskever, 2016) and learning to execute (Zaremba & Sutskever, 2014; Bieber et al., 2020).
Revised: Kurach et al., 2016; Kaiser & Sutskever, 2016) and learning to execute (Zaremba & Sutskever, 2014; Bieber et al., 2020).

Original: The goal of both problems is for the neural network to learn to emulate a function f , which is “algorithmic” in the sense that it can be represented by a short program, such as addition or polynomial evaluation, from input-output behavior.
Revised: The goal of both problems is for the neural network to learn to emulate a function f , which is “algorithmic” in the sense that it can be represented by a short program, such as addition or polynomial evaluation, from input-output behavior.

Original: In neural algorithm induction , the goal is to learn a single algorithm, and each training example gives a single input and desired output represented as strings.
Revised: In neural algorithm induction , the goal is to learn a single algorithm, and each training example gives a single input and desired output represented as strings.

Original: Therefore, the training data is D = { x i , f ( x i ) } Ni =1 .
Revised: Therefore, the training data is D = { x i , f ( x i ) } Ni =1 .

Original: For learning to execute , we want the model to produce the result of a program, represented as source code, on some input.
Revised: For learning to execute , we want the model to produce the result of a program, represented as source code, on some input.

Original: If each π i is the source code of a program f i , then the training data is D = { ( π i , x i , f i ( x i ))
Revised: If each π i is the source code of a program f i , then the training data is D = { ( π i , x i , f i ( x i ))

Original: Ni
Revised: Ni

Original: =
Revised: =

Original: (it is common for each f i to have multiple input-output examples, but we omit this to lighten notation).
Revised: (it is common for each f i to have multiple input-output examples, but we omit this to lighten notation).

Original: The main idea of this paper is that to solve a given algorithmic task, we simply encode the intermediate steps of the algorithm as text and train the model to emit them to a buffer that we call a “scratchpad.”
Revised: The main idea of this paper is that to solve a given algorithmic task, we simply encode the intermediate steps of the algorithm as text and train the model to emit them to a buffer that we call a “scratchpad.”

Original: For example, let us consider the algorithmic induction task of learning long addition.
Revised: For example, let us consider the algorithmic induction task of learning long addition.

Original: To teach a model to add 29 to 57, a training example may look like the text in Figure 2, where the steps of the grade-school long addition algorithm are written out explicitly.
Revised: To teach a model to add 29 to 57, a training example may look like the text in Figure 2, where the steps of the grade-school long addition algorithm are written out explicitly.

Original: [Figure]
Revised: [Figure]

Original:  Learning to execute tasks  be  π i before the input, scratchpad, and desired out put.
Revised: of the target. Learning to execute tasks can be encoded in a similar way, except now we add the source code π i before the input, scratchpad, and desired output.

Original: An example of a training example for a learning to execute task is shown in Figure 1.
Revised: An example of a training example for a learning to execute task is shown in Figure 1.

Original: At training time, the model will be given the input plus target for standard likelihood-based training.
Revised: At training time, the model will be given the input plus target for standard likelihood-based training.

Original: At test time, the model will be given only the input and will be required to predict the target, e.g., by beam search or temperature sampling.
Revised: At test time, the model will be given only the input and will be required to predict the target, e.g., by beam search or temperature sampling.

Original: In principle, any sequence model could be used for this.
Revised: In principle, any sequence model could be used for this.

Original: In this work, we choose to use decoder-only Transformer language models, but other sequence models could be effective, such as encoder-decoder models (Raffel et al., 2019), or recurrent networks.
Revised: In this work, we choose to use decoder-only Transformer language models, but other sequence models could be effective, such as encoder-decoder models (Raffel et al., 2019), or recurrent networks.

Original: Adding a scratchpad has several potential advantages: First, the model has adaptive computation time, that is, it can now process the information for as long as needed, depending on the complexity of the task given the input.
Revised: Adding a scratchpad has several potential advantages: First, the model has adaptive computation time, that is, it can now process the information for as long as needed, depending on the complexity of the task given the input.

Original: Second, the model can store the intermediate state of its computation in the scratch buffer and refer back to it by attending to its context.
Revised: Second, the model can store the intermediate state of its computation in the scratch buffer and refer back to it by attending to its context.

Original: This removes the need to store all intermediate state in activations.
Revised: This removes the need to store all intermediate state in activations.

Original: Third, by forcing the model to output concrete intermediate states by sampling from the generative model, we aim to reduce the propagation and compounding of small errors, because states are quantized to token embeddings.
Revised: Third, by forcing the model to output concrete intermediate states by sampling from the generative model, we aim to reduce the propagation and compounding of small errors, because states are quantized to token embeddings.

Original: Compounded errors can show up in methods—like Neural Turing Machines (Graves et al., 2014)—that use recurrence to support extended computations.
Revised: Compounded errors can show up in methods—like Neural Turing Machines (Graves et al., 2014)—that use recurrence to support extended computations.

Original: Finally, examining a model’s scratchpad output can help us identify common errors and correct them by revising the scratchpad format.
Revised: Finally, examining a model’s scratchpad output can help us identify common errors and correct them by revising the scratchpad format.

Original: We found this ability to interpret errors to be useful in this work.
Revised: We found this ability to interpret errors to be useful in this work.

Original: For all experiments, we used pre-trained dense decoder-only Transformer language models, ranging in size from 2 million to around 100 billion parameters (with the largest denoted 100B+).
Revised: For all experiments, we used pre-trained dense decoder-only Transformer language models, ranging in size from 2 million to around 100 billion parameters (with the largest denoted 100B+).

Original: These models were pre-trained on web documents and dialog data.
Revised: These models were pre-trained on web documents and dialog data.

Original: We omit some details here—including exact model sizes—for double-blind review, but will include speciﬁc details in the published version.
Revised: We omit some details here—including exact model sizes—for double-blind review, but will include speciﬁc details in the published version.

Original:  A DDITION
Revised:  A DDITION

Original: As a ﬁrst task, we consider integer addition.
Revised: As a ﬁrst task, we consider integer addition.

Original: The baseline addition task presents two numbers as the input, and the target is their sum.
Revised: The baseline addition task presents two numbers as the input, and the target is their sum.

Original: For example:
Revised: For example:

Original: [-Input: 2 9 + 5Target: 
-]
Revised: [-Input: 2 9 + 5Target: 
-][Figure]

Original: [Figure]
Revised: [Figure]

Original: We implement the scratchpad by including the intermediate steps of the long addition algorithm in the target, as in Figure 2.
Revised: We implement the scratchpad by including the intermediate steps of the long addition algorithm in the target, as in Figure 2.

Original: We train several models on integer addition problems with inputs that have 1-8 digits.
Revised: We train several models on integer addition problems with inputs that have 1-8 digits.

Original: We then test performance on in-distribution addition problems (with up to 8 digit inputs), and on out-of-distribution problems with 9 and 10 digit inputs.
Revised: We then test performance on in-distribution addition problems (with up to 8 digit inputs), and on out-of-distribution problems with 9 and 10 digit inputs.

Original: The models were ﬁne-tuned on 100k examples for 5k steps with batch size 32.
Revised: The models were ﬁne-tuned on 100k examples for 5k steps with batch size 32.

Original: There are 10k in-distribution test examples, and 1k test examples for each out-of-distribution task.
Revised: There are 10k in-distribution test examples, and 1k test examples for each out-of-distribution task.

Original: We examine the performance as a function of model size, ranging from 2M to 1B parameters.
Revised: We examine the performance as a function of model size, ranging from 2M to 1B parameters.

Original: We compare performance to a baseline which includes the input and target numbers, but no intermediate scratchpad steps.
Revised: We compare performance to a baseline which includes the input and target numbers, but no intermediate scratchpad steps.

Original: Results Figure 3 compares the performance of the scratchpad algorithm with the baseline.
Revised: Results Figure 3 compares the performance of the scratchpad algorithm with the baseline.

Original: We see that beyond a critical model size, models are able to solve the addition task using the scratchpad, while models trained without a scratchpad fail to do so even at the largest tested scale.
Revised: We see that beyond a critical model size, models are able to solve the addition task using the scratchpad, while models trained without a scratchpad fail to do so even at the largest tested scale.

Original: On the out-ofdistribution tasks (9-10 digit addition), we ﬁnd that models trained without scratchpad completely fail, while models trained with scratchpad show consistent improvement as a function of model size.
Revised: On the out-ofdistribution tasks (9-10 digit addition), we ﬁnd that models trained without scratchpad completely fail, while models trained with scratchpad show consistent improvement as a function of model size.

Original:  P OLYNOMIAL E VALUATION
Revised:  P OLYNOMIAL E VALUATION

Original: Next we focus on a slightly higher-level task: evaluating polynomials.
Revised: Next we focus on a slightly higher-level task: evaluating polynomials.

Original: Inspired by the “polynomial evaluation” subproblem in Saxton et al.
Revised: Inspired by the “polynomial evaluation” subproblem in Saxton et al.

Original: (2019), we generate a dataset of polynomials of degree less than or equal to three, with integer coefﬁcients and inputs constrained to the range [ − 10 , 10] .
Revised: (2019), we generate a dataset of polynomials of degree less than or equal to three, with integer coefﬁcients and inputs constrained to the range [ − 10 , 10] .

Original: We also restrict outputs to the range [ − 1000 , 1000] .
Revised: We also restrict outputs to the range [ − 1000 , 1000] .

Original: We generate a training dataset of 10,000 polynomials and a test dataset of size 2,000.
Revised: We generate a training dataset of 10,000 polynomials and a test dataset of size 2,000.

Original: An example scratchpad target for this task is shown in Figure 4, with each term of the polynomial evaluated separately.
Revised: An example scratchpad target for this task is shown in Figure 4, with each term of the polynomial evaluated separately.

Original: As in the previous section, we compare the results of direct execution with the results of using the scratchpad.
Revised: As in the previous section, we compare the results of direct execution with the results of using the scratchpad.

Original: In this experiment, we evaluate in the few-shot regime using a 100B+ parameter pre-trained decoder-only model, as previous work indicates that very large models may be able to perform additions and multiplications with 3 or fewer digits few-shot (Brown et al., 2020).
Revised: In this experiment, we evaluate in the few-shot regime using a 100B+ parameter pre-trained decoder-only model, as previous work indicates that very large models may be able to perform additions and multiplications with 3 or fewer digits few-shot (Brown et al., 2020).

Original: We use n = 4 example problems in the few-shot prompt.
Revised: We use n = 4 example problems in the few-shot prompt.

Original: We also evaluate in the ﬁne-tuning regime with an 8B parameter model ﬁne-tuned for 2000 steps on the training set.
Revised: We also evaluate in the ﬁne-tuning regime with an 8B parameter model ﬁne-tuned for 2000 steps on the training set.

Original: The results of both evaluations are shown in Table 1.
Revised: The results of both evaluations are shown in Table 1.

Original: We ﬁnd that scratchpad execution outperforms direct execution signiﬁcantly in both the few-shot and ﬁne-tuning regimes.
Revised: We ﬁnd that scratchpad execution outperforms direct execution signiﬁcantly in both the few-shot and ﬁne-tuning regimes.

Original:  E XECUTING P YTHON P ROGRAMS
Revised:  E XECUTING P YTHON P ROGRAMS

Original: We have shown that scratchpads can help algorithm induction, that is, they can help models learn to implement a particular algorithm with direct algorithm-speciﬁc supervision.
Revised: We have shown that scratchpads can help algorithm induction, that is, they can help models learn to implement a particular algorithm with direct algorithm-speciﬁc supervision.

Original: But needing to handdesign the intermediate states for every new task is sub-optimal.
Revised: But needing to handdesign the intermediate states for every new task is sub-optimal.

Original: In this section, we evaluate whether a model can learn to implement a new algorithm by executing arbitrary code.
Revised: In this section, we evaluate whether a model can learn to implement a new algorithm by executing arbitrary code.

Original: To test this capability, we follow the problem setup from Austin et al.
Revised: To test this capability, we follow the problem setup from Austin et al.

Original: (2021), in which language models are asked to predict the result of executing a given Python program on a particular input.
Revised: (2021), in which language models are asked to predict the result of executing a given Python program on a particular input.

Original: Language models performed poorly at this task, even on programs which are solutions to a programming tasks the model is able tosolve.
Revised: Language models performed poorly at this task, even on programs which are solutions to a programming tasks the model is able to

Original: Here we show that the scratchpad technique can dramatically improve the ability of language models to execute programs.

Original: [Figure]
Revised: [Figure]

Original: with a scratchpad.
Revised: with a scratchpad.

Original: Each term in the polynomial is computed separately and then added.
Revised: Each term in the polynomial is computed separately and then added.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Revised: solve.

Revised: Here we show that the scratchpad technique can dramatically improve the ability of language models to execute programs.

Original: Direct execution prediction Our main baseline is the direct execution prediction procedure explored in Austin et al.
Revised: Direct execution prediction Our main baseline is the direct execution prediction procedure explored in Austin et al.

Original: Models are shown the source code for a function, and asked to predict the output of running the function on speciﬁc inputs.
Revised: Models are shown the source code for a function, and asked to predict the output of running the function on speciﬁc inputs.

Original: For example, the function in Figure 1 takes as input a string s and a character ch , and removes the ﬁrst and last instances of the character ch from the string s .
Revised: For example, the function in Figure 1 takes as input a string s and a character ch , and removes the ﬁrst and last instances of the character ch from the string s .

Original: The direct execution prompt and target for this task are shown in the “Direct Execution
Revised: The direct execution prompt and target for this task are shown in the “Direct Execution

Original: Prediction” box in Figure 1.
Revised: Prediction” box in Figure 1.

Original: A task is considered solved under this regime if the model correctly outputs the target string.
Revised: A task is considered solved under this regime if the model correctly outputs the target string.

Original: Execution prediction via scratchpad tracing As discussed above, direct execution prediction requires the model to correctly output the result of executing the entire function in a single pass.
Revised: Execution prediction via scratchpad tracing As discussed above, direct execution prediction requires the model to correctly output the result of executing the entire function in a single pass.

Original: Direct execution prediction has been shown to perform poorly on Python programs in Austin et al.
Revised: Direct execution prediction has been shown to perform poorly on Python programs in Austin et al.

Original: We therefore design a scratchpad formulation of the execution task, in which models predict the output of a program by ﬁrst predicting the sequence of intermediate states computed during the program’s execution.
Revised: We therefore design a scratchpad formulation of the execution task, in which models predict the output of a program by ﬁrst predicting the sequence of intermediate states computed during the program’s execution.

Original: Formally, we train models to predict an alternating sequence of 1) the ordered sequence of source code lines executed, and 2) the state of the local variables after each line is executed.
Revised: Formally, we train models to predict an alternating sequence of 1) the ordered sequence of source code lines executed, and 2) the state of the local variables after each line is executed.

Original: We call this object the program’s trace , and it allows us to track both the control ﬂow— the sequence of operations executed—and how the state changes as a result of each operation.
Revised: We call this object the program’s trace , and it allows us to track both the control ﬂow— the sequence of operations executed—and how the state changes as a result of each operation.

Original: We represent the trace as a string, with the line of code reproduced directly, and the state information represented as a JSON dictionary.
Revised: We represent the trace as a string, with the line of code reproduced directly, and the state information represented as a JSON dictionary.

Original: For example, the “Scratchpad Tracing” box in Figure 1 contains the tracing prompt and trace target for the function discussed above.
Revised: For example, the “Scratchpad Tracing” box in Figure 1 contains the tracing prompt and trace target for the function discussed above.

Original: Concretely, for each function to be traced, the prompt is formed by printing the function deﬁnition, followed by a line which calls the function on a particular input: output = fn
Revised: Concretely, for each function to be traced, the prompt is formed by printing the function deﬁnition, followed by a line which calls the function on a particular input: output = fn

Original: name(input value) , where fn name and input value are replaced with the corresponding function name and input value.
Revised: name(input value) , where fn name and input value are replaced with the corresponding function name and input value.

Original: In Figure 1, note how the correct output of remove Occ("PHP","P") is shown in the last line of the trace, assigned to the variable "output" .
Revised: In Figure 1, note how the correct output of remove Occ("PHP","P") is shown in the last line of the trace, assigned to the variable "output" .

Original: A tracing example is considered to have the correct execution output if the encoding of the value assigned to the variable output in the last line is a semantic match with the target output value (here, "output": "P" ).
Revised: A tracing example is considered to have the correct execution output if the encoding of the value assigned to the variable output in the last line is a semantic match with the target output value (here, "output": "P" ).

Original: We consider a task to be executed correctly if all given input-output examples are correctly executed.
Revised: We consider a task to be executed correctly if all given input-output examples are correctly executed.

Original: We can also test whether there is a “trace exact match” between the model prediction and the ground truth trace, by a) semantically comparing each state in the trace to the corresponding state in the ground truth trace, and b) comparing the sequence of source code lines predicted with the ground truth sequence.
Revised: We can also test whether there is a “trace exact match” between the model prediction and the ground truth trace, by a) semantically comparing each state in the trace to the corresponding state in the ground truth trace, and b) comparing the sequence of source code lines predicted with the ground truth sequence.

Original: Experimental setup As a proof-of-concept, we ﬁrst show that scratchpad tracing greatly improves execution performance on synthetic Python data.
Revised: Experimental setup As a proof-of-concept, we ﬁrst show that scratchpad tracing greatly improves execution performance on synthetic Python data.

Original: Then, we compare scratchpad tracing and execution on the human-written Python problems from Austin et al.
Revised: Then, we compare scratchpad tracing and execution on the human-written Python problems from Austin et al.

Original: We ﬁnd that a novel data augmentation technique that uses programs generated by the model as additional training data can signiﬁcantly increase tracing performance on real data, whereas this augmentation technique hurtsperformance for direct execution.
Revised: We ﬁnd that a novel data augmentation technique that uses programs generated by the model as additional training data can signiﬁcantly increase tracing performance on real data, whereas this augmentation technique hurts

Original: Finally, we show that incorporating tracing data from additional sources further improves tracing performance, indicating that the scratchpad tracing technique explored here may scale well with more data.

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Revised: performance for direct execution.

Revised: Finally, we show that incorporating tracing data from additional sources further improves tracing performance, indicating that the scratchpad tracing technique explored here may scale well with more data.

Original: For all experiments on Python code, we use a Transformer model with around 100 billion parameters (100B+), a context window of 1024 tokens and a limit of 512 generation tokens.
Revised: For all experiments on Python code, we use a Transformer model with around 100 billion parameters (100B+), a context window of 1024 tokens and a limit of 512 generation tokens.

Original: Unless otherwise stated, all ﬁne-tuning runs used a batch size of 8192 tokens and a learning rate of 3e-5, and model inference was performed with decoding temperature set to T = 0 , equivalent to greedy decoding.
Revised: Unless otherwise stated, all ﬁne-tuning runs used a batch size of 8192 tokens and a learning rate of 3e-5, and model inference was performed with decoding temperature set to T = 0 , equivalent to greedy decoding.

Original: 5. S CRATCHPAD B EATS D IRECT E XECUTION FOR S YNTHETIC P YTHON P ROGRAMS
Revised: 5. S CRATCHPAD B EATS D IRECT E XECUTION FOR S YNTHETIC P YTHON P ROGRAMS

Original: In our ﬁrst experiment, we test the few-shot and ﬁne-tuned execution capabilities of our models on simple synthetic Python programs.
Revised: In our ﬁrst experiment, we test the few-shot and ﬁne-tuned execution capabilities of our models on simple synthetic Python programs.

Original: This provides a proof-of-concept for our tracing technique.
Revised: This provides a proof-of-concept for our tracing technique.

Original: We use a dataset of synthetic Python programs modiﬁed from Bieber et al.
Revised: We use a dataset of synthetic Python programs modiﬁed from Bieber et al.

Original: These programs include small integers (0, 1, and 2), simple while loops, and if statements.
Revised: These programs include small integers (0, 1, and 2), simple while loops, and if statements.

Original: We construct a corpus of synthetic programs to mimic the size of the MBPP dataset in Austin et al.
Revised: We construct a corpus of synthetic programs to mimic the size of the MBPP dataset in Austin et al.

Original: (2021), with 400 training programs, 100 validation programs, and 200 test programs.
Revised: (2021), with 400 training programs, 100 validation programs, and 200 test programs.

Original: For each program, three random integer inputs are sampled from the range 0 to 9.
Revised: For each program, three random integer inputs are sampled from the range 0 to 9.

Original: We test execution and scratchpad tracing under both few-shot and ﬁne-tuning conditions.
Revised: We test execution and scratchpad tracing under both few-shot and ﬁne-tuning conditions.

Original: For fewshot experiments, the prompt contains three examples of previous tracing problems, as shown in
Revised: For fewshot experiments, the prompt contains three examples of previous tracing problems, as shown in

Original: Appendix A. For ﬁne-tuned experiments, we ﬁne-tune models to convergence on the training split, as judged by validation perplexity.
Revised: Appendix C. For ﬁne-tuned experiments, we ﬁne-tune models to convergence on the training split, as judged by validation perplexity.

Original: For the few-shot scratchpad experiment, we noticed that models would not assign the variable name output to the ﬁnal value in the trace, and would instead continue using v0 (the name of the variable returned in the function f ) as the variable name for the ﬁnal output line.
Revised: For the few-shot scratchpad experiment, we noticed that models would not assign the variable name output to the ﬁnal value in the trace, and would instead continue using v0 (the name of the variable returned in the function f ) as the variable name for the ﬁnal output line.

Original: We therefore modiﬁed the accuracy criterion from checking whether the value of output in the last line of the trace is correct, to checking whether the value of v0 is correct.
Revised: We therefore modiﬁed the accuracy criterion from checking whether the value of output in the last line of the trace is correct, to checking whether the value of v0 is correct.

Original: (Under naive scoring, the few-shot tracing accuracy is roughly zero.)
Revised: (Under naive scoring, the few-shot tracing accuracy is roughly zero.)

Original: An example of this behavior is shown in Appendix B.
Revised: An example of this behavior is shown in Appendix D.

Original: Results Table 2 shows our results on synthetic Python problems.
Revised: Results Table 2 shows our results on synthetic Python problems.

Original: In both few-shot and ﬁne-tuned settings, the scratchpad tracing technique leads to higher overall execution accuracy on the 200 test problems.
Revised: In both few-shot and ﬁne-tuned settings, the scratchpad tracing technique leads to higher overall execution accuracy on the 200 test problems.

Original: Fine-tuning also improves performance more for the scratchpad tracing technique than it does for direct execution.
Revised: Fine-tuning also improves performance more for the scratchpad tracing technique than it does for direct execution.

Original: 5. S CRATCHPAD B EATS D IRECT E XECUTION FOR R EAL P ROGRAMS
Revised: 5. S CRATCHPAD B EATS D IRECT E XECUTION FOR R EAL P ROGRAMS

Original: In our second set of experiments, we explore how well the scratchpad performs compared to execution on real data.
Revised: In our second set of experiments, we explore how well the scratchpad performs compared to execution on real data.

Original: Our main evaluation dataset is the MBPP dataset, introduced in Austin et al.
Revised: Our main evaluation dataset is the MBPP dataset, introduced in Austin et al.

Original: MBPP consists of 1000 programming problems, each of which contains a natural language speciﬁcation, a ground-truth Python program, and three input-output test cases.
Revised: MBPP consists of 1000 programming problems, each of which contains a natural language speciﬁcation, a ground-truth Python program, and three input-output test cases.

Original: These programs involve computation using a large variety of types, including ints, strings, ﬂoats, dictionaries, tuples, and more, and include many language features and control-ﬂow structures, such as loops, comprehensions, library imports, API calls and recursion.
Revised: These programs involve computation using a large variety of types, including ints, strings, ﬂoats, dictionaries, tuples, and more, and include many language features and control-ﬂow structures, such as loops, comprehensions, library imports, API calls and recursion.

Original: The evaluation split of the MBPP dataset contains 500 tasks.
Revised: The evaluation split of the MBPP dataset contains 500 tasks.

Original: In order to separate out effects of the generation window size, we report all evaluation metrics on the subset of these tasks for which the ground-truth trace ﬁts within the generation window of the model for all three of the input-output examples.
Revised: In order to separate out effects of the generation window size, we report all evaluation metrics on the subset of these tasks for which the ground-truth trace ﬁts within the generation window of the model for all three of the input-output examples.

Original: This leaves a subset of 212 testtasks.
Revised: This leaves a subset of 212 test

Original: Increasing generation and context window length is an important issue for Transformer-based models, but we view it as orthogonal and leave it for future work.

Original: We report “per-task” execution and tracing accuracies, which require all examples to be correctly executed/traced.
Revised: We report “per-task” execution and tracing accuracies, which require all examples to be correctly executed/traced.

Original: We additionally report “per-example” accuracies, which correspond to the total fraction of test examples which are executed/traced correctly across the dataset.
Revised: We additionally report “per-example” accuracies, which correspond to the total fraction of test examples which are executed/traced correctly across the dataset.

Original: We ﬁnd that training scratchpad models on an dataset augmented with samples from the model signiﬁcantly improves performance for the scratchpad model, while it harms the direct execution model.
Revised: We ﬁnd that training scratchpad models on an dataset augmented with samples from the model signiﬁcantly improves performance for the scratchpad model, while it harms the direct execution model.

Original: Combining tracing training data from several sources further improves scratchpad model performance. 
Revised: Combining tracing training data from several sources further improves scratchpad model performance. [Table] tasks.

Revised: Increasing generation and context window length is an important issue for Transformer-based models

Original: [Table]
Revised: [Table]

Revised: but we view it as orthogonal and leave it for future work

Original: [Table]
Revised: [Table]

Original:  P ERFORMANCE IS P OOR IN THE V ERY -L OW -D ATA R EGIME
Revised: 5.2. P ERFORMANCE IS P OOR IN THE V ERY -L OW -D ATA R EGIME

Original: In our ﬁrst experiment with the MBPP data
Revised: In our ﬁrst experiment with the MBPP data

Original: [Table]
Revised: [Table]

Original: we train a scratchpad tracing model on the 374 training tasks (3 examples per task, so 1122 overall examples).
Revised: we train a scratchpad tracing model on the 374 training tasks (3 examples per task, so 1122 overall examples).

Original: We discard all training examples which exceed the context window.
Revised: We discard all training examples which exceed the context window.

Original: We compare overall execution results against a model trained on the same 374 training tasks to perform direct execution.
Revised: We compare overall execution results against a model trained on the same 374 training tasks to perform direct execution.

Original: The columns labeled “MBPP” for Direct Execution and Scratchpad in Table 3 show the results of this experiment.
Revised: The columns labeled “MBPP” for Direct Execution and Scratchpad in Table 3 show the results of this experiment.

Original: Neither the scratchpad model or the direct execution model achieve good performance (5% and 10% output accuracy, respectively), and direct execution outperforms the scratchpad model.
Revised: Neither the scratchpad model or the direct execution model achieve good performance (5% and 10% output accuracy, respectively), and direct execution outperforms the scratchpad model.

Original: 5.2. S AMPLED P ROGRAMS M AKE G OOD S CRATCHPAD T RAINING D ATA
Revised: 5.2. S AMPLED P ROGRAMS M AKE G OOD S CRATCHPAD T RAINING D ATA

Original: Next, we employ a data augmentation technique to increase the size of the training dataset: We ﬁrst run few-shot synthesis on the 374 MBPP training tasks using the pre-trained 100B+ model, as described in Austin et al.
Revised: Next, we employ a data augmentation technique to increase the size of the training dataset: We ﬁrst run few-shot synthesis on the 374 MBPP training tasks using the pre-trained 100B+ model, as described in Austin et al.

Original: For each task, we sample and record 80 candidate programs { P s } from the model at temperature T = 0 .
Revised: For each task, we sample and record 80 candidate programs { P s } from the model at temperature T = 0 .

Original: We can then create a new execution datapoint using the candidate program P s , the original three inputs for the task { x i } i =1 , 2 , 3 , and the three new outputs which result from running the candidate program on the original three inputs: { y i new } i =1 , 2 , 3 , where y i new
Revised: We can then create a new execution datapoint using the candidate program P s , the original three inputs for the task { x i } i =1 , 2 , 3 , and the three new outputs which result from running the candidate program on the original three inputs: { y i new } i =1 , 2 , 3 , where y i new

Original: = P s ( x i ) .
Revised: = P s ( x i ) .

Original: We discard any candidate programs for which execution results in an error.
Revised: We discard any candidate programs for which execution results in an error.

Original: Note that the outputs of y i new may or may not be equal to the original outputs, depending on the computation performed by the generated program P s .
Revised: Note that the outputs of y i new may or may not be equal to the original outputs, depending on the computation performed by the generated program P s .

Original: Therefore, this augmented direct execution dataset has both additional new programs and new outputs compared to the original dataset.
Revised: Therefore, this augmented direct execution dataset has both additional new programs and new outputs compared to the original dataset.

Original: We can analogously create a tracing dataset for our scratchpad model by tracing the execution of each candidate program P s on each x i .
Revised: We can analogously create a tracing dataset for our scratchpad model by tracing the execution of each candidate program P s on each x i .

Original: This process produces much larger tracing and execution datasets with 17k new programs, which we refer to as MBPP-aug .
Revised: This process produces much larger tracing and execution datasets with 17k new programs, which we refer to as MBPP-aug .

Original: Conceptually, we have augmented the dataset using a combination of tools already available to us, namely a) the neural model, and b) program execution via a Python interpreter.
Revised: Conceptually, we have augmented the dataset using a combination of tools already available to us, namely a) the neural model, and b) program execution via a Python interpreter.

Original: We ﬁne-tune direct execution and scratchpad models on this new augmented dataset MBPP-aug, using the same process as above.
Revised: We ﬁne-tune direct execution and scratchpad models on this new augmented dataset MBPP-aug, using the same process as above.

Original: The “MBPP-aug” columns in Table 3 show the results of this experiment.
Revised: The “MBPP-aug” columns in Table 3 show the results of this experiment.

Original: While the direct execution approach suffers a decrease in accuracy when trained on this additional data, the performance of the scratchpad model is greatly improved; the model trained on the augmented data solves more than three times the number of tasks as the model trained on only the original MBPP programs.
Revised: While the direct execution approach suffers a decrease in accuracy when trained on this additional data, the performance of the scratchpad model is greatly improved; the model trained on the augmented data solves more than three times the number of tasks as the model trained on only the original MBPP programs.

Original: We also note that if we measure the raw correctness across samples, the model already achieves 26.8% exact trace match, which is surprisingly high.
Revised: We also note that if we measure the raw correctness across samples, the model already achieves 26.8% exact trace match, which is surprisingly high.

Original: 5. S CRATCHPAD T RAINING M AKES G OOD U SE OF L ARGE D ATASETS
Revised: 5. S CRATCHPAD T RAINING M AKES G OOD U SE OF L ARGE D ATASETS

Original: [Figure]
Revised: [Figure]

Original: In this section
Revised: In this section

Original: [Figure]
Revised: [Figure]

Original: we examine whether collecting additional tracing data from human-written programs further improves tracing performance.
Revised: we examine whether collecting additional tracing data from human-written programs further improves tracing performance.

Original: This will allow us to understand whether the tracing procedure here is likely to scale well when slightly out-of-distribution tracing data is added to the ﬁne-tuning set.
Revised: This will allow us to understand whether the tracing procedure here is likely to scale well when slightly out-of-distribution tracing data is added to the ﬁne-tuning set.

Original: We experiment using two datasets:
Revised: We experiment using two datasets:

Original: Single-line programs This dataset consists of roughly 9 million examples of single-line Python transformations.
Revised: Single-line programs This dataset consists of roughly 9 million examples of single-line Python transformations.

Original: Figure 6 (Top) shows examples of these transformations.
Revised: Figure 6 (Top) shows examples of these transformations.

Original: Each transformation consists of an initial set of variables and corresponding values, a single line of Python (together these form the input), and the new set of variables and values which results from running the line (the target).
Revised: Each transformation consists of an initial set of variables and corresponding values, a single line of Python (together these form the input), and the new set of variables and values which results from running the line (the target).

Original: When training on single-line data, we do not introduce intermediate scratchpad steps.
Revised: When training on single-line data, we do not introduce intermediate scratchpad steps.

Original: While this dataset does not provide examples of the high-level, multi-line control ﬂow of a trace, the data provides good supervision for modeling the execution of individual lines of code, which is a key component of tracing.
Revised: While this dataset does not provide examples of the high-level, multi-line control ﬂow of a trace, the data provides good supervision for modeling the execution of individual lines of code, which is a key component of tracing.

Original: This data was collected by Fraser Greenlee, and can be accessed here.
Revised: This data was collected by Fraser Greenlee, and can be accessed here.

Original: CodeNet The Project CodeNet dataset
Revised: CodeNet The Project CodeNet dataset

Original: (Puri et al., 2021) consists of millions of user submissions to approximately 4,000 coding problems.
Revised: (Puri et al., 2021) consists of millions of user submissions to approximately 4,000 coding problems.

Original: These submission include both correct and incorrect solutions to programming problems.
Revised: These submission include both correct and incorrect solutions to programming problems.

Original: However, from the experiment with MBPP-aug above, we know that incorrect or broken programs can still provide a useful training signal.
Revised: However, from the experiment with MBPP-aug above, we know that incorrect or broken programs can still provide a useful training signal.

Original: We additionally improved our tracing technique to allow tracing programs with errors; when an error is reached, the error message is added to the end of the trace text and tracing is stopped.
Revised: We additionally improved our tracing technique to allow tracing programs with errors; when an error is reached, the error message is added to the end of the trace text and tracing is stopped.

Original: We extracted a total of 670, traces from the CodeNet data.
Revised: We extracted a total of 670, traces from the CodeNet data.

Original: For each dataset, we ﬁrst ﬁne-tune the model on these datasets, and then perform a second ﬁne-tuning on MBPP-aug until convergence.
Revised: For each dataset, we ﬁrst ﬁne-tune the model on these datasets, and then perform a second ﬁne-tuning on MBPP-aug until convergence.

Original: Results Results are shown in Table 3.
Revised: Results Results are shown in Table 3.

Original: As above, we report execution accuracy across tasks.
Revised: As above, we report execution accuracy across tasks.

Original: We additionally report trace accuracy across tasks, to understand the extent to which the entire trace is accurately predicted.
Revised: We additionally report trace accuracy across tasks, to understand the extent to which the entire trace is accurately predicted.

Original: We also report the raw execution and trace accuracy across all test examples, as an additional metric to compare models.
Revised: We also report the raw execution and trace accuracy across all test examples, as an additional metric to compare models.

Original: Training on either the single-line dataset or the CodeNet dataset alone seem to both provide gains over MBPP-aug (23.4% and 25.2% tasks executed correctly, respectively).
Revised: Training on either the single-line dataset or the CodeNet dataset alone seem to both provide gains over MBPP-aug (23.4% and 25.2% tasks executed correctly, respectively).

Original: However, combining both CodeNet and the single-line dataset seems lead to the highest performance; tracing produces the correct ﬁnal output for 26.6% of the tasks, and nearly a quarter of the tasks (24.6%) are traced perfectly for all three examples.
Revised: However, combining both CodeNet and the single-line dataset seems lead to the highest performance; tracing produces the correct ﬁnal output for 26.6% of the tasks, and nearly a quarter of the tasks (24.6%) are traced perfectly for all three examples.

Original: These results seem promising: the neural network can often exactly trace programs.
Revised: These results seem promising: the neural network can often exactly trace programs.

Original: In particular, greedily decoding from the best model produces the exact correct trace for almost 42% of all traces.
Revised: In particular, greedily decoding from the best model produces the exact correct trace for almost 42% of all traces.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: The tasks in this paper can be viewed as exploring one criticism of large language models, namely, to what extent do they simply rely on surface-level statistical correlations on text, without learn- ing semantics or world knowledge (Bender & Koller, 2020)?
Revised: The tasks in this paper can be viewed as exploring one criticism of large language models, namely, to what extent do they simply rely on surface-level statistical correlations on text, without learn- ing semantics or world knowledge (Bender & Koller, 2020)?

Original: In response, Li et al.
Revised: In response, Li et al.

Original: (2021) provide evidence that pre-trained language models do indeed construct approximate representations of the semantics of the situations they describe in text.
Revised: (2021) provide evidence that pre-trained language models do indeed construct approximate representations of the semantics of the situations they describe in text.

Original: In the context of programs, Austin et al.
Revised: In the context of programs, Austin et al.

Original: (2021) approach this question by exploring the learning to execute task on MBPP, which we consider in
Revised: (2021) approach this question by exploring the learning to execute task on MBPP, which we consider in

Original: Section 5.2.
Revised: Section 5.2.

Original: The idea behind this task was to explore whether neural models for synthesis that generate code could also execute it.
Revised: The idea behind this task was to explore whether neural models for synthesis that generate code could also execute it.

Original: While that work ﬁnds existing models perform poorly at predicting execution, we show that adding a scratchpad allows these models to perform better.
Revised: While that work ﬁnds existing models perform poorly at predicting execution, we show that adding a scratchpad allows these models to perform better.

Original: Work in learning to execute has considered whether off-the-shelf recurrent neural networks (Zaremba & Sutskever, 2014) or more specialized architectures (Dehghani et al., 2018; Bieber et al.,  2020) have an inductive bias that is sufﬁciently well suited for executing  arbitrary code.
Revised: Work in learning to execute has considered whether off-the-shelf recurrent neural networks (Zaremba & Sutskever, 2014) or more specialized architectures (Dehghani et al., 2018; Bieber et al., 2020; Wang et al., 2020) have an inductive bias that is sufﬁciently well suited for executing and reasoning about arbitrary code.

Original: The related problem of neural algorithm induction has attracted considerable interest (Graves et al., 2014;
Revised: The related problem of neural algorithm induction has attracted considerable interest (Graves et al., 2014;

Original: Kurach et al., 2016; Kaiser & Sutskever, 2016; Graves et al., 2016; Reed & de Freitas, 2016).
Revised: Kurach et al., 2016; Kaiser & Sutskever, 2016; Graves et al., 2016; Reed & de Freitas, 2016; Veliˇckovi´c et al., 2020a;b).

Original: This work proposes new neural architectures, inspired by theoretical models of computation, whose inductive bias allows them to more easily learn algorithm induction tasks.
Revised: This work proposes new neural architectures, inspired by theoretical models of computation, whose inductive bias allows them to more easily learn algorithm induction tasks.

Original: Several methods for algorithm induction speciﬁcally add adaptive computation time to sequence models (Graves, 2016; Dehghani et al., 2018; Banino et al., 2021).
Revised: Several methods for algorithm induction speciﬁcally add adaptive computation time to sequence models (Graves, 2016; Dehghani et al., 2018; Banino et al., 2021).

Original: In particular, universal transformers include adaptive computation time, and are evaluated both on algorithm induction and on learning to execute tasks (Dehghani et al., 2018).
Revised: In particular, universal transformers include adaptive computation time, and are evaluated both on algorithm induction and on learning to execute tasks (Dehghani et al., 2018).

Original: In contrast, a scratchpad is a simple way both to provide a  model with adaptive computation time, and also to provide supervision about how to use that additional computation, without requiring modiﬁcation to the underlying architecture.
Revised: In contrast, a scratchpad is a simple way both to provide a transformer model with adaptive computation time, and also to provide supervision about how to use that additional computation, without requiring modiﬁcation to the underlying architecture.

Original: Algorithm induction has also been connected to pre-trained models.
Revised: Algorithm induction has also been connected to pre-trained models.

Original: Lu et al.
Revised: Lu et al.

Original: (2021) show that
Revised: (2021) show that

Original: Transformers can be used to some extent as universal computation engines, by pre-trainingthem on natural language, and ﬁne-tuning a small fraction of the weights on non-language tasks, including simple algorithm induction tasks.
Revised: Transformers can be used to some extent as universal computation engines, by pre-training on natural language, and ﬁne-tuning a small fraction of the weights on non-language tasks, including simple algorithm induction tasks.

Original: Finally, supervised approaches to semantic parsing (Zelle & Mooney, 1996; Zettlemoyer & Collins, 2005; Kwiatkowksi et al., 2010; Wong & Mooney, 2006) predict the text of a database query, which can then be executed to answer a natural language question.
Revised: Finally, supervised approaches to semantic parsing (Zelle & Mooney, 1996; Zettlemoyer & Collins, 2005; Kwiatkowksi et al., 2010; Wong & Mooney, 2006) predict the text of a database query, which can then be executed to answer a natural language question.

Original:  L IMITATIONS AND F UTURE W ORK
Revised:  L IMITATIONS AND F UTURE W ORK

Original: Context window size In this work, we limit our experiments to problems where the scratchpad text ﬁts within the model generation window.
Revised: Context window size In this work, we limit our experiments to problems where the scratchpad text ﬁts within the model generation window (512 tokens).

Original: However, many problems require very long scratchpad generations.
Revised: However, many problems require very long scratchpad generations.

Original: Therefore, fully realizing the potential of the scratchpad technique may require further improvements in transformer generation window size.
Revised: Therefore, fully realizing the potential of the scratchpad technique may require further improvements in transformer generation window size.

Original: This is an active area of research in NLP (Tay et al., 2020), and improvements would be beneﬁcial for the scratchpad technique.
Revised: This is an active area of research in NLP (Tay et al., 2020), and improvements would be beneﬁcial for the scratchpad technique.

Original: Learning to use the scratchpad without supervision A clear next step is to try to learn to use the scratchpad without direct supervision.
Revised: Learning to use the scratchpad without supervision A clear next step is to try to learn to use the scratchpad without direct supervision.

Original: A simple method would be to use reinforcement learning (RL) techniques: models would be rewarded for correctly answering questions, with reward inversely proportional to the number of scratchpad tokens used.
Revised: A simple method would be to use reinforcement learning (RL) techniques: models would be rewarded for correctly answering questions, with reward inversely proportional to the number of scratchpad tokens used.

Original: We would hope that learning to use the scratchpad would be a transferable skill; for example, a model could potentially use the algorithm it learned to perform long addition to succeed at polynomial evaluation.
Revised: We would hope that learning to use the scratchpad would be a transferable skill; for example, a model could potentially use the algorithm it learned to perform long addition to succeed at polynomial evaluation.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work we showed—through experiments on long addition, polynomial evaluation, and Python code execution—that allowing models to read from and write to a simple scratchpad can improve their performance on algorithmic tasks.
Revised: In this work we showed—through experiments on long addition, polynomial evaluation, and Python code execution—that allowing models to read from and write to a simple scratchpad can improve their performance on algorithmic tasks.

Original: Such models may be a ﬁrst step toward combining the knowledge-compression capabilities of large language models with reasoning capabilities, in order to build models that understand code as well as write it.
Revised: Such models may be a ﬁrst step toward combining the knowledge-compression capabilities of large language models with reasoning capabilities, in order to build models that understand code as well as write it.

Original: This could be useful for a variety of applications that require both working with natural language and reasoning about program semantics, such as program synthesis, neural-guided program analysis, and interactive programming assistants.
Revised: This could be useful for a variety of applications that require both working with natural language and reasoning about program semantics, such as program synthesis, neural-guided program analysis, and interactive programming assistants.

Original: The scratchpad technique presented here might not take us all the way toward that goal, but we hope it is an important step.
Revised: The scratchpad technique presented here might not take us all the way toward that goal, but we hope it is an important step.

Revised: R EPRODUCIBILITY

Revised: All ﬁne-tuning and evaluation datasets used in this work are either open source or synthetic and easily reproducible (this excludes MBPP-Aug, which depends on generations from the pre-trained transformer model used in this work).

Revised: Examples of exact prompts are shown in the Appendix, so that they can be exactly reproduced.

Revised: Although the pre-training details are not open-source, they correspond to the details in Austin et al.

Original: A E XAMPLE FEW - SHOT PROMPT FOR SYNTHETIC P YTHON EXPERIMENTS
Revised: A E FFECTS OF S CRATCHPAD E XECUTION T RAINING ON S YNTHESIS P ERFORMANCE

Revised: To measure the extent to which ﬁne-tuning on the tracing task described above affects the model’s ability to perform program synthesis we ran a few-shot synthesis experiment using the “MBPPaug + CodeNet + single line” model.

Revised: Speciﬁcally, we performed few-shot synthesis on the MBPP dataset, as described in Austin et al.

Revised: For each MBPP synthesis task, 80 candidate programs are sampled from the model ( T = 0 . 5 ), and the task is considered solved if any of the candidate programs satisfy all three test cases.

Revised: For more details, see Austin et al.

Revised: The “MBPP-aug +

Revised: CodeNet + single line” model achieved an overall synthesis accuracy of 54%, compared to the 62% accuracy of the original few-shot model in Austin et al.

Revised: This indicates that the scratchpad execution training does not completely disrup the model’s ability to perform other few-shot tasks.

Revised: B L ONG A DDITION A BLATION S TUDY

Revised: In our long addition experiments in Section we compared a model that was trained to perform “direct execution” (the baseline) vs a model trained to use a scratchpad.

Revised: Since the model trained to use the scratchpad gets an additional signal from all the intermediate steps shown, we also study what happens if the scratchpad model is subsequently trained to perform direct execution (i.e., directly output the target without using the scratchpad).

Revised: The result is shown in Figure 7 where we followed the same training procedure as for the original direct execution baseline and scratchpad models.

Revised: We see no signiﬁcant beneﬁts from doing any intermediate training using a scratchpad.

Revised: This indicates that the extra training-time information seen by the scratchpad model does not seem solely responsible for the scratchpad model’s improved performance.

Revised: C E XAMPLE FEW - SHOT PROMPT FOR SYNTHETIC P YTHON EXPERIMENTS

Original: Below is an example of a prompt for few-shot synthetic Python synthesis problems:
Revised: Below is an example of a prompt for few-shot synthetic Python synthesis problems:

Original: Consider the following Python function: def f(v0): v0 +=v4 = 2 while v4 > 0: v4 -= 1 v0 *=return v output = f(6)
Revised: Consider the following Python function: def f(v0): v0 +=v4 = 2 while v4 > 0: v4 -= 1 v0 *=return v output = f(6)

Original: What is the execution trace?
Revised: What is the execution trace?

Original: [BEGIN] state: {} line: def f(v0): state: {"f": "<callable_object f>"} line: output = f(6) state: {"v0": 6} line: v0 += 0 state: {"v0": 6} line: v4 = 2 state: {"v0": 6, "v4": 2} line: while v4 > 0: state: {"v0": 6, "v4": 2} line: v4 -= 1 state: {"v0": 6, "v4": 1} line: v0 *= 2 state: {"v0": 12, "v4": 1} line: while v4 > 0: state: {"v0": 12, "v4": 1} line: v4 -= 1 state: {"v0": 12, "v4": 0} line: v0 *= 2 state: {"v0": 24, "v4": 0} line: while v4 > 0: state: {"v0": 24, "v4": 0} line: return v0 state: {"f": "<callable_object f>", "output": 24} [DONE]
Revised: [BEGIN] state: {} line: def f ): state callable_object line: output = f(6) state: {"v0": 6} line: v0 += 0 state: {"v0": 6} line: v4 = 2 state: {"v0": 6, "v4": 2} line: while v4 > 0: state: {"v0": 6, "v4": 2} line: v4 -= 1 state: {"v0": 6, "v4": 1} line: v0 *= 2 state: {"v0": 12, "v4": 1} line: while v4 > 0: state: {"v0": 12, "v4": 1} line: v4 -= 1 state: {"v0": 12, "v4": 0} line: v0 *= 2 state: {"v0": 24, "v4": 0} line: while v4 > 0: state: {"v0": 24, "v4": 0} line: return v0 state: {"f": "<callable_object f>", "output": 24} [DONE]

Original: Consider the following Python function: def f(v0): v0 -=v0 += 2 v0 -=return v output = f(4)
Revised: Consider the following Python function: def f(v0): v0 -=v0 += 2 v0 -=return v output = f(4)

Original: What is the execution trace?
Revised: What is the execution trace?

Original: [BEGIN] state: {} line: def f(v0): state: {"f": "<callable_object f>"} line: output = f(4) state: {"v0": 4} line: v state: {" } line: state 6}line } line: return v0 state output
Revised: [BEGIN] state: {} line: def f(v0): state: {"f": "<callable_object f>"} line: output = f(4) state: {"v0": 4} line: v0 -= 0 state: {"v0": 4} line: v0 += 2 state: {"v0": 6} line: v0 -= 0 state: {"v0": 6} line: return v0 state: {"f": "<callable_object f>", "output": 6} [DONE]

Original: } [DONE] Consider the following Python function: def f(v0): v0 -=v8 = 2 while v8 > 0:v8 -= 1 v0 *=return voutput = f(4) What is the execution trace? [BEGIN] state: {} line: def f(v0): state: {"f": "<callable_object f>"} line: output = f(4) state: {"v0": 4} line: v0 -= 0 state: {"v0": 4} line: v8 = 2 state: {"v0": 4, "v8": 2} line: while v8 > 0: state: {"v0": 4, "v8": 2} line: v8 -= 1 state: {"v0": 4, "v8": 1} line: v0 *= 1 state: {"v0": 4, "v8": 1} line: while v8 > 0: state: {"v0": 4, "v8": 1} line: v8 -= 1 state: {"v0": 4, "v8": 0} line: v0 *= 1 state: {"v0": 4, "v8": 0} line: while v8 > 0: state: {"v0": 4, "v8": 0} line: return v0 state: {"f": "<callable_object f>", "output": 4} [DONE]
Revised: Consider the following Python function: def f(v0): v0 -=v8 = 2 while v8 > 0: v state: {"v0": 4, "v8": 2} line: v8 -= 1 state: {"v0": 4, "v8": 1} line: v0 *= 1 state: {"v0": 4, "v8": 1} line: while v8 > 0: state: {"v0": 4, "v8": 1} line: v8 -= 1 state: {"v0": 4, "v8": 0} line: v0 *= 1 state: {"v0": 4, "v8": 0} line: while v8 > 0: state: {"v0": 4, "v8": 0} line: return v0 state: {"f": "<callable_object f>", "output": 4} [DONE]

Original: Consider the following Python function: def f(v0): v0 +=v0 *= 2 v0 *=return v output = f(8)
Revised: Consider the following Python function: def f(v0): v0 +=v0 *= 2 v0 *=return v output = f(8)

Original: What is the execution trace?
Revised: What is the execution trace?

Original: [BEGIN]
Revised: [BEGIN]

Original: B I NCORRECT OUTPUT FOR SYNTHETIC P YTHON PROBLEMS
Revised: D I NCORRECT OUTPUT FOR SYNTHETIC P YTHON PROBLEMS

Original: Below is an example of the failure of the few-shot synthetic experiments discussed in Section 5.1.
Revised: Below is an example of the failure of the few-shot synthetic experiments discussed in Section 5.1.

Original: Given a problem such as:
Revised: Given a problem such as:

Original: Consider the following Python function: def f(v0): v0 +=v0 *= 2 v0 *=return v output = f(8)
Revised: Consider the following Python function: def f(v0): v0 +=v0 *= 2 v0 *=return v output = f(8)

Original: What is the execution trace?
Revised: What is the execution trace?

Original: The model produces the following trace
Revised: The model produces the following trace: state

Original: [Table]
Revised: [Equation]

Original: [Table]
Revised: [Equation] line [Equation] def f(

Original: [Table]
Revised: [Equation] ): state [Equation] v

Original: [Table]
Revised: [Equation]

Revised: at 0x00BFC1C

Original: [Table]
Revised: [Equation] [Equation] line [Equation]

Original: [Table]
Revised: v [Equation] =state [Equation]

Original: [Table]
Revised: v [Equation] [Equation] line: v0 *=state: {

Original: [Table]
Revised: [Equation] v [Equation] : 20} line: v0 *=state: {"v0": 20} line: return v state: {"f": "

Original: [Table]
Revised: [Equation]

Original:  Note that this trace is exactly correct, except for the last line: where the model predicted "v0": 20 , the correct output is "output": 20 .
Revised: function f at 0x00BFC1C0>", "v0": 20} Note that this trace is exactly correct, except for the last line: where the model predicted "v0": 20 , the correct output is "output": 20 .

Original: Because this type of error consistently occurs in the few-shot synthetic Python experiments, we modiﬁed the evaluation script slightly to consider this output to be correct. 
Revised: Because this type of error consistently occurs in the few-shot synthetic Python experiments, we modiﬁed the evaluation script slightly to consider this output to be correct. [Figure] [Figure]