Venue: NAACL 2024

# LM-Infinite: Zero-Shot Extreme Length Generalization

for Large Language Models

Chi Han\({}^{1}\) , Qifan Wang\({}^{2}\), Hao Peng\({}^{1}\), Wenhan Xiong\({}^{3}\), Yu Chen\({}^{1}\)+, Heng Ji\({}^{1}\), Sinong Wang\({}^{3}\)

\({}^{1}\) University of Illinois Urbana-Champaign, \({}^{2}\) Meta, \({}^{3}\) GenAI Meta, \({}^{4}\) Anytime AI

\({}^{1}\){chihana3,haopeng,hengji@illinois.edu,

\({}^{23}\){wqfcr,xwhan,sinongwang}@meta.com,

\({}^{4}\)ychen@anytime-ai.com

Work performed as an intern in Meta GenAI.This work was done while the author was at Meta.

###### 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 encoding 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\(\times\) decoding speed up and 7.5\(\times\) memory saving over the original model. Our code will be publicly available upon publication.

## 1 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 [68, 9], 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 [68, 9]. 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 [23] and Alibi [26] 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 [18, 19, 17], they are not directly applicable to pre-trained models for length generalization without further training. Through both theoretical analysis and empirical investigation, SS3 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 contexts_without parameter updates_. LM-Infinite consists of two major components designed to alleviate the three factors above. (1) a \(\mathbf{\Lambda}\)-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 Mohathsami and Jaggi (2023) and Qasper Dasigi et al. (2021), 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\(\times\) decoding speed up and 7.5\(\times\) GPU memory saving over the original LLMs.

## 2 Background and Related Work

### Relative Positional Encodings

The traditional absolute positional encodings provide the absolute position information, usually with the help of a sequence of vectors called _position embeddings_Vaswani et al. (2017); Kenton and Toutanova (2019); Ke et al. (2020). They, however, have trouble when the model encounters unseen positions in inputs longer than the training length. Relative positional encodings aim to address the limitations of previous-generation positional encoding methods and consider the relative distances between tokens instead of the absolute positions. Examples include a learned attention logit bias in T5Raffel et al. (2020), Transformer-XL Dai et al. (2019), Skyformer Chen et al. (2021), SketchingChen et al. (2022) and Sandwich Chi et al. (2023), a fixed linear attention decay Press et al. (2021), and rotating query and key sequences based on distances such as RoPE Su et al. (2021); Li et al. (2023), CAPE Likhomanenko et al. (2021) and XPos Sun et al. (2022); Ding et al. (2023). Despite some promising empirical evidence, length generalization failures are still widely observed when directly applied to large language models Kaiockendev (2023). In what follows, we briefly discuss two widely used relative positional encoding methods. They lay out the necessary context for our onward discussion and experiments.

Rotary Position Embedding (RoPE; Su et al. (2021))It rotates the key and query vectors based on positions before computing the inner product. Specifically, each vector \({\bf x}\) (either key or query) is split into pairs of elements \(\left\{(x_{0},x_{1}),(x_{2},x_{3}),\cdots\right\}\), with each pair interpreted as a 2-dimensional vector. RoPE then rotates the vector \((x_{a},x_{a+1})\) of token \(i\) with angle \(\theta_{a,i}=i\omega_{a}\), where \(\omega_{a}\) is the rotating speed associated with dimension pair \((a,a+1)\). After rotation, the 2-D vector becomes \(\begin{pmatrix}\cos i\omega_{a}&-\sin i\omega_{a}\\ sin i\omega_{a}&\cos i\omega_{a}\end{pmatrix}\begin{pmatrix}x_{i}\\ x_{i+1}\end{pmatrix}\). They show that the inner product between rotated query \({\bf q}_{i}\) and rotated key \({\bf k}_{j}\) is solely determined by \({\bf q}_{i}\), \({\bf k}_{j}\), and their relative distance \(i-j\). We always have \(i\geq j\) due to the causal attention mask.

AlBiBress et al. (2021)It offsets all attention logits between tokens \(i,j\) by a linear term \(-m(i-j)\) and become \({\bf q}^{\top}_{i}{\bf k}_{j}-m(i-j)\). To this end, the MPT-7B codes implement an offset matrix as an additive term in attention logits.

### Efforts Towards Length Generalization

In light of generalization failures observed in LLMs, one straightforward solution is to finetune LLMs on longer text sequences Chen et al. (2023); Tworkowski et al. (2023); Tao et al. (2023); Kiyono et al. (2021); Anil et al. (2022). These approaches do not address the underlying causes of length generalization failures and require massive training resources. Other solutions propose to grant LLMs access to longer contexts without really reading them in full Zhou et al. (2023); Bueno et al. (2022); Mohathsami and Jaggi (2023); Yang et al. (2023). Augmenting LLMs with retrieval-based memories Wu et al. (2021); Guu et al. (2020); Borgeaud et al. (2022); Khandelwal et al. (2019); Kaiser et al. (2016); Yogatama et al. (2021) also make LLMs applicable to a large database. These designs, however, usually need finetuning and are not directly compatible with the existing LLMs. Our work, in contrast, facilitates zero-shot length generalization. Another similar work Ratner et al. (2023) increases con 

text length with attention patterns without further training. However, it is limited to the in-context learning setting.

## 3 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 Touvron et al. (2023b); Team (2023). We use Llama-2Touvron et al. (2023b), 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 tokensWith 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 tokensOn 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

Figure 1: We identify three factors underlying the length generalization failure in LLMs in §3. (a) Factor 1: Unseen distances between tokens cause attention logits to explode. (b) Factor 2: An unseen number of tokens can cause attention entropy to increase beyond the training range as the length increases. (c) Factor 3: Starting few tokens occupy a distinct feature region and should not be discarded. The two blue regions at the upper center and lower right correspond to the initial tokens that are highly concentrated but also very far from later tokens. The lower-left region contains the thousands of overlapping dots corresponding to the later tokens.

 

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 SS5.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 Sun et al. (2022) and Longformer Beltagy et al. (2020). 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 spacePerhaps counter-intuitively:

Observation: 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 Kazemnejad et al. (2023), 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 Kazemnejad et al. (2023) (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 SSE. 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\(\sim\)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 Beltagy et al. (2020); Ding et al. (2023); Zaheer et al. (2020). 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

Figure 2: (a) LM-Infinite is a plug-and-play solution for various LLMs, consisting of a \(\Lambda\)-shaped mask and a distance ceiling during attention. For clarity, this figure shows a toy scenario where \(L_{\text{pre-train}}\) and \(n_{\text{starting}}\) are both 2. (b) We also provide a conceptual model for understanding how relative position encoding works.

 

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.**

## 4 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 \({\bf\Lambda}\)-shaped attention mask and a distance ceiling. Besides, re-introducing the middle top-\(k\) tokens is optional for enhanced downstream performance.

\({\bf\Lambda}\)-shaped attention maskIt contains two attention spans: the _starting_ one allows each token to attend to the first \(n_{\text{starting}}\) tokens if they come before the current one; the ending one allows each token to attend to most recent \(L_{\text{pretrain}}\) tokens. \(L_{\text{pretrain}}\) is the maximum length during training. Other tokens are ignored. In ablation studies in SSA, we find that choosing \(n_{\text{starting}}\in[5,100]\) generally achieves equally good performance. Note that \(n_{\text{starting}}=0\) (i.e., attending only to the most recent tokens) substantially hurts the performance. This resolves Factors 2 and 3 in SS3 by both limiting the number of tokens under attention and ensuring the starting few tokens are attended.

Distance ceilingLM-Infinite further bounds the "effective distance" to \(L_{\text{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(\mathbf{q},\mathbf{k},d)\), where \(d\) is the distance between two tokens. Now we modify it as \(w(\mathbf{q},\mathbf{k},d^{\prime}))\) where \(d^{\prime}=\min(d,L_{\text{pretrain}})\). Figure 2(a) shows an illustrative example where the distance ceiling is \(L_{\text{pretrain}}=2\). This addresses Factor 1 in SS3 by bounding the distance value in attention calculation.

Optionally attending to top-\(k\) tokens in the middleLM-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 (SS5.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 \(d=\frac{1}{2}L_{\text{pre-train}}\). These hyperparameters are selected based on a held-out Passkey Retrieval validation set, where we set \(k=5\) and \(h=5\), 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 SS5.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 \({\bf\Lambda}\)-shaped mask is conceptually similar to the attention patterns derived from heuristics Beltagy et al. (2020); Ding et al. (2023); Zaheer et al. (2020). However, we formally show in SS3 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 detailsLM-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 \(\Lambda\)-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 \(\mathbf{k}\) vectors unrotated and rotate all \(\mathbf{q}\) vectors to a fixed distance \(L_{\text{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_{\text{pretrain}}|\) and apply the \({\bf\Lambda}\)-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 SS3, 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 

\({\bf v}_{i}\) vectors, without the starting few tokens, the attention output can not reach regions that \({\bf 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" Peysakhovich and Lerer (2023) 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.

## 5 Evaluation

We evaluate LM-Infinite with LLaMA-7B Touvron et al. (2023), Llama-2-7b Touvron et al. (2023), MPT-7B Team (2023), and GPT-J-6B Wang and Komatsuzaki (2021). 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.

### Language Modeling with Extremely Long Context

We use ArXiv and OpenWebText2 corpora from the Pile dataset Gao et al. (2020), which contain preprint papers from ArXiv and Reddit submissions, respectively. We evaluate with negative log-likelihood (NLL) and perplexity (exponential of NLL). Figure 3 plots the NLL curves on the ArXiv

Figure 3: LM-Infinite flattens the negative log-likelihood (NLL) curves of various LLMs on ArXiv dataset without any parameter updates. The trends are similar to MPT-7B-Storywriter, an explicitly fine-tuned LLM. Llama-2 outputs NaN values on long sequences so the curve is relatively shorter.

 

dataset. Here, we break down the models' perplexity performance by positions so that the curve shows the NLL that the model achieves around that specific position, averaged across all evaluated sequences. Llama-2 outputs NaN probabilities on sequences that are slightly longer than 10K, thereby its shorter curve. All vanilla models run out of memory at \({\sim}32\)K lengths.1 The baselines' NLL quickly blows up when the tested sequences are longer than what they train on. With LM-Infinite, all models can generalize to sequences that are substantially longer than the lengths they are trained on, retaining the NLL performance. This validates our length failure factor analyses in SS1. The longer ends of curves have larger variances because of fewer documents of those lengths. In Figure 4, we further evaluate LM-Infinite + Llama-2 on a sequence of **200M** tokens, which is constructed by sampling with replacement from the ArXiv dataset and concatenating all data. LM-Infinite shows the ability to remain stably low log-perplexity level over extreme lengths.

Footnote 1: We run on a single A100 GPU with 80GB GPU memory.

Table 1 summarizes the perplexity performance at a few milestone lengths (2K, 4K, 8K, 16K, and 32K) on ArXiv and OpenWebText2, which shows a similar trend. OpenWebText2 has very few data instances over a length of 32K, so we omit the column. With LM-Infinite, all models can generalize to unseen lengths, and LM-Infinite achieves best perplexity in 7 out of 9 cases. On LLaMA + LM-Infinite, the perplexity decreases as length increases and position becomes larger. Surprisingly, _without any parameter update_, LM-Infinite outperforms many strong baselines that are trained on substantially longer text segments. As a direct comparison, MPT-7B+LM-Infinite achieves only slightly worse performance than its fine-tuned counterpart, MPT-7B-Storywriter. This confirms that LM-Infinite is a promising alternative to resource-consuming fine-tuning.

### Ablation Study

Figure 5 provides an ablation study with the LLaMA model on the ArXiv dataset about why both components in LM-Infinite are essential for maintaining LLM functionality over the length of 8K. We compare LM-Infinite with its variants to show the efficacy of the design and also to validate the factors in SS3. Among all curves, only LM-Infinite has relatively stable log-perplexity, meaning that components in LM-Infinite are all essential for successful length generalization. The vanilla LLM model (the "vanilla" curve) fails immediately with exploding NLL. If we only apply \(\mathbf{\Lambda}\)-shaped mask (the "\(\Lambda\)" curve) and do not bound inter-token distance (Factor 1), the NLL still explodes immediately after pre-training lengths. The "ceiling" curve only applies the distance ceiling technique but not the \(\Lambda\)-shaped mask to limit the number of attended tokens. The performance still degenerates (evidenced by an ever-increasing NLL). This confirms that the existence of Factor 2, too many to

Figure 4: LM-Infinite generalizes Llama-2 to an extreme length of 200M. The dashed line is the NLL at 8K length of the vanilla Llama-2 model.

Figure 5: Ablation study on LLaMA in §5.2. x-axis is token position and y-axis is negative log-likelihood (NLL). The vanilla model (vanilla), using a windowed attention (window), using only a \(\Lambda\)-shaped attention mask (\(\Lambda\)), and only the ceiling the distance value (ceiling) all more or less suffer from perplexity explosion. Only LM-Infinite can retrain the performance while generalizing to unseen lengths.

 

kens, is still harming the LLMs' performance. The "window" curve shows a baseline with the sliding-window attention pattern, which only attends to the most recent tokens in a sliding window without altering the input text. It produces the second worst NLL values, which indicates a significant performance and fluency degradation. This confirms our theoretical analysis of factor 3. Due to its visibly much worse performance, we exclude it from other evaluations.

Another similar baseline to "window" is the truncation baseline, which drops excessive tokens altogether when the context is longer than the model's pre-training length, keeping only the most recent ones. This truncation process happens before the forward pass starts and essentially removes the truncated text from the input to the model without changing the attention mechanism. We compared this baseline in two places in the paper. In SS5.3 and Table 2, LM-infinite outperforms this baseline on downstream tasks. In Section 5.4 and Figure 6, LM-infinite achieves a better trade-off between computation complexity and generation quality than this baseline.

### Downstream Evaluation

As LLMs are often deployed for downstream tasks, we evaluate how LM-Infinite performs on two long-input tasks under the zero-shot setting: Passkey Retrieval Mohtashami and Jaggi (2023) and Qapser Dasigi et al. (2021). Passkey Retrieval buries a passkey at a random position in a long distraction text and, in the end, asks what the passkey is. Qasper is a question-answering dataset on scientific papers with a total of 1.5K testing questionnaire pairs. We evaluate Llama-2-7b-chat, as its instruction tuning enables good task-solving ability Bai et al. (2023), with middle top-5 tokens enabled on higher than 5-th layer (see SS4 for definition and Appendix A for hyperparameter selection). Results are listed in Table 2. LM-Infinite consistently outperforms the baselines on both tasks, with

\begin{table}
\begin{tabular}{l r r r r r r r r} \hline \hline  & \multicolumn{8}{c}{**Passkey Retrieval**} & \multicolumn{3}{c}{**Qasper**} \\ \cline{2-9} Model & 6K & 8K & 10K & 12K & 16K & average & \\ \hline
**Original** & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.2 \\
**Truncated** & 66.0 & 55.3 & 38.8 & 32.8 & 27.3 & 44.0 & 30.1 \\
**LM-Infinite** & **70.3** & **90.8** & **86.5** & **79.3** & **79.1** & **81.2** & **31.3** \\ \hline \hline \end{tabular}
\end{table}
Table 2: Downstream evaluation on Passkey Retrieval and Qasper. LM-Infinite enables Llama-2 to consistently outperform both the original model and the baseline that truncates inputs to 4K. The truncation baseline drops excessive tokens altogether when the context is longer than the model’s pretraining length, keeping only the most recent ones, which happens before the forward pass starts without changing the attention mechanism.

\begin{table}
\begin{tabular}{l l c r r r r r r r r r} \hline \hline  & \multicolumn{4}{c}{**BLEU**} & \multicolumn{4}{c}{**ROUGE**} \\ \cline{3-13}
**Model** & \multirow{2}{*}{\(L_{\text{pretrain}}\)} & \multirow{2}{*}{2K} & \multirow{2}{*}{4K} & \multirow{2}{*}{8K} & \multirow{2}{*}{16K} & \multirow{2}{*}{32K} & \multirow{2}{*}{2K} & \multirow{2}{*}{4K} & \multirow{2}{*}{8K} & \multirow{2}{*}{16K} & \multirow{2}{*}{32K} \\ \cline{1-1}

 

a 37.2 percentage gain on Passkey Retrieval and a 1.2 percentage gain on the Qasper task. Passkey retrieval locates useful information uniformly in a sequence, so the performance of the truncated baseline largely depends on whether the remaining part covers the passkey. On Qasper, the top-\(k\) attention is necessary for achieving good performance, which indicates that similarly important information in the middle needs to be attended to. This suggests that it can improve downstream task performance on long inputs without fine-tuning while the vanilla model immediately fails.

### Generation Quality

We further evaluate LM-Infinite's generation quality on ArXiv and OpenWebText2 test sets, with BLEU Papineni et al. (2002) and ROUGE Lin (2004) (ROUGE-L). We let the LLMs generate 100 tokens after each milestone length and use the following 100 tokens in original texts as references. As the generation is time-consuming, we sample 100 long sequences for evaluation for each dataset. The results are summarized in Table 3. The trend is similar to that in the last section: _without_ parameter updates, LM-Infinite successfully allows LLMs to retain their generation quality while generating sequences longer than training, comparable to the fine-tuned baselines such as MPT-7B-SW. The generation results from the vanilla LLMs are poor and contain mostly nonsensical texts, resulting in many close-to-zero scores. For some BLEU scores, it yields zero {2,3,4}-gram overlaps with the reference texts. As BLEU is weighted geometric mean over {1,2,3,4}-gram precisions, the final BLEU scores for those columns are 0. Appendix Table 8 presents some generation output examples that can provide a good picture of the generation quality. We also evaluate the efficiency in Appendix G: with 32K-long sequences, LM-Infinite achieves 2.7\(\times\) decoding speedup and 7.5\(\times\) GPU memory saving.

A few example generations are shown in Appendix H. We also compare LM-Infinite with a simple truncation-based baseline by repeatedly truncating excessive contexts. However, as the generation length increases, frequent truncations and re-encoding of new contexts are required. The larger the truncation window is, the more context is kept, but the larger the computational overhead. We let the models generate 10k tokens on ArXiv. In Figure 6, it is clear that LM-Infinite achieves a substantially better quality-efficiency tradeoff. With similar computation, LM-Infinite outperforms the baseline by about 5 BLEU. To achieve a similar BLEU, LM-Infinite incurs only \(<\)25% computational overhead than the truncation baseline.

## 6 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\(\times\) decoding speed up and a 7.5\(\times\) 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.

Figure 6: LM-Infinite achieves a better trade-off between computation complexity with generation quality than simple truncation.

 



### Limitations

This work evaluates a wide range of open-domain LLMs. However, without access to the source code of proprietary LLMs such as ChatGPT, the proposed method could not be evaluated on them. Furthermore, due to limited computational resources and time, the proposed method has not been evaluated on texts with even larger lengths, such as 1G. The model is designed on relative positional encoding Transformer models, which is the mainstream backbone for most modern LLMs. The question of how LM-Infinite can enable more efficient fine-tuning or pre-training can also be explored in future work.

## Acknowledgement

This research is partially supported by U.S. DARPA KAIROS Program No. FA8750-19-2-1004, and DARPA INCAS Program No. HR001121C0165. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies, either expressed or implied, of DARPA, or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein.

## References

* A. Anil, Y. Wu, A. Andreassen, A. Lewkowycz, V. Misra, V. Ramasesh, A. Slone, G. Gur-Ari, E. Dyer, and B. Neyshabur (2022)Exploring length generalization in large language models. Advances in Neural Information Processing Systems35, pp. 38546-38556. Cited by: SS1.
* S. Bai, X. Lv, J. Zhang, Y. He, J. Qi, L. Hou, J. Tang, Y. Dong, and J. Li (2024)Longalign: a recipe for long context alignment of large language models. arXiv preprint arXiv:2401.18058. Cited by: SS1.
* Y. Bai, X. Lv, J. Zhang, H. Lyu, J. Tang, Z. Huang, Z. Du, X. Liu, A. Zeng, L. Hou, Y. Dong, J. Tang, and J. Li (2023)Longbench: a bilingual, multitask benchmark for long context understanding. External Links: Link Cited by: SS1.
* I. Beltagy, M. E. Peters, and A. Cohan (2020)Longformer: the long-document transformer. arXiv preprint arXiv:2004.05150. Cited by: SS1.
* S. Borgeaud, A. Mensch, J. Hoffmann, T. Cai, E. Rutherford, K. Millican, G. Bm Van Den Driessche, J. Lespiau, B. Damoc, A. Clark, et al. (2022)Improving language models by retrieving from trilions of tokens. In International conference on machine learning, pp. 2206-2240. Cited by: SS1.
* M. C. Bueno, C. Gemmell, J. Dalton, R. Lottofo, and R. Nogueira (2022)Induced natural language rationales and interleaved markup tokens enable extrapolation in large language models. In Proceedings of the 1st Workshop on Mathematical Natural Language Processing (MathNLP), pp. 17-24. Cited by: SS1.
* S. Chen, S. Wong, L. Chen, and Y. Tian (2023a)Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306.15595. Cited by: SS1.
* Y. Chen, Q. Zeng, H. Ji, and Y. Yang (2021)Skyformer: remodel self-attention with gaussian kernel and nystrom method. In Proc. Thirty-fifth Annual Conference on Neural Information Processing Systems (NeurIPS2021), Cited by: SS1.
* Human Language Technologies (NAACL-HLT2022), Cited by: SS1.
* Y. Chen, S. Qian, H. Tang, X. Lai, Z. Liu, S. Han, and J. Jia (2023b)Longlora: efficient fine-tuning of long-context large language models. In The Twelfth International Conference on Learning Representations, Cited by: SS1.
* T. Chi, T. Fan, A. Rudnicky, and P. Ramadge (2023)Dissecting transformer length extrapolation via the lens of receptive field analysis. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 13522-13537. Cited by: SS1.
* Z. Dai, Z. Yang, Y. Yang, J. G. Carbonell, Q. Le, and R. Salakhutdinov (2019)Transformer-xl: attentive language models beyond a fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 2978-2988. Cited by: SS1.
* P. Dasigi, K. Lo, I. Beltagy, A. Cohan, N. A Smith, and M. Gardner (2021)A dataset of information-seeking questions and answers anchored in research papers. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 4599-4610. Cited by: SS1.
* J. Ding, S. Ma, L. Dong, X. Zhang, S. Huang, W. Wang, and F. Wei (2023)Longnet: scaling transformers to 1,000,000,000 tokens. arXiv preprint arXiv:2307.02486. Cited by: SS1.
* Y. Ding, L. Zhang, C. Zhang, Y. Xu, N. Shang, J. Xu, F. Yang,

and Mao Yang. 2024. Longrope: Extending llm context window beyond 2 million tokens. _arXiv preprint arXiv:2402.13753_.
* Dong et al. (2024) Harry Dong, Xinyu Yang, Zhenyu Zhang, Zhangyang Wang, Yuejie Chi, and Beidi Chen. 2024. Get more with less: Synthesizing recurrence with kv cache compression for efficient llm inference. _arXiv preprint arXiv:2402.09398_.
* Gao et al. (2020) Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. 2020. The Pile: An 800gb dataset of diverse text for language modeling. _arXiv preprint arXiv:2101.00027_.
* Guu et al. (2020) Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. 2020. Retrieval augmented language model pre-training. In _International conference on machine learning_, pages 3929-3938. PMLR.
* Haussler (2018) David Haussler. 2018. Decision theoretic generalizations of the pac model for neural net and other learning applications. In _The Mathematics of Generalization_, pages 37-116. CRC Press.
* Huang et al. (2023) Xijie Huang, Li Lyna Zhang, Kwang-Ting Cheng, and Mao Yang. 2023. Boosting llm reasoning: Push the limits of few-shot learning with reinforced in-context pruning. _arXiv preprint arXiv:2312.08901_.
* Jiang et al. (2023) Huiqiang Jiang, Qianhui Wu, Xufang Luo, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, and Lili Qiu. 2023. Longllmingua: Accelerating and enhancing llms in long context scenarios via prompt compression. _arXiv preprint arXiv:2310.06839_.
* Kaiokendev (2023) Kaiokendev. 2023. Things iin learning while training superhot. [https://kaiokendev.github.io/til#extending-context-to-8k](https://kaiokendev.github.io/til#extending-context-to-8k).
* Kaiser et al. (2016) Lukasz Kaiser, Ofir Nachum, Aurko Roy, and Samy Bengio. 2016. Learning to remember rare events. In _International Conference on Learning Representations_.
* Kazemnejad et al. (2023) Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. 2023. The impact of positional encoding on length generalization in transformers. _arXiv preprint arXiv:2305.19466_.
* Ke et al. (2020) Guolin Ke, Di He, and Tie-Yan Liu. 2020. Rethinking positional encoding in language pre-training. In _International Conference on Learning Representations_.
* Meng-Wei Chang Kenton and Kristina Toutanova (2019) Jacob Devlin Ming-Wei Chang Kenton and Lee Kristina Toutanova. 2019. Bert: Pre-training of deep bidirectional transformers for language understanding. In _Proceedings of NAACL-HLT_, pages 4171-4186.
* Khandelwal et al. (2019) Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. 2019. Generalization through memorization: Nearest neighbor language models. In _International Conference on Learning Representations_.
* Kiyono et al. (2021) Shun Kiyono, Sosuke Kobayashi, Jun Suzuki, and Kentaro Inui. 2021. Shape: Shifted absolute position embedding for transformers. In _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing_, pages 3309-3321.
* Lee et al. (2024) Kuang-Huei Lee, Xinyun Chen, Hiroki Furuta, John Canny, and Ian Fischer. 2024. A human-inspired reading agent with gist memory of very long contexts. _arXiv preprint arXiv:2402.09727_.
* Li et al. (2023) Dacheng Li, Rulin Shao, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph E. Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang. 2023. How long can open-source llims truly promise on context length?
* Likhomanenko et al. (2021) Taitana Likhomanenko, Qiantong Xu, Gabriel Synnaeve, Ronan Collobert, and Alex Rogozhnikov. 2021. Cape: Encoding relative positions with continuous augmented positional embeddings. _Advances in Neural Information Processing Systems_, 34:16079-16092.
* Lin (2004) Chin-Yew Lin. 2004. Rouge: A package for automatic evaluation of summaries. In _Text summarization branches out_, pages 74-81.
* Liu et al. (2023) Xiaoran Liu, Hang Yan, Chenxin An, Xipeng Qiu, and Dahua Lin. 2023. Scaling laws of rope-based extrapolation. In _The Twelfth International Conference on Learning Representations_.
* Luo et al. (2024) Kun Luo, Zheng Liu, Shitao Xiao, and Kang Liu. 2024. Bge landmark embedding: A chunking-free embedding method for retrieval augmented longest-context large language models. _arXiv preprint arXiv:2402.11573_.
* Lv et al. (2024) Kai Lv, Xiaoran Liu, Qipeng Guo, Hang Yan, Conghui He, Xipeng Qiu, and Dahua Lin. 2024. Longwannuan: Towards systematic measurement for long text quality. _arXiv preprint arXiv:2402.13583_.
* Mohtashami and Jaggi (2023) Amirkeivan Mohtashami and Martin Jaggi. 2023. Landmark attention: Random-access infinite context length for transformers. _arXiv preprint arXiv:2305.16300_.
* Oren et al. (2024) MataneI Oren, Michael Hassid, Yossi Adi, and Roy Schwartz. 2024. Transformers are multi-state rnns. _arXiv preprint arXiv:2401.06104_.
* Papineni et al. (2002) Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. 2002. Bleu: a method for automatic evaluation of machine translation. In _Proceedings of the 40th annual meeting of the Association for Computational Linguistics_, pages 311-318.
* Peng et al. (2023) Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. 2023. Yarn: Efficient context window extension of large language models. In _The Twelfth International Conference on Learning Representations_.
* Peysakhovich and Lerer (2023) Alexander Peysakhovich and Adam Lerer. 2023. Attention sorting combats recency bias in long context language models. _arXiv preprint arXiv:2310.01427_.

 

David Pollard. 1990. Empirical processes: Theory and applications. In NSF-CBMS Regional Conference Series in Probability and Statistics, JSTOR. Cited by: SS2.
* O. Press, N. Smith, and M. Lewis (2021)Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference on Learning Representations, External Links: Link Cited by: SS4.3.
* Z. Qiu, J. Li, S. Huang, W. Zhong, and I. King (2024)Cloingeval: a chinese benchmark for evaluating long-context large language models. arXiv preprint arXiv:2403.03514. Cited by: SS2.
* Z. Qu (2023)Gpt rotational position embedding for length extrapolation. In Proceedings of the 2023 6th International Conference on Machine Learning and Natural Language Processing, pp. 166-170. Cited by: SS2.
* C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Y. Zhou, W. Li, and P. J. Liu (2020)Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research21 (1), pp. 5485-5551. Cited by: SS2.
* J. Rasley, S. Rajbhandari, Q. Ruwase, and Y. He (2020)DeepSpeed: system optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 3505-3506. Cited by: SS2.
* N. Ratner, Y. Levine, Y. Belinkov, O. Ram, I. Magar, O. Abend, E. Karpas, A. Shashua, K. Leyton-Brown, and Y. Shoham (2023)Parallel context windows for large language models. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 6383-6402. Cited by: SS2.
* S. Ren and K. Q. Zhu (2024)On the efficacy of eviction policy for key-value constrained generative language model inference. arXiv preprint arXiv:2402.06262. Cited by: SS2.
* N. Shao, S. Xiao, Z. Liu, and P. Zhang (2024)Flexibly scaling large language models contexts through extensible tokenization. arXiv preprint arXiv:2401.07793. Cited by: SS2.
* K. Song, X. Wang, S. Cho, X. Pan, and D. Yu (2023)Zebra: extending context window with layerwise grouped local-global attention. arXiv preprint arXiv:2312.08618. Cited by: SS2.
* J. Su, Y. Lu, S. Pan, A. Murtadha, B. Wen, and Y. Liu (2021)Roformer: enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864. Cited by: SS2.
* Y. Sun, L. Dong, B. Patra, S. Ma, S. Huang, A. Benhaim, V. Chaudhary, X. Song, and F. Wei (2022)A length-extrapolatable transformer. arXiv preprint arXiv:2212.10554. Cited by: SS2.
* M. Tao, Y. Feng, and D. Zhao (2023)A frustratingly easy improvement for position embeddings via random padding. arXiv preprint arXiv:2305.04859. Cited by: SS2.
* M. NLP Team (2023)Introducing mpt-?b: a new standard for open-source, commercially usable llms. Note: Accessed: 2023-05-05 Cited by: SS2.
* H. Touvron, T. Lavril, G. Izacard, X. Martinet, M. Lacchaux, T. Lacroix, B. Roziere, N. Goyal, E. Hambro, F. Azhar, et al. (2023a)LIMA: open and efficient foundation language models. arXiv preprint arXiv:2302.13971. Cited by: SS2.
* H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashkykov, S. Batra, P. Bhagava, S. Bhosale, et al. (2023b)Llama 2: open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288. Cited by: SS2.
* S. Tworkowski, K. Staniszewski, M. Pacek, Y. Wu, H. Michalewski, and P. Milos (2023)Focused transformer: contrastive training for context scaling. arXiv preprint arXiv:2307.03170. Cited by: SS2.
* A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin (2017)Attention is all you need. Advances in neural information processing systems30, pp. . Cited by: SS2.
* B. Wang and A. Komatsuzaki (2021)GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model. Note: [https://github.com/kingofloz/mesh-transformer-jax](https://github.com/kingofloz/mesh-transformer-jax) Cited by: SS2.
* C. Wang, R. Ning, B. Pan, T. Wu, Q. Guo, C. Deng, G. Bao, Q. Wang, and Y. Zhang (2024)Novelqa: a benchmark for long-range novel question answering. arXiv preprint arXiv:2403.12766. Cited by: SS2.
* Y. Wu, M. Norman Rabe, D. Hutchins, and C. Szegedy (2021)Memorizing transformers. In International Conference on Learning Representations, Cited by: SS2.
* G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis (2024)Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, Cited by: SS2.
* K. Yang, D. Klein, N. Peng, and Y. Tian (2023)DOC: improving long story coherence with detailed outline control. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), Toronto, Canada, pp. 3378-3465. Cited by: SS2.

 

Zi Yang and Nan Hua. 2024. Attendre: Wait to attend by retrieval with evicted queries in memory-based transformers for long context processing. _arXiv preprint arXiv:2401.04881_.
* Yogatama et al. (2021) Dani Yogatama, Cyprien de Masson d'Autume, and Lingpeng Kong. 2021. Adaptive semiparametric language models. _Transactions of the Association for Computational Linguistics_, 9:362-373.
* Yuan et al. (2024) Tao Yuan, Xuefei Ning, Dong Zhou, Zhijie Yang, Shiyao Li, Minghui Zhuang, Zheyue Tan, Zhuyu Yao, Dahua Lin, Boxun Li, et al. 2024. Lv-eval: A balanced long-context benchmark with 5 length levels up to 256k. _arXiv preprint arXiv:2402.05136_.
* Zaheer et al. (2020) Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. 2020. Big bird: Transformers for longer sequences. _Advances in neural information processing systems_, 33:17283-17297.
* Zhang et al. (2024a) Peitian Zhang, Zheng Liu, Shitao Xiao, Ninglu Shao, Qiwei Ye, and Zhicheng Dou. 2024a. Soaring from 4k to 400k: Extending llm's context with activation beacon. _arXiv preprint arXiv:2401.03462_.
* Zhang et al. (2024b) Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, et al. 2024b. \(\infty\) bench: Extending long context evaluation beyond 100k tokens. _arXiv preprint arXiv:2402.13718_.
* Zhang et al. (2024c) Zhenyu Zhang, Runjin Chen, Shiwei Liu, Zhewei Yao, Olaturji Ruwase, Beidi Chen, Xiaoxia Wu, and Zhangyang Wang. 2024c. Found in the middle: How language models use long contexts better via plug-and-play positional encoding. _arXiv preprint arXiv:2403.04797_.
* Zhang et al. (2024d) Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Liannim Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Re, Clark Barrett, et al. 2024d. H2o: Heavy-hitter oracle for efficient generative inference of large language models. _Advances in Neural Information Processing Systems_, 36.
* Zhou et al. (2023) Wangchunshu Zhou, Yuchen Eleanor Jiang, Peng Cui, Tiannan Wang, Zhenxin Xiao, Yifan Hou, Ryan Cotterell, and Mrinmaya Sachan. 2023. Recurrentgpt: Interactive generation of (arbitrarily) long text. _arXiv preprint arXiv:2305.13304_.
* Zhu et al. (2023) Dawei Zhu, Nan Yang, Liang Wang, Yifan Song, Wenhao Wu, Furu Wei, and Sujian Li. 2023. Pose: Efficient context window extension of llms via positional skip-wise training. In _The Twelfth International Conference on Learning Representations_ 

Implementation Details

In this section, we introduce some implementation details of LM-Infinite as well as the hyper-parameter selection.

### Computational resources

All experiments are run on single A100 GPUs with 80GB GPU memory each. The 200M length generalization runs for 20 hours. The downstream tasks take 3\(\sim\)7 hours to evaluate each. Our work does not involve any training or fine-tuning. All models are loaded with Huggingface2 code repository. Rouge and BLEU scores are loaded from evaluate3 package. Datasets and models are used with permission from their licenses.

Footnote 2: [https://huggingface.cohttps:/huggingface.co](https://huggingface.cohttps:/huggingface.co)

Footnote 3: [https://huggingface.co/docs/evaluate/index](https://huggingface.co/docs/evaluate/index)

### The size of starting attention span

We vary the value of \(n_{\text{starting}}\) and find LM-Infinite to be tolerant with it taking a wide range of values without sacrificing the NLL values. Speicifically, we evaluate it on sequences of 16k length in the ArXiv validation set and calculate the average NLL.

### Reintroducing Top-\(k\) Middle Tokens

This optional technique involves optionally attending to \(k\) tokens in the middle with the largest attention logits. Here, the \(k\) tokens are selected independently for each attention head and only apply to layers higher than \(h\)-th layer. These tokens have an attention distance of \(d=\frac{1}{2}L_{\text{pre-train}}\). We select these hyper-parameters based on a held-out validation set of Passkey Retrieval. On Llama-2, we use \(k=5\) and \(h=5\). As an ablation study, we vary each hyper-parameter and observe its effects on Passkey Retrieval accuracy.

On the Qasper dataset, for both vanilla models and LM-Infinite, we use 6K sub-sequence of inputs as prompts and use a systematic prompt format described in Llama-2 paper Touvron et al. (2023b).

## Appendix B Additional Related Work

After our preprint, there have been papers that cite our work and investigate zero-shot or few-shot length generalization of LLMs. As many absolute or relative position encoding methods are based on periodic functions, Qu (2023); Ding et al. (2024); Liu et al. (2023); Jiang et al. (2023) propose to apply LLMs (fine-tuned or not) on decreased period frequencies (which is equivalent to interpolating position indices) to adapt LLMs to longer sequences. Some other papers finetune LLMs with designed attention patterns Oren et al. (2024); Zhang et al. (2024) on long contexts, using neural tangent kernel Peng et al. (2023), or with low-rank adaptationLoRA Chen et al. (2023b). Yang and Hua (2024) instead proposes a wait-to-attend mechanism to extend length limits for memory-based Transformers. Other ways of key-value cache selection/eviction methods are investigated in Ren and Zhu (2024); Dong et al. (2024); Zhang et al. (2024). Similarly, Huang et al. (2023); Lee et al. (2024) tackles long context by learning to dynamically prune, select, or summarize contexts. Alternatively, context compression methods Shao et al. (2024) propose to learn to compress long contexts into shorter embedding sequences. Some work proposes alternative position encodings Song et al. (2023); Zhang et al. (2024c); Zhu et al. (2023) or landmark token embeddings Luo et al. (2024) that enable extendable context limits. Xiao et al. (2024) is a later concurrent work to ours with a similar approach to LLM length generalization. Unlike our work, they feed a sequence to an LLM token-by-token which limits their extreme length generalization (4M v.s. 200M of ours), and more importantly, they do not show improvements on downstream tasks without pre-training an LLM from scratch.

\begin{table}
\begin{tabular}{l l l l l l l} \hline \hline  & \multicolumn{9}{c}{\(n_{\text{starting}}\)} \\
0 & 1 & 2 & 10 & 100 & 1000 & 2000 \\ \hline
6.43 & 1.03 & 1.03 & 1.02 & 1.02 & 1.81 & 4.96 \\ \hline \hline \end{tabular}
\end{table}
Table 4: Effect on LM-Infinite’s NLL by varying \(n_{\text{starting}}\).

\begin{table}
\begin{tabular}{l l l l l l l} \hline \hline  & \multicolumn{9}{c}{\(k\)} \\
1 & 3 & 5 & 10 & 20 & 50 & 200 \\ \hline
0.69 & 0.81 & 0.81 & 0.79 & 0.8 & 0.79 & 0.73 \\ \hline \hline \end{tabular}
\end{table}
Table 5: Effect of varying \(k\).

 

Finally, there is a lot of new benchmarks Qiu et al. (2024); Zhang et al. (2024b); Yuan et al. (2024); Wang et al. (2024); Lv et al. (2024); Bai et al. (2024) proposed to evaluate the long-context performance of LLMs.

## Appendix C Formal Statement of Theorem 1

Let us denote the logit function with relative position encoding as \(w({\bf q},{\bf k},d)\in\mathbb{R}\). It maps the query \({\bf q}\), key \({\bf k}\), and their distance \(d\), to an attention logit. The final attention weights are usually calculated by a softmax operation. For example, given \(n\) tokens with indices \((1,\cdots,n)\), the attention by the last token on a preceding token at position \(i\) is:

\[\text{Attn}(\text{token}_{n},\text{token}_{i})=\frac{e^{w({\bf q}_{n},{\bf k }_{i},n-i)}}{\sum_{j=1}^{n}e^{w({\bf q}_{n},{\bf k}_{j},n-j)}} \tag{1}\]

Then the formal theorem of Theorem 1 is as follows:

**Theorem 2**.: _(Formal) Let \({\bf q}\) and \({\bf k}\) be random vectors sampled from training distributions \(\mathcal{D}_{\mathbf{q}}\) and \(\mathcal{D}_{\mathbf{k}}\), respectively, where \(\mathcal{D}_{\mathbf{q}}\) and \(\mathcal{D}_{\mathbf{k}}\) are the trained distributions for \({\bf q}\) and \({\bf k}\), respectively. We use the pseudo-dimension \(\dim_{P}(\cdot)\) defined in Pollard (1990), which measures the representation capacity of a function family. Assume that the set of distances based logit functions \(\mathcal{H}=\{w(\cdot,\cdot,d)|d\in\mathbb{N}\}\) has bounded pseudo-dimension \(\dim_{P}(\mathcal{H})=r\)4. Let us also define the distinguish-ability of two distances \(d\) and \(d^{\prime}\) under \(w\) as their expected squared difference: \(\mu_{w}(d,d^{\prime})=\operatorname*{\mathbb{E}}_{{\bf q}\sim\mathcal{D}_{ \mathbf{q}},{\bf k}\sim\mathcal{D}_{\mathbf{k}}}(w({\bf q},{\bf k},d)-w({\bf q },{\bf k},d^{\prime}))^{2}\). We assume that \(w\) is not limited to recognizing only a finite group of distances, otherwise, all distances longer than a threshold will become almost the same as shorter distances. Formally, for any \(n\), there is a partition of \([0..n]\) into \(\alpha(n)\) groups so that, \(\mu_{w}(d,d^{\prime})\leq\epsilon\) for any \(d,d^{\prime}\) from the same group. \(\alpha(n)\in\mathbb{N}\) is non-decreasing and unbounded function. Then we have:_

Footnote 4: This is true for most current techniques. See discussions in Appendix F

We first borrow a lemma from Haussler (2018), which we paste below. Note that a cover size \(\mathcal{N}(\epsilon,\mathcal{H},\mu)\) is defined as the minimum cardinal of a cover-set \(\mathcal{H}^{\prime}\) so that any element of \(h\in\mathcal{H}\) will be within \(\epsilon\) distance to at least one element \(h^{\prime}\in\mathcal{H}^{\prime}\).

**Lemma 3**.: _Let \(\mathcal{H}\) be a function family mapping from space \(X\) to range \([0,B]\), and its pseudomension \(\dim_{P}(\mathcal{H})=r\). Then for any probabilistic measure \(P\) on \(X\), and \(\epsilon\in[0,B]\), we have that the \(\epsilon\) cover of \(\mathcal{H}\) under metric \(\mu(h_{1},h_{2})=\operatorname*{\mathbb{E}}_{x\sim P}(h_{1}(x)-h_{2}(x))^{2}\) is bounded by:_

\[\mathcal{N}_{P}(\epsilon,\mathcal{H},\mu)\leq 2\left(\frac{2eB}{\epsilon} \ln\frac{2eB}{\epsilon}\right)^{r}\]

With this lemma, we can go on to prove Theorem 2 as follows.

Proof.: We prove by contradiction. Assume that \(\sup_{{\bf q},{\bf k},d\leq n}|w({\bf q},{\bf k},d)|<a=\left(\frac{\alpha(n)} {2}\right)^{\frac{1}{2r}}\frac{\epsilon}{4e}\). Without loss of generality, we can shift all the values to range \([0,2a]\). Then the function family \(\mathcal{H}=\{w(\cdot,\cdot,d)|d\in\mathbb{N}\}\) will have cover size \(\mathcal{N}_{P}(\epsilon,\mathcal{H},\mu)\leq 2\left(\frac{4ea}{\epsilon}\ln \frac{4ea}{\epsilon e}\right)^{r}<\alpha(n)\).

However, this is smaller than the number of different distances and relative weight attentions \(\mathcal{H}\), which means that at least 2 functions will be close to each other \((w(\cdot,\cdot,d),w(\cdot,\cdot,d^{\prime}))^{2}<\epsilon\). This constitutes a contradiction with the distinguishability condition.

## Appendix D Formal Statement and Proof of Proposition 1

The formal statement of Proposition 1 is the following:

**Proposition 2**.: _(Attention Entropy Explosion) Let \(w_{1},w_{2},\cdots,w_{n}\in[-B,B]\) be a sequence of attention logits. Then the entropy of the attention distribution they span is asymptotically lower bounded by \(\ln n\):_

\[H\left(\left(\frac{e^{w_{i}}}{\sum_{j=1}^{n}e^{w_{j}}}|1\leq i\leq n\right) \right)=\Omega(\ln n)\]

The entropy approaches \(+\infty\) as \(n\) grows larger.

Proof.: Note that entropy on a discrete distribution is defined as \(\text{Entropy}(P)=-\sum_{i}p_{i}\ln p_{i}\). Then the attention entropy determined by attention logits

\begin{table}
\begin{tabular}{c c c c c c c} \hline  & & & \multicolumn{5}{c}{\(h\)} \\
0 & 4 & 5 & 6 & 8 & 16 & 24 \\ \hline
0.81 & 0.94 & 0.94 & 0.91 & 0.46 & 0.45 & 0.46 \\ \hline \end{tabular}
\end{table}
Table 7: Effect of varying \(h\).

 

\(\{w_{i}|1\leq i\leq n\}\) is

Entropy (Attention) \[=-\sum_{i}\frac{e^{w_{i}}}{\sum_{j}e^{w_{j}}}\ln\frac{e^{w_{i}}}{ \sum_{j}e^{w_{j}}}\] \[=-\sum_{i}\frac{e^{w_{i}}}{\sum_{j}e^{w_{j}}}\left(w_{i}-\ln\sum_{ j}e^{w_{j}}\right)\] \[=-\sum_{i}\frac{e^{w_{i}}}{\sum_{j}e^{w_{j}}}w_{i}+\ln\sum_{j}e^{ w_{j}}\] \[\geq-\max_{i}w_{i}+\ln(ne^{-B})\] \[\geq\ln n-2B\] \[= \Omega(\ln n)\]

## Appendix E More on Implicitly Encoded Positions

We also plot the token features of more layers with PCA projection to the 2D plane in Figure 7. We see that from layer 2 to higher layers, the initial few tokens occupy a distinct region with later tokens. Therefore, if these tokens are discarded by window attention during attention, the attention output, which is a weighted sum of \(v_{i}\) vectors, will reside in a different region. This can explain why windowed attention does not work and why the first few tokens need to be kept by our \(\boldsymbol{\Lambda}\)-shaped attention.

## Appendix F Pseudo-Dimension Assumption on Attention Logit Functions

In Theorem 3.4, we assumed that the family of distance-based logit functions \(\mathcal{H}=\{w(\cdot,\cdot,d)|d\in\mathbb{N}\}\) has a finite pseud-dimension. In this section, we demonstrate that most current implementations of relative positional encodings do have a finite pseudo-dimension. Let us discuss a few examples in the following:

T5-Bias and AlibiIt is easy to see that, the difference between any two logit functions is uniform: \(w(\cdot,\cdot,d_{1})-w(\cdot,\cdot,d_{2})=\text{bias}(d_{1})-\text{bias}(d_{2})\) regardless of the input. Therefore this family cannot shatter any set larger than 2, so the pseudo-dimension is 1.

Windowed AttentionThis operation is equivalent to limiting the family to a finite size \(|\mathcal{H}|=d_{\mathrm{max}}+1\), where \(d_{\mathrm{max}}\) is the size of the window. Therefore \(\dim_{P}(\mathcal{H})\leq d_{\mathrm{max}}+1\).

NoPEAs there is no explicit positional encoding implemented, all distances are equivalent. The pseudo-dimension is 1.

RoPE, CAPE, and XPosFor RoPE, the logit function \(w(\mathbf{q},\mathbf{k},d)\) is the weighted sum of finite fixed sinusoidal functions \(\{\sin(\omega_{i}d),\cos(\omega_{i}d)\}\). The size of this set is equivalent to the feature dimension number \(k\). We know that \(\dim_{P}(\mathcal{H}_{1}+\mathcal{H}_{1})\leq\dim_{P}(\mathcal{H}_{1})+\dim_{ P}(\mathcal{H}_{2})\). Also, the scaling of a single function can only have a pseudo-dimension of 2. Therefore, the whole family has a bounded pseudo-dimension \(\dim_{P}(\mathcal{H})\leq 2k\). The analysis on CAPE and XPos is similar.

## Appendix G Computational Efficiency Evaluation

To evaluate the computational efficiency of LM-Infinite, we run the Llama-2-7B model on 100 sequences of 32k length in the ArXiv dataset. The hardware is a single A100 GPU with 80GB GPU memory. As the memory is not enough to host the whole computation graph during decoding, we use DeepSpeed Rasley et al. (2020) with Zero3 optimization. We also have to modify the computation code to further reduce GPU memory usage and prevent out-of-memory errors. With that in mind, the vanilla Llama-2-7B model encodes with an average speed of 48.19s per sequence, while LM-Infinite encodes with an average of 15.26s per sequence, a 3.16x speedup. The vanilla Llama-2-7B model decodes with 7.34s per token, while LM-Infinite decodes with 2.70s per token, a 2.72x speedup. We also evaluate the GPU memory usage on 32k-length inputs, the statistics of which are profiled with PyTorch Profiler. The vanilla model uses 33.2Gb memory per sequence, while LM-Infinite uses 4.41Gb per sequence, a 7.53\(\times\) memory saving.
 

* [25] M. C. Bader, A. D. Krizhevsky, and S. Ghahramani. Learning from a single image to predict the best performance of the model. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_, pages 1-4, 2016.
* [26] J. L. Friedman, R. Girshick, and Y. Bengio. Generative adversarial nets for large scale data. In _Advances in neural information processing systems_, volume 389 of _Lecture Notes in Artificial Intelligence_, pages 77-80, 2005.
* [27] P. F. Donoho, N. Srivastava, and I. Sutskever. Dropout: a deep learning approach to stochastic optimization. _arXiv preprint arXiv:1412.6098_, 2014.
* [28] H. Dong, X. Li, and Z. Sun. Deep residual learning with convolutional networks. In _NIPS conference on computer vision_, pages 1338-1346. PMLR, 2014.
* [29] T. Lin, W. Liu, and Z. Sun. Deep residual learning with convolutions. In _Neural machine translation (ICLR)_, pages 1099-1109. MIT Press, Cambridge, MA, USA, 2014.
* [30] E. Maaten, O. Hinton, and J. Salakhutdinov. Imagenet classification by jointly connected layers. In _International conference on Machine Learning_, pages 1099-1110. AAAI, 2014.
* [31] V. Mnih, I. Reed, and S. Erhan. Intriliguing deep neural networks. In _Advances in Neural Information Processing Systems (IJCNN)_, pages 1-10, 2006.
* [32] J. R. Socher. _The Annals of Statistics_, 2nd edn. of _Proc. IEEE International Congress on Computer Vision and Pattern Recognition_, pages 1-10, 2006.

[MISSING_PAGE_POST]

 



\begin{table}
\begin{tabular}{l l} \hline \hline  & **ArXiv Dataset** \\ \hline
**Context** & _Our proof of the Kottwitz formula imitates these arguments of Langlands and Rapoport of deriving the Kottwitz formula from their conjecture. As such, their theory of Galois gerbs is a major ingredient in this work. Another essential ingredient is a recent work of Kisin I@Kisni17] on the aforementioned conjecture of Langlands and Rapoport. Previously, Kottwitz I@Kontwitz92] proved the formula in PEL-type cases (of simple Lie type A or C) by a method which is based on the Honda-Tate theory. This method however cannot be applied in general Hodge-type situations, and indeed_ \\ \hline
**Vanilla Liama-2** & _of this (of over-equary-her, and [t and, in the...cister ’ and an of the model to by. by –, this, by the. It, and it, 7.. (and at equist (oft me. to the [[I WNE (to. and for the (((de in the (for the anddistile-c.:[de for in an inc rt