A Length-Extrapolatable Transformer (2023)

Abstract: Position modeling plays a critical role in Transformers. In this paper, we focus on length extrapolation, i.e., training on short texts while evaluating longer sequences. We define attention resolution as an indicator of extrapolation. Then we propose two designs to improve the above metric of Transformers. Specifically, we introduce a relative position embedding to explicitly maximize attention resolution. Moreover, we use blockwise causal attention during inference for better resolution. We evaluate different Transformer variants with language modeling. Experimental results show that our model achieves strong performance in both interpolation and extrapolation settings. The code will be available at https://aka.ms/LeX-Transformer .

## introduction transformer @xcite has shown strong performance in nlp and become a de-facto backbone @xcite @xcite . however, most of them have a crucial shortcoming: they can only deal with the in-distribution size of inputs. figure 1 shows that the perplexity of previous transformers increases rapidly when the input sequence is getting longer. it is usually infeasible to train a model with all possible input lengths. therefore, a length-extrapolatable transformer is essential for wider usage. in sequence modeling, position information plays a crucial role in building the correct representation and understanding of the latent meaning. for recurrent neural networks such as lstm @xcite , the calculation is done along the sequence order in o(n) time. however, the parallel attention module makes it hard to encode position effectively. first, @xcite proposes absolute sinusoidal position embedding, and @xcite adjusts it to a learnable one. the absolute design is computation-efficient, but not comparable with subsequent relative ones @xcite @xcite . among many relative position embeddings, rope @xcite shows better performance and is used to many plms such as palm @xcite . however, it can't deal with sequences with exceeding length. alibi @xcite mitigates the extrapolation problem but sacrifices the general performance. since different strategies concentrate on some part of the position feature, it is essential to build a comprehensive view and guide the transformer's design systematically. first, a transformer should be sensitive to order. otherwise, it will degenerate into a bag-of-word model which confuses the whole meaning. then, position translation can't hurt the representation, especially, when a prefix is added to the target sentence, the representation should stay the same with an attention mask on the prefix. after that, a good sequence model needs to deal with any input length. as illustrated before, the length problem is not universal but special for transformer. especially, when a transformer is pre-trained under a maximal length, it is not affordable to re-train for applying to tasks with longer sequences. finally, when a transformer satisfies the principles above, we evaluate its performance, which requires thorough experiments and empirical analysis. considering all the properties above, we propose extrapolatable position embedding (xpos), which is a universal-good design for transformers. based on rope's design, we propose attention resolution as a metric to measure position monotonicity. then, we generalize its mathematical form, where an exponential decay is added to the rotation matrix. xpos preserves the advantage of rope, and behaves stably at long-term dependency. besides, inspired by sparse attention methods @xcite @xcite , we choose blockwise causal attention to increase attention resolution, which improves the performance of length extrapolation for language modeling. we train different transformers from scratch. we evaluate models on pg22 and qmsum @xcite with various input lengths. on the interpolation experiments, lex transformer reaches minimal perplexity. in the extrapolation experiments, our methods can continue decreasing the perplexity while other methods either can't extrapolate (i.e., perplexity increases) when the input length is very long. figure 1 shows clearly that lex transformer has an opposite tendency compared with others. we summarize our contributions as follows: ## a length-extrapolatable transformer we define attention resolution as the indicator of the transformer's capability on encoding position in section 3.1. then we propose two ways to maximize the resolution metric, i.e., improve the length interpolation and extrapolation of transformers. first, we introduce a relative position encoding method (section 3.2) to explicitly maximize attention resolution. second, we propose to use blockwise causal masking (section 3.3) during extrapolation inference for improved resolution. in the following section, we denote d as the hidden dimension and l as the input length. for each attention layer, query, key, and value are calculated by input ## conclusion we propose lex transformer to accurately capture position information for transformers. we define attention resolution as the metric of length extrapolation and design a solution to improve the modeling. extensive experiments on language modeling show that our method achieves lower perplexity on longer sequences while training on short texts. the simplicity also makes the method a go-to augmentation for transformer-based language models. in addition, attention resolution provides a more principled view for position modeling, which sheds light on future architecture design.