Source-linked AI summary

Randomized Positional Encodings Boost Length Generalization of Transformers

Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Bennani, Shane Legg, Joel Veness

arXiv:2305.16843v1cs.LGcs.AIcs.CLstat.ML

TL;DR

Transformers struggle to generalize to longer sequences, while training directly on longer sequences is computationally costly. This paper randomizes ordered positional encodings over an extended position range and reports higher test accuracy across algorithmic reasoning tasks. Its evaluation is confined to synthetic tasks, leaving generalization to other domains open.

  • Problem

    Positional encodings for positions beyond the training length are out of distribution on longer test sequences, limiting evidence about Transformer length generalization.

    Source: §1 · p. 1, §3 · p. 3
  • Method

    The method samples an ordered subset of positions from a range extending beyond training and test positions, then uses those positions for token encodings.

    Source: §1 · p. 1, §3 · p. 3
  • Results

    12.0% average test-accuracy increase is reported for the randomized method across algorithmic reasoning tasks.

    Source: §1 · p. 2
  • Takeaways & Limitations

    The evaluation reports improved length generalization and computational performance over naively training on longer sequences.

    Source: §5 · p. 5

    Evaluation is confined to synthetic algorithmic reasoning tasks, so generalizability to other tasks and domains remains an open question.

    Source: Limitations · p. 5

Abstract

from arXiv · show

Transformers have impressive generalization capabilities on tasks with a fixed context length. However, they fail to generalize to sequences of arbitrary length, even for seemingly simple tasks such as duplicating a string. Moreover, simply training on longer sequences is inefficient due to the quadratic computation complexity of the global attention mechanism. In this work, we demonstrate that this failure mode is linked to positional encodings being out-of-distribution for longer sequences (even for relative encodings) and introduce a novel family of positional encodings that can overcome this problem. Concretely, our randomized positional encoding scheme simulates the positions of longer sequences and randomly selects an ordered subset to fit the sequence's length. Our large-scale empirical evaluation of 6000 models across 15 algorithmic reasoning tasks shows that our method allows Transformers to generalize to sequences of unseen length (increasing test accuracy by 12.0% on average).

1 Introduction

Transformers need positional encodings to represent token order, but positions beyond the training length are out of distribution at test time. The paper introduces randomized positional encodings and reports improved length generalization on algorithmic reasoning tasks.

1 Introduction

  • Positions beyond the maximum training length N are out of distribution when evaluating on longer sequences.The introduction frames this as a barrier to length generalization, including on tasks such as binary addition.Source: §1 · p. 1
  • Randomized positional encodings assign each training example an ordered subset of positions drawn from a range extending beyond positions seen in training or evaluation.The approach is compatible with existing positional encoding schemes and is presented as more efficient than training directly on longer sequences.Source: §1 · pp. 1–2
  • 12.0% average and up to 43.5% higher test accuracy are reported across algorithmic reasoning tasks.The introduction describes the method as leaving in-domain generalization performance unaffected.Source: §1 · p. 2

2 Related Work

Related work spans positional encoding methods and approaches to systematic length generalization. The paper positions its evaluation on algorithmic reasoning tasks as a test of methods that prior positional encodings do not solve in this setting.

2 Related Work

  • Earlier positional encodings include scaled sinusoids, learned embeddings, relative distances, RoPE (Su et al., 2021), and distance-proportional attention bias.The paper says these approaches fail at length generalization on algorithmic reasoning tasks.Source: §2 · p. 2
  • Randomized learned positional encodings are described as a special case of the paper's family (Li and McClelland, 2022).The paper also cites feature and position randomization in graph neural networks and continuous dynamical systems for longer sequences.Source: §2 · p. 2
  • Other systematic length-generalization approaches include embedding scaling, early stopping, adaptive computation time, geometric attention, and hierarchical reinforcement learning.The paper evaluates on the benchmark of Delétang et al. (2023), which unifies work on Transformer recognition of formal languages.Source: §2 · p. 2

3 Randomized Positional Encodings

The method addresses positional encodings shifting out of distribution as test sequences exceed the training length. It samples ordered positions from an extended range, while requiring the maximum test length to be known when choosing that range.

3 Randomized Positional Encodings

  • For training to maximum length N, conventional positional encodings shift away from their training distribution on test sequences of length M > N.The proposed scheme relies on order information and is expected to generalize up to M ≤ L, where L is configurable.Source: §3 · p. 3
  • At each training step, the method samples a length n from 1 to N, then samples n indices without replacement from 1 to L, sorts them, and uses their positional encodings in order.At test time it applies the same procedure to all M token positions.Source: §3 · p. 3
  • The sampled positions preserve token order but are resampled at every step, so encodings are not directly tied to exact token positions.The paper reports an ablation validating the necessity of ordering the sampled positions.Source: §3 · p. 3
  • Training on short sequences can match prior approaches trained on longer sequences when longer sequences share the same structure and their positions are observed during training.The paper notes global attention has O(ℓ^2) running time, making its scheme faster than direct long-sequence training; in-domain performance is largely unaffected.Source: §3 · p. 3
  • The maximum test length M must be known in advance to choose L ≫ M; choosing L much larger than N or M may leave too few unique indices observed and hurt both in- and out-of-distribution performance.The paper says performance is largely unaffected by the precise value of L in its experiments, but practitioners may need to tune it for their domains.Source: §3 · p. 3

4 Experimental Evaluation

The experiments compare positional encodings on algorithmic reasoning tasks and evaluate unseen sequence lengths. Randomized encodings improve average test accuracy, and a short-sequence randomized-relative model reaches high accuracy faster than a long-sequence classical-relative model.

4 Experimental Evaluation

  • The benchmark covers modular arithmetic, string reversal and duplication, binary addition and multiplication, and bucket sort.The paper follows the setup of Delétang et al. (2023), evaluating an encoder-only model trained to length 40 and tested on lengths 41–500.Source: §4 · pp. 3–4
  • 12.0% average test-accuracy increase is reported for the randomized method across the benchmark; the randomized relative encoding solves tasks previously out of reach, including REVERSE STRING and MISSING DUPLICATE.Table 1 reports accuracy averaged over test lengths and maximized over 10 seeds and 3 learning rates.Source: Table 1 · p. 4, §4 · p. 4
  • The method is instantiated with a range of positional encodings, including none, sinusoidal, relative, ALiBi, RoPE, learned, and label-based encodings.The paper states label-based encodings (Li and McClelland, 2022) are equivalent to randomized learned encodings and are subsumed by its method.Source: §4 · p. 4
  • A model trained on sequences up to length 40 reaches test accuracy above 90% roughly 35.4 times faster than the naive model trained up to length 500.Figure 2 compares average accuracy over lengths 41–500 against training time; the randomized model runs at 168.4 steps per second versus 22.1 for the naive approach on an NVIDIA V100 GPU.Source: §4 · pp. 4–5, Figure 2 · p. 5

5 Conclusion

The conclusion presents randomized positional encodings as a way to improve length generalization by sampling ordered positions from a wider range than training or evaluation positions. The paper reports stronger length-generalization performance than prior work and better computational performance than training on longer sequences.

5 Conclusion

  • The paper introduces a family of positional encodings that improves Transformers’ length generalization capabilities.Source: §5 · p. 5
  • Randomized positional encodings sample positions from a wider range than training or test lengths while preserving their order.The approach is intended to address conventional encodings becoming out-of-distribution at longer sequence lengths.Source: §5 · p. 5
  • The large-scale evaluation reports that the method significantly outperforms prior work in length generalization.The conclusion also reports superior computational performance over the naive approach of training on longer sequences.Source: §5 · p. 5

Limitations

The evaluation is confined to synthetic algorithmic reasoning tasks, leaving generalization to other tasks and domains open. The paper also notes that the approach addresses only one factor in poor length generalization and introduces a maximum-position hyperparameter.

Limitations

  • Evaluation is confined to synthetic algorithmic reasoning tasks, so generalization to other tasks and domains remains an open question.The authors say further evaluation is needed to understand the method’s potential in real-world applications.Source: Limitations · p. 5
  • The method introduces a maximum-position hyperparameter L that practitioners may need to tune for their problem domains.The authors report performance is largely unaffected by L’s precise value in their Appendix B.1 experiments.Source: Limitations · p. 5
  • The approach ameliorates only one failure mode of length generalization; attention becoming less peaked on longer sequences is another cited factor.Source: Limitations · p. 5

A Experimental Details

The experiments use a 15-task algorithmic reasoning suite with sampled training and test distributions. They evaluate an encoder-only Transformer across positional encodings and training settings, reporting negligible computational overhead for the method.

A Experimental Details

  • The experiment suite consists of 15 algorithmic reasoning tasks with continuously sampled training and testing distributions.Source: A Experimental Details · p. 8
  • Models train for 2,000,000 steps with batch size 128 and are tested across sequence lengths 41–500.Testing uses a batch of 500 for each length; each task and positional encoding uses 10 initialization seeds.Source: A Experimental Details · p. 8
  • The study uses a five-block, eight-head encoder-only Transformer and reports negligible computational overhead from the method.The model has d_model = 64; relative and randomized relative encodings use 270,146 parameters.Source: A Experimental Details · p. 8

B.1 Ablation Study

The ablations examine the maximum sampling position L and whether subsampled positions are sorted. Accuracy is largely stable across reasonable L values, while sorting the randomized sin/cos positions substantially improves results.

B.1 Ablation Study

  • The maximum sampling position and sorting of subsampled positions are the two components examined in the ablation study.Source: §B.1 · p. 8
  • Test accuracy on REVERSE STRING and MISSING DUPLICATE is largely unaffected by L across the tested sweep.The study considers L values of 1024, 2048, 4096, and 8192; the reported stability applies when L exceeds the maximum evaluation length M but is not unreasonably large.Source: §B.1 · p. 8, Figure 3 · p. 8
  • 15.7% average test-accuracy increase, up to 76.3% on certain tasks, is reported when subsampled randomized sin/cos positions are sorted.Without sorting, the approach fails to beat random accuracy on all but permutation-invariant CYCLE NAVIGATION.Source: §B.1 · pp. 8–9, Table 3 · p. 9
  • The authors identify sorting as crucial because Transformers need positional encodings’ relative order to correspond to token positions.The unsorted comparison is described as a sanity check because the authors do not expect strong performance without order information.Source: §B.1 · pp. 8–9

B.2 Comparison to Prior Work

Randomized positional encodings outperform geometric attention overall, though geometric attention performs better on a task with inherent locality bias. Across the broader comparison, randomized variants also improve over their original positional encodings on average.

B.2 Comparison to Prior Work

  • Randomized positional encodings significantly outperform their original counterparts on average in mean test accuracy across random seeds.Table 4 reports mean and standard deviation for scores averaged over all test lengths.Source: §B.2 · p. 9, Table 4 · p. 10
  • Randomized learned encodings achieve non-trivial to strong length generalization on most tasks, unlike learned encodings that fail to beat random accuracy.The randomized version considers all possible embedding columns during training; standard learned encodings leave positions beyond training length unlearned.Source: §B.2 · p. 9
  • The comparison covers common positional encodings and adds a Neural Data Router variant using geometric attention and directional encoding, without gating or shared layers.The paper also compares with none, sin/cos, relative, ALiBi, RoPE, learned, and label-based encodings.Source: §B.2 · p. 9
  • 9.7% higher maximum test accuracy on average is achieved by randomized positional encodings over geometric attention, though not on every task.Geometric attention performs substantially better on MODULAR ARITHMETIC (SIMPLE), which has an inherent locality bias.Source: Table 5 · p. 10, §B.2 · p. 10

B.3 Analysis

On REVERSE STRING sequences beyond the training length, randomized relative encodings retain activation and attention patterns associated with the task, while standard relative encodings do not. The analysis also reports near-perfect accuracy for the randomized model on longer sequences compared with a substantial drop for the standard model.

B.3 Analysis

  • On 30 REVERSE STRING sequences, randomized relative encoding accuracy is 0.994 at length 150, versus 0.596 for standard relative encoding.At length 40, both models attain 1.0 average accuracy.Source: §B.3 · p. 10
  • Figure 4 plots performance curves across tasks and encodings, distinguishing in-domain lengths from unseen lengths at the training-range boundary.Sequences to the right of the dashed red line are beyond the training range.Source: Figure 4 · p. 11
  • At length 150, randomized relative activations overlap with length-40 activations across layers, while standard relative activations are out of distribution in layers 3 and 4.The comparison uses 2D PCA projections; qualitatively similar results were obtained for sin/cos and learned encodings.Source: §B.3 · p. 10, Figure 5 · p. 12
  • At length 150, randomized relative encodings retain the reversal-related X pattern in attention matrices, while it breaks down for standard relative encodings.Both encodings show the pattern at the maximum training length of 40.Source: §B.3 · p. 12, Figure 6 · p. 13
Loading 2305.16843v1…