Source-linked AI summary

Show Your Work: Scratchpads for Intermediate Computation with Language Models

Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, Augustus Odena

arXiv:2112.00114v1cs.LGcs.NE

TL;DR

Language models struggle with multi-step algorithmic computation despite strong one-pass generation abilities. The paper trains Transformers to emit intermediate computation steps in scratchpads, finding improved performance across long addition, polynomial evaluation, and Python code execution. The authors conclude that scratchpads are a step toward combining language-model knowledge with reasoning about program semantics.

  • Problem

    Language models struggle with precise multi-step computations such as addition and program execution, even when they can generate code for programming tasks.

  • Method

    The paper encodes intermediate algorithmic steps as text and trains Transformers to emit them into a scratchpad before producing the final answer.

  • Results

    Scratchpads improve performance across long addition, polynomial evaluation, and Python code execution, including addition generalization to larger instances.

  • Takeaways & Limitations

    Scratchpads may help combine large language models’ knowledge-compression capabilities with reasoning about program semantics.

  • Takeaways & Limitations

    Experiments are limited to scratchpad texts fitting within a 512-token generation window, while many problems require much longer generations.

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper argues that one-pass Transformer computation limits algorithmic reasoning and proposes supervised scratchpads containing intermediate steps. Across addition, polynomial evaluation, and program execution, scratchpads improve multi-step task performance.

  • Motivation: Transformers struggle with precise, unbounded algorithmic computations despite generating code for simple programming problems.The introduction cites failures on multi-digit addition, Python execution, and programs with loops.
  • Motivation: One-pass computation prevents models from adapting computation time to problem difficulty before producing an answer.The proposed response changes task design rather than the model or training procedure.
  • Approach: A scratchpad lets models emit arbitrary intermediate tokens before the final answer, trained as text with standard supervised learning.For addition, the intermediate tokens encode the steps of the standard long-addition algorithm.
  • Results: Scratchpads improve long-addition performance and out-of-distribution generalization to larger problem instances.The addition experiments are conducted in the fine-tuning regime.
  • Results: Scratchpads help Transformers evaluate polynomials in both few-shot and fine-tuning regimes.The paper presents polynomial evaluation as a higher-level task than addition.
  • Results: Predicting full program traces line by line with local-variable annotations dramatically improves execution prediction.The paper presents this as its most general application of scratchpads.

2 METHOD

The method encodes algorithmic intermediate states as text and trains sequence models to emit them in a scratchpad before the desired output. This exposes computation steps while retaining standard Transformer-based sequence prediction.

  • Task formulation: Algorithmic tasks are framed as learning functions represented by short programs, including addition and polynomial evaluation.The paper distinguishes algorithm induction from learning to execute, where source code is evaluated on an input.
  • Scratchpad construction: The method encodes an algorithm’s intermediate steps as text and trains the model to emit them into a scratchpad buffer.For long addition, the scratchpad explicitly writes the grade-school addition steps.
  • Scratchpad construction: Learning-to-execute examples prepend source code to the input, scratchpad, and desired output.This applies the same intermediate-step representation to program execution.
  • Training and inference: Training uses standard likelihood-based supervision, while testing requires predicting the target from the input alone.The paper uses decoder-only Transformers, though other sequence models could also be used.
  • Motivation: Scratchpads can provide adaptive computation time and persistent intermediate state through the generated context.The authors also propose that emitting concrete states may reduce propagation and compounding of small errors.
  • Models: Addition experiments train pre-trained dense decoder-only Transformers ranging from 2 million to 137 billion parameters.The models were pretrained on web documents and dialogue data.

3 ADDITION

For integer addition, the scratchpad target contains the intermediate steps of long addition, enabling evaluation both within the 1–8 digit training range and beyond it.

  • Setup: The addition scratchpad records the intermediate steps of the long-addition algorithm.Models are trained on inputs containing 1–8 digits.
  • Setup: The evaluation covers in-distribution problems up to 8 digits and out-of-distribution problems with 9- and 10-digit inputs.The baseline omits intermediate scratchpad steps.
  • Results: Beyond a critical model size, scratchpad-trained models solve addition while no-scratchpad models fail even at the largest tested scale.On 9–10 digit addition, scratchpad models improve consistently with model size, whereas baseline models completely fail.

4 POLYNOMIAL EVALUATION

The polynomial task evaluates low-degree integer-coefficient polynomials, using scratchpads that compute each term separately before summing the results.

  • Task setup: The dataset contains polynomials of degree at most three with integer coefficients and inputs constrained to [−10, 10].Outputs are restricted to [−1000, 1000], with 10,000 training and 2,000 test polynomials.
  • Implementation: Digit spacing is used so each digit maps to a separate token in the models’ representation.
  • Scratchpad computation: The scratchpad evaluates each polynomial term separately and then adds the terms.Figure 4 illustrates this decomposition of polynomial evaluation.

5 EXECUTING PYTHON PROGRAMS

The paper evaluates scratchpad tracing for executing synthetic and real Python programs, representing intermediate control-flow and variable states before producing outputs. Tracing improves synthetic execution and benefits from augmented and combined training data on real programs, although very-low-data MBPP performance remains poor.

  • Execution prediction via scratchpad tracing: Scratchpad tracing represents execution as alternating source-code lines and local-variable states, exposing both control flow and state changes.The trace is serialized with executed lines and JSON-like local-variable states before the final output.
  • 5.1 SCRATCHPAD BEATS DIRECT EXECUTION FOR SYNTHETIC PYTHON PROGRAMS: Synthetic programs show higher overall execution accuracy with scratchpad tracing than with direct execution in both few-shot and fine-tuned settings.Fine-tuning improves tracing more than direct execution on the synthetic benchmark.
  • 5.2.1 PERFORMANCE IS POOR IN THE VERY-LOW-DATA REGIME: On the original MBPP training setup, direct execution reaches 10% output accuracy versus 5% for scratchpad tracing, indicating poor performance in the very-low-data regime.The comparison uses 374 training tasks and evaluates a 212-task subset whose traces fit within the model’s generation window.
  • 5.2.2 SAMPLED PROGRAMS MAKE GOOD SCRATCHPAD TRAINING DATA: Model-generated program augmentation improves scratchpad performance while reducing direct-execution accuracy on MBPP.The augmentation procedure creates additional programs and outputs by executing sampled candidate programs, then traces them for scratchpad training.
  • 5.3 SCRATCHPAD TRAINING MAKES GOOD USE OF LARGE DATASETS: Combining CodeNet and single-line training data yields 26.6% task execution accuracy and 24.6% perfect traces across all three examples.Greedy decoding from the best model produces the exact correct trace for almost 42% of all traces.

6 RELATED WORK

The paper situates scratchpad-based algorithmic reasoning within work on semantic understanding, program execution, neural algorithm induction, and pretrained Transformers as computation engines.

  • Semantic understanding: Prior work examines whether language models learn semantic representations rather than relying only on surface-level statistical correlations.The paper connects this question to evidence that pretrained models construct approximate situation semantics.
  • Program execution: Learning-to-execute research studies recurrent and specialized architectures for executing and reasoning about arbitrary code.Related neural algorithm induction work proposes architectures with inductive biases suited to these tasks.
  • Pretrained computation: Pretrained Transformers have also been fine-tuned on non-language tasks, including simple algorithm induction problems, as universal computation engines.This work extends the broader connection between pretrained models and algorithmic computation.

7 LIMITATIONS AND FUTURE WORK

The paper identifies context-window length as a practical boundary for scratchpads and proposes learning scratchpad use without direct supervision as future work.

  • Context window size: Experiments are limited to problems whose scratchpad text fits within a 512-token generation window.Many problems require longer generations, so larger Transformer generation windows may be needed to realize the technique’s potential.
  • Learning without supervision: A proposed next step is to learn scratchpad use without direct supervision, potentially through reinforcement learning.The suggested reward favors correct answers while penalizing the number of scratchpad tokens used.
  • Transfer: The authors hope scratchpad use could transfer across algorithms, such as from long addition to polynomial evaluation.This transfer possibility is presented as a future expectation rather than an established result.

8 CONCLUSION

Experiments on long addition, polynomial evaluation, and Python execution show that scratchpads can improve language-model performance on algorithmic tasks. The authors frame this as progress toward models that combine language knowledge with program reasoning, while noting that the technique is not yet sufficient for that broader goal.

  • Conclusion: Scratchpads improve performance across long addition, polynomial evaluation, and Python code execution.The conclusion summarizes evidence across increasingly complex algorithmic tasks.
  • Implication: The technique may help combine large language models’ knowledge-compression capabilities with reasoning capabilities.The authors connect this direction to models that understand code as well as write it.
  • Applications: Potential applications include program synthesis, neural-guided program analysis, and interactive programming assistants.These applications require both natural-language interaction and reasoning about program semantics.
  • Scope: The authors characterize scratchpads as an important step rather than a complete solution to combining language and program reasoning.The conclusion explicitly limits the scope of the contribution.

A EFFECTS OF SCRATCHPAD EXECUTION TRAINING ON SYNTHESIS PERFORMANCE

This appendix evaluates whether fine-tuning on execution tracing affects few-shot program synthesis using the MBPP-aug + CodeNet + single line model.

  • Experimental setup: The study measures the effect of fine-tuning on tracing performance in a few-shot MBPP program-synthesis experiment.For each task, 80 candidate programs are sampled at T = 0.5, and a task is solved if any candidate satisfies the specification.

B LONG ADDITION ABLATION STUDY

The experiments compare direct execution with scratchpad-trained models later trained to output targets directly, alongside synthetic Python execution-trace examples. The traces represent execution through successive variable updates and loop iterations.

  • Long addition ablation study: The ablation compares a direct-execution baseline with a scratchpad model subsequently fine-tuned to perform direct execution.The comparison tests whether scratchpad performance comes solely from the additional intermediate-step training signal.
  • Synthetic Python execution traces: An execution-trace prompt asks the model to report program states as each function line executes.The examples include function calls, assignments, loop conditions, updates, and the returned output.
  • Synthetic Python execution traces: A second trace applies v0 *= 2 twice, changing the initial value 6 to the returned output 24.The trace records v4 decreasing from 2 to 0 between the two multiplications.
  • Synthetic Python execution traces: A shorter trace applies v0 -= 0, v0 += 2, and v0 -= 0, returning output 6 from input 4.The intermediate states show the value changing from 4 to 6 before return.

D INCORRECT OUTPUT FOR SYNTHETIC PYTHON PROBLEMS

The synthetic Python example contains a trace that is nearly correct but formats its final state differently from the expected output. The evaluation was adjusted to accept this recurring formatting error, while the ablation compares scratchpad and direct-execution training.

  • Incorrect output for synthetic Python problems: The model’s trace correctly computes v0 as 20 after applying the listed updates to input 8.The predicted intermediate values are 10, 20, and 20.
  • Incorrect output for synthetic Python problems: The only reported error is the final field name: the model predicts v0 instead of output.The passage states that the trace is otherwise exactly correct.
  • Incorrect output for synthetic Python problems: Because this formatting error recurs in few-shot synthetic Python experiments, the evaluation script counts the output as correct.This changes evaluation treatment rather than the model’s generated trace.
  • Long addition ablation study: The long-addition ablation compares baseline direct execution with scratchpad training followed by direct execution.The intermediate scratchpad-training stage is described as providing an additional signal from shown intermediate steps.
Loading 2112.00114v1…