LM -Infinite: Zero-Shot Extreme Length Generalization for Large Language Models (2024)

Abstract: Today’s large language models (LLMs) typically train on short text segments (e.g., <4K tokens) due to the quadratic complexity of their Transformer architectures. As a result, their performance suffers drastically on inputs longer than those encountered during training, substantially limiting their applications in real-world tasks involving long contexts such as encod- ing scientific articles, code repositories, or long dialogues. Through both theoretical analysis and empirical investigation, this work identifies three major factors contributing to this length generalization failure. Our theoretical analysis reveals that commonly used techniques like using a sliding-window attention pattern or relative positional encodings are inadequate to address them. Answering these challenges, we propose LM-Infinite, a simple and effective method for enhancing LLMs’ capabilities of handling long contexts. LM-Infinite is highly flexible and can be used with most modern LLMs off-the-shelf. Without any parameter updates, it allows LLMs pre-trained with 2K or 4K-long segments to generalize to up to 200M length inputs while retaining perplexity. It also improves performance on downstream tasks such as Passkey Retrieval and Qasper in the zero-shot setting. LM-Infinite brings substantial efficiency improvements: it achieves 2.7× decoding speed up and 7.5× memory saving over the original model. Our code will be publicly available upon publication.

## introduction large language models (llms) have recently advanced the state-of-the-art across various natural language processing tasks. they typically train on text segments of fewer than 4k tokens @xcite , primarily due to the computational overhead quadratic in the input lengths of their transformer architectures. as a result, they face challenges in generalization to inputs that are excessively longer than what they are trained on and suffer substantial deterioration in their performance @xcite . this limits their applicability in tasks that require long-range contexts, such as encoding scientific articles, source code repository generation, or long-context dialogues. extensive efforts have been devoted to addressing this length generalization challenge. relative positional encodings such as rope @xcite and alibi @xcite have been widely adopted by state-of-the-art llms, which calculate attention based on inter-token distance instead of absolute positions, hoping to avoid model failures due to unseen absolute position embeddings. moreover, although applying a sliding-window attention pattern on the transformer architecture can reduce the memory overhead @xcite @xcite , they are not directly applicable to pre-trained models for length generalization without further training. through both theoretical analysis and empirical investigation, §3 pinpoints three primary factors underlying the length generalization failures: (1) the challenge of handling unseen distances among tokens, (2) the difficulty of attending to unseen numbers of tokens, and (3) implicitly encoded absolute positional information in initial tokens. these challenges can make llms' computational features, such as attention logits and hidden vectors, deviate from the training distribution, leading to failures of length generalization. existing techniques fall short of addressing these underlying issues. answering these challenges, we propose lm-infinite, a simple and effective method to enhance transformer llms' capabilities for modeling long contextswithout parameter updates. lm-infinite consists of two major components designed to alleviate the three factors above. (1) a λ-shaped attention mask and (2) a ceiling on attention dis-tances. the former forces the model to attend to only the beginning of the sequence and the most recent tokens within a pre-defined window, ignoring the rest. the latter component caps the relative distance values to the maximum the model has seen during training. it can also optionally re-introduce top-k tokens in the middle to achieve better performance in some downstream tasks. lm-infinite is highly flexible and applies to any off-the-shelf llms that use relative positional encoding and does not require any finetuning. our experiments thoroughly evaluate lm-infinite on a variety of tasks and llms. on arxiv (academic papers) and openwebtext2 (reddit posts) lm-infinite facilitates zero-shot generalization for a wide range of llms to texts up to 200m tokens, retaining the language modeling perplexity and generation quality. without any parameter updates, lm-infinite improves scores compared with the original model and truncation baselines on downstream tasks including passkey retrieval (mohtashami and jaggi, 2023) and qasper @xcite , which are two established benchmarks for long-context evaluation. we observe a 37.2% gain on passkey retrieval and a 1.2% gain on qasper in the zero-shot setting. lm-infinite also brings substantial efficiency improvements: it achieves 2.7× decoding speed up and 7.5× gpu memory saving over the original llms. ## why do transformer llms fail to generalize to long contexts? through a series of theoretical and experimental investigations, this section aims to identify the potential causes underlying current llms' failure in length generalization. our discussion assumes transformer-based llms that use relative positional encodings, as this design is widely adopted in today's llms @xcite . we use llama-2 @xcite , which is pre-trained with 4k-length segments, for investigation. on sequences longer than the training length, we will show that the unseen inter-token distances, the increasing number of attended tokens, and the implicitly encoded position information of the starting tokens can all make certain computational features out of the training distribution. as deep models can be sensitive to input distribution shifts, these factors need to be handled for llms to generalize to unseen lengths. factor 1: challenges in handling unseen distances among tokens with relative positional encoding, the impact of positions on the attention weight between two tokens depends solely on their relative distance. as the sequence length grows exceedingly long, some distance values will surpass those seen during training. we make the following informal theoretical claim: theorem 1. (informal) for an attention mechanism using relative positional encoding, the attention logits must explode to infinities to differentiate previously unseen distances apart as the sequence length increases. the formal theorem and its proof can be found in appendix c. we also empirically verify this on llama-2 on the arxiv dataset truncated down to 8k length. we extract the attention logits of all attention heads and their maximum attention logits on different sequence lengths in figure 1(a) . it shows the average and variance among attention heads. we see that the attention logits increase to substantially larger values when the sequence length exceeds the training length of 4k. to mitigate this issue, we conjecture that it may help to cap the relative distance values to the maximum that the model has seen during training (i.e., a distance ceiling). however, as we will see from the proposition below, addressing logit explosion leads to another challenge. factor 2: attending to unseen numbers of tokens on longer sequences, tokens at later positions must distribute attention weights across a larger context. we then make the following claim that, if attention logits are bounded but the number of tokens to attend to is not limited, it can cause the attention entropy to increase beyond the training range: proposition 1. if the attention logits are bounded, as the sequence becomes longer, the attention entropy grows to infinity. a formal statement as well as the proof can be found in appendix d. this conclusion is further empirically verified by plotting the attention entropy against context lengths in figure 1(b) . the curve shows an ever-increasing attention entropy. this trend, although increasing logarithmically, is still harming the llms' performance, as we will illustrate in the ablation study in §5.2 and figure 5 . this suggests that we should bound the attention context size to ensure that the attention entropy stays in seen ranges during pre-training and avoid degenerated outputs. a simple windowed attention, where each token only attends to the nearest tokens within a distance, might handle factors 1 and 2. this is similar to the block-diagonal attention mask used in xpos @xcite and longformer @xcite . however, as we will show in the next paragraph, this introduces another factor that can also fail llms. factor 3: starting tokens occupy a distinct feature space perhaps counter-intuitively: observation 1. even without explicit absolute positional embeddings, attention outputs of the first few tokens can occupy a distinct representational space compared to other positions. therefore, when passed to later layers, these starting tokens have distinct value vectors from their lower-layer outputs. this follows from theorem 1 in @xcite , which proves that the absolute positions can be implicitly encoded in the outputs of tokens of a single attention layer, even without positional encodings. in their construction, the starting tokens' signals are the strongest and easiest to distinguish from other tokens. as relative positional encoding is strictly more expressive than no positional encoding setting in @xcite (e.g., by letting all distances have the same attention function), the same conclusion applies to relative positional encoding as well. as an empirical verification, we take the hidden states output by the second layer of llama-2 and plot a principal component analysis (pca) projection into a 2-d plane in figure 1(c ). more figures for other layers can be found in §e. the dots correspond to the first 4096 tokens in 32 sequences, with blue ones corresponding to the initial tokens and red tokens being the tail ones. the two blue regions at the upper center and lower right correspond to the initial highly concentrated tokens (whose positions are around 0∼25) and are very far from later tokens. the lower-left region contains the remaining overlapping tokens in a sequence (zoomed in to another box). the plot shows that the vector representations of the initial tokens concentrate on regions in the feature space that are distinct from the remaining tokens. this fresh finding reveals a fundamental flaw of the sliding-window attention pattern, which restricts the attention to the most recent tokens within a predefined window size, a widely adopted baseline recently @xcite @xcite . as attention is essentially a weighted average over the value vectors, sliding-window attention discards the initial tokens, keeping the attention output from reaching the regions that value vectors of the initial tokens occupy. this enforces the model to handle a different region during the computation, introducing additional generalization challenges. as a straightforward solution to this issue, the initial tokens need to be kept in the attention computation. ## our proposal: lm-infinite inspired by the analyses and take-away messages in the previous section, we propose lm-infinite to achieve zero-shot length generalization for llms. an overview of lm-infinite is illustrated in figure 2(a ). this simple solution consists of two components: a λ-shaped attention mask and a distance ceiling. besides, re-introducing the middle top-k tokens is optional for enhanced downstream performance. λ-shaped attention mask it contains two attention spans: the starting one allows each token to attend to the first n starting tokens if they come before the current one; the ending one allows each token to attend to most recent l pretrain tokens. l pretrain is the maximum length during training. other tokens are ignored. in ablation studies in §a, we find that choosing n @xmath0, 100] generally achieves equally good performance. note that n @xmath1.e., attending only to the most recent tokens) substantially hurts the performance. this resolves factors 2 and 3 in §3 by both limiting the number of tokens under attention and ensuring the starting few tokens are attended. distance ceiling lm-infinite further bounds the "effective distance" to l pretrain . this only affects the starting few tokens when attended by tokens at later positions. specifically, in relative positional encoding, the original attention logit is w(q, k, d), where d is the distance between two tokens. now we modify it as w(q, k, d ′ )) where figure 2 (a) shows an illustrative example where the distance ceiling is l @xmath2. this addresses factor 1 in §3 by bounding the distance value in attention calculation. optionally attending to top-k tokens in the middle lm-infinite can optionally attend to k tokens in the middle with the largest attention logits. this is particularly useful in downstream tasks where information in the middle tokens matters ( §5.3). here the k tokens are selected independently for each attention head in layers higher than h-th layer, and have an attention distance of @xmath3. these hyperparameters are selected based on a held-out passkey retrieval validation set, where we set @xmath4, with more details in appendix a. our selection of k and h does not depend on task-specific tuning, and in our experiments, we apply this same set of hyperparameters in all other downstream tasks and achieve consistent improvements. these intermediate tokens do not hurt performance. rather, in the evaluation of downstream tasks in §5.3, intermediate tokens are more useful and selectively attending to top-k tokens brings substantial performance improvements with little impact on the efficiency. for llm generation and inference, however, we find the intermediate tokens unnecessary to attend to for lm-infinite to achieve good perplexity or generation quality. lm-infinite's λ-shaped mask is conceptually similar to the attention patterns derived from heuristics @xcite @xcite . however, we formally show in §3 factor 3 that these previous approaches theoretically cannot generalize to unseen lengths but require parameter updates. this inherent limitation motivates the other two components in lm-infinite to achieve zero-shot length generalization. implementation details lm-infinite is applicable in all transformer models with relative positional encoding. one only needs to replace the attention function in each transformer layer with lm-infinite without any parameter updates. the λ-shaped attention mask is relatively straightforward to implement. in rope, attention logits in the ending attention span follow the original calculation. in the starting attention span (excluding its overlap with the ending span), we keep all k vectors unrotated and rotate all q vectors to a fixed distance l pretrain . then the logits in two spans can be composed. augmenting alibi with lm-infinite is also straightforward. we simply clip the offset matrix with a minimum value of -|ml pretrain | and apply the λ-shaped attention mask. discussion. in figure 2 (b), we show a conceptual model of how relative positional encoding functions. this conceptual model reflects the design choices of lm-infinite. in this conceptual model, a long context can be roughly partitioned into 3 parts: the starting tokens encode strong absolute position information (factor 3). as explained in §3, they are essential to attention to because their features occupy a distinct region in the feature space. as attention is essentially a weighted average over v i vectors, without the starting few tokens, the attention output can not reach regions that v i vectors of the initial tokens occupy. the rear tokens provides primarily their relative positions to the final tokens. their importance probably arises from the "recency bias" @xcite learned by llms during pre-training. the middle tokens encode less position-sensitive information. as analyzed in factor 2, including too many intermediate tokens does more harm than good to length generalization. ## evaluation we evaluate lm-infinite with llama-7b (touvron et al., 2023a), llama-2-7b @xcite , mpt-7b (team, 2023), and gpt-j-6b @xcite . llama-7b and gpt-j-6b are pre-trained with 2k lengths and the other models are pre-trained with 4k lengths. llama, llama-2, and gpt-j use rope encoding, and mpt-7b uses alibi encoding. mpt-7b-storywriter (fine-tuned on long sequences) is used as one of the baselines. ## conclusions and future work this work proposes a zero-shot length generalization method for various off-the-shelf llms without parameter updates. through theoretical analysis and empirical investigation, this work identifies three major factors contributing to this length generalization failure. our theoretical analysis further reveals why truncating the attention window and relative positional encodings are inadequate to address them. our solution, lm-infinite, is a simple and effective method for enhancing llms' capabilities of handling long contexts. it allows llms pre-trained with 2k or 4k-long segments to generalize to up to 200m length inputs while retaining perplexity. it also improves performance on downstream tasks such as passkey retrieval and qasper in the zero-shot setting. it brings substantial efficiency improvements: a 2.7× decoding speed up and a 7.5× memory saving over the original model. lm-infinite's computational efficiency and ease of use allow researchers without enormous computational resources to use llms on long sequences. future work can investigate if these techniques allow for more efficient and effective llm pre-training and fine-tuning. another direction is to apply lm-infinite to applications such as long reasoning, long-dialogue, retrieval-augmented generation, or long literature generation.