PaperlayerView exact revision on arXiv

Source-linked AI summary

Transformers Can Achieve Length Generalization But Not Robustly

Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, Denny Zhou

arXiv:2402.09371v1cs.LGcs.AIcs.CL

TL;DR

Transformers often struggle to generalize to longer sequences, including on decimal addition. This paper evaluates combinations of data format and position encoding, finding that an appropriate combination supports almost-perfect generalization to 2.5× the training length.

  • Problem

    Transformers often struggle to generalize to longer sequences, and prior work found limited length generalization on decimal addition.

    Source: §1 · pp. 1–2

  • Method

    The paper evaluates data-format and position-encoding combinations for decimal addition, with a recipe using FIRE, randomized positions, reversed formatting, and index hints.

    Source: §1 · p. 2

  • Results

    With appropriate data formatting and position encoding, Transformers achieve almost-perfect generalization on sequences up to 2.5× the training length.

    Source: §7 · p. 11

  • Takeaways & Limitations

    For decimal addition, the paper concludes that effective length generalization depends on the right combination of data format and position encoding, rather than an inherent Transformer design barrier.

    Source: §7 · p. 11

    Robust length generalization remains elusive, even with carefully tuned regularization hyperparameters.

    Source: §7 · p. 11

Abstract

from arXiv · show

Length generalization, defined as the ability to extrapolate from shorter training sequences to longer test ones, is a significant challenge for language models. This issue persists even with large-scale Transformers handling relatively straightforward tasks. In this paper, we test the Transformer's ability of length generalization using the task of addition of two integers. We show that the success of length generalization is intricately linked to the data format and the type of position encoding. Using the right combination of data format and position encodings, we show for the first time that standard Transformers can extrapolate to a sequence length that is 2.5x the input length. Nevertheless, unlike in-distribution generalization, length generalization remains fragile, significantly influenced by factors like random weight initialization and training data order, leading to large variances across different random seeds.

1. Introduction

The paper tests whether Transformers can generalize beyond training lengths on decimal addition and finds that performance depends on the combination of data format and position encoding. The selected combination reaches 2.5× the training length, but this generalization remains fragile.

1. Introduction

  • The study evaluates whether Transformers learn scalable algorithms or rely on shortcuts, using decimal addition, where prior work found limited length generalization.The authors frame N-digit addition as synthetic language learning that can reveal whether Transformers internalize algorithms.Source: §1 · pp. 1–2
  • Training on 40-digit addition yields over 98% accuracy on 100-digit addition (2.5× extension), exceeding the prior reported extension ratios.The best recipe combines FIRE position encodings, randomized positions, reversed formatting, and index hints.Source: Figure 1 · p. 1, §1 · p. 2
  • The successful recipe combines FIRE position encodings with randomized positions, reversed formatting, and index hints.The paper reports that success depends on the choice of position encoding as well as data format.Source: §1 · p. 2
  • The observed 2.5× length generalization is fragile and significantly affected by random weight initialization and training-data order.The authors report substantial variation across these factors.Source: §1 · p. 2

2. Position Encoding and Data Formats

The section surveys positional encodings and data formats proposed for length generalization, noting that neither architectural changes nor scaling alone necessarily resolve extrapolation. It reviews how encoding choices and addition-specific formats affect the task.

2. Position Encoding and Data Formats

  • Prior work explores positional encodings, attention mechanisms, and data formatting, while some reported modifications are ad hoc or perform poorly on longer sequences and scaling alone may be insufficient.The section frames positional encoding and data formats as approaches to improving the learning efficacy of existing Transformers.Source: §2 · p. 2

2.1. Position Encoding for Length Generalization

  • APE often struggles on longer sequences, while RoPE has limited length generalization and context-length extensions do not necessarily help algorithmic tasks.APE attaches position vectors to token embeddings; RoPE encodes relative positions through rotations of query and key vectors.Source: §2.1 · p. 3
  • Additive RPEs modify attention through a position-dependent bias, with T5, Alibi, and KerpleLog using different functions of key- and query-position differences.The bias is added to attention logits; the listed methods use learned scalars or a distance-based function.Source: §2.1 · p. 3
  • Decoder-only Transformers can acquire positional understanding without explicit PE, and one study found NoPE outperforming specialized PEs on simple algorithmic tasks.Encoder-only Transformers are permutation equivariant without positional encodings.Source: §2.1 · p. 3
  • Randomized PE samples positional encodings from beyond test-time lengths while preserving their order, helping models adapt to larger encodings at inference.The cited approach is intended to eliminate out-of-distribution position encodings during testing.Source: §2.1 · pp. 3–4

2.2. Data Formats

  • Reversing addition places least-significant digits first, aligning autoregressive output with digit-by-digit computation and reducing the task to local steps involving operand digits and carry.The standard printed order instead presents most-significant digits first.Source: §2.2 · p. 4
  • Index hints label operand and answer digits, enabling Transformers to perform indexing via induction heads.For example, 42 + 39 = 81 is represented as a4b2 + a3b9 = a8b1.Source: §2.2 · p. 4
  • Figure 2 contrasts PE and format choices; Shen et al. (2023)'s random spacing generalizes 10-to-11-digit addition but falters on longer sequences, while Zhou et al. (2023) and this work use index hints.The figure caption says all compared studies use reversed format; this work combines FIRE and Randomized PE.Source: §2.2 · p. 4, Figure 2 · p. 4

3. A Recipe for Length Generalization in Decimal Addition

The proposed recipe combines FIRE and randomized position encodings with reversed addition and index hints. Each choice is motivated by a distinct aspect of extrapolating decimal addition, and the authors ablate these decisions and alternatives.

3. A Recipe for Length Generalization in Decimal Addition

  • The best model combines FIRE and randomized position encodings with reversed format and index hints.The authors report this combination as the one used by the model behind Figure 1.Source: §3 · p. 5
  • FIRE is intended to help length generalization because its position encodings are more expressive than other PEs.The authors attribute the expressiveness comparison to Li et al. (2023).Source: §3 · p. 5
  • Randomized position encodings are intended to prevent overfitting to position indices and index differences seen during training.The authors describe randomization as crucial for avoiding that overfitting.Source: §3 · p. 5
  • Reversed format helps decompose the computation into local, Markovian steps that depend only on the preceding step.Source: §3 · p. 5
  • Index hints are intended to ease operand identification by helping match the operands to add at each step.The authors link the hints to the challenge of identifying the right operands.Source: §3 · p. 5
  • The authors ablate each recipe decision and test some alternative choices in Section 4.Source: §3 · p. 5

4. Experiments

The experiments evaluate how data format and positional encoding affect addition length generalization, and show that strong extrapolation can coexist with substantial sensitivity to random seeds and training-data order.

4.1. Setup

  • Training used 30M examples at lengths 1–40 with reversed format and index hints sampled from 102 ordered symbols; the base Transformer has 25M parameters.Source: §4.1 · p. 5
  • Experiments were repeated five times because performance variance was high, and hyperparameters for randomized positional encoding and random space augmentation were tailored to each positional encoding.Source: §4.1 · p. 5

4.2. Results

  • 100-digit addition shows near-perfect generalization in the best FIRE trial trained only on lengths 1–40, while other positional encodings degrade visibly beyond length 60.Source: §4.2 · pp. 5–6, Figure 3 · p. 6
  • Index hints improve length generalization across positional encodings; without them, NoPE and FIRE also lose 40-digit in-distribution accuracy, while RoPE and KerpleLog fail out of distribution.Source: Figure 4 · p. 6, §4.2 · p. 6
  • Reversed format outperforms standard format across positional encodings; FIRE remains strong in standard format but declines beyond 60 digits, likely from carry-propagation limits.Source: §4.2 · p. 6, Figure 5 · p. 7
  • Random space augmentation helps RoPE and KerpleLog but harms NoPE and FIRE, while randomized positional encoding boosts FIRE and degrades KerpleLog.Source: §4.2 · pp. 6–7, Figure 7 · p. 7, Figure 8 · p. 7
  • Reversed-format training shows a sharp, grokking-like transition, unlike the gradual improvement in standard format; the transition marks learning the addition algorithm.Source: §4.2 · p. 6, Figure 6 · p. 7
  • Out-of-distribution length generalization varies substantially across weight initializations and data orders despite similar training losses and perfect in-distribution generalization.Some weight initializations are resilient to changes in data order, suggesting “fortunate” configurations for robust length generalization.Source: §4.2 · pp. 7–8, Figure 10 · p. 8, Figure 11 · p. 8

5. Analysis

The analysis finds that training length strongly affects extrapolation, while model size has mixed effects; errors and training trajectories also reveal substantial variation across methods and runs.

Error analysis

  • Errors are distributed roughly uniformly across output indices, and—excluding two near-zero-accuracy runs—over 90% of errors in incorrect 100-digit examples are single-digit mistakes.The analysis also finds no significant difference between examples with and without carry, so carry propagation does not majorly impede length generalization.Source: §5 · p. 8

Performance evolution during training

  • FIRE’s out-of-distribution accuracy generally rises steadily during training, whereas other position encodings are more volatile and some show a sudden accuracy surge without in-distribution change.Models reach near-perfect in-distribution accuracy early, while out-of-distribution accuracy follows architecture-dependent trajectories.Source: §5 · p. 8, Figure 12 · p. 9

Sequence length during training

  • Training on 40 digits generalizes to 100 digits (2.5×), compared with 45 digits (1.5×) after training to 30, 25 digits (1.25×) after 20, and no extrapolation after 10.Increasing training length improves generalization across position encodings.Source: §5 · p. 9, Figure 13 · p. 9

Scaling model size

  • Model-size effects are mixed: larger models help at shorter lengths, but longer-regime outcomes vary by position encoding, and the 2M model deteriorates on longer sequences.The 5M model reaches 80% accuracy on 100-digit addition after training on lengths 1–40, surpassing the 268M model.Source: §5 · pp. 9–10, Figure 14 · p. 9, Figure 15 · p. 10

Does stronger regularization reduce variance?

  • Higher weight decay slightly increases the likelihood of effective generalization, but neither weight decay nor dropout ensures robustness; initialization and data order still strongly influence variance.With weight decay of 1e-6 or 1.0, the models attain approximately 80% accuracy on 100-digit addition after training on 40-digit sequences; dropout of 0.2 severely impairs generalization.Source: §5 · p. 10, Figure 16 · p. 10

6. Related Work

Prior length-generalization work has explored position encodings and data formats, while this paper emphasizes that their combination matters. Its experiments find that FIRE and coordinated data design support stronger generalization.

Data format for Length Generalization

  • Earlier approaches target either positional encodings or data formats, and prior work also reports effects from data distributions and sampling strategies.The paper contrasts those approaches with good performance under naive random sampling, suggesting sophisticated sampling may be redundant in its setting.Source: §6 · pp. 10–11

Position Encoding for Length Generalization

  • Prior findings differ on positional encoding: one study found better algorithmic-task results without it, whereas these experiments identify FIRE as crucial for optimal length generalization.The paper also reports that position encoding and data design work synergistically.Source: §6 · pp. 10–11

Data format for Length Generalization

  • For Transformers trained from scratch, prior work reports gains from reversed formatting and scratchpads, while data distribution and sampling can also influence generalization.Related work additionally studies heuristic formats for pretrained models and auxiliary tasks or train-set priming in other settings.Source: §6 · p. 11

7. Conclusion

The paper concludes that length generalization is achievable with an appropriate pairing of data format and position encoding, while robust generalization remains unresolved. It also describes additive relative positional encoding and its data-generation setup.

7. Conclusion

  • Transformers achieve almost perfect generalization up to 2.5× training length with appropriate data formatting and position encoding, but robust generalization remains elusive.The analysis finds a significant dependency between position-encoding type and data format.Source: §7 · p. 11

A.1. Additive Relative Positional Encoding (RPE)

  • Additive relative positional encoding places positional information in attention logits as a bias matrix induced by a function of token positions.Variants differ in how that position-dependent bias is formulated and parameterized.Source: §A.1 · p. 15
  • T5’s relative-position bias divides distances into logarithmically scaled buckets, each associated with a parameter.The described setup uses K+1 buckets and a predefined distance L1.Source: §A.1 · p. 15
  • The appendix describes Alibi’s distance-proportional bias, Kerple’s logarithmic or power variants, and FIRE’s learnable function-based bias.FIRE uses a learnable MLP and a monotonically increasing position transformation.Source: §A.1 · p. 15

B.1. Data Generation

  • The default format reverses the addition sequence and includes index hints; training samples consecutive hints from an ordered set of 102 symbols, and inference uses the same hint-sampling strategy.Examples use uniformly selected operand lengths up to the training maximum, with independently sampled operands and zero padding for possible carry digits.Source: §B.1 · p. 16

B.2. Training Details

The appendix specifies the model and training setup, then provides figures covering training behavior, extrapolation across position encodings, index hints, and run-to-run variance.

B.2. Training Details

  • The base model is a 25M-parameter Transformer with six blocks, 512 hidden units, a 2048-unit feedforward layer, RMSNorm, and 50,000 AdamW training steps.Training uses weight decay 0.1, no dropout, a 500-step warm-up, cosine decay from 3e-4, sequence packing, and greedy evaluation.Source: §B.2 · p. 16, Table B.1 · p. 17

C.1. Training Loss and Sequence Exact Match Accuracy of Reverse Format with Index Hint trained up to 40

  • Figure C.1 reports exact-match accuracy from 20 to 100 digits across ten trials trained to 40 digits with index hints and reverse format; Figure C.2 shows varied extrapolation despite near-zero training loss.The loss comparison is across ten reverse-format trials after 10K training steps.Source: Figure C.1 · p. 18, Figure C.2 · p. 18

C.2. Training Loss and Next-token Prediction Accuracy of Standard and Reverse Format

  • Figure C.3 compares training loss and next-token accuracy for standard and reverse formats across four position encodings; reverse format shows distinct grokking during training.Standard format instead shows gradual enhancement.Source: Figure C.3 · p. 19

C.3. The evolution of EM Accuracy during training in reverse format using 4 PEs

C.4. Effect of Index Hint

  • Without index hints, all tested position encodings fail beyond trivial lengths and struggle even on in-distribution generalization.Figure C.8 reports best sequence exact match across five trials trained up to length 10.Source: Figure C.8 · p. 22

C.5. Source of Variance

  • Figure C.9 plots 100-digit exact-match accuracy across ten FIRE trials, comprising two weight seeds and five data seeds, after training to 40 digits.The figure’s setup separates weight-seed and data-seed variation.Source: Figure C.9 · p. 22

C.6. Error Analysis

Errors are spread broadly across positions, although successful generalizations mostly make only single-digit errors; error counts exhibit an exponential distribution.

C.6. Error Analysis

  • Across 10 runs, errors are broadly distributed across positions, while specific checkpoints show positional error propensities.In successful generalizations, more than 90% of errors are single-digit inaccuracies.Source: Figure C.10 · p. 23
  • In successful generalizations, more than 90% of errors are confined to single-digit inaccuracies.Source: Figure C.10 · p. 23
  • Error counts exhibit an exponential distribution.Source: §C.6 · p. 23, Figure C.11 · p. 23

C.7. Training Digit Length Effect

The appendix examines training digit length, model size, FIRE scaling, and hyperparameters; model-size effects vary by regime and encoding, while higher dropout impedes length generalization.

C.7. Training Digit Length Effect

  • The training-digit-length study compares best sequence exact-match accuracy across five trials, two model sizes, and training lengths up to 10, 20, 30, and 40 digits using four position encodings.Source: Figure C.13 · p. 24

C.8. Model Size Effect

  • Model-size scaling consistently improves shorter regimes (1-10, 1-20), but effects in larger regimes vary by position encoding and do not noticeably reduce variance.Larger models outperform smaller ones with RoPE and KerpleLog, but underperform with NoPE and FIRE.Source: Figure C.14 · p. 25

C.9. FIRE Related Scaling

  • The FIRE scaling study compares best sequence exact-match accuracy across five trials, four model sizes, and training lengths up to 10, 20, 30, and 40 digits.Source: Figure C.15 · p. 26

C.10. Hyperparameter Study

  • Higher dropout markedly impedes length generalization, whereas a lower rate shows negligible impact.The hyperparameter study reports EM Accuracy at Length 100 and identifies 3e-4 as the seemingly optimal learning rate.Source: Figure C.17 · p. 27, Figure C.16 · p. 27, §C.10 · p. 27

D. Training Loss and Sequence Exact Match Accuracy

Across the appendix's addition experiments, near-zero training losses do not ensure similar length generalization across runs; the figures compare accuracy across lengths, formats, and augmentation or position-encoding conditions.

D.1. Reverse Format without Index Hint trained up to 40-digit addition

  • Despite nearly zero log perplexity losses after 10K steps, runs show very different length generalization in the reverse-format experiment trained up to 40 digits.Source: Figure D.2 · p. 28

D.2. Reverse Format without Index Hint trained up to 10-digit addition

  • The reverse-format experiment trained up to 10 digits compares exact-match accuracy on 20- to 100-digit addition across four position encodings and plots training loss over 10 trials.Despite nearly zero log perplexity losses after 10K steps, runs show different length generalization.Source: Figure D.3 · p. 29, Figure D.4 · p. 29

D.3. Standard Format with Index Hint trained up to 40

  • The standard-format section presents exact-match accuracy on 20- to 100-digit addition and training loss across 10 trials.The accuracy figure describes training up to 40 digits and four position encodings.Source: Figure D.5 · p. 30, Figure D.6 · p. 30

D.4. Random Space Augmentation with Reverse Format with Index Hint trained up to 40-digit addition

  • The random-space-augmentation figures compare accuracy and training loss with versus without augmentation across 10 trials in the reverse-format, index-hint setting.The accuracy comparison covers 20- to 100-digit addition after training up to 40 digits with four position encodings.Source: Figure D.7 · p. 31, Figure D.8 · p. 32

D.5. Randomized Position Encoding with Reverse Format with Index Hint trained up to 40-digit addition

  • The randomized-position-encoding figures compare accuracy and training loss with versus without randomized position encoding across 10 trials.The accuracy comparison covers 20- to 100-digit addition after training up to 40 digits with four position encodings.Source: Figure D.9 · p. 33, Figure D.10 · p. 33

Related papers

Loading 2402.09371v1…