Randomized Positional Encodings Boost Length Generalization of Transformers (2023)

Abstract: 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).

## introduction transformers are emerging as the new workhorse of machine learning as they underpin many recent breakthroughs, including sequence-to-sequence modeling @xcite , image recognition @xcite , and multi-task learning @xcite . however, recent work @xcite demonstrated that transformers fail to generalize to longer sequences on seemingly simple tasks such as binary addition. thus, while certain problems can be solved without length generalization, algorithmic reasoning generally requires this ability, similar to many real-world settings such as online or continual learning. while the transformer's attention mechanism can recognize complex relationships amongst to-figure 1 : test-time evaluation with longer inputs. the standard positional encoding vector has values larger than those observed during training. our approach avoids this problem by assigning a random (ordered) positional encoding vector using the full range of possible test positions to each training example. ## randomized positional encodings unlike rnns @xcite , which are unrolled over tokens one step at a time, transformers process large chunks of the input sequence in parallel via global attention @xcite . as a result, transformers do not need to "remember" previous tokens, but they do have to break the permutation-invariance of the attention mechanism. to that end, the embeddings of the input sequence are generally augmented with positional encodings. for example, the vanilla transformer adds the following positional encodings to the embedded input sequence before passing it to the attention layers: while positional encodings generally succeed at inducing the required positional information for sequences of fixed length, they are one of the main failure modes preventing length generalization. concretely, for a transformer with standard positional encodings trained on a curriculum of sequences of maximum length n , test sequences of length @xmath0, with the shift getting increasingly large as m grows. to address this, we propose a randomized encoding scheme, which relies only on order information, and can be expected to generalize up to sequences of length m , where @xmath1, with a configurable hyperparameter l. ## experimental evaluation problem setup we closely follow the experiment setup of @xcite and evaluate our method on a wide range of algorithmic reasoning tasks such as modular arithmetic, reversing/duplicating a string, binary addition/multiplication, and bucket sort. the tasks are derived from formal language recognition and thus grouped according to the chomsky hierarchy @xcite , which partitions languages into regular (r), context-free, context-sensitive (cs), and recursively enumerable. regular tasks can be solved by a finite-state automaton (fsa), deterministic context-free (dcf) tasks can be solved by an fsa with access to a deterministic stack, and we consider the encoder-only model of the original seq-to-seq transformer @xcite , as used in popular pre-trained language models such as bert @xcite or gopher @xcite . thus, for tasks that require a multitoken output sequence y (e.g., duplicating a string), we pad the input sequence with |y| empty tokens and compute the entire transformer output from the padded sequence (i.e., we do not use autoregressive sampling). we train the model on sequences of length sampled uniformly from u(1, n ), with @xmath2, and evaluate it on sequences of length {n + 1, . . . , m }, with @xmath3. we set the maximum position @xmath4.1). we report the accuracy averaged over all unseen sequence lengths, i.e., n + 1, . . . , m , for the best-performing model out of 10 different parameter initialization seeds and three learning rates 1 × 10 -4 , 3 × 10 -4 , 5 × 10 -4 . we use the same hyperparameters as @xcite and provide the full experiment setup in appendix a. we make our code publicly available at https://github.com/deepmind/ randomized_positional_encodings . comparison to prior work we compare our method to a wide range of positional encodings: none, sin / cos @xcite , relative @xcite , alibi (press et al., 2022), rope @xcite , learned @xcite , and label-based @xcite . note that the label encodings proposed by @xcite are equivalent to randomized learned positional encodings and thus subsumed by our method. we instantiate our randomized positional encoding scheme with all the above encodings and show the average test accuracy in ## conclusion we introduced a novel family of positional encodings that significantly improves the length generalization capabilities of transformers. our positional encodings are based on the insight that conventional positional encodings will be out-ofdistribution when increasing the sequence length. thus, to overcome this issue, we randomly sample our encodings from a wider range than the lengths seen at test time while keeping the order. our largescale empirical evaluation demonstrates that our method significantly outperforms prior work in terms of length generalization while offering superior computational performance over the naive approach of training the model on longer sequences.