PaperlayerView exact revision on arXiv

Source-linked AI summary

Length Generalization in Arithmetic Transformers

Samy Jelassi, Stéphane d'Ascoli, Carles Domingo-Enrich, Yuhuai Wu, Yuanzhi Li, François Charton

arXiv:2306.15400v1cs.LG

TL;DR

Transformers struggle with integer arithmetic on numbers larger than those seen in training. This paper tests relative position embeddings and train set priming, finding that their length-generalization results differ across arithmetic tasks.

  • Problem

    Large integers are rare in text corpora, and their absence limits language models’ mathematical ability to extrapolate arithmetic to larger numbers.

    Source: §1 · p. 1

  • Method

    The paper evaluates relative position embeddings and train set priming for length generalization across four arithmetic tasks.

    Source: §1 · p. 2

  • Results

    Relative position embeddings support addition generalization up to 20 digits, while priming 5-digit × 3-digit multiplication with 50 of 5000 examples supports generalization to 35-digit operands.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results show that transformers can generalize arithmetic to longer operands, with the approach depending on the task.

    Source: §1 · p. 2

    Multiplication experiments fixed the second operand at 3 digits and selected 35-digit extrapolation arbitrarily, though the authors expect results to hold for reasonable lengths.

    Source: §1 · p. 3

Abstract

from arXiv · show

We examine how transformers cope with two challenges: learning basic integer arithmetic, and generalizing to longer sequences than seen during training. We find that relative position embeddings enable length generalization for simple tasks, such as addition: models trained on $5$-digit numbers can perform $15$-digit sums. However, this method fails for multiplication, and we propose train set priming: adding a few ($10$ to $50$) long sequences to the training set. We show that priming allows models trained on $5$-digit $\times$ $3$-digit multiplications to generalize to $35\times 3$ examples. We also show that models can be primed for different generalization lengths, and that the priming sample size scales as the logarithm of the training set size. Finally, we discuss potential applications of priming beyond arithmetic.

1 Introduction

The paper studies why transformers struggle to extrapolate arithmetic to larger numbers, a gap left by prior work focused on in-distribution evaluation. It finds relative position embeddings support length generalization for addition, while multiplication requires train set priming.

Motivation

  • Prior arithmetic-transformer studies largely used in-distribution evaluation, while existing out-of-distribution extrapolation results were disappointing.Large integers are rare in text corpora, motivating models that extrapolate arithmetic learned from smaller numbers.Source: §1 · p. 1

Contributions

  • For 5 × 3 multiplication, adding 50 target-distribution examples out of 5,000 enables generalization to 35-digit operands.The paper reports that this uses far fewer long examples than fine-tuning.Source: §1 · p. 2
  • Relative position embeddings support addition length generalization: models trained on 5-digit numbers can generalize to 20-digit operands.The paper also reports generalization to 15-digit sums.Source: §1 · p. 2
  • Relative position embeddings and other simple techniques fail to provide length generalization for multiplication.Fine-tuning on long sequences helps, but requires many target-distribution samples and causes catastrophic forgetting.Source: §1 · p. 2
  • The priming sample size scales logarithmically with training-set size, and primed models can target multiple extrapolation lengths.The paper frames 35-digit extrapolation as an illustrative setting and fixes the second multiplication operand at 3 digits.Source: §1 · pp. 2–3
  • With 500 priming examples, a model trained on 5-digit multiplication generalizes across 6- to 35-digit operands.The paper contrasts this with 500 examples being insufficient to train a model directly for those lengths.Source: §1 · p. 3

Related work

Prior transformer work applies to symbolic mathematics, arithmetic, and longer-sequence processing, but length generalization in arithmetic remains challenging. This paper uses relative position embeddings for addition and introduces train set priming for multiplication.

Transformers for mathematics

  • Transformers have been applied to symbolic computation, theorem proving, formula simplification, and numerical calculations, while arithmetic studies document limitations.Modular arithmetic and modular inversion are among the tasks explored in this literature.Source: Related work · p. 3

Length generalization

  • Length-generalization methods include alternative position embeddings, new tokens, and new attention mechanisms, alongside approaches for mathematical sequences.Prior work includes recurrent models and hand-crafted arithmetic modules; autoregressive approaches also use fine-tuning or scratchpads.Source: Related work · p. 3
  • This paper applies relative position embeddings to addition and introduces train set priming for multiplication.Its model overview uses token and position embeddings, a transformer encoder, and per-position output classifiers.Source: Related work · p. 3, Figure 1 · p. 4

2 Experimental setup

The experiments frame four arithmetic operations as symbol-sequence prediction and compare encoder-only transformer variants. Training, evaluation, and output formatting are designed to test generalization from 5-digit training examples to longer operands.

2.1 Problems and encodings

  • The setup covers addition, modular addition, multiplication, and modular multiplication, with models predicting the result from two positive integer operands.Inputs and outputs are digit sequences with operation tokens; multiplication outputs use twice the operand length, while addition outputs use one extra position.Source: §2.1 · pp. 4–5
  • Padding inputs and outputs to fixed lengths is important for the model to handle carries.The setup uses <PAD> and distinguishes in-distribution testing at training length from out-of-distribution testing at greater lengths.Source: §2.1 · p. 5
  • Standard training uses ntrain-digit examples; fine-tuning adds target-length examples, while priming mixes a small fixed share of long examples into training.The first operands come from a fixed training set, while second operands are sampled online.Source: §2.1 · p. 5

2.2 Model and training

  • The models are encoder-only Transformers or UTransformers, trained from scratch as supervised per-position symbol classifiers with cross-entropy loss.The encoder output is truncated to nout positions and a linear layer predicts each output symbol.Source: §2.2 · pp. 5–6
  • Universal Transformers use shared layers, which the authors motivate as fitting recursive arithmetic algorithms and repeated processing of long operands.The authors choose encoder-only models as a simpler alternative to autoregressive or sequence-to-sequence architectures for these tasks.Source: §2.2 · p. 6

3 Addition: relative position embeddings enable length generalization

Relative position embeddings enable substantial addition length generalization, unlike absolute position embeddings. Generalization depends on model configuration, including a minimum hidden size and, for UTransformers, sufficient depth.

Results

  • All compared models achieve 100% in-domain accuracy despite their different out-of-distribution behavior.The experiments train on up to 5-digit operands and test addition at lengths from 6 to 20 digits.Source: §3 · p. 7
  • Absolute-position models achieve only 3.1% accuracy on 6-digit tests and 0% at all longer tested lengths.The authors attribute this to position-specific rules that leave unseen digit positions unhandled.Source: §3 · pp. 7–8
  • Relative position models reach 99.9% accuracy on 10-digit and 98.3% on 15-digit addition tests, while accuracy falls to 21.3% on 20-digit numbers.The key-relative-position variant is crucial for extrapolation.Source: §3 · p. 8

Depth and dimension for longer extrapolation

  • At least 512 hidden dimensions for Transformers and 256 for UTransformers are needed for extrapolation; 1024-dimensional models reach 17-digit extrapolation.UTransformers need 6 layers to extrapolate, while 2-layer Transformers reach 10-digit numbers.Source: §3 · p. 8, Figure 2 · p. 8

4 Modular arithmetic

The difficulty of modular arithmetic depends on the modulus: powers of 10 yield fixed-length outputs, while other moduli add integer division. Generalization varies across these cases and between modular addition and multiplication.

Modular arithmetic

  • Modular operations with c = 10^k use only the k last operand digits and produce a constant-length k-digit result; other moduli add integer division.The fixed number of relevant digits makes powers-of-10 cases easier to learn and generalize, according to the paper.Source: §4 · pp. 8–9

Modular addition

  • For c ∈ {100, 1000}, modular addition generalizes to large numbers with RPE, while APE reaches 73.3% accuracy on 10-digit numbers.The reported 73.3% is for APE-based models on 10-digit numbers.Source: §4 · p. 9
  • For c ∈ {101, 128}, no model learns 5-digit modular addition in-domain, and scaling architectures to 14 layers and 1280 dimensions brings no improvement.The authors relate this to the difficulty of modular arithmetic in the general case.Source: §4 · p. 9

Modular multiplication

  • For modular multiplication, both APE and RPE achieve 100% on 35-digit numbers when c = 100; for c = 1000, RPE reaches 83% on 20-digit and 55% on 30-digit numbers.For c = 1000, APE reaches 43% on 20-digit numbers.Source: §4 · p. 9

Modular arithmetic

  • APE-based models generalize when an operation uses a fixed number of input tokens and produces a constant-length output.This is the authors’ interpretation of their modular-arithmetic experiments.Source: §4 · p. 9

5 Multiplication: train set priming for length generalization

Relative position embeddings alone do not yield length generalization for 5 × 3 multiplication, and fine-tuning on long examples can forget the original task. Train set priming uses a small number of long examples to support extrapolation, including across lengths when the priming data cover those lengths.

5.1 Relative position embeddings and fine-tuning

  • Relative position embeddings fail to generalize on 5 × 3 multiplication, despite near-100% in-domain accuracy; simpler 5 × 1 and 5 × 2 tasks reach only 12.2% and 16.9% at 6 digits.Those simpler tasks also fail at longer sequence lengths.Source: §5.1 · p. 10
  • Fine-tuning on 1,000 examples of 35 × 3 multiplication learns the long task but leaves the model unable to perform 5 × 3 multiplication.The paper identifies this loss of the original ability as catastrophic forgetting.Source: §5.1 · p. 10

5.2 Priming for length generalization in multiplication.

  • 50 priming examples enable close to 100% accuracy on both 5 × 3 and 35 × 3 multiplication, using 20 times fewer examples than fine-tuning.For 35 × 3, priming requires 75,000 examples over training (1,500 passes over 50 fixed examples), compared with 7.5 million for fine-tuning.Source: §5.2 · pp. 10–11
  • Fewer than 25 priming examples prevent generalization, and accuracy rises with priming rate above that threshold.The 25-example threshold corresponds to a 0.5% priming rate.Source: §5.2 · p. 11
  • Priming sample size grows logarithmically with training-set size: 30, 70, and 100 examples are needed for training sets of 10^3, 10^4, and 10^5, respectively.These priming rates are 3%, 0.7%, and 0.1%, respectively.Source: §5.2 · p. 11
  • Priming sample size scales linearly with extrapolation length: 10 examples suffice for 6-digit generalization, compared with 50 for 35-digit generalization.Source: §5.2 · p. 11
  • Most mixed-length curricula fail to extrapolate to 35 × 3, but priming with a mixture of 34- and 35-digit numbers learns faster and reaches higher extrapolation accuracy.Source: §5.2 · p. 11

5.3 Priming for extrapolation at all lengths

  • Priming on lengths from 6 to 35 yields over 95% accuracy across those lengths with 500 examples, while priming only even lengths does not generalize to odd lengths.Priming with 35-digit numbers alone generalizes only to 35-digit operands.Source: §5.3 · pp. 11–12

6 Discussion

The discussion attributes length generalization to how models learn output digits, while noting that varying output lengths may challenge APEs. It also analyzes addition failures, priming, and possible extensions beyond arithmetic.

6.2 Failure cases in addition

  • Almost all 20-digit addition failures involve at least three carries and two consecutive carries; larger carry counts have no further impact.Figure 6 analyzes accuracy by total and consecutive carries.Source: Figure 6 · p. 13, §6.2 · p. 13
  • Wrong 20-digit sums usually contain only a few incorrect digits, concentrated in the first and second positions.These positions correspond to the largest powers of ten in the sum.Source: Figure 6 · p. 13, §6.2 · p. 13

6.1 Why do RPEs extrapolate better than APEs?

  • APE models fail on carry-free element-wise addition while RPE models generalize, disproving carry propagation as the explanation for RPE success.When APE models generalize, they almost always predict the five leftmost result digits—their in-domain positions.Source: §6.1 · p. 13
  • APE models length-generalize in modular tasks when the modulus is a power of 10 and output length stays constant, suggesting varying output lengths matter.The authors identify varying output lengths as an important factor in APE extrapolation failures.Source: §6.1 · p. 13
  • RPE models learn most output digits simultaneously, unlike APE models, suggesting they learn one algorithm across positions that helps generalization.The first and last digits are learned first; the other digits are learned simultaneously by RPE models, while APE models appear to learn positions independently.Source: §6.1 · p. 13

6.3 More about priming

  • Priming also enables APE models to length-generalize, but their required priming rate is 10 times larger, at ε = 10%.The authors say this higher rate means primed APE models are not a viable proposition.Source: §6.3 · p. 14
  • With priming, multiplication models learn several output digits concurrently; the authors identify simultaneous learning of training positions as a key determinant of length generalization.Without priming, RPE multiplication models learn one digit at a time over 1,000 epochs, while primed RPE and APE models learn digits concurrently.Source: §6.3 · p. 14

6.4 Priming beyond arithmetic

  • The authors conclude that train set priming can improve arithmetic length generalization with fewer target-distribution samples than fine-tuning and without catastrophic forgetting.They propose investigating priming for other mathematical problems, compositionality, theory, and adaptation of pretrained language models to new tasks.Source: §6.4 · p. 14

A Additional experiments

The appendix collects additional modular-arithmetic results and complementary experiments on digitwise addition and priming.

A Additional experiments

  • The appendix first gives complete modular addition and multiplication results referenced in Section 4.These are presented in Subsection A.1.Source: A Additional experiments · p. 19
  • It then reports APE and RPE results on digitwise addition and shows that APE models can be primed to length-generalize at a higher priming rate.The APE priming results appear in Subsection A.3.Source: A Additional experiments · p. 19
  • The appendix also plots the order in which RPE and APE models make correct predictions across digit positions.These plots appear in Subsection A.5.Source: A Additional experiments · p. 19

A.1 Additional experiments on modular arithmetic

The additional modular-arithmetic experiments report different extrapolation outcomes across moduli for addition and multiplication.

A.1 Additional experiments on modular arithmetic

  • The reported modular addition and multiplication results cover moduli 100, 1000, 128, and 101, with accuracy evaluated on 100,000-example test sets.The appendix captions identify the respective results in Tables 4 and 5.Source: Table 4 · p. 19, Table 5 · p. 20
  • For modular addition, models extrapolate when the modulus is a power of 10 but fail for moduli 128 and 101.Table 4 reports accuracy on 100,000-example test sets for moduli 100, 1000, 128, and 101.Source: §A.1 · p. 19, Table 4 · p. 19
  • For modular multiplication, models extrapolate for power-of-10 moduli; they generalize non-trivially for 128 but fail for 101.The authors offer, but do not confirm, the hypothesis that the difference relates to 101 being prime and 128 being a power of 2.Source: §A.1 · p. 20, Table 5 · p. 20

A.2 Element-wise addition experiments

The element-wise addition experiments test extrapolation from five-digit training examples to 20-digit sums and compare APE with RPE models.

A.2 Element-wise addition experiments

  • The experiment trains a UTransformer on five-digit element-wise addition using 50,000 examples and tests extrapolation on 10,000 20-digit examples.Table 6 reports extrapolation accuracy for APE and RPE models.Source: §A.2 · p. 21, Table 6 · p. 21
  • RPE models generalize from five-digit element-wise addition training to 20-digit test examples, whereas APE models fail.Training uses 50,000 examples and evaluation uses 10,000 test examples.Source: §A.2 · p. 21, Table 6 · p. 21
  • When APE models predict element-wise sums, they perform well on the leftmost five digits seen during training but fail on rightmost digits.Figure 8 plots the APE model's digitwise test accuracy.Source: §A.2 · p. 21, Figure 8 · p. 21

A.3 Multiplication experiments using APEs

With train set priming, an APE model extrapolates from 5 × 3 to 35 × 3 multiplication with 95% test accuracy; priming also changes how it learns digit positions.

A.3 Multiplication experiments using APEs

  • Figure 9a reports an APE model’s accuracy on 5 × 3 and 35 × 3 multiplications; a standard UTransformer was trained on 5 × 3, tested on 35 × 3, with 5,000 training examples.Source: Figure 9 · p. 21
  • Without priming, the model learns digit positions independently, while 10% priming makes it learn the positions together.The independent-position pattern is described as a sign of memorization.Source: §A.3 · p. 22

A.4 Training accuracy for varying priming rates

  • Figure 10 reports digitwise accuracy on training examples across varying priming rates.Source: Figure 10 · p. 22

A.5 Test accuracy for varying priming rates

  • Figure 11 reports digitwise predictions on test examples across varying priming rates.Source: Figure 11 · p. 23
Loading 2306.15400v1…