[Paper Review] Show Your Work: Scratchpads for Intermediate Computation with Language Models

Maxwell Nye, Anders Andreassen+10 more|arXiv (Cornell University)|Nov 30, 2021

Topic ModelingComputer Science130 citations

TL;DR

The paper shows that adding a scratchpad of intermediate computation steps enables transformers to perform complex multi-step computations (like long addition, polynomial evaluation, and Python program execution) far more effectively than direct prediction, including better generalization to larger inputs.

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". 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.

Motivation & Objective

• Motivate and enable large language models to perform unbounded multi-step computations by emitting intermediate steps (scratchpads).
• Demonstrate that scratchpads improve performance on long addition, polynomial evaluation, and Python program execution tasks.
• Show that scratchpads offer adaptive computation time and traceable intermediate state without changing model architecture.
• Explore data-augmentation and multi-source tracing data to scale scratchpad benefits.

Proposed method

• Introduce scratchpad as an intermediate text buffer where the model emits step-by-step computation results.
• Train decoder-only Transformer models to predict both the input/output and the scratchpad content using supervised learning.
• Apply scratchpad to: (i) long integer addition with digit-wise steps, (ii) polynomial evaluation, (iii) execution of Python programs via tracing of execution traces.
• Compare scratchpad-based execution with direct execution prediction across few-shot and fine-tuning regimes.
• Use various model sizes (from 2M to 137B parameters) and standard training setups.

Experimental results

Research questions

• RQ1 Can a scratchpad enable Transformers to perform multi-step algorithmic tasks more accurately than direct single-pass prediction?
• RQ2 Does emitting intermediate computation traces improve out-of-distribution generalization in tasks like long addition and polynomial evaluation?
• RQ3 Can scratchpad tracing improve the accuracy of predicting Python program execution on given inputs?
• RQ4 How does scratchpad performance scale with model size and data regimes (few-shot vs fine-tuning)?

Key findings

• Scratchpads dramatically improve multi-step computation in Transformers across addition, polynomial evaluation, and program execution tasks.
• Scratchpad-trained models generalize better to larger input sizes than models trained without scratchpads.
• Scratchpad-based execution improves both few-shot and fine-tuning performance on polynomial evaluation.
• Tracing-based execution (predicting program traces) significantly enhances Python program execution predictions compared to direct execution, given sufficient training data.
• Data augmentation and multi-source tracing data (Single-line, CodeNet) further boost scratchpad performance, enabling higher trace accuracy and more tasks solved.