Shortformer: Better Language Modeling using Shorter Inputs (2021)

Abstract: Increasing the input length has been a driver of progress in language modeling with transformers. We identify conditions where shorter inputs are not harmful, and achieve perplexity and efficiency improvements through two new methods that decrease input length. First, we show that initially training a model on short subsequences before moving on to longer ones both reduces overall training time and, surprisingly, substantially improves perplexity. Second, we show how to improve the efficiency of recurrence methods in transformers, which let models condition on previously processed tokens when generating sequences that exceed the maximal length the transformer can handle at once. Existing methods require computationally expensive relative position embeddings; we introduce a simple alternative of adding absolute position embeddings to queries and keys instead of to word embeddings, which efficiently produces superior results. We show that these recurrent models also benefit from short input lengths. Combining these techniques speeds up training by a factor of 1.65, reduces memory usage, and substantially improves perplexity on WikiText-103, without adding any parameters.

## introduction scaling up transformer @xcite language models @xcite @xcite has been an important driver of progress in nlp. language models require data to be segmented into subsequences for both training and inference: memory constraints limit a language model to handling at most a few thousand tokens at once, while many training and evaluation datasets are much longer. recent work focuses on increasing the length of input subsequences, which determines the maximum number of tokens a model can attend to @xcite @xcite . we challenge the assumption that longer input subsequences are always better by showing that existing transformers do not always effectively use them. we then introduce new methods based on shorter input subsequences that improve runtime, memory efficiency, and perplexity. we first investigate how input subsequence length affects transformer language models ( §3). naïve evaluation-where we split a large evaluation set into multiple nonoverlapping subsequences, each evaluated independently-initially supports the commonly-held belief that models that train and do inference on longer subsequences achieve better perplexity ( however, when we evaluate each model with a sliding window @xcite , outputting one token at a time using the maximal amount of context, we find-surprisingly-that models using subsequences exceeding 1,024 tokens do not further improve performance ( we conclude that the performance gains (using naïve evaluation) of models that use longer subsequences occur not only because of their better modeling ability, but partly because they divide the evaluation set into longer subsequences. this division helps because of an issue we call the early token curse: by default, early tokens in a subsequence will have short histories to attend to. using longer subsequences means fewer tokens will suffer from the early token curse. for example, when using inputs of length 1,024, about 94% of tokens get to attend to more than 64 preceding tokens. if we use inputs of length 128, only 50% of tokens get to attend to 64 or more preceding tokens. based on this analysis, we explore how to improve models by using shorter inputs. we introduce two techniques. staged training ( §4) first, we show that initially training on shorter subsequences (before moving to longer ones) leads not only to much faster and more memory-efficient training, but it surprisingly also greatly improves perplexity, suggesting that longer inputs are harmful early in training. position-infused attention ( §5) second, we consider a natural way to avoid the early token curse during training and inference: attending to cached representations from the previously evaluated subsequence @xcite . this approach interferes with conventional absolute position embeddings in a way that forced @xcite to use relative position embeddings, which are computationally expensive. we introduce a fast, simple alternative: instead of adding absolute position embeddings to word embeddings-thereby entangling a word's content and positional information-we add them to the keys and queries in the self-attention mechanism (but not to the values). this does not increase parameter count or runtime. token representations can then be cached and reused in subsequent computations. we show that when using this method, shorter subsequence models outperform longer ones. finally, we show additive gains from combining staged training and position-infused attention (shortformer, §6), resulting in a model that trains much quicker and achieves better perplexity on wikitext-103. we also show that these results transfer to language modeling on the toronto book corpus ( §a.5, appendix). ## background and experimental setup transformer language models map a list of tokens x n-l:n-1 to a probability distribution over the next token x n . we refer to the list of tokens as the current input subsequence (whose length is l). causal masking lets us make l predictions at once, with the prediction for token i + 1 conditioned on the ith token and all previous inputs x n-l:i-1 , but not on future inputs. we define the number of tokens the model can attend to at each timestep as its effective context window. note that l is not to be confused with the (typically much greater) length of a training or evaluation dataset. during inference, language models can be used for two distinct tasks: generation and evaluation. in order to define these tasks, we first define nonoverlapping and sliding window inference. nonoverlapping inference to evaluate a string longer than l, we can evaluate each subsequence of l tokens independently. this fast approach is commonly used during training; if used, tokens in one subsequence cannot condition on those in the previous subsequence, giving rise to the early token curse discussed in §1. see figure 1(a) . ## how does context window size affect transformers? segmenting a corpus into subsequences results in different effective context windows for different timesteps depending on where they fall in a segment. subsequence length l is an upper bound on the effective context window at each timestep. when making the first prediction, the model attends only to the first input token. when making the second prediction, the model attends to the first two inputs, and so on, up to the lth timestep where the model can attend to all input tokens when making the lth prediction. ## repositioning position embeddings sliding window inference substantially improves performance by increasing the minimum effective context window size. but it is very slow. we could solve this by letting the model attend to representations of the previous subsequence during inference on the current one. in this case, the same token representations would be used in different positions since a token generated near the end of one subsequence would be cached and reused near the start of the next one. however, transformer model representations entangle positional and content information, so a cached token representation would encode an incorrect position when reused in a new position. transformerxl @xcite uses relative position embeddings to solve this problem. however, that approach is slower and uses more parameters and memory than the baseline transformer. 9 we solve this using no extra parameters, memory, or runtime. we also show that our method can use much shorter input subsequences and still achieve superior performance. transformer language models the baseline transformer lm, given a token list t of length l and a tensor p containing the first l position embeddings, produces l next-token predictions using the following procedure: ## shortformer results to assess whether the gains from staged training, pia and caching are additive, we take our best caching pia model, with subsequence length 512, and apply staged training to it, training it with a subsequence length of between 32 to 256 for the first half of training. foot_7 the best model (with initial subsequence length 128), which we call shortformer, achieves 17.47 dev. set perplexity and trains 65% faster than the baseline. since its attention matrices are of dimension 512 • 1,024 (the baseline's are 3,072 • 3,072), our model uses less memory ( §a.4, appendix). it has the same number of parameters as the baseline. figure 3 (appendix) compares our best models using each method we presented (and their combination) to the baseline. it shows that combining caching, pia and staged training (shortformer) yields the quickest training and best perplexity when using nonoverlapping evaluation. evaluation speed is similar for all of these models. finally, since it uses a cache, sliding window evaluation does not increase shortformer's performance. by training the baseline with staged training (and no pia or caching), we obtain a model (our best model from §4.2) that, with sliding window eval., obtains even better results, but that model is much slower than shortformer ( shortformer outperforms the baseline's perplexity and performs within a standard deviation of the sandwich transformer @xcite and transformerxl. it does not outperform the compressive transformer @xcite , routing transformer @xcite and knn-lm @xcite , which make orthogonal improvements that can be applied to any language model, at the price of slower decoding. combining them with our approach may yield further gains. these results are similar to those we obtain on the toronto book corpus ( §a.5 in the appendix). many recent papers have explored improving transformer efficiency by reducing the quadratic cost of self-attention, motivated by scaling to longer sequences @xcite @xcite . we instead demonstrate improved results with shorter sequences, which naturally also improve efficiency. one way to reduce transformer memory usage is to sparsify the attention matrix by letting the model attend only to a subset of nearby tokens at each timestep @xcite @xcite . training on shorter subsequence lengths is much more efficient: we use multiple, but much smaller, attention matrices. since attention uses memory and computation in a way that scales quadratically with input size, splitting the inputs into multiple subsequences each processed independently lets us use less memory and run faster. like our method, @xcite attend at each timestep to a growing number of neighbors as training progresses, but they use five stages, which we found not to be superior to our two-staged method. the adaptive attention span model of @xcite learns the maximum effective context window sizes for each head at each layer independently. like in our method, context window sizes are smaller at the start of training and lengthen as training progresses. we show that a simple approach of manually choosing two subsequence lengths is highly effective. in addition, keeping subsequence lengths equal across all heads and layers lets us save memory and runtime. ## conclusion our results challenge the conventional wisdom that longer subsequences are always better. by first training on shorter subsequences and then progressing to longer ones via staged training, we improve perplexity and reduce training time. we additionally propose position-infused attention, which enables caching and efficiently attending to previous outputs; we show that models using this method do not require large input subsequences. we finally show that these two methods can be combined to produce a speedier and more accurate model.