PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Chi Han, Qifan Wang, Hao Peng, Wenhan Xiong, Yu Chen, Heng Ji, Sinong Wang

arXiv:2308.16137v7cs.CLcs.AI

TL;DR

LLMs trained on short segments can struggle with longer contexts, limiting their use on long documents and conversations. LM-Infinite combines a Λ-shaped attention mask with a distance ceiling to extend context without parameter updates. It supports zero-shot generalization to inputs up to 200M tokens while retaining perplexity and improves downstream performance and efficiency.

  • Problem

    LLMs typically train on segments shorter than 4K tokens, and their performance deteriorates on longer inputs needed for tasks involving long-range context.

    Source: §1 · p. 1

  • Method

    LM-Infinite uses a Λ-shaped mask to attend to beginning and recent tokens and a distance ceiling, without parameter updates.

    Source: §1 · pp. 1–2, §4 · p. 5

  • Results

    LM-Infinite generalizes to texts up to 200M tokens while retaining perplexity and generation quality, improves zero-shot Passkey Retrieval and Qasper scores, and speeds decoding 2.7× while saving 7.5× GPU memory over original LLMs.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    The method offers zero-shot long-context handling for off-the-shelf LLMs, including models pretrained on 2K- or 4K-length segments.

    Source: §6 · p. 9

    The method is designed for relative-positional-encoding Transformers, was not evaluated on proprietary models such as ChatGPT, and was not tested on inputs as long as 1G tokens.

    Source: Limitations · p. 10

Abstract

from arXiv · show

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 theoretical analysis and empirical investigation, this work identifies three major factors contributing to this length generalization failure. Our theoretical analysis further reveals that commonly used techniques like truncating the attention window 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.7x decoding speed up and 7.5x memory saving over the original model. Our codes are released at \url{https://github.com/Glaciohound/LM-Infinite}.

1 Introduction

LM-Infinite targets the sharp loss of performance LLMs experience on contexts longer than their short training segments. It combines a Λ-shaped attention mask with a distance ceiling, improving long-context benchmarks and efficiency without parameter updates.

1 Introduction

  • LLMs trained on segments shorter than 4K tokens struggle to generalize to much longer inputs, limiting tasks that require long-range context.Examples include scientific articles, source code repositories, and long-context dialogues.Source: §1 · p. 1
  • Relative positional encodings and sliding-window attention do not resolve the underlying length-generalization problems, and sliding-window methods are not directly applicable to pretrained models without further training.The paper identifies unseen token distances, unseen numbers of attended tokens, and position information encoded in initial tokens as the central factors.Source: §1 · p. 1
  • LM-Infinite combines a Λ-shaped mask that retains beginning and recent tokens with a distance ceiling that caps relative distances at the training maximum.The method can optionally reintroduce middle top-k tokens for some downstream tasks.Source: §1 · pp. 1–2
  • 37.2% gain on Passkey Retrieval and 1.2% gain on Qasper accompany zero-shot generalization to texts up to 200M tokens, with 2.7× decoding speedup and 7.5× GPU memory saving over original LLMs.The paper reports retained language-modeling perplexity and generation quality on ArXiv and OpenWebText2.Source: §1 · p. 2

2 Background and Related Work

The section reviews relative positional encodings and prior approaches to extending context. It motivates LM-Infinite by describing the limitations of finetuning, retrieval-based methods, and attention patterns that do not address the identified length-generalization factors.

2.1 Relative Positional Encodings

  • Relative positional encodings use token-to-token distances rather than absolute positions; RoPE rotates query and key vectors, while AliBi adds a linear distance offset to attention logits.The paper notes that length-generalization failures remain despite promising evidence for these encodings.Source: §2.1 · p. 2

2.2 Efforts Towards Length Generalization

  • Prior approaches include finetuning on longer sequences, retrieval-based memories, and attention patterns; they may require substantial training, lack compatibility with existing LLMs, or be limited to in-context learning.The paper contrasts these with its zero-shot approach.Source: §2.2 · pp. 2–3
  • Figure 1 summarizes three factors behind length-generalization failures: unseen distances, an unseen number of attended tokens, and initial tokens occupying a distinct feature region.The figure depicts the initial-token regions as concentrated and distant from later-token representations.Source: Figure 1 · p. 3

3 Why do Transformer LLMs Fail to Generalize to Long Contexts?

The analysis attributes length-generalization failures to shifts in computational features caused by unseen distances, larger attended contexts, and position information concentrated in starting tokens. It argues these factors must be addressed for relative-position Transformer LLMs to generalize to longer inputs.

3 Why do Transformer LLMs Fail to Generalize to Long Contexts?

  • The analysis assumes Transformer LLMs with relative positional encodings and investigates Llama-2, pretrained on 4K-length segments.It argues the three factors can shift computational features out of the training distribution.Source: §3 · p. 3
  • For relative positional encodings, unseen inter-token distances can make attention logits explode; Llama-2 experiments show substantially larger logits beyond its 4K training length.The analysis proposes capping distances at the maximum seen during training, while noting that this exposes another challenge.Source: §3 · p. 3

Factor 2: attending to unseen numbers of tokens

  • When logits are bounded but the attended context keeps growing, attention entropy can exceed its training range and grow without bound as sequences lengthen.The paper reports an empirically increasing entropy trend and suggests bounding context size to keep entropy within seen ranges.Source: Factor 2: attending to unseen numbers… · pp. 3–4
  • Starting tokens can occupy a distinct representational space, and sliding-window attention that discards them prevents later attention outputs from reaching their feature region.The analysis attributes this implicit absolute-position signal to starting tokens even without explicit absolute positional embeddings.Source: Factor 2: attending to unseen numbers… · pp. 4–5

4 Our proposal: LM-Infinite

LM-Infinite targets zero-shot length generalization with a Λ-shaped attention mask and a distance ceiling. It can optionally restore top-k middle tokens for downstream tasks.

4 Our proposal: LM-Infinite

  • The Λ-shaped mask attends to the first nstarting tokens and the most recent Lpretrain tokens, ignoring other tokens.Values of nstarting from 5 to 100 generally perform equally well; nstarting = 0 substantially hurts performance.Source: §4 · p. 5
  • The distance ceiling replaces relative distance d with min(d, Lpretrain), bounding the effective attention distance.It affects starting tokens when later-positioned tokens attend to them.Source: §4 · p. 5
  • For downstream tasks, LM-Infinite can selectively attend to the middle tokens with the largest attention logits.The selected top-k tokens are chosen independently per attention head; the same k and h settings are used across downstream tasks, while middle tokens are unnecessary for generation and inference.Source: §4 · p. 5
  • Unlike conceptually similar heuristic attention patterns, LM-Infinite is presented as avoiding their theoretical limitation on unseen-length generalization without parameter updates.The method can be implemented by replacing the attention function in Transformer models with relative positional encoding, including RoPE and AliBi.Source: §4 · p. 5
  • The mask preserves access to starting tokens, whose features occupy a distinct region, while limiting intermediate tokens that can harm length generalization.The paper relates rear-token importance to recency bias and describes middle-token information as less position-sensitive.Source: §4 · pp. 5–6
  • 7 of 9 columns show LM-Infinite achieving the highest perplexity on ArXiv and OpenWebText2 without parameter updates.Lpretrain denotes the training segment length.Source: Table 1 · p. 6

5 Evaluation

Across language modeling, zero-shot downstream tasks, and generation evaluations, LM-Infinite extends models beyond their training lengths while retaining quality. Ablations show that both its attention mask and distance ceiling are needed for length generalization.

5.1 Language Modeling with Extremely Long Context

  • At 200M tokens, LM-Infinite keeps Llama-2’s log-perplexity stably low and achieves the best perplexity in 7 of 9 milestone cases.It also outperforms many baselines trained on longer segments and is only slightly worse than fine-tuned MPT-7B-Storywriter.Source: §5.1 · p. 7, Figure 4 · p. 7

5.2 Ablation Study

  • Only the full method keeps log-perplexity relatively stable; vanilla, windowed, Λ-only, and distance-ceiling-only variants degrade beyond training lengths.Windowed attention has the second-worst NLL and is excluded from later evaluations. The truncation baseline drops excess tokens before the forward pass, retaining the most recent ones.Source: §5.2 · pp. 7–8, Figure 5 · p. 7

5.4 Generation Quality

  • Across ArXiv and OpenWebText2, LM-Infinite retains generation quality beyond training lengths, comparable to or better than fine-tuned MPT-7B-Storywriter.With 32K-long sequences, it achieves 2.7× decoding speedup and 7.5× GPU memory saving.Source: Table 3 · p. 8, §5.4 · p. 9

5.3 Downstream Evaluation

  • In zero-shot evaluation, LM-Infinite gains 37.2 percentage on Passkey Retrieval and 1.2 percentage on Qasper over the reported baselines.Passkey truncation performance depends on whether the retained context contains the passkey; Qasper benefits from attending to top-k middle tokens.Source: §5.3 · pp. 8–9, Table 2 · p. 8

5.4 Generation Quality

  • At similar computation, LM-Infinite exceeds truncation by about 5 BLEU; at similar BLEU, it incurs <25% computational overhead.Figure 6 compares computation complexity with generation quality.Source: §5.4 · p. 9, Figure 6 · p. 9

6 Conclusions and Future Work

The paper concludes that LM-Infinite enables zero-shot length generalization without parameter updates, extending models trained on 2K or 4K segments to inputs up to 200M tokens while retaining perplexity. It also reports downstream-task improvements and efficiency gains, and identifies further applications and training-efficiency questions for future work.

6 Conclusions and Future Work

  • LM-Infinite lets models trained on 2K or 4K segments handle inputs up to 200M tokens while retaining perplexity.The method requires no parameter updates.Source: §6 · p. 9
  • The method improves zero-shot performance on Passkey Retrieval and Qasper and delivers 2.7× decoding speedup and 7.5× memory savings over the original model.The paper attributes practical access to long sequences to its computational efficiency and ease of use.Source: §6 · p. 9
  • Future work could test whether these techniques make LLM pre-training and fine-tuning more efficient and effective.The paper also proposes exploring long reasoning, long dialogue, retrieval-augmented generation, and long literature generation.Source: §6 · p. 9

Limitations

The evaluation covers open-domain LLMs but not proprietary models such as ChatGPT, whose source code was unavailable. Resource and time constraints also limited testing beyond the evaluated lengths, while the method is designed for Transformers with relative positional encodings.

Limitations

  • The method was not evaluated on proprietary models such as ChatGPT because their source code was unavailable.The paper says it evaluated a wide range of open-domain LLMs.Source: Limitations · p. 10
  • Resource and time constraints left inputs as long as 1G tokens unevaluated.Source: Limitations · p. 10
  • LM-Infinite is designed for Transformer models with relative positional encodings, and more efficient fine-tuning or pre-training remains future work.Source: Limitations · p. 10

A Implementation Details

The implementation-details section reports the experimental resources and software used. Experiments ran on single A100 GPUs, and the work involved no training or fine-tuning.

A.1 Computational resources

  • Experiments used single A100 GPUs with 80GB of memory each.Source: §A.1 · p. 14
  • The 200M-length generalization runs took 20 hours, while each downstream task took 3–7 hours to evaluate.Source: §A.1 · p. 14
  • The experiments involved no training or fine-tuning; models were loaded using Hugging Face code, and Rouge and BLEU scores came from the evaluate package.Datasets and models were used with permission from their licenses.Source: §A.1 · p. 14

A.2 The size of starting attention span

This material reports an attention-span sensitivity check, details the optional top-k middle-token technique and its evaluation settings, and surveys later related work and benchmarks.

A.2 The size of starting attention span

  • On 16k-token ArXiv validation sequences, LM-Infinite tolerates a wide range of starting attention-span values without sacrificing average NLL.Source: §A.2 · p. 14, Table 4 · p. 14

A.3 Reintroducing Top-k Middle Tokens

  • The optional top-k technique selects middle tokens with the largest attention logits independently per head, for layers above h; Llama-2 uses k = 5 and h = 5.The hyperparameters were selected using held-out Passkey Retrieval validation; Qasper prompts use 6K-token input subsequences.Source: §A.3 · p. 14
  • Tables 5, 6, and 7 examine varying k, middle-token attention distance, and h, respectively.Their captions specify the ablation dimensions but not their outcomes.Source: Table 5 · p. 14, Table 6 · p. 14, Table 7 · p. 15

B Additional Related Work

C Formal Statement of Theorem 1

Theorem 2 examines whether relative-position attention can distinguish increasing token distances under bounded representational capacity. Its proof uses a covering-number argument to show a conflict between that capacity and an unbounded number of distinguishable distance groups.

C Formal Statement of Theorem 1

  • The theorem assumes a finite pseudo-dimension for the family of distance-based logit functions and an unbounded, non-decreasing number of groups of distances that are distinguishable beyond tolerance ϵ.Distances within each group have expected squared logit difference at most ϵ under the training distributions for queries and keys.Source: C Formal Statement of Theorem 1 · p. 15
  • The proof invokes a covering-size bound for function families with finite pseudo-dimension, using the logit range and pseudo-dimension to bound the number of functions needed to cover the family.The cited lemma defines cover size through expected squared distance under a probability measure.Source: C Formal Statement of Theorem 1 · p. 15
  • The proof argues by contradiction: if the number of distance groups grows at least logarithmically with n, the cover-size bound forces two distinct distance logits to be within ϵ.The resulting similarity contradicts the assumed distinguishability of those distances.Source: C Formal Statement of Theorem 1 · p. 15

D Formal Statement and Proof of Proposition 1

The appendix formalizes an attention-entropy lower bound, links early-token features to the effect of windowed attention, and checks the finite pseudo-dimension assumption for several positional-encoding schemes.

D Formal Statement and Proof of Proposition 1

  • The attention distribution's entropy is asymptotically lower-bounded by ln n, so it approaches infinity as the number of logits grows.The proof starts from the discrete entropy definition and expresses attention probabilities through normalized exponentiated logits.Source: D Formal Statement and Proof of… · pp. 15–16

E More on Implicitly Encoded Positions

  • The appendix illustrates token-feature separation across layers with a two-dimensional PCA projection in Figure 7.The projection is presented as further visualization of the initial tokens' distinct region in higher layers.Source: E More on Implicitly Encoded Positions · p. 16
  • In layers 2 and above, the initial few tokens occupy a distinct feature region; discarding them can move the attention output into a different region.The authors use this to explain why windowed attention may not work and why their Λ-shaped attention keeps those tokens.Source: E More on Implicitly Encoded Positions · p. 16

F Pseudo-Dimension Assumption on Attention Logit Functions

  • For T5-Bias and Alibi, the difference between any two distance-logit functions is input-independent, yielding pseudo-dimension 1.The passage states that this family cannot shatter any set larger than two.Source: F Pseudo-Dimension Assumption on… · p. 16
  • Windowed attention restricts the distance-logit family to dmax + 1 functions, so its pseudo-dimension is at most dmax + 1.Source: F Pseudo-Dimension Assumption on… · p. 16
  • For NoPE, all distances are equivalent because no explicit positional encoding is used, and the pseudo-dimension is 1.Source: F Pseudo-Dimension Assumption on… · p. 16
  • For RoPE, the logit family has pseudo-dimension at most 2k; the appendix says the CAPE and XPos analyses are similar.The bound follows from a finite set of fixed sinusoidal functions and a bound on the pseudo-dimension of their sums and scalings.Source: F Pseudo-Dimension Assumption on… · p. 16

G Computational Efficiency Evaluation

On 32k ArXiv sequences, LM-Infinite reduces Llama-2-7B encoding and decoding time and GPU memory use. The appendix also relates early-token features to windowed attention and provides example-generation material after 8k contexts.

G Computational Efficiency Evaluation

  • 2.72x decoding speedup and 7.53× memory saving: LM-Infinite decodes at 2.70s per token versus 7.34s for vanilla Llama-2-7B and uses 4.41Gb per sequence versus 33.2Gb.On 100 ArXiv sequences of 32k length, encoding averages 15.26s per sequence versus 48.19s for vanilla, a 3.16x speedup; evaluation used one 80GB A100 and DeepSpeed Zero3 with code modifications to avoid out-of-memory errors.Source: G Computational Efficiency Evaluation · p. 16

H Example Generation

  • In Llama, at second or higher layers, the initial few tokens encode a strong position signal and occupy a distinct feature region; abandoning them might move the attention output vector out of the pre-training distribution.Source: Figure 7 · p. 17
  • Table 8 presents example text generations on ArXiv and OpenWebText2 after 8k context lengths.Source: Table 8 · p. 18
Loading 2308.16137v7…