Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation

Since the introduction of the transformer model by Vaswani et al . ( 2017 ) , a fundamental question has yet to be answered : how does a model achieve extrapolation at inference time for sequences that are longer than it saw during training ? We first show that extrapolation can be enabled by simply changing the position representation method , though we find that current methods do not allow for efficient extrapolation . We therefore introduce a simpler and more efficient position method , Attention with Linear Biases ( ALiBi ) . ALiBi does not add positional embeddings to word embeddings ; instead , it biases query-key attention scores with a penalty that is proportional to their distance . We show that this method trains a 1.3 billion parameter model on input sequences of length 1024 that extrapolates to input sequences of length 2048 , achieving the same perplexity as a sinusoidal position embedding model trained on inputs of length 2048 but training 11 % faster and using 11 % less memory . ALiBi ’ s inductive bias towards recency also leads it to outperform multiple strong position methods on the WikiText-103 benchmark . 1 1 INTRODUCTION . When constructing a transformer-based language model , a major design decision is the length of training sequences , denoted L herein , which has to date been equivalent to the length of inference sequences . More context , achieved by larger L , improves predictions at inference time . But longer sequences are more expensive to train on.2 Before transformers , RNN language models were trained on shorter-L sequences and assumed to generalize to longer contexts at inference time ( Mikolov et al. , 2010 ; Mikolov & Zweig , 2012 ; Zaremba et al. , 2014 ) . Vaswani et al . ( 2017 ) , introducing the transformer , speculated that it “ may [ ... ] extrapolate to sequence lengths longer than the ones encountered during training. ” We define extrapolation as a model ’ s ability to continue performing well as the number of input tokens during validation increases beyond the number of tokens on which the the model was trained . We find that transformer language models ( LMs ) that use sinusoidal position embeddings have very weak extrapolation abilities ; see Figure 1 . We demonstrate that this failure to extrapolate is caused by the position embedding method . As shown in Figure 1 , recent alternatives to the original sinusoidal position method ( Su et al. , 2021 ; Raffel et al. , 2020 ) have improved extrapolation . However , the better of these , the T5 bias , is considerably slower than the sinusoidal approach and uses extra memory and parameters ( Figure 2 ) . We therefore introduce Attention with Linear Biases ( ALiBi ) to facilitate efficient extrapolation . ALiBi negatively biases attention scores with a linearly decreasing penalty proportional to the distance between the relevant key and query . Our simple approach eliminates position embeddings . 1Code & models : https : //github.com/ofirpress/attention_with_linear_biases 2Figure 7 in the appendix plots training speed , in words per second , against L. Compared to a sinusoidal model trained on the same input length , our method requires no additional runtime or parameters and incurs a negligible ( 0–0.7 % ) memory increase . ALiBi can be implemented by changing only a few lines of existing transformer code . Using ALiBi , a transformer LM can be trained on short-L sequences and therefore at much lower cost , and it can still be reliably applied to long sequences at runtime . For example , a 1.3 billion parameter LM trained on L = 1024 tokens with ALiBi achieves the same perplexity as a sinusoidal model trained on L = 2048 when both are tested on sequences of 2048 tokens , even though our model is 11 % faster and uses 11 % less memory . Though performance peaks at around two times the number of tokens that the model was trained on , ALiBi maintains strong performance even on sequences of length 10,000 . In recently explored settings where NLP training examples are given as context to an LM ( Brown et al. , 2020 ) , our approach will allow exposure to more examples . Additionally , it enables generation of longer outputs . 2 CURRENT APPROACHES DO NOT EXTRAPOLATE EFFICIENTLY . We show for the first time that the sinusoidal position method , which technically should be able to extrapolate , in practice has very limited extrapolation capabilities . Though the rotary position method improves over the sinusoidal one , it still does not achieve satisfying results . Holding everything else constant , we are the first to observe that the T5 bias position method leads to better extrapolates than either of these , and so we conclude that extrapolation ability depends heavily on the position embedding . Unfortunately , the T5 bias is computationally costly ( Figure 2 ) . 2.1 BACKGROUND AND EXPERIMENTAL SETUP . A transformer LM receives a list of tokens and outputs a probability distribution representing its prediction for the next token . We call the input list the current input subsequence since the inputs to language models are typically subsequences from ( much longer ) training or evaluation sequences . During both training and perplexity evaluation ( i.e. , scoring a fixed sequence ) , many predictions can be calculated at once ; this is done using a “ causal mask ” that ensures each position ’ s prediction is influenced only by tokens to its left . Let L be the length of each input subsequence during training ; it includes L predictions , which on average have access to L+12 tokens of ( left ) context . To explore a model ’ s extrapolation abilities , we are interested in cases where sequences of length Lvalid > L are considered at evaluation time . When L differs between inference and training , we use L to refer to the length of subsequences during training and Lvalid to refer to their length at validation . Nonoverlapping Inference To train on or evaluate a sequence longer than L tokens , it is typical to segment the sequence into L-length subsequences and train on or evaluate them independently . Unless otherwise stated , we use nonoverlapping inference to report perplexity scores . Extrapolation During Inference Formally , the functions that define a transformer layer are agnostic to input length ; 3 they map from some arbitrary , unfixed number of input vectors to the same number of output vectors . When transformers are applied to data that is inherently sequential , like text , positional information is injected into the inputs in various ways . Vaswani et al . ( 2017 ) discussed two options for embedding positions into vectors to be added to word embeddings : learning embeddings for specific positions and unlearned sinusoidal embeddings . They observed similar performance between these two but preferred the sinusoidal approach , which they argued might extrapolate to longer input sequences during inference . We find that this model can not extrapolate to more than a few dozen tokens beyond L.4 Experiment Setup We first test the extrapolation abilities of various position methods on the WikiText-103 corpus ( Merity et al. , 2016 ) using the transformer language model of Baevski & Auli ( 2018 ) . We use this model because of its prominent role in recent language modeling developments ( Khandelwal et al. , 2020 ; Press et al. , 2021 ) . The training set is about 103 million tokens from English Wikipedia ( half a gigabyte ) . The model has 16 transformer layers of dimension 1024 , with 8 heads , and a feedforward inner dimension of 4096 . This model ties the word embedding and softmax matrices ( Press & Wolf , 2017 ; Inan et al. , 2017 ) . In our experiments , other than varying the position method and training subsequence length , we modify no other hyperparameters , including the random seed and number of training epochs ( 205 ) . 2.2 MEASURING EXTRAPOLATION . Sinusoidal Position Embeddings Sinusoidal position embeddings ( Vaswani et al . ( 2017 ) ; §3.5 ) are constant , non-learned vectors that are added to token embeddings on input to the first layer of the transformer . They are frequently used in transformer language modeling ( Baevski & Auli , 2018 ; Lewis et al. , 2021 ) and machine translation ( Vaswani et al. , 2017 ; Ott et al. , 2018 ) models . We first consider the unmodified model of Baevski & Auli ( 2018 ) , which uses sinusoidal position embeddings , and train it on L = 512 tokens ; we then run inference with it on the validation set on L + k tokens , with k ranging from 0 to 15,000 . Figure 1 ( left ) and the corresponding Table 2 ( in the appendix ) show that while the model improves perplexity up to k = 20 , performance stops improving and stays steady from k = 20 to k = 50 and then begins degrading . Similar results are obtained for a model trained with L = 1024 tokens ( Figure 1 ( right ) and Table 3 in the appendix ) . That model improves for up to Lvalid = L+ 50 tokens , after which performance declines . 3These include the embedding lookup , feedforward sublayer , and softmax layer , which act independently on vector inputs , as well as the attention sublayers , whose parameters do not depend on input length ( and which must handle variable-length inputs , e.g. , due to causal masking ) . 4The learned positional embedding approach does not have a way to encode positions greater than L ; it therefore has no ability to extrapolate . Rotary Position Embeddings The rotary method was introduced by Su et al . ( 2021 ) and has recently been popularized by the open source GPT-3 ( Brown et al. , 2020 ) implementation GPTJ ( Wang & Komatsuzaki , 2021 ) . Instead of adding sinusoidal embeddings at the bottom of the transformer , they multiply the keys and queries of every attention layer by sinusoidal embeddings . Unlike the sinusoidal or learned positional embedding approach , the rotary method injects position information into the model at every layer , not just at the initial one . In addition , it adds no position information to the values of the self-attention sublayer . The output of a self-attention sublayer is a linearly transformed , weighted sum of the input value vectors ; therefore , by not inserting position information into the values , the outputs of each transformer-layer contain no explicit position information . We suspect that this segregation of position information may be beneficial for extrapolation , and we draw inspiration from it in the design of our method ( §3 ) . We apply the rotary position embedding method to our Baevski & Auli baseline.5 The perplexity results ( Figure 1 and Appendix Tables 2 and 3 ) are better than the sinusoidal approach : the model with L = 512 ( L = 1024 ) improves perplexity with up to k = 200 ( k = 100 ) more tokens than it saw during training , but this comes at the cost of slower training and inference ( Figure 2 ) . T5 Bias Though most models use trained or sinusoidal position embeddings , the T5 model of Raffel et al . ( 2020 ) uses a relative position method ( Shaw et al. , 2018 ; Huang et al. , 2019 ) that adds no position information to word embeddings ( as in the previous method ) . Instead , it modifies the way attention values are computed . We refer to this as the “ T5 bias ” method.6 To compute attention values in the unmodified transformer , we compute the dot product of every query with every relevant key and then softmax these attention values . In this method , we compute the attention values as before , but then we add a learned , shared bias to each query-key score that is dependent on just the distance between the query and key . Therefore , all query-key scores where the query and key distance are zero ( i.e. , the query and key represent the same token ) get a specific learned bias , all scores where the query and key are one word away get a different learned bias , and so on , up to a certain point , from where multiple different distances share the same learned bias ( which might be beneficial for extrapolation ) . As in the rotary method , the T5 bias injects position information into the model at every layer and integrates no explicit position information into the self-attention value vectors . Raffel et al . ( 2020 ) propose that the T5 bias may allow extrapolation , but they did not report experiments testing this . Here , we show that the T5 bias does allow language models to extrapolate . We do this by again modifying the Baevski & Auli model , this time to insert the T5 bias into it.7 As Figure 1 shows , the T5 bias improves perplexity with longer sequences than the ones it was trained on , i.e. , k = 600 ( k = 800 ) extra tokens for a model trained on L = 512 ( L = 1024 ) input tokens . Unfortunately , this impressive performance comes at a cost : training is at least twice as slow as with the sinusoidal model . Therefore , this model ’ s extrapolation ability provides no efficiency advantage . For example , to do inference on 1024 tokens , we could either train the sinusoidal model with L = 1024 or train the T5 bias model on L = 512 tokens and extrapolate to 1024 for inference . However , the L = 1024 sinusoidal model runs at 28.5k words per second ( WPS ) , while the L = 512 T5 bias model runs at 14.4k WPS ( Appendix Table 1 ) , so there is no speedup when training on shorter sequences with this method.8 5Our rotary method implementation is based on the code in https : //github.com/JunnYu/ RoFormer_pytorch , which is linked to from the official repository of Su et al . ( 2021 ) : ( https : //github.com/ZhuiyiTechnology/roformer ) . After we finished running our experiments with the rotary method , we were informed that the runtime of the code linked above could be optimized , making it only 2 % slower than the sinusoidal approach . This optimization would not change extrapolation performance . 6This method is similar to the one used in Parikh et al . ( 2016 , Equation 7 ) . 7Our T5 bias implementation is based on the one used in HuggingFace Transformers ( Wolf et al. , 2020 ) , which in turn is based on the official Mesh Tensorflow T5 code . 8Narang et al . ( 2021 ) benchmarked the T5 bias as being just 8.7 % slower than the sinusoidal approach ; thus , while always incurring a runtime penalty , this method ’ s runtime could be faster depending on the choice of hardware and software frameworks used . Narang et al . used the Tensorflow T5 library running on TPUs , while we used the PyTorch Fairseq library running on GPUs .

Summary: This paper studies input length extrapolation for Transformer language models; i.e., how Transformer LMs perform on test sequences that are longer than training sequences. The paper finds that how positions are encoded plays a crucial role for input length extrapolation. Models with sinusoidal and rotary position embeddings do not extrapolate well, while T5’s position-dependent attention mechanism (dubbed T5 bias) enables better extrapolation. The paper then proposes ALiBi, another attention mechanism that also allows extrapolation while being computationally more efficient than T5 bias. These results are empirically confirmed on two datasets.