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 investigates the extrapolation capability of transformer-based language models. The authors observed that existing positional encoding methods (e.g., sinusoidal embedding, relative positional embedding) fail to generalize to longer sequences in language modeling tasks. Therefore, they introduce a new positional encoding method called ALiBi, which adds temporal bias to the multi-head attention to penalize attention score proportional to token distances. Experimental results show that ALiBi has significantly stronger extrapolation capability compared to other positional encoding methods.