A Length-Extrapolatable Transformer

Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers, pages 14590–14604 July 9-14, 2023 ©2023 Association for Computational Linguistics A Length-Extrapolatable Transformer Yutao Sun1∗, Li Dong2, Barun Patra2, Shuming Ma2, Shaohan Huang2 Alon Benhaim2, Vishrav Chaudhary2, Xia Song2, Furu Wei2 Tsinghua University1 Microsoft2 https://github.com/microsoft/torchscale Abstract Position modeling plays a critical role in Trans- formers. In this paper, we focus on length ex- trapolation, i.e., training on short texts while evaluating longer sequences. We define at- tention resolution as an indicator of extrapo- lation. Then we propose two designs to im- prove the above metric of Transformers. Specif- ically, we introduce a relative position embed- ding to explicitly maximize attention resolu- tion. Moreover, we use blockwise causal at- tention during inference for better efficiency. The proposed architecture is named Length- Extrapolatable (LEX) Transformer. We evalu- ate different Transformer variants on language modeling. Experimental results show that our model achieves better performance in both in- terpolation and extrapolation settings. The code will be available at https://aka.ms/ LeX-Transformer. 1 Introduction Transformer (Vaswani et al., 2017) has shown strong performance in NLP and become a de-facto backbone (Dosovitskiy et al., 2020; Radford et al., 2021; Wang et al., 2022). 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 Transform- ers 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 repre- sentation and understanding of the latent mean- ing. For Recurrent Neural Networks such as LSTM (Hochreiter and Schmidhuber, 1997), the calculation is done along the sequence order in O(N) time. However, the parallel attention module ∗Work done during internship at Microsoft Research. Figure 1: The perplexity of different Transformer de- signs with various input lengths. makes it hard to encode position effectively. First, Vaswani et al. (2017) proposes absolute sinusoidal position embedding, and Devlin et al. (2019) ad- justs it to a learnable one. The absolute design is computation-efficient, but not comparable with sub- sequent relative ones (Shaw et al., 2018; Su et al., 2021; Press et al., 2021). Among many relative po- sition embeddings, ROPE (Su et al., 2021) shows better performance and is used to many PLMs such as PaLM (Chowdhery et al., 2022). However, it can’t deal with sequences with exceeding length. Alibi (Press et al., 2021) 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 degener- ate 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 …

Content truncated to 500 words