PaperlayerView exact revision on arXiv

Source-linked AI summary

Teaching Arithmetic to Small Transformers

Nayoung Lee, Kartik Sreenivasan, Jason D. Lee, Kangwook Lee, Dimitris Papailiopoulos

arXiv:2307.03381v1cs.LG

TL;DR

Prior work leaves it challenging to disentangle which data and model factors elicit arithmetic in next-token predictors. This study trains small transformers from scratch and finds that changing arithmetic formats, including reversing outputs and adding intermediate steps, improves learning.

  • Problem

    The factors that elicit arithmetic capabilities in next-token predictors remain difficult to untangle because prior studies vary in data complexity and tasks.

    Source: §1 · p. 3

  • Method

    The study trains small transformers from random initialization on arithmetic with next-token prediction and ablates data formats, sampling, scale, and other training conditions.

    Source: §1 · p. 4

  • Results

    Reversing outputs and adding intermediate steps improve arithmetic accuracy and sample complexity, with sharp data-scale transitions in some settings.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Detailed instructive data and even output reversal improve arithmetic learning in the studied transformers.

    Source: §14 · p. 30

    The models did not predict beyond the number of digits on which they were trained.

    Source: §13 · p. 30

Abstract

from arXiv · show

Large language models like GPT-4 exhibit emergent capabilities across general-purpose tasks, such as basic arithmetic, when trained on extensive text data, even though these tasks are not explicitly encoded by the unsupervised, next-token prediction objective. This study investigates how small transformers, trained from random initialization, can efficiently learn arithmetic operations such as addition, multiplication, and elementary functions like square root, using the next-token prediction objective. We first demonstrate that conventional training data is not the most effective for arithmetic learning, and simple formatting changes can significantly improve accuracy. This leads to sharp phase transitions as a function of training data scale, which, in some cases, can be explained through connections to low-rank matrix completion. Building on prior work, we then train on chain-of-thought style data that includes intermediate step results. Even in the complete absence of pretraining, this approach significantly and simultaneously improves accuracy, sample complexity, and convergence speed. We also study the interplay between arithmetic and text data during training and examine the effects of few-shot prompting, pretraining, and model scale. Additionally, we discuss length generalization challenges. Our work highlights the importance of high-quality, instructive data that considers the particular characteristics of the next-word prediction objective for rapidly eliciting arithmetic capabilities.

1 Introduction

The study isolates what helps small, randomly initialized transformers learn arithmetic under next-token prediction, finding that data format and sampling shape learning efficiency and accuracy.

3 Preliminaries and Experimental Setup

  • Controlled experiments train small transformers from scratch on arithmetic tasks, using held-out operand pairs and data formats that progressively add information.NanoGPT has approximately 10.6 million parameters and uses the conventional next-token prediction objective.Source: §1 · p. 3, §3 · pp. 5–6

1 Introduction

  • Reversed outputs and balanced sampling improve arithmetic learning; accuracy can shift sharply with training-data scale, sometimes consistent with low-rank matrix completion.The figure contrasts plain, reverse, simplified-scratchpad, and detailed-scratchpad formats; plain formatting never reaches 100% accuracy, while the other formats require fewer samples as detail increases.Source: §1 · p. 3, Figure 1 · p. 4
  • Training with chain-of-thought data significantly improves accuracy and sample complexity, even without language pretraining.The authors conjecture that decomposing a compositional task into components creates a higher-dimensional but easier-to-learn function map.Source: §1 · p. 3
  • Learning arithmetic operations together improves each task, while one-shot prompting boosts accuracy and additional examples bring no significant improvement.Source: §1 · p. 3
  • Pretraining supports reasonable performance on some basic arithmetic tasks with few fine-tuning samples, but reverse formatting can hurt pretrained models trained on standard formats; scale helps but is not necessary.Source: §1 · p. 4
  • Models achieve high accuracy at trained digit lengths but struggle on omitted lengths and beyond the trained range.The authors characterize the learned arithmetic as a mapping constrained to trained lengths, rather than a flexible algorithm.Source: §1 · p. 4
  • The study’s distinction from prior instructive-data work is its focus on randomly initialized models, extensive ablations, and theoretical explanations of observed phenomena.Source: §1 · p. 4

2 Related Works

Different data formatting methods for addition

For addition, output order and sample composition shape how efficiently small transformers learn: reversing results makes the digitwise task simpler, while balancing digits and carry patterns improves accuracy. Scratchpad formats add intermediate information and improve sample efficiency.

Different data formatting methods for addition

4.2 Reversing the Output

  • Plain-format accuracy plateaus at slightly over 85% with 10,000 samples, while reverse format reaches perfect 3-digit addition after a sharp transition between 1,000 and 4,000 samples.Source: Figure 4 · p. 7, §4.2 · p. 8

4.1 Training on Conventional Data

  • For additions up to three digits, MSB-first output requires access to all operand digits, implying a more complex global algorithm than digitwise addition.Source: §4 · p. 7, §4.1 · p. 8

4.2 Reversing the Output

  • Reversing the output order lets the model generate the LSB first, using each operand’s corresponding digit and the carry from the previous digit.Source: §4.2 · p. 8

5 Connection to Low-Rank Matrix Completion

The paper connects addition learning to rank-2 matrix completion, while showing that transformers generalize beyond that framework. It then finds that stepwise training data can improve sample efficiency, with the design of intermediate steps materially affecting performance.

5.1 Addition Tables are Rank-2 Matrices

  • NanoGPT and low-rank matrix completion show a similar sharp transition near 1,500 samples on 2-digit addition, linking the model’s rapid learning to matrix completion.Addition can be represented as a rank-2 matrix-completion problem, with a theoretical transition at O(n) observed entries.Source: §5.1 · p. 9, Figure 5 · p. 9

5.2 NanoGPT Generalizes better than Matrix Completion solutions

  • NanoGPT retains near-perfect addition accuracy when many operands are excluded, unlike matrix completion, which cannot recover even a single empty row or column.The authors report that excluding numbers sometimes improves performance and conjecture this may reflect regularization.Source: §5.2 · pp. 9–10, Table 1 · p. 10

6.1 Training on Chain-of-Thought Data

  • Detailed scratchpad data reaches perfect addition accuracy with 1,000 samples, while simplified scratchpad needs 2,000; the authors attribute the gains to decomposing the task into easier components.The reverse format requires more than twice the simplified scratchpad sample count to reach 100% accuracy.Source: §6.1 · p. 11

6.2 The Importance of Intermediate Step Design: Subtraction

  • For subtraction, the scratchpad version that compares operands performs considerably worse than the version using digit-wise subtraction; most Version 2 errors misidentify the larger operand.The authors use Version 1 in detailed scratchpad experiments unless otherwise specified.Source: §6.2 · p. 12, Figure 8 · p. 13
  • With noisy intermediate labels, addition sample efficiency declines, but sufficient training still yields full accuracy.The study replaces intermediate steps, not the final correct answer.Source: §6.2 · p. 13, Figure 9 · p. 13

7 Extending to Longer Digit Addition

For longer additions, reverse and scratchpad formats maintain more consistent sample requirements than plain formatting. During fine-tuning, detailed scratchpads can preserve earlier digit-length skills, while format consistency affects learning speed.

7.1 Training from Random Initialization

  • As digit length increases, plain-format addition needs more training examples, while reverse and scratchpad formats have more consistent sample requirements.The performance gap between modified and plain formats becomes more pronounced at higher digit lengths.Source: §7.1 · p. 15, §7 · p. 15, Figure 10 · p. 15

7.2 Fine-Tuning from Pretrained Models

  • Fine-tuning on 4-digit additions in plain format immediately reduces accuracy on 1–3-digit additions, whereas detailed scratchpad fine-tuning learns 4-digit addition while maintaining earlier performance.The detailed scratchpad model initially dips on 3-digit additions, then recovers as training proceeds.Source: §7.2 · p. 16, Figure 11 · p. 16
  • For pretrained k-digit models, learning k + 1-digit addition takes 1,000–5,000 examples in reverse and scratchpad formats; plain format requires increasingly more examples as digit length grows.For reverse format specifically, processing the i-th operand digit and carry from the previous position remains consistent regardless of the number of digits.Source: Figure 12 · p. 17, §7.2 · p. 17

7.3 Impact of Formats on Fine-Tuning

  • Fine-tuning is faster when the pretrained and fine-tuning formats match; detailed scratchpad models are especially difficult to adapt to other formats.The authors attribute the latter difficulty to unlearning patterns specific to the verbose scratchpad format.Source: §7.3 · p. 18, Figure 13 · p. 18

8 Teaching Arithmetic Operations Beyond Addition

The study extends its formatting experiments beyond addition to subtraction, multiplication, sine, and square root. Its arithmetic coverage remains bounded, and some proposed intermediate steps for unary operations involve computations that may be no simpler than the target task.

8 Teaching Arithmetic Operations Beyond Addition

  • The experiments cover subtraction and multiplication on positive integers, plus sine and square root on restricted decimal ranges and precision.Subtraction uses up to 3 digits, multiplication up to 2-digit operands, and sine and square root are evaluated at up to four digits of precision.Source: §8 · p. 18, §8.1 · p. 19

8.1 Extended Arithmetic Operations

  • Scratchpads decompose subtraction and multiplication into digit-level steps, but sine and square root use Taylor-series and Newton-method iterations that include exponentiation or division.The paper notes those unary-operation intermediate computations may not be inherently simpler than the original operation; its scope also excludes some arithmetic operations and broader floating-point arithmetic.Source: §8 · p. 18, §8.1 · p. 19

Detailed scratchpad formatting for different arithmetic tasks

Detailed scratchpads break some arithmetic tasks into intermediate steps, but their benefits vary by operation. Joint training on five arithmetic tasks improves zero-shot multiplication and square-root performance, while slightly reducing subtraction performance.

Detailed scratchpad formatting for different arithmetic tasks

Sine

  • Detailed scratchpads significantly improve few-sample performance over plain or reverse formats for subtraction and multiplication.Reverse formatting is not particularly effective for multiplication.Source: Sine · p. 20, Figure 15 · p. 21
  • Sine and square-root scratchpads yield less improvement because their intermediate steps are more intricate than those for addition, subtraction, and multiplication.The described scratchpads use Taylor series for sine and Newton’s method for square root.Source: Sine · p. 20, Detailed scratchpad formatting for… · p. 20

8.2 Jointly Training on All Five Arithmetic Tasks

  • The joint dataset randomizes the sequence of addition, subtraction, multiplication, sine, and square-root tasks.It includes 10,000 examples each for addition, subtraction, sine, and square root, and 3,000 for multiplication.Source: §8.2 · p. 20
  • For sine and square root, scratchpad formatting provides limited improvement, unlike the stronger gains reported for subtraction and multiplication.The figure caption attributes the difference to the complexity of the detailed intermediate steps.Source: Sine · p. 20, Figure 15 · p. 21
  • Joint-training evaluation compares zero-shot and few-shot results with models trained separately on the same task datasets.Few-shot prompts may use exemplars from any of the five tasks or text prompts.Source: §8.2 · p. 21
  • Joint training significantly improves zero-shot multiplication and square-root performance, but slightly reduces subtraction performance.Few-shot prompting generally improves performance, with similar results from task-specific and unrelated-task examples.Source: §8.2 · p. 21

9 Mixing Shakespeare with Arithmetic Data

The experiments mix Shakespeare text with addition data and test whether prompts help models select the desired output context. Few-shot addition prompts generally improve performance, while text prompts help in mixed-data settings but can hurt when training uses addition alone.

9 Mixing Shakespeare with Arithmetic Data

  • The mixed-data experiment combines Shakespeare text with plain and detailed-scratchpad addition examples, varying the number of each example type.The dataset contains 1,115,394 Shakespeare tokens, 10,000 plain addition examples, and 3,000 detailed scratchpad examples; addition samples are interspersed among dialogue chunks.Source: §9 · p. 22
  • Few-shot evaluations average 1-, 2-, and 3-shot prompts over fixed exemplar sets, using format-matched examples for plain and scratchpad addition.The study also tests text prompts before addition queries.Source: §9 · p. 22, Figure 16 · p. 22
  • Few-shot prompting improves addition performance, with plain addition approaching near-perfect performance at 40,000 training samples.Text prompts also maintain high plain-addition performance when many addition examples are included in the mixed dataset.Source: §9 · p. 22, Figure 17 · p. 23
  • In mixed training, more addition examples increase how often addition follows Shakespeare text, which the authors link to fewer inconsistencies under text prompts.This is the authors’ proposed explanation for why text prompts can remain effective when the model has seen both text and addition.Source: §9 · pp. 22–23
  • With addition-only training, text prompts degrade performance, while correct or noisy addition exemplars improve it.Even noisy exemplars that retain the A + B = C format perform nearly identically to few-shot prompting.Source: §9 · p. 23, Figure 18 · p. 23

10 Fine-tuning, Scaling, and Pretraining in Larger Models

Larger models and pretrained initialization improve arithmetic performance in the tested settings, while formatting and tokenization choices affect results. Detailed scratchpads can improve sample efficiency but require substantially more tokens, so their training and inference costs matter.

10 Fine-tuning, Scaling, and Pretraining in Larger Models

  • 85-million-parameter GPT-2 outperforms NanoGPT on plain and detailed-scratchpad addition evaluations.Figure 19 compares the models trained from scratch jointly on text and addition.Source: Figure 19 · p. 24, §10 · p. 24
  • Fine-tuning pretrained GPT-2 performs better than training from scratch, and GPT-2 performs highly with character-level tokenization or Tiktoken when spaces separate digits.The authors report that spaces support consistent number tokenization with Tiktoken.Source: §10 · p. 24, Figure 20 · p. 25
  • Pretrained GPT-3 improves addition and subtraction performance over randomly initialized NanoGPT or GPT-2 trained with 1,000 samples, but reverse and simplified-scratchpad formats help addition while hurting subtraction.The authors suggest adapting from plain-format arithmetic may contribute to this format-dependent difference.Source: §10 · p. 25, Table 5 · p. 25, Table 6 · p. 25
  • For sine and square root, 1,000 training samples are insufficient for exact answers at epsilon = 0.The authors call for larger training datasets to better assess these tasks.Source: §10 · p. 26, Table 7 · p. 26
  • Few-shot prompting improves all three GPT-3 models, but becomes less effective after fine-tuning; their zero-shot performance is reported as quite poor.The authors therefore focus their comparisons primarily on zero-shot performance.Source: §10 · p. 26

11 Token Efficiency Across Data Formats

  • Detailed scratchpad training is more sample-efficient but uses more tokens per sample, whereas reverse formatting is the most token-efficient approach.The authors note that token demand affects both training and inference costs; their cost analysis counts tokens per sample rather than unique tokens across samples.Source: §11 · pp. 26–27, Figure 21 · p. 26, Table 8 · p. 27

12 Length Generalization

Across the reported addition tests, models do not generalize to digit lengths absent from training, even when they handle unseen examples at trained lengths. The paper limits its conclusions to elementary arithmetic and small-scale GPT-type models, leaving generalization to practical LLM scales uncertain.

12 Length Generalization

  • NanoGPT trained on 1- and 3-digit reverse-formatted additions cannot perform 2- or 4-digit additions, and models trained up to 7 digits cannot generalize to 8-digit additions.The figure also reports failure to generalize from 10-digit to 11-digit binary addition.Source: §12 · p. 27, Figure 22 · p. 27

Examples for length generalization prompts

13 Limitations

  • The experiments observe no prediction beyond the number of digits used in training, consistent with the reported difficulty of length generalization.The discussion notes prior reports of length generalization under some conditions, including models with more than 10^8 parameters in Nye et al. (2021).Source: §13 · p. 30
  • The experiments use GPT-type decoder-only models and primarily character-level tokenization; whether the findings generalize to practical LLM scales remains uncertain.The authors describe scaling and BPE results as preliminary.Source: §13 · p. 30
  • The study focuses on simple arithmetic to isolate factors in capability emergence, so it does not establish findings for more complicated abilities.The authors distinguish this focus from existing literature demonstrating complicated abilities in practice.Source: §13 · p. 30

14 Conclusion

  • The conclusion identifies higher token costs for detailed scratchpads and difficult length generalization as important constraints on arithmetic training.It also leaves open how to curate data so models learn an algorithm rather than an approximate function map.Source: §14 · p. 30

Appendix

The appendix proves why least-significant-digit-first addition needs only local digits and a carry, and examines formatting and a low-rank matrix-completion connection. Zero-padding and symbol wrapping improve performance, while the matrix-completion experiment exhibits a sharp transition as more entries are revealed.

A Proofs

  • For addition output in decreasing significance, an algorithm must access all input digits from the first output digit; increasing-significance output needs only the current digits and prior carry.The local procedure computes each output digit and requires only the two input digits at that position and the carry from the previous position.Source: A Proofs · p. 35

B.1 Zero-Padding and Symbol Wrapping

  • Zero-padding and ‘$’-wrapping significantly improve plain-format addition, but even with the delimiter it remains below reverse formatting in accuracy and sample efficiency.The appendix adopts ‘$’-wrapping for its modified reverse format while retaining unmodified plain format as the baseline.Source: §B.1 · p. 36, Table 9 · p. 36, Figure 24 · p. 36

B.2 Low-Rank Matrix Completion

  • The rank-2 addition-matrix completion algorithm fills an unknown entry when three entries of a 2 × 2 submatrix are known, enforcing a zero determinant.It searches iteratively for such submatrices and reports success or failure in completing the matrix.Source: §B.2 · p. 37
  • A sharp phase transition in completion success appears when O(n) entries are revealed in the n × n addition matrix.The experiment varies the number of randomly revealed entries and averages success across 100 trials for each setting.Source: §B.2 · p. 37

B.3 Prompting with Text

The appendix tests whether pure-text few-shot prompts affect addition performance, expecting out-of-context text exemplars to hurt accuracy. It uses five types of text prompts.

B.3 Prompting with Text

  • The experiment tests five pure-text exemplar types, including short or long Shakespeare text, text outside the dataset, and text containing numbers.The expectation is that out-of-context text prompts will hurt accuracy, unlike few-shot addition examples.Source: §B.3 · p. 38
  • The five prompt types vary in length, whether they appear in the Shakespeare dataset, and whether they include numbers.Prompt 1 is short and outside the dataset; Prompt 2 is short and within it; Prompt 3 is longer and within it; Prompt 4 includes numbers; Prompt 5 is long and outside it.Source: §B.3 · p. 38
  • The text-prompt examples are detailed in Figure 25.Source: §B.3 · p. 38

Text prompts for few-shot experiments

Text-prompt effects on addition vary with the exemplars, while relevant few-shot examples are beneficial and text prompts often hinder performance. The appendix also examines sine accuracy around zero and ±1.

Text prompts for few-shot experiments

  • Figure 26 compares five prompt types; its solid lines show mean performance across five exemplars and shaded regions show standard deviation.Prompt types range from short or long Shakespeare text to text outside that dataset and prompts with numbers.Source: Figure 26 · p. 39
  • Text prompts followed by addition queries generally harm performance, whereas relevant 1-, 2-, or 3-shot exemplars are beneficial.The text-prompt comparison averages five few-shot exemplars and 25 text exemplars.Source: Text prompts for few-shot experiments · p. 40, Figure 27 · p. 40
  • Long text prompts, especially Prompt 5, typically produce a larger performance decline, and addition accuracy varies notably across text-prompt choices.Prompt 5 is a long prompt not present in the Shakespeare dataset.Source: Figure 26 · p. 39, Text prompts for few-shot experiments · p. 40
  • Figure 27 compares NanoGPT and GPT-2 trained on Shakespeare with varying plain-addition or scratchpad sample counts, evaluated on plain and detailed-scratchpad test prompts.The figure distinguishes text prompts from few-shot exemplars and reports variability with shaded standard-deviation regions.Source: Figure 27 · p. 40

B.4 Analyzing the results on Sine/Sqrt

  • Across data formats, sine performs especially well near sin(x) = 0, where the authors conjecture sin(x) ≈ x makes the function easier to learn.Accuracy also improves near ±1, potentially for similar reasons.Source: §B.4 · p. 41, Figure 28 · p. 41
  • Figure 28 analyzes sine and square-root errors across tolerance thresholds; the reported sine pattern includes improved accuracy near ±1.The passage attributes the near-zero sine behavior to the approximation sin(x) ≈ x and conjectures a similar reason near ±1.Source: Figure 28 · p. 41, §B.4 · p. 41

C Experimental Setup

The experiments use custom arithmetic datasets and Shakespeare text, with addition data balanced by digit length and carry count. They compare plain, reverse, and scratchpad formats, including task-specific reasoning steps.

C Experimental Setup

C.1 Dataset

  • Arithmetic datasets are constructed by operation, and test operand pairs are sampled so they do not appear in training.The text dataset is Shakespeare, comprising 40,000 dialogue lines and 1,115,394 characters.Source: §C.1 · p. 41

C.1.1 Data Balancing

  • Balancing addition examples by digit length and carry count gives a clear advantage over the sampling alternatives.Table 10 compares random, balanced-digit, and balanced-carry sampling for zero-padded addition.Source: §C.1.1 · p. 42, Table 10 · p. 42
  • The 10,000-example addition set contains all 100 one-digit cases, 900 two-digit cases, and 9,000 three-digit cases, with equal counts for 0–3 carries.The test set has 10,000 non-overlapping examples; larger and smaller training sets are assembled from the balanced set while maintaining carry balance.Source: §C.1.1 · p. 42

C.1.2 Data Formatting

  • Formatting conditions include plain output, reversed output digits, and simplified or detailed scratchpads that provide algorithmic reasoning steps.For addition scratchpads, digit-wise sums and carries are given from least to most significant digit; zero-padding is used only in the sampling comparison.Source: §C.1.2 · pp. 42–43
  • Detailed scratchpads use task-specific intermediate steps, including digit-wise subtraction, shifted partial products for multiplication, sine’s Taylor series, and Newton iterations for square root.The sine steps involve exponentiation, while the square-root steps involve division, which the text notes may be as complex as the target operation.Source: §C.1.2 · p. 43

C.2 Model

The appendix specifies the model variants, training configurations, and arithmetic prompt examples used in the experiments. In the tandem-training curves, addition accuracy continues to improve with more iterations, while Shakespeare perplexity tends to increase.

C.2 Model

  • Experiments primarily use decoder-only NanoGPT with character-level tokens, randomly initialized learnable absolute positions, and temperature 0.8; “carry-on” refers to borrow.NanoGPT is a scaled-down GPT-2 variant with half the self-attention layers, heads, and embedding dimension.Source: §C.2 · p. 43
  • NanoGPT context length is 256 for plain and reverse formatting and 1024 for scratchpads; GPT-2 uses 1024 throughout.Train examples are approximately 13–15 characters in plain/reverse formats, 64 in simplified scratchpad, and 281 in detailed scratchpad.Source: §C.2 · p. 44
  • Pretrained GPT-3 fine-tuning experiments use the Ada, Curie, and Davinci models.Source: §C.2 · p. 44

C.3 Hyperparameter Configurations

  • Flash attention is used for memory efficiency and speed; bfloat16 is typical, with float16 on Nvidia 2080 GPUs and no significant performance differences observed.GPT-2 uses batch size 8 with gradient accumulation to reach an effective batch size of 64.Source: §C.3 · p. 44
  • Addition accuracy continued to improve with iterations, while Shakespeare perplexity tended to rise, indicating some overfitting despite reasonable prompted text.The curves compare NanoGPT and GPT-2 trained with varying plain addition and detailed scratchpad samples alongside Shakespeare data.Source: Figure 30 · p. 45

D Prompt Examples

  • The prompt appendix provides examples of plain, reverse, simplified-scratchpad, and detailed-scratchpad formats across five arithmetic operations.Source: D Prompt Examples · p. 46

D.1 Addition

The noisy simplified-scratchpad appendix illustrates variants of digit-sum and carry information, while the GPT-3 fine-tuning example specifies prompt-completion formatting.

D.6 Noisy Simple Scratchpad

  • The appendix supplies one example for each noise case used in the simplified-scratchpad experiments.Source: §D.6 · p. 49

Noisy Simple Scratchpad Examples

D.7 Example data for GPT-3 fine-tuning

  • The GPT-3 fine-tuning example separates prompt and completion, spaces digits in plain and reverse formats for consistent tokenization, and uses “###” as the stop sequence.Source: §D.7 · p. 50
Loading 2307.03381v1…