Keywords: program synthesis, transformers, language models, pre-training, program induction
TL;DR: We train very large pre-trained language models to execute algorithms and Python programs by predicting the intermediate states line-by-line.
Abstract: Large pre-trained language models perform remarkably well on tasks that can be done "in one pass", such as generating realistic text or synthesizing computer programs. However, they struggle with tasks that require unbounded multi-step computation, such as adding integers or executing programs. Surprisingly, we find that these same models are able to perform complex multi-step computations - even in the few-shot regime - when asked to perform the operation "step by step", showing the results of intermediate computations. In particular, we train Transformers to perform multi-step computations by asking them to emit intermediate computation steps into a "scratchpad". We hypothesize that by providing supervision on the intermediate computation steps, the model gains additional learning signal on how to systematically generalize from small computations to larger ones. 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, even when we care only about the final result. Even though the model is required to predict many more tokens, it is still better at predicting the final results, because the individual prediction steps are easier. We believe that this result provides an early indication of the potential power of intermediate computation within language models.