Extending Input Contexts of Language Models through Training on Segmented Sequences (2024)

Abstract: Effectively training language models on longinputs poses many technical challenges. As acost consideration, languages models are pre-trained on a fixed sequence length before beingadapted to longer sequences. We explore var-ious methods for adapting models to longerinputs by training on segmented sequences andan interpolation-based method for extendingabsolute positional embeddings. We developa training procedure to extend the input con-text size of pretrained models with no architec-tural changes and no additional memory coststhan training on the original input lengths. Bysub-sampling segments from long inputs whilemaintaining their original position the model isable to learn new positional interactions. Ourmethod benefits both models trained with abso-lute positional embeddings, by extending theirinput contexts, as well as popular relative posi-tional embedding methods showing a reducedperplexity on sequences longer than they weretrained on. We demonstrate our method canextend input contexts by a factor of 4× whileimproving perplexity.

## introduction transformer-based models @xcite capture sequence information through positional embeddings (pe). there are two types of pes: absolute and relative. absolute positional embeddings (ape) learn a separate embedding for each position in a sequence; these embeddings are added to the input of the first layer. relative positional embeddings (rpe) encode the relative distance between positions, often by weighting attention score of positions further away less. the ability for models to process long sequences efficiently is of growing importance as models become more capable. increased input context allows for more complex in-context learning examples @xcite . additionally, they allow for question answering and summarization over scientific papers and patents @xcite @xcite . due to rpe's positional information only being a function of relative distance these methods can be applied to any input sequence length. in practice, popular rpe methods fail to generalize to sequences longer than they were trained on. furthermore, self-attention's memory cost is quadratic meaning training on long sequences becomes prohibitively expensive as the sequence length grows. in this work, we study the problem of extending the input context of pre-trained decoder-only transformer-based models, considering those that use either absolute or relative positional embeddings. we show that an interpolation-based approach allows ape models to extrapolate to sequence lengths longer then they were trained onmatching or outperforming the extrapolation ability of rpe methods like alibi @xcite and rope @xcite . to further improve the ability of these models to take advantage of the longer input context, we present resource-efficient methods that continuously pre-train ape-and rpebased models on carefully sampled segmented subsequences of long sequences. doing so simulates training on long sequences while remaining within a fixed input length. this allows the models to efficiently learn the embeddings of the newly created absolute positions or the relative embeddings associated with the longer pairwise distances. we experiment with models trained with apes, rope, and alibi to verify our method improves the extrapolation performance independent of the choice of positional embeddings. results show that interpolating the embedding matrix of absolute positional embeddings without any additional training allows for extrapolation to sequences 5× the original input context. furthermore, our segment-based methods are able to increase the extrapolation ability of all positional embedding approaches. when applied to apes this method achieves 87% the performance of training on sequences twice as long at no extra memory footprint. the paper is organized as follows: first, we conduct a review of various existing literature that motivated our approach. second, we formally define the problem of length extrapolation and propose our methods for efficiently extending a model's input context. third, we provide a detailed breakdown of our experimental setup and methodology to enable reproducibility. finally, we present our results along with a thorough discussion and analysis. ## methods there are three reasons that motivate this work. first, there exist numerous high-quality pre-trained models whose input context is limited to 1k-2k tokens. extending the input context of these models will further increase their applicability. second, even though methods that rely on relative positional embeddings can operate on input contexts that are longer than what they were trained on, their out-of-the-box extrapolation performance is not good @xcite . third, due to the quadratic complexity of self-attention and the linear compute/memory complexity of transformers w.r.t. sequence length, direct training on long input contexts is resource intensive. this limits the input context that we can directly train on. ## results we conduct our experiments and present results in such a way to answer the following questions: ## analysis our results demonstrate that segmented training is a viable approach to extend the input context size of language models. it is not immediately intuitive why, especially given that the relative positional embeddings methods are not learned. for absolute positional embeddings the reasoning is fairly straightforward. first, in section 5.1 we demonstrated interpolating the embedding matrix led to reasonable extrapolation without any training. before any training occurs the model already has some extrapolation ability. the segmented sequences allow for positions further away than the input size normally allows to interact and learn how to incorporate information. in the case of relative positional embedding methods these results are less intuitive. both rpe methods penalize the attention scores of positions as a function of their relative distance, meaning that initially there is not much attention across chunk boundaries. we hypothesize that through training on segmented sequences the model learns to at- ## conclusion in this work we proposed a simple and memory efficient approach to extend the effective input context size of models through training on sequences created by sampling segments from long documents. we demonstrated our method is robust to the choice of positional embeddings and allows models to be trained on sequences at least 4× their original input length. furthermore, our results on extending absolute positional embeddings through interpolation demonstrated they can extrapolate better than