One-sentence Summary

University of Illinois Urbana-Champaign and Meta propose LM-Infinite, a zero-shot method that identifies three key factors behind LLM length generalization failure and addresses them without parameter updates, enabling models pre-trained on $2\text{K}\text{-}4\text{K}$2K-4K token segments to handle inputs up to $200\text{M}$200M tokens while achieving $2.7\times$2.7× decoding speedup and $7.5\times$7.5× memory reduction.

Key Contributions

• The paper provides a theoretical and empirical account of why short-context LLMs fail on longer inputs, identifying three length-generalization factors and showing that sliding-window attention and relative positional encodings alone are insufficient.

• The paper introduces LM-Infinite, an inference-time length-generalization method for off-the-shelf LLMs that keeps attention behavior closer to the training regime, enabling models pre-trained on 2K or 4K segments to process inputs up to 200M tokens without parameter updates while retaining perplexity.

• Evaluations show that LM-Infinite improves zero-shot performance on Passkey Retrieval and Qasper and provides efficiency gains, including a 2.7x decoding speedup and a 7.5x memory saving over the original model.

Introduction

Large language models often train on sequences under about 4K tokens because Transformer attention scales poorly with input length. This creates a practical gap for applications that need long-range context, such as scientific paper understanding, code repositories, and extended dialogue.

Prior work has used relative positional encodings, sliding windows, retrieval memories, or long-sequence finetuning, but these approaches still struggle to explain and fix length generalization failures. The authors identify three core issues: unseen token distances, unseen attention context sizes, and hidden absolute-position effects near sequence starts.

The authors propose LM-Infinite, a training-free method that uses a Λ-shaped attention mask and a ceiling on attention distances. This lets off-the-shelf LLMs with relative positional encodings generalize to very long inputs, up to 200M tokens in their experiments, while improving speed and reducing GPU memory use.

Method

The authors leverage an important observation about Transformer hidden states: even without explicit absolute positional embeddings, the first few tokens can occupy a distinct region in the feature space. A PCA projection of second-layer hidden states for Llama-2 shows that initial-token representations form clusters that are separate from later-token representations. Since attention aggregates value vectors, removing these starting tokens can prevent attention outputs from reaching the feature regions associated with sequence beginnings. This motivates explicitly preserving a small starting span during attention. As shown in the visualization below:

Building on this observation, the authors propose LM-Infinite, a parameter-free modification of the attention mechanism for Transformer models with relative positional encodings. The method is designed for zero-shot length generalization: it changes how attention logits and attention masks are computed, but does not require parameter updates or retraining. Its core architecture consists of a $\Lambda$Λ-shaped attention mask, a distance ceiling, and an optional top-$k$k intermediate-token retrieval mechanism. Refer to the framework diagram below:

The $\Lambda$Λ-shaped attention mask restricts each token to two attention spans. The starting span allows a query token to attend to the first $n_{\mathrm{starting}}$nstarting​ tokens, provided those tokens appear before the current position. The ending span allows the query token to attend only to recent tokens within the pretraining context length. For a query at position $i$i, the attended set can be written as

$$\mathcal{A}_i = \{j \le i : j < n_{\mathrm{starting}}\} \cup \{j \le i : i-j \le L_{\mathrm{pretrain}}\}.$$Ai​={j≤i:j<nstarting​}∪{j≤i:i−j≤Lpretrain​}.

All other middle tokens are masked out. This limits the number of tokens participating in each attention computation, while still guaranteeing that the starting tokens remain visible. If $n_{\mathrm{starting}} = 0$nstarting​=0, the model reduces to a recent-window attention pattern and loses access to the initial-token feature space, which harms performance.

LM-Infinite also introduces a distance ceiling. For two tokens with relative distance

$$d_{ij} = i-j,$$dij​=i−j,

the original relative-positional attention logit can be viewed as a function

$$w(\mathbf{q}_i, \mathbf{k}_j, d_{ij}).$$w(qi​,kj​,dij​).

The method replaces the raw distance with a bounded effective distance:

$$d'_{ij} = \min(d_{ij}, L_{\mathrm{pretrain}}),$$dij′​=min(dij​,Lpretrain​),

and computes the attention logit as

$$w(\mathbf{q}_i, \mathbf{k}_j, d'_{ij}).$$w(qi​,kj​,dij′​).

This modification only matters when distant query tokens attend to the starting span. It prevents unseen, excessively large distance values from entering the attention logit function, thereby keeping the positional input within the range observed during pretraining.

The optional top-$k$k component extends the mask by re-introducing a small number of middle tokens. For each attention head, the method can select the $k$k intermediate tokens with the largest attention logits and add them to the attended set. The selected tokens are chosen independently per head in layers deeper than a threshold layer $h$h. This is useful for downstream tasks where information located in the middle of a long context must be retrieved, while still limiting the overall attention entropy introduced by a growing sequence length.

The design can be summarized conceptually by partitioning a long context into three functional regions. The starting tokens carry strong position-sensitive features and should remain visible. The rear tokens supply recent relative positional information around the current query. The middle tokens encode information that is less dependent on exact position, and attending to too many of them can increase attention entropy and destabilize inference at unseen lengths.

Implementation is lightweight because LM-Infinite only changes the attention computation. For RoPE-based models, tokens in the recent window follow the original rotated attention calculation. For the starting span, excluding its overlap with the recent window, key vectors are kept unrotated and query vectors are rotated to a fixed effective distance $L_{\mathrm{pretrain}}$Lpretrain​ before composing the final logits. For AliBi-style models, the positional offset matrix is clipped so that effective distances respect the same ceiling, and the $\Lambda$Λ-shaped mask is applied afterward. This preserves the original relative positional encoding interface while adapting it for extended inference lengths.

Experiment

Experiments evaluate LM-Infinite on several long-context language models using ArXiv and OpenWebText2 for language modeling, zero-shot passkey retrieval and Qasper for downstream tasks, and long text generation with BLEU and ROUGE. The results show that it allows models to generalize far beyond their pretraining context while maintaining stable perplexity, and ablations confirm that both bounded token distance and selective attention are necessary for length generalization. It also improves long-input retrieval and question answering, preserves generation quality beyond training lengths, and offers a promising alternative to expensive fine-tuning.

LM-Infinite improves long-context language modeling on ArXiv and OpenWebText2 without parameter updates. It usually achieves the lowest perplexity among compared models at evaluated lengths, especially when test sequences exceed the training segment lengths. It also remains close to a fine-tuned long-context baseline despite requiring no finetuning. LM-Infinite obtains best perplexity in most evaluated length columns. Vanilla baselines degrade sharply beyond their training length, while LM-Infinite preserves stable perplexity. MPT-7B with LM-Infinite nearly matches the fine-tuned Storywriter model.

LM-Infinite lets Llama-2 handle long-input retrieval and question answering better than both the unmodified model and a truncation baseline. Across increasing context lengths, it maintains comparatively high Passkey Retrieval accuracy, while the baseline loses useful information when it discards earlier context. The original model fails to retrieve passkeys and shows only minimal Qasper accuracy. LM-Infinite consistently improves Passkey Retrieval over truncation as context length grows. LM-Infinite gives a small average gain on Qasper over truncation, indicating better use of information that appears outside the recent context.

Across the reported long text generation results on ArXiv, LM-Infinite helps pretrained models maintain usable generation quality as context lengths increase beyond their original training length. Vanilla MPT-7B and Llama-2 often degrade to near-zero scores at longer milestones, while LM-Infinite remains much closer to the strong fine-tuned MPT-7B-Storywriter baseline. The BLEU and ROUGE trends indicate better long-range generation quality without parameter updates. Vanilla models show major generation failures at longer lengths, with many scores dropping close to zero. LM-Infinite preserves generation quality across longer sequences and stays comparable to the fine-tuned MPT-7B-Storywriter baseline. Adding LM-Infinite to Llama-2 produces a clear improvement over the vanilla model at longer generation milestones.

The experiments evaluate LM-Infinite on long-context language modeling, long-input retrieval and question answering, and long text generation using ArXiv, OpenWebText2, Passkey, and Qasper benchmarks. The results show that it usually achieves lower perplexity and more stable generation and retrieval performance as context length increases beyond the original training limit, while unmodified or truncated baselines degrade sharply. Overall, LM-Infinite improves long-context behavior without parameter updates and remains close to strong fine-tuned long-context baselines.