LM-Infinite: Simple On-the-Fly Length Generalization for Large Language Models

ArXiv: 2308.16137

🎯 Pitch

Transformers collapse on long inputs not just due to quadratic costs, but because their attention mechanisms break mathematically on unseen distances and token counts. LM-Infinite fixes this zero-shot—without any fine-tuning—by forcing attention into a Λ-shape and capping distance values, unlocking sequences up to 200 million tokens while preserving perplexity.

---

1. Executive Summary

This paper analyzes why Transformer-based large language models fail to generalize to input sequences longer than those seen during training, identifying three root causes—attention logit explosion from unseen inter-token distances, attention entropy growth from attending to an unbounded number of tokens, and distinct representational space occupied by initial tokens that sliding-window attention discards—and proposes LM-Infinite, a zero-shot method combining a Λ-shaped attention mask (attending only to initial tokens and a recent window) with a distance ceiling (capping relative position values to the pretraining maximum). Without any parameter updates, LM-Infinite enables LLMs pretrained on 2K–4K tokens to generalize to sequences up to 200M tokens while retaining perplexity, achieving 2.7× decoding speedup and 7.5× memory savings, with a 37.2 percentage-point gain on Passkey Retrieval—establishing that the length generalization failure is addressable through principled attention manipulation alone, not requiring expensive fine-tuning on long sequences.

2. Context and Motivation

The Core Problem: Transformers Break When They See Longer Sequences Than They Trained On

Today's large language models are almost universally trained on short text segments—typically 2K or 4K tokens (Section 1). This is not a deliberate design choice about what constitutes an ideal context length; it is a hard constraint imposed by the quadratic complexity of the standard Transformer self-attention mechanism. When you double the sequence length, you quadruple the computation and memory required for attention. Training on 32K-token segments would be 64× more expensive than training on 4K-token segments for the attention layers alone, making it economically prohibitive to simply "train on longer texts" as a universal solution.

Yet the world is full of long documents. Scientific papers routinely span tens of thousands of words. Code repositories contain millions of tokens spread across thousands of files. Long conversational histories in chatbots or multi-turn planning systems can accumulate context far beyond what the model was designed to handle. When an LLM pretrained on 4K segments is asked to process an 8K, 16K, or 100K-token input, it fails catastrophically—perplexity explodes, outputs become nonsensical, and the model may even produce NaN values (as the authors observe with Llama-2 on sequences slightly longer than 10K tokens; Figure 3, Table 1). This is the length generalization failure: the model cannot extrapolate its learned behaviors to sequence lengths it never encountered during training.

The gap this paper addresses is therefore both practical and fundamental. Practically, it limits what LLMs can be used for out-of-the-box without expensive retraining. Fundamentally, it reveals that something about the Transformer architecture—even with techniques explicitly designed to handle arbitrary lengths—prevents smooth extrapolation to unseen sequence dimensions. The authors frame a "paradox" that is central to their motivation: why should LLMs with relative positional encodings—which compute attention weights based on token-to-token distances rather than absolute positions—still fail on long inputs? After all, distance-based attention should, in principle, be invariant to the absolute length of the sequence. The fact that it is not means that there are deeper, less obvious failure modes at play.

The Real-World Stakes

The importance of solving this problem extends far beyond academic curiosity about Transformer internals. The authors highlight several concrete application domains in Section 1:

• Scientific literature processing: Encoding full research papers (which frequently exceed 8K tokens) to answer questions, generate summaries, or extract structured knowledge.
• Code repository understanding: Large codebases contain millions of tokens across files, and a model that can only attend to the most recent 4K tokens of a repository cannot reason about cross-file dependencies or architectural patterns that span the entire codebase.
• Long-form dialogue systems: Customer support conversations or multi-turn planning sessions can accumulate context that far exceeds standard training lengths.
• Retrieval-augmented generation (RAG): When LLMs are given retrieved documents as context, the combined prompt can easily exceed pretraining lengths if multiple documents or long passages are included.

Without addressing length generalization, practitioners are forced into one of several unsatisfactory workarounds: (1) truncate inputs, losing potentially critical information; (2) fine-tune the model on longer sequences at substantial computational cost; or (3) use retrieval-based systems that only expose the model to small chunks at a time, sacrificing the ability to reason across chunks. None of these are ideal, and all represent compromises that a genuine solution to length generalization would eliminate.

There is also a computational equity dimension to this problem. Fine-tuning on long sequences requires enormous computational resources—clusters of GPUs processing millions of long-context tokens. This locks out smaller research labs and practitioners who lack access to such infrastructure. A zero-shot method that works on off-the-shelf pretrained models, as LM-Infinite promises, democratizes access to long-context capabilities. The authors make this point explicitly in Section 6, noting that LM-Infinite enables "researchers without enormous computational resources to use LLMs on long sequences."

How Prior Work Approached This (And Why It Falls Short)

The paper situates itself within a rich landscape of prior attempts to extend LLMs' effective context length. These approaches can be broadly categorized, and the authors systematically identify where each falls short.

Relative Positional Encodings: Necessary but Not Sufficient

The most fundamental architectural advance was the shift from absolute to relative positional encodings (Section 2.1). Earlier Transformers like the original BERT and GPT used learned absolute position embeddings—each position from 0 to L_max had its own embedding vector. When a model trained with position 0–4095 encountered position 4096 at test time, it had no embedding for that position and would either fail or behave unpredictably.

Relative encodings solve this specific problem by making attention a function of the distance between tokens, not their absolute positions. Rotary Position Embedding (RoPE; Su et al., 2021) rotates the query and key vectors before computing their inner product, such that the resulting attention logit depends only on the relative distance i - j. Alibi (Press et al., 2021) adds a linear bias -m(i - j) to the attention logit, again making it distance-dependent. These techniques are now standard in state-of-the-art LLMs: LLaMA, Llama-2, and GPT-J use RoPE; MPT-7B uses Alibi.

The intuition was that relative encodings should naturally generalize to any length—since distances are just integers, and the model can compute w(q, k, d) for any d, there should be no architectural barrier to processing sequences of arbitrary length. As the authors note in Section 2.1, "despite some promising empirical evidence, length generalization failures are still widely observed when directly applied to large language models." The models still break. This means that the problem is not (solely) about the encoding scheme being unable to represent long distances—it is about something subtler in how the model's computations distribute when exposed to unseen distances and unseen numbers of tokens.

Sliding-Window Attention: Efficient but Destructive

Another family of approaches restricts each token's attention to a local window of nearby tokens, rather than attending to the full sequence. This is the core idea behind Longformer (Beltagy et al., 2020), BigBird (Zaheer et al., 2020), and LongNet (Ding et al., 2023). The motivation is twofold: it reduces the quadratic complexity of attention to linear, and it prevents the model from having to process distances longer than the window size.

However, the authors identify a fundamental problem with this approach that prior work had not articulated (Section 3, Factor 3). Sliding-window attention discards the first few tokens of the sequence entirely once the window has moved past them. If these initial tokens encode crucial information—either semantic content from the prompt's beginning or implicit position signals—the model loses access to it. The authors demonstrate empirically that the vector representations of the first few tokens (positions 0–~25) occupy a distinct region of the feature space compared to all later tokens (Figure 1c). When sliding-window attention excludes these initial tokens, "the attention output will reside in a different region" (Appendix E), forcing the model to operate on a distribution of value vectors it was never trained to handle.

More critically for this paper's positioning: sliding-window attention is typically introduced during pretraining, not applied post-hoc to already-trained models. An off-the-shelf Llama-2 model was trained with full causal attention over 4K tokens; simply substituting a sliding-window mask at inference time changes the attention pattern the model was optimized for, introducing a distribution shift that degrades performance. The authors show this concretely in Figure 5: the "window" baseline (sliding-window attention applied to a pretrained LLaMA) produces "the second worst NLL values" among all ablation conditions.

Fine-Tuning on Longer Sequences: Expensive and Lacks Mechanistic Insight

Perhaps the most straightforward solution is to simply continue training (fine-tune) the model on longer text segments. This is the approach taken by several works: MPT-7B-Storywriter (Team, 2023) was fine-tuned on 65K-token sequences; LongLLaMA (which extends LLaMA to 8K) and various other fine-tuning approaches (Chen et al., 2023a; Tworkowski et al., 2023) similarly invest substantial compute into teaching the model to handle longer contexts through additional training.

The authors acknowledge that this approach works (MPT-7B-Storywriter serves as a strong baseline in their experiments), but they identify two key limitations:

1. Computational cost: Fine-tuning on long sequences requires processing orders of magnitude more tokens than the original pretraining, and the quadratic attention cost means that each training step is substantially more expensive. This puts long-context fine-tuning out of reach for most practitioners.

2. Lack of mechanistic understanding: Fine-tuning treats the length generalization failure as a black-box problem—the model sees longer texts and eventually learns to handle them through gradient updates. But it does not explain why the failure occurs in the first place. Without this understanding, we cannot know whether fine-tuning truly addresses the root causes or merely patches symptoms, nor can we design more efficient interventions. As the authors put it (Section 1), these approaches "do not address the underlying causes of length generalization failures."

Retrieval-Based and Memory-Augmented Methods: Circumvention, Not Solution

Some approaches avoid processing long contexts in full by augmenting LLMs with retrieval mechanisms (Wu et al., 2021; Guu et al., 2020; Borgeaud et al., 2022) or external memories (Khandelwal et al., 2019; Yogatama et al., 2021). These methods allow the model to access information from a large database without attending to all of it simultaneously, effectively sidestepping the length generalization problem.

The authors note that these designs "usually need finetuning and are not directly compatible with the existing LLMs" (Section 2.2). More fundamentally, retrieval-augmented approaches solve a different problem than what LM-Infinite addresses: they enable access to a large corpus through selective retrieval, but they do not allow the model to simultaneously attend to all parts of a single very long sequence for holistic reasoning. For tasks requiring integration of information across distant parts of a document—comparing the methods section of a paper with its results, or tracking a variable's definition to its use 100 pages later—the ability to attend across arbitrary distances within a single coherent context is essential.

How This Paper Positions Itself: Diagnosis Before Treatment

The central intellectual move of this paper is to diagnose before treating. Rather than proposing yet another architectural modification or training recipe, the authors first ask: what exactly goes wrong inside a Transformer when the input gets too long? Their answer—the three factors explicated in Section 3—provides a mechanistic account of the failure that prior work had not articulated.

This diagnostic approach leads naturally to their proposed solution. Each component of LM-Infinite maps directly to one of the identified failure modes:

• The Λ-shaped attention mask resolves Factor 2 (too many tokens to attend to causes entropy explosion) by limiting the attention context to a fixed-size window, and resolves Factor 3 (discarded initial tokens) by explicitly preserving the first n_starting tokens in the attention span.
• The distance ceiling resolves Factor 1 (unseen distances cause attention logit explosion) by capping relative distances to L_pretrain, preventing the model from ever encountering a distance value it was not trained on.

This mapping from diagnosis to treatment is what distinguishes LM-Infinite from prior heuristic approaches. A sliding-window attention mask (Beltagy et al., 2020) also limits the attention context, but it does so without understanding why the initial tokens matter—and as a result, it discards them and fails. The Λ-shape is not arbitrary; it is precisely the minimal modification to a windowed attention pattern that preserves the essential representational properties the model learned during pretraining.

The paper also positions itself as zero-shot and compatible with off-the-shelf models. Unlike fine-tuning approaches that require access to training infrastructure and long-sequence data, LM-Infinite is a plug-and-play modification to the attention mechanism that requires no parameter updates. This is not just a convenience—it is a commitment to the idea that the length generalization failure is not about missing knowledge (which would require training to acquire) but about computational features drifting out of distribution (which can be corrected by constraining the computation itself). The fact that the method works across four model families (LLaMA, Llama-2, GPT-J, MPT) with two different positional encoding schemes (RoPE and Alibi) substantiates this claim: the underlying failure modes are architectural, not model-specific.

The Conceptual Framework: Relative Positional Attention as a Partitioned Space

The paper provides a conceptual model in Figure 2(b) that organizes the long context into three functional regions:

• Starting tokens (initial ~few dozen): These encode strong absolute position information. Their value vectors occupy a distinct region of the feature space. Removing them would shift the expected attention output to a different region than what the model was optimized for during training. This is why the Λ-shaped mask must include them.

• Rear tokens (within L_pretrain of the current position): These provide primarily relative position information and benefit from the "recency bias" that LLMs learn during pretraining—tokens closer to the current position tend to be more relevant. This is why a local window around the current position is essential.

• Middle tokens (everything between the starting tokens and the rear window): These encode less position-sensitive information. Including too many of them does more harm than good because it dilutes the attention weights and causes entropy explosion (Factor 2). However, for certain downstream tasks where key information may be located in the middle (e.g., Passkey Retrieval buries the answer at a random position), the optional top-k reintroduction mechanism allows the model to selectively attend to promising middle tokens.

This three-part partition is not presented as a theoretical guarantee but as a conceptual model that rationalizes the design choices in LM-Infinite. It explains why the Λ-shape works rather than merely reporting that it does. It also makes a testable prediction: if the middle tokens are truly "less position-sensitive," then reintroducing them with a fixed distance ceiling of L_pretrain/2 should allow the model to attend to their content without suffering from the distance-related or entropy-related failure modes—which is exactly what the top-k mechanism does.

The Gap This Paper Fills

In summary, the paper addresses a specific, well-defined gap: no prior method enables zero-shot, plug-and-play length generalization for off-the-shelf LLMs while maintaining both perplexity and downstream task performance. Relative positional encodings were theoretically motivated but empirically insufficient. Sliding-window attention requires retraining and discards critical initial-token information. Fine-tuning on long sequences is expensive and lacks mechanistic insight. Retrieval-based methods circumvent rather than solve the core attention problem.

The paper's contribution is not just a new method but a framework for understanding why length generalization fails—and from that understanding, a minimal, principled intervention that addresses each failure mode without requiring any modification to the model's parameters. As Figure 2 illustrates, LM-Infinite is essentially a constraint on the attention computation that keeps all its intermediate quantities (logits, entropies, value-vector regions) within the ranges the model observed during pretraining, allowing it to continue operating normally on sequences orders of magnitude longer than it was designed for.

3. Technical Approach

3.1 Reader Orientation

LM-Infinite is a plug-and-play attention-masking and distance-capping method applied at inference time. It solves the problem that LLMs break catastrophically on sequences longer than their pretraining length by keeping the model's internal computational features—attention logits, entropy, and the set of vector regions reachable through attention—within the distribution the model was trained on.

3.2 Big-Picture Architecture (Diagram in Words)

The system modifies the standard Transformer self-attention computation at every layer, replacing the full causal attention mask with a constrained attention pattern and modifying the distance argument passed to the relative positional encoding function. There are three components:

1. Λ-shaped attention mask — restricts each token to attend only to: (a) the first n_starting tokens of the sequence, and (b) the most recent L_pretrain tokens within a local window. All other tokens are masked out, receiving zero attention weight before the softmax. This controls the number of tokens entering the attention distribution (solving Factor 2) and preserves access to the distinct vector space occupied by initial tokens (solving Factor 3).

2. Distance ceiling — when computing the attention logit w(q, k, d) for relative positional encodings (RoPE or Alibi), the true distance d is replaced with d' = min(d, L_pretrain). This prevents the model from ever processing a distance larger than what it saw during pretraining (solving Factor 1).

3. Optional top-k middle token reintroduction — for downstream tasks where critical information may be located in the masked middle region, each attention head in layers above a threshold h can independently select the k middle tokens with the highest attention logits and attend to them with a fixed distance of L_pretrain / 2. This is not needed for perplexity or generation quality; it is a task-specific enhancement.

Information flows as follows: input tokens are embedded and passed into the first Transformer layer → the standard attention score computation q^T k is performed → the Λ-shaped mask zeroes out disallowed token pairs → the distance ceiling modifies the relative position argument before positional biases are added → a standard softmax produces the attention distribution → value vectors are aggregated via the attention weights → the output proceeds through the rest of the Transformer layer normally → this repeats identically at every layer.

3.3 Roadmap for the Deep Dive

• First, the three failure modes from Section 3, recapped with their mechanistic implications—because every design choice in LM-Infinite is a direct response to a specific failure mode, and understanding the cause-effect mapping is essential.
• Second, the Λ-shaped attention mask in full detail—its structure, the n_starting parameter, how the two attention spans are composed, and why the Λ-shape specifically works when pure windowed attention fails.
• Third, the distance ceiling mechanism—how it is implemented for RoPE and Alibi separately, and the theoretical justification (Theorem 1) explaining why capping distances prevents attention logit explosion.
• Fourth, the optional top-k middle token mechanism—when it is needed, how k and h are selected, and why reintroducing middle tokens with a fixed distance works without reintroducing the original failure modes.
• Fifth, the conceptual model in Figure 2(b) that unifies these design choices—the three-part partition of a long sequence and what each region encodes.
• Sixth, implementation specifics for RoPE and Alibi—the concrete code-level transformations that realize LM-Infinite in practice.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a diagnosis-and-intervention paper whose core idea is that three identifiable computational distribution shifts cause length generalization failure, and a minimal constraint on the attention mechanism—restricting which tokens can be attended and capping the distance values—is sufficient to prevent all three shifts simultaneously.

---

The Three Failure Modes Recapped (Mechanistic Root Causes)

Before describing LM-Infinite's components, it is essential to understand exactly what goes wrong inside a Transformer when the input exceeds pretraining length, because each component of LM-Infinite is a direct response to one of these three failure modes. The analysis in Section 3 is the intellectual foundation of the entire method.

Factor 1: Unseen distances cause attention logits to explode. With relative positional encodings, the attention logit between two tokens is computed as w(q, k, d), where d = i - j is the integer distance between their positions. During pretraining on sequences of length L_pretrain, the model only ever processes distances d ∈ {1, 2, ..., L_pretrain - 1}. When a test sequence exceeds L_pretrain, distances d > L_pretrain appear for the first time.

Theorem 1 formalizes why this is catastrophic. The logit function w(·, ·, d) must be able to distinguish different distances—otherwise all long-range token pairs would be treated identically regardless of whether they are 10K or 100K apart. If the function family H = {w(·, ·, d) | d ∈ ℕ} has bounded pseudo-dimension r (a measure of representational capacity; all common relative encoding schemes satisfy this, as shown in Appendix F), and if it can distinguish α(n) distinct distance groups within [0, n], then the magnitude of the logits must grow at least as:

$\sup_{q,k,d \leq n} |w(q, k, d)| \geq \left(\frac{\alpha(n)}{2}\right)^{\frac{1}{2r}} \frac{\epsilon}{4e}$

where α(n) is the number of distinguishable distance groups for sequence length n, r is the pseudo-dimension of the logit function family, and ϵ is the distinguishability threshold between groups.

What it computes: a lower bound on the maximum absolute value of any attention logit as a function of sequence length n. As n grows and the number of distinct distance groups α(n) increases (since longer sequences contain more distinct distances), the term (α(n)/2)^{1/(2r)} grows, forcing the logits to take on increasingly extreme values.

Why this form: the proof (Appendix C) proceeds by contradiction using Haussler's covering lemma. If logits were bounded, the function family H would have a covering number smaller than the number of distances it needs to distinguish, contradicting the assumption that w can tell distances apart. The pseudo-dimension r appears in the exponent because it governs how many distinct functions can be packed into a bounded range—the smaller r is, the faster the logits must grow to accommodate new distances. RoPE, for example, has r = 2k where k is the feature dimension, giving it substantial capacity but still finite, meaning logit explosion is inevitable.

The empirical evidence is shown in Figure 1(a): the maximum attention logit across all heads jumps from roughly 5–10 at 4K length to 15–20 at 8K length—an explosion in magnitude. When these inflated logits pass through the softmax, they produce near-one-hot attention distributions focused on a tiny subset of available tokens, breaking the model's ability to integrate information across the full context.

Factor 2: Attending to an unbounded number of tokens causes attention entropy to grow without bound. At position i, the model distributes its attention over all i preceding tokens. The entropy of this attention distribution measures how evenly the model spreads its attention:

$H\left(\left\{\frac{e^{w_i}}{\sum_{j=1}^n e^{w_j}}\right\}_{1 \leq i \leq n}\right) = \Omega(\ln n)$

where w_i are the attention logits bounded in [-B, B] and n is the number of attended tokens.

What it computes: a lower bound on the entropy of a softmax attention distribution over n tokens when logits are bounded. The entropy grows at least logarithmically with n.

Why this form: the proof (Appendix D, Proposition 2) is a direct algebraic manipulation. Entropy is defined as -Σ p_i ln p_i. Substituting the softmax expression and simplifying yields -Σ (e^{w_i} / Σ e^{w_j}) w_i + ln Σ e^{w_j} ≥ -max_i w_i + ln(n e^{-B}) = ln n - 2B. The ln n term emerges from the sum inside the logarithm: even if all logits were equal (producing uniform attention), the normalization constant Σ e^{w_j} grows with n, and the entropy must increase to accommodate the larger support. The lower bound -2B reflects the worst case where one token dominates attention, but the ln n term is unavoidable regardless of the logit pattern.

At pretraining length n = 4096, the entropy is approximately ln(4096) ≈ 8.3 nats. At test length n = 32768, it is ln(32768) ≈ 10.4 nats—outside the training distribution. Figure 1(b) empirically confirms this ever-increasing entropy curve. High entropy means the model's attention is diluted across many tokens rather than focused on the most relevant ones, degrading the quality of the aggregated value vector. The model was optimized under the assumption that attention would be distributed over roughly 4K tokens; distributing it over 32K tokens means each token receives a smaller fraction of the total attention mass, making it harder for the model to select and amplify important signals.

Factor 3: Initial tokens occupy a distinct representational space and cannot be discarded. The authors make a counterintuitive observation (Observation 1): even without explicit absolute position embeddings, the first few tokens of a sequence have attention outputs that occupy a different region of the feature space compared to all later tokens. This follows from Kazemnejad et al. (2023)'s Theorem 1, which proves that absolute position can be implicitly encoded in the outputs of even a single attention layer without any positional encoding at all—the starting tokens' signals are the strongest and easiest to distinguish from other tokens because of their unique attention pattern (they attend to few or no preceding tokens). Since relative positional encoding is strictly more expressive than no encoding, the same conclusion applies.

Figure 1(c) visualizes this: a PCA projection of the hidden states after layer 2 of Llama-2 shows the first ~0–25 tokens (blue) concentrated in two distinct clusters (upper center and lower right) that are far from the thousands of overlapping later tokens (red, in the lower-left region). Additional plots in Appendix E (Figure 7) show this pattern persists across layers 1 through 20.

The implication is severe for sliding-window attention. Attention is a weighted average of value vectors: output = Σ α_i v_i. If the initial tokens' value vectors v_0, v_1, ..., v_25 live in a different part of the vector space than all other tokens, then any attention output that includes contributions from these initial tokens can reach that region of space. If a sliding-window attention mask excludes these tokens (because the window has moved past them), the attention output is constrained to a weighted average of only the later tokens' value vectors—and cannot reach the regions the model learned to access during pretraining. The model is forced to operate on a subspace of the full value-vector space that it never encountered during training, introducing a distribution shift in every subsequent layer's input.

This is why pure windowed attention applied to a pretrained model fails so badly (the "window" curve in Figure 5 produces the second-worst NLL): it isn't just losing semantic information from the prompt's beginning; it is fundamentally altering the representational geometry that the model expects.

---

The Λ-Shaped Attention Mask

The Λ-shaped attention mask is LM-Infinite's mechanism for resolving Factors 2 and 3 simultaneously. It is a binary mask applied to the attention logits before the softmax, setting the logits for disallowed token pairs to -∞ (practically, a very large negative number) so that after softmax, those positions receive zero attention weight.

The mask is defined by two spans for each query position i:

• Starting span: token i attends to tokens 0, 1, ..., n_starting - 1 if i > n_starting. If i ≤ n_starting, the starting span covers all tokens up to i - 1 (standard causal masking for the very beginning of the sequence).

• Ending span: token i attends to tokens i - L_pretrain + 1, i - L_pretrain + 2, ..., i - 1 (the most recent L_pretrain tokens, with standard causal masking ensuring no future tokens are included).

• Overlap: if the two spans overlap (which happens when i is within L_pretrain of the start), the union is taken—no token is double-counted.

All tokens not in either span receive an attention logit of -∞.

The name "Λ-shaped" comes from the shape of the mask when visualized as an attention matrix: for a sufficiently long sequence, the attention pattern forms two bands—a thin horizontal band at the top (all later tokens attending to the first n_starting tokens, forming the left leg of the Λ) and a diagonal band along the main diagonal (tokens attending to the most recent L_pretrain tokens, forming the right leg of the Λ). Between them is a triangular region of zero attention.

The two spans are not symmetric in function. The ending span resolves Factor 2 by bounding the number of attended tokens: each query attends to at most n_starting + L_pretrain tokens, where L_pretrain is 2048 or 4096 depending on the base model. This keeps the attention entropy within the range seen during training—the model distributes attention over roughly the same number of tokens regardless of the total sequence length.

The starting span resolves Factor 3 by ensuring that the distinctive value vectors of the initial tokens are always available to every later token's attention distribution. Even a token at position 100,000 can attend to token 5 (with a distance-ceiling-modified distance, as discussed below), allowing its attention output to reach the regions of feature space that only the initial tokens' value vectors can provide.

The n_starting parameter. The authors sweep this value (Appendix A, Table 4) and find the method is "tolerant with it taking a wide range of values without sacrificing the NLL values." Specifically, on 16K-length sequences from the ArXiv validation set:

Setting n_starting = 0 (i.e., pure windowed attention with no starting token preservation) produces dramatically worse perplexity (6.43 vs. ~1.02), confirming that the starting span is essential. Values from 1 to 100 all produce effectively equivalent performance. The authors "recommend choosing n_starting ∈ [5, 100]" (Section 4). The reason values up to 100 work well is that the PCA analysis (Figure 1c) shows the initial ~25 tokens are the ones occupying the distinct feature region; including a few dozen more does not hurt because those tokens are still within the attention budget and do not shift the distribution meaningfully.

The degradation at n_starting = 1000 and 2000 is because the total number of attended tokens (n_starting + L_pretrain for tokens far from the end) exceeds the pretraining attention budget, bringing Factor 2 (entropy growth) back into play.

Comparison with sliding-window attention. A pure sliding-window mask (the "window" ablation in Figure 5) corresponds to n_starting = 0 and an ending span of L_pretrain tokens. The ablation study shows this produces NLL values that are substantially worse than LM-Infinite, second only to the completely unmodified vanilla model. This is the key evidence that Factor 3 (discarded initial tokens) is a real and severe failure mode independent of Factors 1 and 2—even with the attention context bounded to a fixed size, removing the initial tokens causes a catastrophic distribution shift.

The authors make this point explicitly in Section 4: "As attention is essentially a weighted average over the value vectors, sliding-window attention discards the initial tokens, keeping the attention output from reaching the regions that value vectors of the initial tokens occupy. This enforces the model to handle a different region during the computation, introducing additional generalization challenges."

---

The Distance Ceiling

The distance ceiling resolves Factor 1 by preventing the relative positional encoding function from ever processing a distance larger than what it saw during pretraining. It is a simple numerical transformation applied to the distance argument before the positional bias is computed.

Formally, the original attention logit with relative positional encoding is:

$w(q_i, k_j, d)$

where q_i is the query vector for token i, k_j is the key vector for token j, and d = i - j is their relative distance (always non-negative due to causal masking).

LM-Infinite replaces d with a capped version:

$d' = \min(d, L_{\text{pretrain}})$

producing the modified logit:

$w(q_i, k_j, d')$

What it computes: a clipped distance value. If the true distance d is less than or equal to L_pretrain, it passes through unchanged. If d exceeds L_pretrain, it is replaced by L_pretrain. The positional encoding function w then operates on this capped value as if the tokens were at most L_pretrain positions apart.

Why this form: the alternative—letting d grow without bound—leads to the attention logit explosion proven in Theorem 1. Simply normalizing or scaling the distances (e.g., dividing by a factor) would change the effective distances for all token pairs, including those within L_pretrain, altering the model's behavior even on sequences of training length. The min(d, L_pretrain) operation is identity for all distances the model was trained on and only modifies unseen distances. This ensures that the model's behavior on sequences up to L_pretrain is completely unchanged, maintaining backward compatibility.

There is a subtlety: the distance ceiling is a "hard" cap that makes all distances beyond L_pretrain indistinguishable from each other. Tokens that are 5K positions apart and tokens that are 100K positions apart both receive the same positional bias as if they were L_pretrain positions apart. Theorem 1's premise is that the logit function must distinguish these distances to avoid treating them identically—but LM-Infinite deliberately chooses to treat them identically. This is acceptable because:

1. For tokens within the ending span (≤ L_pretrain distance), the true distance is used, so local structure is preserved exactly.
2. For the starting span tokens, the distance is always capped to L_pretrain, meaning all later tokens see the initial tokens as being effectively at distance L_pretrain. This is a deliberate information loss—the model cannot distinguish whether it is at position 5K or 100K relative to the start—but the authors' empirical results show this information is not critical for maintaining perplexity or generation quality. What matters is that the model can access the initial tokens' value vectors; their exact distance appears to be less important than their content.

Implementation for RoPE. In Rotary Position Embedding, the attention logit between query q_i and key k_j is not computed as an explicit function w(q, k, d) with a separate distance argument. Instead, RoPE rotates the query and key vectors before the dot product, such that the resulting logit implicitly depends on the relative distance. Specifically, each pair of dimensions (a, a+1) in the query and key vectors is rotated by angles θ_a,i = i·ω_a and θ_a,j = j·ω_a, respectively, where ω_a is the rotation frequency for that dimension pair. The dot product ⟨rotated_q_i, rotated_k_j⟩ then depends only on i - j through trigonometric identities.

To apply the distance ceiling in RoPE, the authors keep the key vectors for the initial tokens unrotated (or equivalently, rotated as if they were at position 0) and rotate the query vectors for all later tokens as if they were at a fixed distance L_pretrain from the start. Specifically (Section 4, "Implementation details"):

"In RoPE, attention logits in the ending attention span follow the original calculation. In the starting attention span (excluding its overlap with the ending span), we keep all k vectors unrotated and rotate all q vectors to a fixed distance L_pretrain."

This means that for a token at position i attending to an initial token at position j < n_starting, the effective distance used in the RoPE rotation is not i - j (which could be arbitrarily large), but rather a fixed value L_pretrain. The rotation applied to the query vector is cos(L_pretrain · ω_a) and sin(L_pretrain · ω_a) in the appropriate dimension pairs, regardless of the actual i position.

For tokens in the ending span (the overlap region where a token attends to recent tokens within L_pretrain distance), standard RoPE rotation is used without modification.

Implementation for Alibi. In Attention with Linear Biases, an additive bias -m · (i - j) is added to each attention logit, where m is a head-specific slope. To apply the distance ceiling, the offset matrix is simply clipped:

"We simply clip the offset matrix with a minimum value of -|m * L_pretrain|"

This means that for any token pair with distance d > L_pretrain, the bias is capped at -|m · L_pretrain| rather than continuing to grow linearly with d. This is a direct implementation of d' = min(d, L_pretrain) in the additive bias space.

Both implementations ensure that the positional signal for very long distances is identical to the positional signal for distance L_pretrain, preventing the attention logits from growing to unseen magnitudes.

---

Optional Top-k Middle Token Reintroduction

The Λ-shaped mask eliminates all middle tokens—those between the initial n_starting tokens and the ending window of L_pretrain tokens—from the attention computation. This is necessary to bound the attention entropy (Factor 2), but it has an obvious downside: if important information happens to be located in the middle of a long document, the model cannot access it.

For language modeling (perplexity) and text generation, this is acceptable because token prediction is primarily driven by local context and broad document-level signals from the beginning. The empirical results in Section 5.1 and 5.4 confirm that LM-Infinite without middle tokens achieves strong perplexity and generation quality.

However, for downstream tasks like Passkey Retrieval—where a critical piece of information (a passkey number) is deliberately buried at a random position in a long distractor text, and the model must find and reproduce it—the model must attend to the middle. Similarly, for question answering on long scientific papers (Qasper), relevant passages may be located anywhere in the document.

The optional top-k reintroduction mechanism addresses this (Section 4, "Optionally attending to top-k tokens in the middle"):

"LM-Infinite can optionally attend to k tokens in the middle with the largest attention logits. This is particularly useful in downstream tasks where information in the middle tokens matters."

How it works. For each attention head in layers above the h-th layer, the model computes the attention logits between the current query and all middle tokens (those in neither the starting span nor the ending span) using a fixed distance of d = L_pretrain / 2. This fixed distance is used because:

1. It is within the pretraining distance range, so the model has seen positional signals at this distance during training.
2. It is a "neutral" middle distance that does not encode strong recency or primacy biases, appropriate for tokens whose true distance is unknown and variable.
3. It allows the content-based attention scores (q^T k) to dominate the token selection, with the positional bias being a constant factor.

Among all middle tokens, the k tokens with the highest resulting attention logits are then included in the attention computation for that head, along with the tokens from the starting and ending spans. The remaining middle tokens remain masked out.

Hyperparameter selection. The authors select k and h based on a held-out validation set of Passkey Retrieval (Appendix A.3). They report:

• k = 5 and h = 5 as the chosen values (Section 4).
• Varying k from 1 to 200 shows that k = 3 or k = 5 work best, with accuracy around 0.81; k = 1 drops to 0.69; k = 200 drops to 0.73 (Table 5).
• Varying the attention distance for the middle tokens shows that d = 2048 (= L_pretrain / 2 for Llama-2) works best (0.81 accuracy), with both shorter (512) and longer (4096) distances decreasing performance (Table 6).
• Varying h shows that only applying top-k to layers 5 and above works well (0.94 accuracy), while applying it to all layers (h = 0, accuracy 0.81) or only very high layers (h = 24, accuracy 0.46) degrades performance (Table 7). The authors hypothesize that lower layers encode more position-sensitive features where introducing middle tokens with a fixed distance would be disruptive, while higher layers are more content-focused and can benefit from the additional information.

Why top-k per head independently. The selection is done "independently for each attention head" because different heads specialize in different types of attention patterns. One head might learn to attend to noun phrases, another to syntactic dependencies, another to topic-relevant content. Allowing each head to select its own top-k middle tokens means that head-specific relevance can be captured—a token that is highly relevant for one attention function may be irrelevant for another, and vice versa.

Why this does not reintroduce Factor 2. The total number of attended tokens with top-k reintroduction is n_starting + L_pretrain + k per head. With n_starting ≤ 100, L_pretrain = 4096, and k = 5, this is roughly 4201 tokens—still within the pretraining attention budget. The entropy does not grow unboundedly because the attention distribution is still over a bounded set of tokens. Crucially, the k tokens are not a random sample of middle tokens; they are the highest-scoring tokens, meaning they represent concentrated, high-value information rather than diffuse noise. Adding 5 high-signal tokens to a distribution over 4100 tokens has negligible impact on the overall entropy while potentially providing critical information.

No quality degradation. The authors note (Section 4): "These intermediate tokens do not hurt performance." When top-k is enabled, the model's perplexity and generation quality on standard language modeling tasks are not degraded, because the model can simply assign low attention weights to the reintroduced tokens if they are not useful. The mechanism is strictly additive in terms of what the model can attend to, not what it must attend to.

---

The Conceptual Model: Three Regions of a Long Sequence

The design of LM-Infinite is not ad hoc; it follows from a conceptual model of how relative positional attention partitions a sequence into functionally distinct regions (Figure 2b, Section 4 "Discussion"). Understanding this model explains why the Λ-shape and the distance ceiling work together.

The model divides a long context into three parts:

Starting tokens (initial ~few dozen). These encode strong absolute position information. Theorem 1 from Kazemnejad et al. (2023) proves that the first few tokens in a causal attention architecture inevitably have attention outputs that are distinguishable from all later tokens—their causal mask pattern (attending to very few preceding tokens) makes them unique. In the PCA visualization (Figure 1c), these tokens occupy two distinct clusters far from the main mass of token representations. The authors interpret this as meaning that the initial tokens' value vectors v_0, v_1, ... are necessary for the attention output to reach certain regions of the feature space. Without them, the model operates on a subspace of value vectors, creating a distribution shift in every subsequent layer.

This is why the Λ-shaped mask must include the starting span—not primarily for the semantic content of the initial tokens (though that may also matter), but for their role in the representational geometry of attention. Even if the initial tokens were random noise, their presence in the attention computation might still be necessary because the model was trained to expect attention outputs to be linear combinations that include these specific vectors.

Rear tokens (within L_pretrain of the current position). These provide primarily relative position information and benefit from "recency bias." During pretraining, the model learns that nearby tokens are more relevant for next-token prediction than distant ones, and allocates more attention mass to recent positions. The ending span of the Λ-shaped mask preserves this local structure exactly—tokens within L_pretrain distance are attended with their true relative distances, maintaining the model's learned recency preferences.

The size of this span is set to L_pretrain because this is the maximum local context the model was trained to use. Making it smaller would deprive the model of useful local context; making it larger would include distances beyond pretraining, reintroducing Factor 1.

Middle tokens (everything else). These encode less position-sensitive information. As the distance from both the start and the current position grows, the positional signal becomes less informative—a token at position 5000 in a 10000-token document is neither "recent" nor "initial," and its exact distance matters less for prediction than its content. The authors' empirical finding that middle tokens can be reintroduced with a fixed distance of L_pretrain/2 without hurting performance supports this: the model does not need precise distance information for middle tokens; it primarily needs their content vectors.

This explains why the Λ-shape masks out most middle tokens by default—they contribute to Factor 2's entropy growth without providing sufficient positional benefit to justify the dilution of attention. And it explains why selectively reintroducing the top-k by content-based attention (with a fixed neutral distance) works: it recovers the content benefit without the entropy cost of attending to all middle tokens.

---

Implementation Details: RoPE

For models using Rotary Position Embedding (LLaMA, Llama-2, GPT-J-6B), the implementation requires modifying how the query and key rotations are applied (Section 4, "Implementation details").

The standard RoPE attention logit for a query q_i at position i attending to a key k_j at position j is:

$\langle R_{\Theta,i} \cdot q_i, \; R_{\Theta,j} \cdot k_j \rangle$

where R_{Θ,i} is the rotation matrix applied to the query based on position i, and R_{Θ,j} is the rotation matrix applied to the key based on position j. The rotation for a vector pair (x_a, x_{a+1}) at position p is given by rotating it by angle θ_a,p = p · ω_a.

LM-Infinite partitions this computation into the two attention spans:

Ending span (recent tokens within L_pretrain distance). Tokens in this span are attended with their true distances. The standard RoPE computation applies without modification. For each query-key pair (i, j) where i - j ≤ L_pretrain and j ≥ n_starting, the logit is:

$\langle R_{\Theta,i} \cdot q_i, \; R_{\Theta,j} \cdot k_j \rangle$

which depends on i - j through the trigonometric identity for the difference of rotated vectors.

Starting span (initial n_starting tokens). For query-key pairs (i, j) where j < n_starting and i > L_pretrain (i.e., the query is far enough from the start that the true distance exceeds L_pretrain), the key vectors are kept unrotated (equivalent to applying R_{Θ,0}, the identity rotation for position 0), and the query vectors are rotated as if they were at position L_pretrain:

$\langle R_{\Theta,L_{\text{pretrain}}} \cdot q_i, \; R_{\Theta,0} \cdot k_j \rangle = \langle R_{\Theta,L_{\text{pretrain}}} \cdot q_i, \; k_j \rangle$

This produces an effective distance of L_pretrain for all such pairs, regardless of the actual i - j distance.

Overlap region. When i ≤ L_pretrain, the starting and ending spans overlap. In this region, standard RoPE applies to all tokens (since all distances are within L_pretrain anyway), and the Λ-shaped mask simply includes all tokens up to i - 1 (standard causal attention).

Implementation note. The authors phrase this as: "keep all k vectors unrotated and rotate all q vectors to a fixed distance L_pretrain." This is an efficient implementation strategy—rather than computing different rotations for each initial-token key based on its position and then capping the effective distance in the dot product, they simply do not rotate the keys of initial tokens at all and rotate the queries to a fixed distance. The trigonometric properties of RoPE ensure that the resulting dot product encodes a fixed relative distance.

---

Implementation Details: Alibi

For models using Alibi (MPT-7B), the implementation is simpler because Alibi uses an additive bias rather than vector rotation (Section 4, "Implementation details").

The standard Alibi attention logit is:

$q_i^T k_j - m \cdot (i - j)$

where m > 0 is a head-specific slope. The term -m · (i - j) is an additive penalty that grows linearly with distance, creating a recency bias—tokens that are further apart receive more negative bias.

LM-Infinite modifies this by capping the bias term:

$q_i^T k_j - \min(m \cdot (i - j), \; |m| \cdot L_{\text{pretrain}})$

Equivalently, the offset matrix (which stores the -m · (i - j) values for all token pairs) is "clipped with a minimum value of -|m · L_pretrain|."

What this computes: for any distance d ≤ L_pretrain, the bias is -m · d exactly as in the original Alibi. For any distance d > L_pretrain, the bias is capped at -|m| · L_pretrain. The attention logit does not continue to decrease (become more negative) with increasing distance, preventing the logits from becoming extremely negative (which, through the softmax, would cause them to be exponentially suppressed, creating the same effective problem as logit explosion—extreme concentration of attention mass on a few tokens).

Combined with the Λ-shaped mask. The Alibi implementation is then combined with the same Λ-shaped attention mask: middle tokens (those in neither the starting span nor the ending span) have their logits set to -∞, completely excluding them from the softmax. The bias clipping is applied to the tokens that are in the attention spans—primarily the starting span tokens, whose true distances exceed L_pretrain.

---

Putting It All Together: The Complete LM-Infinite Attention Computation

Assembling the components, the LM-Infinite attention computation at each layer for a query at position i proceeds as follows:

1. Compute raw attention scores: For all keys k_j where j < i (causal masking), compute the base attention logit. For RoPE models, this is q_i^T k_j (the dot product of rotated vectors based on the rules described above). For Alibi models, this is q_i^T k_j (without any positional bias yet).

2. Apply distance ceiling: For RoPE, the rotations for the starting span already encode the min(d, L_pretrain) ceiling. For Alibi, the additive bias -min(m · (i - j), |m| · L_pretrain) is added to the raw logit.

3. Apply the Λ-shaped mask: Identify which token positions j are in the starting span (j < n_starting) or the ending span (i - j ≤ L_pretrain). Set the logits for all other positions to -∞.

4. Optionally reintroduce top-k middle tokens: For layers above h, for each head independently, compute attention scores for all middle tokens using a fixed distance of L_pretrain/2. Select the k middle tokens with the highest such scores and reset their logits from -∞ to these computed values.

5. Softmax and value aggregation: Compute softmax(logits) to obtain attention weights, then compute the weighted sum of value vectors Σ α_j · v_j. This output proceeds through the standard Transformer layer (residual connection, layer norm, feed-forward network) without any further modification.

What this produces. The attention output for each position i is a weighted combination of:

• Value vectors from the first n_starting tokens (with positional signal corresponding to distance L_pretrain).
• Value vectors from the most recent L_pretrain tokens (with true positional distances).
• Optionally, value vectors from up to k selected middle tokens (with positional signal corresponding to distance L_pretrain/2).

All these value vectors are combined using attention weights that sum to 1, with the softmax operating over a bounded number of logits (approximately n_starting + L_pretrain + k), keeping the attention entropy within the range the model was trained for. No parameter updates are performed; the model's weights remain exactly as they were after pretraining on short sequences.

4. Key Insights and Innovations

Innovation 1: The Length Generalization Problem Is Reformulated from a Representational Issue to a Computational Distribution-Shift Problem

The dominant prior assumption in the field—implicit in the widespread adoption of relative positional encodings like RoPE and Alibi—was that length generalization fails because the model cannot represent positions or distances it has not seen during training. Absolute position embeddings failed because position 4097 had no embedding vector; the natural fix was to make attention depend only on relative distances, which are just integers and therefore unbounded. The expectation was that a model trained on distances up to 4096 should gracefully extrapolate to distance 4097, 10000, or 1000000, because the same mathematical function w(q, k, d) can be evaluated at any integer d.

This paper's central conceptual move is to demonstrate that the failure is not about representation at all—it is about distribution shift in the intermediate computational quantities that the Transformer produces internally. The model can in principle compute w(q_i, k_j, 100000) using RoPE's sinusoidal functions, but doing so produces an attention logit of a magnitude the model never encountered during training. That inflated logit then propagates through the softmax, producing an attention distribution with entropy far outside the training range, which in turn produces value-vector aggregates in regions of the feature space the subsequent layers were never optimized to process.

This reframing is significant because it fundamentally changes what a solution must do. If the problem were representational—"the model hasn't learned how to encode position 100K"—the solution would require training on longer sequences to teach the model new positional representations. But if the problem is distributional—"the model's internal quantities drift out of range when the input gets too long"—the solution can be a constraint on the computation that keeps those quantities in-range, requiring no learning at all. LM-Infinite is the logical endpoint of this reframing: three simple mathematical operations (masking, capping, and optional selective reintroduction) that collectively bound attention logits, attention entropy, and the accessible region of the value-vector space to the ranges observed during pretraining.

The diagnostic evidence for this reframing is in Section 3 and Figure 1. Figure 1(a) shows attention logits jumping from ~5–10 at 4K to ~15–20 at 8K—a quantitative demonstration that the logits, not the positional embeddings, are drifting. Figure 1(b) shows entropy growing logarithmically with sequence length—exactly the pattern predicted by Proposition 1, confirming that the problem is the number of attended tokens, not the encoding of their positions. And Figure 1(c) shows the initial tokens occupying a distinct PCA region—evidence that the issue with sliding-window attention is not lost semantic content but lost access to a region of the feature space.

This is a fundamental reframing, not an incremental improvement. Prior work treated length generalization as a learning problem (fine-tune on longer sequences, design better position encodings); this paper treats it as a constraint-satisfaction problem (identify what drifts and cap it). The distinction matters because constraint-based solutions are zero-shot and architecture-agnostic, while learning-based solutions require training resources and are model-specific.

---

Innovation 2: The "Λ-Shape" Is Derived from a Functional Partition of the Sequence, Not Heuristic Sparsification

Sparse attention patterns—where each token attends to only a subset of preceding tokens rather than the full history—have a long history in efficient Transformer research. Longformer (Beltagy et al., 2020) used a sliding window plus global attention on a few pre-selected tokens. BigBird (Zaheer et al., 2020) combined random, windowed, and global attention. These patterns were designed with a single objective: reduce the quadratic complexity of attention while preserving as much useful information as possible. They were heuristic sparsifications—plausible patterns that seemed reasonable and worked empirically when trained from scratch or fine-tuned.

LM-Infinite's Λ-shaped mask looks superficially similar: a combination of a global span (the first n_starting tokens) and a local window (the most recent L_pretrain tokens). But the derivation is fundamentally different, and this difference has practical consequences that pure sparsification approaches miss.

The paper derives the Λ-shape from a functional analysis of what different regions of a sequence encode in a relative-positional Transformer (the conceptual model in Figure 2b). This analysis reveals that:

• The starting tokens encode strong absolute position information and occupy a distinct representational space (Factor 3, Figure 1c). They are not just "globally important tokens" that a heuristic might select based on content; they are structurally distinctive regardless of content, because every sequence's first few tokens have a unique causal attention pattern. A heuristic that selected "the most important tokens" based on attention scores might or might not include them; the Λ-shape includes them by architectural necessity.

• The rear tokens encode relative position and benefit from recency bias. This aligns with the sliding-window intuition, but the justification is different: the window size must be exactly L_pretrain, not because smaller windows would lose information (they would), but because larger windows would introduce distances beyond the pretraining range, triggering Factor 1's logit explosion.

• The middle tokens are functionally "less position-sensitive"—their exact distance matters less than their content. This is not an assumption but an empirical inference from two observations: (1) reintroducing middle tokens with a fixed distance of L_pretrain/2 works well for downstream tasks without hurting perplexity, and (2) discarding most middle tokens entirely (the default Λ-shape) does not degrade language modeling performance. If middle tokens' positions were critical, both of these findings would be false.

This functional derivation explains a result that would be puzzling under a pure sparsification view: why does a sliding-window attention mask (n_starting = 0) fail so catastrophically on a pretrained model (the "window" ablation in Figure 5, which produces the second-worst NLL values)? From a sparsification perspective, a sliding window is a reasonable approximation to full attention—nearby tokens matter most, and the model was trained with a causal mask anyway. But from the functional partition perspective, removing the starting tokens doesn't just remove some potentially useful information; it shifts the entire attention output into a subspace of the value-vector space that the model was never trained to operate in. The failure is not about missing content; it is about operating in the wrong representational regime. This is why n_starting as low as 1 dramatically improves performance over 0 (Table 4: NLL drops from 6.43 to 1.03)—even a single initial token provides access to the necessary region of the feature space.

This is a fundamental conceptual advance over prior sparse attention work. It provides a principled criterion for which tokens must be included (those occupying functionally distinct representational regions) versus which can be discarded or selectively reintroduced (those in the position-insensitive middle). The Λ-shape is not one of many possible sparsification patterns; it is the minimal pattern that preserves all three functional regions required for the model to remain in-distribution.

---

Innovation 3: The Identification of Verifier-Style Distribution Shifts in Attention Itself—Attention Logit Explosion and Entropy Growth as First-Class Failure Modes

Prior work on Transformers' limitations has largely focused on what the model learns—its knowledge, its representations, its ability to generalize across tasks. Length generalization failures were typically attributed to the model not having learned to use long-range dependencies, or to positional encodings not generalizing. The idea that the attention mechanism itself could produce numerically unstable computations on valid, in-domain inputs that are simply longer than training examples is a distinct and underappreciated class of failure.

This paper identifies two such numerical failure modes with formal precision:

Attention logit explosion (Factor 1, Theorem 1). The theorem proves that if a relative-positional attention function can distinguish arbitrarily many distances (which it must, to be useful), then the magnitude of its logits must grow as sequence length increases. This is not a statement about the model's learned behavior; it is a statement about the function class that the attention mechanism belongs to. Given bounded pseudo-dimension r, the logits must grow at least as Ω(α(n)^{1/(2r)}). For RoPE with feature dimension k, r = 2k, so the growth is Ω(α(n)^{1/(4k)})—sublinear but unbounded.

What makes this insight distinctive is that it identifies a failure mode that is architectural and mathematical, not empirical or data-dependent. Even a perfectly trained model with optimal weights would face this issue. The practical implication is that no amount of training on long sequences can eliminate the logit explosion—at best, fine-tuning can teach the model to be robust to larger logit magnitudes, but the underlying growth continues, and at some sufficiently long sequence length, the logits will again exceed whatever range the model was trained to handle. This explains why even models explicitly fine-tuned on long sequences (like MPT-7B-Storywriter) still have finite effective context limits rather than truly unbounded generalization.

Attention entropy growth (Factor 2, Proposition 1). The proof that H(softmax(w)) = Ω(ln n) for bounded logits is deceptively simple (the algebraic derivation in Appendix D is only a few lines) but has profound implications. It means that as the number of attended tokens grows, the attention distribution inevitably flattens—the model cannot maintain focused attention on a few tokens while also having the capacity to attend to many. This is not a softmax bug; it is an information-theoretic consequence of normalized attention over an expanding support.

The connection to length generalization failure is that the model was trained with attention entropy around ln(4096) ≈ 8.3 nats. When that entropy increases to ln(32768) ≈ 10.4 nats, the downstream layers receive value-vector aggregates that are more uniformly mixed than anything they were optimized to process. The model's learned computations—which assume a certain concentration of attention mass—break because the input statistics have shifted.

Identifying these as first-class failure modes matters because it redirects research effort. If the problem were "the model hasn't learned to use long contexts," the solution would be more training. If the problem is "the attention computation itself becomes numerically pathological at scale," the solution is to constrain the computation—exactly what LM-Infinite's distance ceiling and Λ-shaped mask do. The paper's ablation study (Figure 5) is particularly incisive here: using only the distance ceiling (the "ceiling" curve) addresses Factor 1 but not Factor 2, and the NLL still degrades. Using only the Λ-shaped mask (the "Λ" curve) addresses Factors 2 and 3 but not Factor 1, and the NLL still explodes. Both constraints are necessary because both failure modes are real and independent.

This is a fundamental diagnostic contribution, not an incremental empirical finding. The paper doesn't just show that certain techniques work; it provides a mechanistic vocabulary—logit explosion, entropy growth, representational space shift—that future work can use to analyze and address length generalization in new architectures.

---

Innovation 4: Establishing That Zero-Shot Length Generalization to 200M Tokens Is Possible Through Attention Constraints Alone

This innovation is the empirical proof of concept that validates the reframing in Innovation 1. The paper's headline result—that a model pretrained on 4K-token segments can process a 200M-token sequence while maintaining stable perplexity (Figure 4)—is not just a large number. It demonstrates something conceptually important: the length generalization failure is entirely attributable to the three identified distribution shifts, because addressing those three shifts (and nothing else) enables generalization across five orders of magnitude of sequence length.

Prior work had demonstrated long-context processing through various means: MPT-7B-Storywriter fine-tuned on 65K sequences achieved strong performance at those lengths; retrieval-augmented methods could handle arbitrarily long corpora by feeding chunks to the model; recurrent architectures could theoretically process unbounded sequences. But no prior method had shown that an off-the-shelf, unmodified pretrained LLM—with its original weights frozen—could handle sequences 50,000× longer than its training length while maintaining the same per-token prediction quality.

The comparison with MPT-7B-Storywriter in Table 1 is telling: LM-Infinite applied to the base MPT-7B (pretrained on 4K) achieves "only slightly worse performance than its fine-tuned counterpart" at 32K lengths. This means that the fine-tuning on 65K sequences bought relatively little over simply constraining the attention pattern—the base model already had the necessary knowledge; it just couldn't access it stably on long sequences. LM-Infinite removes the computational obstacles without adding any new knowledge.

The 200M-token result (Figure 4) is particularly significant because it shows that the per-token NLL does not drift upward even at extreme lengths. The dashed reference line is the vanilla Llama-2's NLL at 8K—already past its training length and beginning to degrade. LM-Infinite stays roughly at or below this line for 200M tokens, demonstrating that the constraints are stable—they do not slowly accumulate errors or gradually lose performance as the sequence grows. This is consistent with the theoretical framing: the constraints bound the computational quantities to fixed ranges, so no matter how long the sequence becomes, the model processes each token using logit magnitudes, attention entropies, and value-vector regions indistinguishable from those it saw during pretraining.

This is a fundamental empirical result rather than an incremental scaling demonstration. It establishes an existence proof: unbounded length generalization is possible for Transformer LLMs without any weight modification, without any architectural change, and without any training—provided the attention computation is constrained to prevent distribution shift. The practical ceiling of 200M tokens is likely limited only by the authors' computational resources and patience, not by any inherent limitation of the method.

---

Innovation 5: The Demonstration That "Positional Encoding" Is Not Just About Position—Starting Tokens Encode Functional Prerequisites for Attention

The finding that the first few tokens occupy a distinct representational space (Factor 3, Figure 1c, Appendix E Figure 7) and that removing them causes catastrophic degradation even when all other failure modes are addressed is a conceptual contribution that extends beyond length generalization. It reveals something fundamental about how causal Transformers use position information implicitly: the initial tokens serve as functional anchors for the attention mechanism, providing value vectors that span regions of the feature space no later tokens can reach.

This is counterintuitive because positional encodings (absolute or relative) are typically understood as providing information about position—where a token is in the sequence. The paper shows that the initial tokens encode something more structural: they provide access to certain representational regions that the attention output must be able to reach for downstream layers to function correctly. When those tokens are removed, it is not just that the model loses information about the beginning of the sequence; it is that the very geometry of the attention output space is constrained to a subspace the model was never optimized to operate in.

The evidence for this interpretation comes from multiple angles:

• The PCA visualization (Figure 1c) shows the initial tokens' hidden states are genuinely in different regions of the vector space, not merely at the edges of the same distribution.
• The ablation on n_starting (Table 4) shows that the jump from n_starting = 0 (NLL 6.43) to n_starting = 1 (NLL 1.03) is enormous—a single initial token recovers nearly all the benefit of including 100. This is hard to explain under a "semantic content" hypothesis (one token cannot carry much semantic content) but natural under a "representational anchor" hypothesis (one token's value vector is sufficient to span the necessary region when combined with the rest of the attention distribution).
• The fact that n_starting values from 1 to 100 all perform similarly suggests that as long as some initial tokens are present, the model can access the necessary regions; the exact number beyond the minimum is not critical.

This finding has implications beyond LM-Infinite. It suggests that any attention-modification technique applied to pretrained Transformers—not just length-extension methods but also pruning, distillation, or KV-cache compression approaches—must preserve access to the initial tokens' value vectors or risk inducing a representational distribution shift. It also raises questions about why Transformers learn to use the initial tokens this way during pretraining, and whether this emergent property could be deliberately controlled or exploited in future architectures.

This is a fundamental discovery about Transformer internals, not an incremental observation. It explains a failure mode (why sliding-window attention fails on pretrained models) that prior work had not articulated, and it provides a diagnostic criterion (check whether initial-token value vectors are preserved in the attention computation) that can be applied to evaluate any proposed modification to the attention mechanism.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper uses two corpora from the Pile dataset (Gao et al., 2020) for language modeling: ArXiv (academic preprint papers) and OpenWebText2 (Reddit submissions). For downstream evaluation, Passkey Retrieval (Mohtashami and Jaggi, 2023) buries a random passkey in a long distractor text and asks the model to retrieve it, and Qasper (Dasigi et al., 2021) contains 1.5K question-answer pairs over scientific papers. All datasets are used in their standard test splits as provided by the original sources.

• Base model(s). Four model families are evaluated: LLaMA-7B (Touvron et al., 2023a; pretrained on 2K segments, RoPE encoding), Llama-2-7b (Touvron et al., 2023b; pretrained on 4K, RoPE), MPT-7B (Team, 2023; pretrained on 4K, Alibi encoding), and GPT-J-6B (Wang and Komatsuzaki, 2021; pretrained on 2K, RoPE). For downstream tasks, Llama-2-7b-chat is used because its instruction tuning enables task-solving ability in the zero-shot setting (Section 5.3). The authors argue these models are "representative of the capabilities of many contemporary LLMs" and span both major relative encoding schemes (RoPE and Alibi), allowing evaluation of LM-Infinite's generality.

• Metrics. For language modeling, the primary metric is negative log-likelihood (NLL) computed per token position—specifically, the NLL the model achieves around each position, averaged across all evaluated sequences at that length (Section 5.1). Perplexity (exponential of NLL) is reported in summary tables. For generation quality, BLEU (Papineni et al., 2002) and ROUGE-L (Lin, 2004) are computed by letting models generate 100 tokens after each milestone length and using the following 100 tokens in the original text as references (Section 5.4). For downstream tasks, standard accuracy is used (exact match for Passkey Retrieval, answer correctness for Qasper). Computational efficiency is measured in seconds per sequence (encoding), seconds per token (decoding), and GPU memory usage per sequence (Appendix G).

• Baselines. Several categories of baselines are compared:

  • Vanilla models: the unmodified pretrained LLMs (LLaMA-7B, Llama-2-7b, MPT-7B, GPT-J-6B) evaluated directly on long sequences without any length-generalization intervention.
  • Fine-tuned long-context models: MPT-7B-Storywriter (Team, 2023; fine-tuned on 65K-length sequences using Alibi) and LongLLaMA (fine-tuned on 8K sequences). These represent the "train on longer sequences" approach.
  • Pretrained sparse-attention models: Sandwich (Chi et al., 2023; pretrained with 512-length segments) and XPos (Sun et al., 2022; pretrained with 1K-length segments). These represent architectures designed for length extrapolation from the start.
  • Truncation baseline (Section 5.2-5.4): inputs longer than the model's pretraining length are truncated to keep only the most recent tokens within the training length, dropping excessive tokens before the forward pass. This is a common practical workaround.
  • Windowed attention ("window" ablation; Section 5.2): a sliding-window attention pattern applied to the pretrained model without any other modifications—each token attends only to the most recent L_pretrain tokens. This isolates the effect of the Λ-shape's starting span.

• Generation budget / compute accounting. The paper does not use a FLOPs-based budget framework. Instead, sequence length is the primary independent variable—models are evaluated at milestone lengths (2K, 4K, 8K, 16K, 32K, and in one case 200M tokens) and performance is reported as a function of position within those sequences. For efficiency comparisons (Appendix G, Figure 6), computational cost is measured in TFLOPs per token (a hardware-independent metric) and wall-clock time/memory on a single A100 GPU with 80GB memory. All vanilla models run out of memory at ~32K lengths on this hardware, establishing a practical ceiling for full-attention baselines.

• Cross-validation / statistical protocol. No cross-validation or statistical significance testing is reported. The method is zero-shot and requires no training, so there is no risk of overfitting hyperparameters to test data. Hyperparameters for the optional top-k mechanism are selected on a held-out Passkey Retrieval validation set (Appendix A.3) and then applied identically to Qasper without further tuning. Language modeling evaluations use the standard test splits of ArXiv and OpenWebText2. For generation quality, 100 long sequences are sampled per dataset due to the time-consuming nature of generation (Section 5.4). All hardware experiments use single A100 GPUs.

Main Quantitative Results

Language Modeling on Extremely Long Contexts

The central finding is that LM-Infinite flattens the NLL curves of all tested models on sequences far exceeding their pretraining length. Figure 3 plots NLL as a function of token position on the ArXiv dataset for all four model families.

Vanilla models fail catastrophically. Without LM-Infinite, all baseline models exhibit exploding NLL once the sequence length exceeds their pretraining length (L_pretrain). Llama-2-7b outputs NaN values on sequences slightly longer than 10K, causing its curve to terminate early in Figure 3. LLaMA-7B (2K pretraining) jumps from roughly 3.8 NLL at 2K to ~1.0×10⁴ at 4K—a massive explosion. MPT-7B (4K pretraining) increases from ~5.5 at 4K to ~1.1×10³ at 8K. GPT-J-6B (2K pretraining) goes from ~3.9 at 2K to ~1.3×10³ at 4K. These failures are consistent across both RoPE-based models (LLaMA, Llama-2, GPT-J) and the Alibi-based model (MPT-7B), confirming that the problem is not specific to one encoding scheme.

LM-Infinite stabilizes all models. With LM-Infinite applied, all four model families show essentially flat NLL curves that extend stably to 32K tokens (the maximum length evaluable on the hardware for vanilla models) and beyond. Specifically:

• Llama-2 + LM-Infinite: NLL stays between roughly 3.3 and 6.5 across all positions from 2K to 32K. No NaN values occur.
• LLaMA + LM-Infinite: NLL decreases slightly with position (from ~4.4 at 2K to ~1.0 at 32K), which the authors note as "surprisingly" improved performance at longer lengths.
• GPT-J-6B + LM-Infinite: NLL is stable in the 2.1–3.8 range across all lengths.
• MPT-7B + LM-Infinite: NLL is roughly 4.6–6.8 across all lengths.

Comparison with fine-tuned baselines. The trends are "similar to MPT-7B-Storywriter, an explicitly fine-tuned LLM" (Figure 3 caption). Table 1 quantifies this at milestone lengths for both ArXiv and OpenWebText2. On ArXiv at 32K, MPT-7B + LM-Infinite achieves perplexity 4.6, compared to MPT-7B-Storywriter's 3.6—a gap of only 1.0 perplexity points despite LM-Infinite requiring no training. On OpenWebText2 at 16K, MPT-7B + LM-Infinite achieves 8.9 vs. MPT-7B-Storywriter's 5.1. The fine-tuned model consistently outperforms LM-Infinite, but the margins are modest given the enormous difference in training investment.

LM-Infinite outperforms other zero-shot and pretrained baselines. Table 1 shows LM-Infinite achieving the best perplexity in 7 out of 9 comparisons (3 model families × 3 length milestones where all methods have data). For example, at 8K on ArXiv, Llama-2 + LM-Infinite achieves 3.3 vs. MPT-7B-SW's 4.3, Sandwich's 5.3, and XPos's 20.7. The method is competitive even with models explicitly designed or trained for long contexts.

Extreme length generalization to 200M tokens. Figure 4 evaluates Llama-2 + LM-Infinite on a single concatenated sequence of 200M tokens (constructed by sampling with replacement from ArXiv and concatenating). The NLL remains stably low across all positions, oscillating around the dashed reference line representing vanilla Llama-2's NLL at 8K (roughly 3–8 NLL, with periodic spikes). The authors state that "LM-Infinite shows the ability to remain stably low log-perplexity level over extreme lengths." This is the headline empirical result—the method scales across five orders of magnitude of sequence length without degradation.

Position-dependent trends. On LLaMA + LM-Infinite, the perplexity decreases as length and position increase (Table 1: perplexity drops from 4.4 at 2K to 1.0 at 32K on ArXiv). This is counterintuitive—one would expect prediction to become harder at later positions—but the authors do not provide a detailed explanation. It may reflect that later positions benefit from more accumulated context within the ending window, or that the ArXiv dataset's structure provides stronger local predictability in later sections of papers.

Ablation Study: Why Both Components Are Essential

Figure 5 presents the critical ablation study using LLaMA-7B on ArXiv at 8K length, comparing five conditions:

• Vanilla (unmodified LLaMA): NLL immediately explodes after 2K (the pretraining length), rising sharply from ~4 to >12 by 4K. This is the baseline failure mode.

• Window (pure sliding-window attention, n_starting = 0): NLL is substantially elevated across all positions (~4–8 range) and trends upward with position. This is the "second worst" performer, confirming that simply bounding the attention context without preserving initial tokens is insufficient. The degradation is attributed to Factor 3—discarding the initial tokens' distinct feature space.

• Ceiling (distance ceiling only, no Λ-shaped mask): NLL starts at normal levels (~3.5 at position 2K) but rises steadily with position, reaching ~12 by 8K. The ever-increasing NLL confirms Factor 2—without limiting the number of attended tokens, attention entropy grows and degrades performance even though the distance ceiling prevents logit explosion.

• Λ (Λ-shaped mask only, no distance ceiling): NLL explodes immediately after 2K, similar to the vanilla model. This confirms Factor 1—without capping distances, the attention logits explode for long-distance token pairs even though the attention context is bounded.

• LM-Infinite (both Λ-shaped mask and distance ceiling): NLL remains flat and low (~3.5–4.5) across all positions from 2K to 8K. Only the full method with both components preserves train-length performance.

This ablation is the empirical lynchpin of the paper's diagnostic claims. Each component alone addresses a subset of the failure modes, and each alone fails. Together, they succeed. The results map cleanly onto the three factors from Section 3:

• Ceiling only: addresses Factor 1 but not 2 or 3 → NLL still degrades (Factor 2 active).
• Λ only: addresses Factors 2 and 3 but not 1 → NLL still explodes (Factor 1 active).
• Window only: addresses Factor 2 but not 1 or 3 → NLL still poor (Factors 1 and 3 active).
• LM-Infinite: addresses all three factors → NLL stable.

Downstream Task Evaluation

Table 2 reports zero-shot performance on Passkey Retrieval and Qasper using Llama-2-7b-chat. For LM-Infinite, the optional top-k mechanism is enabled with k = 5, h = 5 (selected on a held-out Passkey Retrieval validation set, Section 4).

Passkey Retrieval. This task places a random passkey at a uniform random position in long distractor text and asks the model to retrieve it. Performance is evaluated at 6K, 8K, 10K, 12K, and 16K lengths.

• Original model: 0.0% at all lengths. The unmodified Llama-2-7b-chat cannot perform the task on any sequence longer than its 4K pretraining length—it fails completely.
• Truncation baseline: 66.0% at 6K, declining monotonically to 27.3% at 16K, averaging 44.0%. This baseline keeps only the most recent 4K tokens of the input. Its performance depends entirely on whether the truncated portion happens to contain the passkey—when the passkey is in the first half of a 6K sequence, truncation loses it; when it is near the end, truncation keeps it. This explains the declining accuracy as sequences get longer (the passkey is uniformly placed, so longer sequences have lower probability that the passkey falls in the last 4K).
• LM-Infinite: 70.3% at 6K, rising to 90.8% at 8K, then declining slightly to 79.1% at 16K, averaging 81.2%. This represents a 37.2 percentage-point gain over the truncation baseline's average (81.2% vs. 44.0%). The non-monotonic pattern (peak at 8K, decline at longer lengths) is not explained in detail, but the 6K result being lower than 8K may reflect boundary effects where the sequence is only slightly longer than the pretraining length and the attention pattern has not fully stabilized.

The key insight is that LM-Infinite's top-k mechanism allows the model to search the middle of the sequence for the passkey—something the truncation baseline cannot do at all. The high accuracy (79–91%) across lengths where truncation achieves 27–66% demonstrates that the model can effectively use its content-based attention to locate relevant information even when it is far from both the start and end of the sequence.

Qasper. This is a question-answering dataset over scientific papers (1.5K test pairs). The task is substantially harder than Passkey Retrieval because it requires reading comprehension, not just pattern matching.

• Original model: 1.2% accuracy. The vanilla Llama-2-7b-chat essentially cannot perform the task on sequences exceeding its pretraining length.
• Truncation baseline: 30.1% accuracy. Truncating to 4K keeps some relevant context but loses information from earlier portions of the papers.
• LM-Infinite: 31.3% accuracy. A 1.2 percentage-point gain over the truncation baseline. This is a more modest improvement than on Passkey Retrieval, reflecting the greater difficulty of the task and the fact that Qasper questions are often answerable from local context within a paper section rather than requiring access to distant parts of the document.

The authors note that "top-k attention is necessary for achieving good performance" on Qasper (Section 5.3), suggesting that the ability to attend to selected middle tokens contributes to the improvement. However, the gain is small, and the paper does not provide an ablation of Qasper performance with and without top-k to quantify this contribution.

Generation Quality

Table 3 reports BLEU and ROUGE-L scores for text generation on ArXiv and OpenWebText2. Models generate 100 tokens after each milestone length, and the generated text is compared against the actual next 100 tokens in the original documents.

Vanilla models fail at generation. For both MPT-7B and Llama-2, generation quality collapses to near-zero BLEU scores at any length beyond L_pretrain. For example, Llama-2 achieves 26.6 BLEU at 2K (within training length) but 0.0 at 4K, 8K, 16K, and 32K on ArXiv. The zero BLEU scores indicate that the generated text contains no {2,3,4}-gram overlaps with the reference—the model produces "mostly nonsensical texts" (Section 5.4). ROUGE-L shows the same pattern, dropping from 31.4 at 2K to 0.2 or 0.0 at all longer lengths.

LM-Infinite preserves generation quality. With LM-Infinite applied:

• Llama-2 + LM-Infinite: BLEU is 26.9 at 2K, 23.6 at 4K, 23.9 at 8K, 24.8 at 16K, and 18.4 at 32K on ArXiv. ROUGE-L is 31.8, 30.9, 28.8, 29.2, and 20.4 respectively. The decline at 32K is notable but the scores remain far above the vanilla model's zero.
• MPT-7B + LM-Infinite: BLEU ranges from 12.6 to 20.2 across lengths on ArXiv, and 2.8 to 5.1 on OpenWebText2. These are lower than Llama-2's scores, reflecting MPT-7B's generally weaker generation capabilities, but they are consistent and do not collapse.

Comparison with fine-tuned MPT-7B-Storywriter. On ArXiv, MPT-7B + LM-Infinite achieves results comparable to MPT-7B-Storywriter. For instance, at 32K, BLEU is 19.7 (LM-Infinite) vs. 14.8 (Storywriter), and ROUGE-L is 26.6 vs. 27.0. In some cases, LM-Infinite slightly outperforms the fine-tuned model. This is a strong result: LM-Infinite achieves generation quality that is competitive with or better than a model explicitly trained on 65K-length sequences, despite requiring no parameter updates.

Zero BLEU scores. Several vanilla model entries show 0 BLEU at longer lengths. The paper explains that "BLEU is weighted geometric mean over {1,2,3,4}-gram precisions," so if the model generates no overlapping {2,3,4}-grams with the reference, the final BLEU is zero even if some unigrams match. This reflects genuinely catastrophic generation degradation rather than a metric artifact.

Computational Efficiency

Appendix G reports efficiency measurements using Llama-2-7B on 100 sequences of 32K length from ArXiv, running on a single A100 GPU with 80GB memory using DeepSpeed ZeRO-3 optimization.

Encoding speed. The vanilla Llama-2-7B encodes at 48.19 seconds per 32K sequence. LM-Infinite encodes at 15.26 seconds per sequence—a 3.16× speedup. This is because LM-Infinite's sparse attention pattern dramatically reduces the number of dot-product operations in the attention layers.

Decoding speed. The vanilla model decodes at 7.34 seconds per generated token. LM-Infinite decodes at 2.70 seconds per token—a 2.72× speedup. This is particularly important for generation tasks where decoding dominates runtime.

GPU memory. The vanilla model uses 33.2 GB of GPU memory per 32K sequence. LM-Infinite uses 4.41 GB per sequence—a 7.53× memory saving. This is the difference between being able to run on a single consumer GPU versus requiring a data-center GPU, and it is what enables the extreme 200M-token sequence processing in Figure 4.

Quality-efficiency tradeoff. Figure 6 compares LM-Infinite with the truncation baseline on a generation task where LLaMA generates 10K tokens on ArXiv. Both methods are evaluated at varying window sizes w (the amount of recent context retained). For truncation, larger w means better quality (more context preserved) but higher computational cost (more tokens to re-encode as the window slides). LM-Infinite achieves a "substantially better quality-efficiency tradeoff": at similar TFLOPs/token (~1.35), LM-Infinite achieves roughly 5 BLEU points higher than truncation. To achieve similar BLEU (~25), LM-Infinite incurs only <25% of the computational overhead of truncation. This demonstrates that the efficiency gains are not merely from doing less computation but from doing smarter computation—LM-Infinite's structured sparsity preserves the functionally essential tokens while truncation blindly keeps only the most recent ones.

Ablation Studies and Robustness Checks

• n_starting parameter sweep (Appendix A, Table 4). On 16K-length ArXiv sequences, n_starting values from 1 to 100 all produce NLL around 1.02–1.03—essentially identical. n_starting = 0 (pure windowed attention) produces NLL 6.43, confirming that including initial tokens is critical, but the exact number beyond the minimum matters little. n_starting = 1000 degrades to 1.81 and n_starting = 2000 to 4.96, confirming that including too many initial tokens reintroduces Factor 2's entropy growth. The authors recommend n_starting ∈ [5, 100].

• Top-k k parameter sweep (Appendix A, Table 5). On Passkey Retrieval, k = 3 and k = 5 both achieve 0.81 accuracy. k = 1 drops to 0.69—insufficient retrievable tokens. k = 200 drops to 0.73—too many reintroduced tokens dilute attention. The non-monotonic relationship shows a clear optimum: reintroducing a small number of high-scoring middle tokens helps, but reintroducing many hurts.

• Middle token attention distance (Appendix A, Table 6). The distance at which reintroduced middle tokens are attended affects performance: d = L_pretrain/2 = 2048 achieves 0.81 Passkey Retrieval accuracy. d = 512 drops to 0.78; d = 4096 drops to 0.63. Using the full pretraining distance for middle tokens (which are not actually at that distance) degrades performance more than using a shorter distance, likely because the positional signal at L_pretrain is associated with "very far" semantics during training and is misleading for middle tokens that are at intermediate distances.

• Layer threshold h for top-k (Appendix A, Table 7). Applying top-k to all layers (h = 0) achieves 0.81 Passkey Retrieval accuracy. Applying it only to layers 5 and above (h = 4 or h = 5) achieves 0.94—a substantial improvement. Applying it only to very high layers (h = 24) drops to 0.46. The authors hypothesize that "lower layers encode more position-sensitive features where introducing middle tokens with a fixed distance would be disruptive, while higher layers are more content-focused." This is a non-obvious finding: the optimal layer to begin selective middle-token attention is neither the very bottom nor the very top, but an intermediate layer.

• Different model families and encoding schemes (Figures 3, Table 1). LM-Infinite works with LLaMA, Llama-2, GPT-J (all RoPE) and MPT-7B (Alibi), demonstrating that the method is not tied to a specific positional encoding scheme. The authors note that "augmenting AliBi with LM-Infinite is also straightforward: we simply clip the offset matrix with a minimum value of -|m * L_pretrain|" (Section 4). The consistent success across encoding schemes supports the claim that the three failure modes are architectural rather than encoding-specific.

• Different datasets (Table 1, ArXiv vs. OpenWebText2). LM-Infinite's performance is consistent across academic papers and social media posts—two domains with different structural properties. On OpenWebText2 at 16K, Llama-2 + LM-Infinite achieves perplexity 8.2 vs. the vanilla model's NaN. The method does not depend on document structure (e.g., section headers in ArXiv) to function.

• Optional top-k doesn't hurt language modeling. The authors state that "for LLM generation and inference, we find the intermediate tokens unnecessary to attend to for LM-Infinite to achieve good perplexity or generation quality" (Section 4). The default LM-Infinite (without top-k) is sufficient for language modeling; top-k is an enhancement specifically for downstream retrieval tasks. This is implicitly an ablation: adding top-k does not degrade perplexity, meaning the middle-token capacity can be added without cost when needed.

• The "ceiling" and "Λ" ablations confirm the independence of failure modes (Figure 5). The ceiling-only ablation (distance cap without attention bounding) shows NLL rising steadily with position—this isolates Factor 2 (entropy growth) because Factor 1 is addressed but the attention context is unbounded so entropy increases. The Λ-only ablation (attention bounding without distance cap) shows NLL exploding immediately after pretraining length—this isolates Factor 1 (logit explosion) because the context is bounded but unseen distances produce explosive logits. The fact that each partial solution fails in a different way (gradual degradation vs. immediate explosion) provides strong evidence that Factors 1 and 2 are genuinely distinct failure modes with different temporal signatures.

• Comparison of logit and entropy behavior (Figure 1). The paper does not present explicit ablations showing that LM-Infinite successfully contains logit magnitudes and attention entropy within training ranges—this would be a valuable supplement. The evidence that it does so is indirect: the NLL curves in Figures 3 and 4 remain flat, which would not be possible if logits were exploding or entropy was growing. However, a direct plot of max attention logit and attention entropy under LM-Infinite (analogous to Figures 1a and 1b) would strengthen the mechanistic claims.

• Replication across model scales. All experiments use 6–7B parameter models. The paper does not evaluate LM-Infinite on smaller models (e.g., 1B parameters) or larger models (e.g., 13B, 70B). The failure modes identified in Section 3 are architectural (logit explosion depends on the pseudo-dimension of the positional encoding function, which scales with feature dimension; entropy growth depends purely on sequence length, not model size), so they should apply across scales, but this is an untested assumption.

Critical Assessment

The experiments provide strong evidence that LM-Infinite enables stable language modeling and generation on sequences far exceeding pretraining length, with the 200M-token result (Figure 4) serving as a compelling existence proof for extreme length generalization through attention constraints. The ablation study (Figure 5) cleanly demonstrates that both the Λ-shaped mask and the distance ceiling are necessary, and that each addresses a distinct failure mode—this is the most rigorous and convincing part of the evaluation. The efficiency gains (2.7× decoding speedup, 7.5× memory saving at 32K; Appendix G) are well-documented and practically significant.

However, several claims from the paper's framing deserve closer scrutiny against the experimental evidence:

Claim: LM-Infinite enables zero-shot length generalization "up to 200M length inputs while retaining perplexity." The 200M-token result is shown for exactly one model (Llama-2) on exactly one constructed sequence (sampled with replacement from ArXiv and concatenated). This is a proof-of-concept rather than a systematic evaluation. The paper does not report: (a) variance across multiple 200M-token sequences, (b) performance at intermediate lengths between 32K and 200M to characterize the scaling behavior, (c) whether the NLL would remain stable indefinitely or eventually degrades. The flat NLL curve in Figure 4 is reassuring, but the single-sequence nature of the evaluation limits the strength of the "up to 200M" claim. A rigorous demonstration would evaluate on multiple independently constructed long sequences and report confidence intervals.

Claim: LM-Infinite "improves performance on downstream tasks such as Passkey Retrieval and Qasper." This claim holds for Passkey Retrieval with a large margin (37.2 percentage points over the truncation baseline) but is less convincing for Qasper (1.2 percentage-point gain). The Qasper result is from a single model (Llama-2-7b-chat) without reported variance or statistical testing. A 1.2-point improvement on a 1.5K-question test set could potentially be within sampling noise, though the paper does not provide error bars. More concerning, the paper does not evaluate LM-Infinite on Qasper without the top-k mechanism, making it unclear whether the gain comes from the Λ-shape and distance ceiling or primarily from the selective middle-token retrieval. An ablation showing Qasper performance with and without top-k would substantially strengthen (or qualify) this claim.

Claim: LM-Infinite is "highly flexible and can be used with most modern LLMs off-the-shelf." The evaluation covers four model families spanning 6–7B parameters and two encoding schemes (RoPE and Alibi). This is broader coverage than most length-generalization papers, which typically evaluate on one or two models. However, all models are in the same size class (6–7B parameters), all are decoder-only Transformers, and none are instruction-tuned (except the chat variant used for downstream tasks). The method has not been demonstrated on encoder-decoder models (T5), mixture-of-experts architectures, or models with absolute positional encodings (though the latter would not benefit since the distance ceiling mechanism assumes relative encodings). The claim of broad applicability is reasonable for current open-source decoder-only LLMs with relative positional encodings, but the bounds of "most modern LLMs" are not empirically established.

Missing experiment: direct measurement of Factor 1 and Factor 2 under LM-Infinite. The paper's core diagnostic claim is that LM-Infinite works by preventing attention logit explosion and entropy growth. Figure 1 shows these phenomena in the vanilla model, and the ablation (Figure 5) shows that each component is necessary, but the paper never directly plots logit magnitudes or attention entropy under LM-Infinite. Showing that logits remain bounded and entropy remains stable under LM-Infinite on long sequences would directly validate the mechanistic story. Without this, the connection between the theoretical analysis (Theorem 1, Proposition 1) and the empirical success remains inferential rather than demonstrated.

Missing baseline: a comparison of LM-Infinite with a model that has been fine-tuned with a similar sparse attention pattern. The paper compares against MPT-7B-Storywriter (fine-tuned on 65K with full Alibi attention) and several pretrained sparse-attention models (Longformer-style, Sandwich, XPos). But it does not compare against the obvious hybrid: taking the same base model and fine-tuning it on long sequences with the Λ-shaped mask, which might combine the benefits of both approaches. This would help disentangle how much of LM-Infinite's performance comes from the attention pattern itself versus from the fact that it constrains the model to stay in-distribution without modifying weights. If fine-tuning with the Λ-shape substantially outperforms zero-shot LM-Infinite, it would suggest that some length-generalization capability still requires learning.

Missing evaluation: the effect of n_starting on downstream tasks. The n_starting parameter is swept for language modeling (Table 4) and found robust across a wide range, but its effect on Passkey Retrieval and Qasper is not reported. If the initial tokens contain critical task information (e.g., instructions in the prompt), the optimal n_starting for downstream tasks might differ from the optimal value for language modeling.

The Qasper improvement deserves qualification. Qasper questions are anchored in specific passages of scientific papers. The 1.2-point gain over truncation could arise from (a) LM-Infinite allowing the model to attend to relevant passages in the middle of the paper that truncation misses, or (b) better processing of the local context due to the stabilized attention distribution. Without a breakdown of performance by question position (do questions about the paper's introduction benefit more or less than questions about the methods section?), the source of the gain is unclear.

Efficiency claims and practical deployment. The 7.5× memory saving and 2.7× decoding speedup are measured on a single A100 with 80GB memory. For the vanilla model, DeepSpeed ZeRO-3 was required to fit the 32K-length computation graph, meaning the baseline is already using memory optimization. LM-Infinite's memory savings would be even more dramatic relative to a vanilla model without DeepSpeed (which would run out of memory entirely). However, the paper does not compare LM-Infinite against other KV-cache compression methods (e.g., H2O, Zhang et al., 2024d; or streaming-LLM-style approaches, Xiao et al., 2024) that also achieve memory savings on long sequences. The claim that LM-Infinite brings "substantial efficiency improvements" is true relative to the full-attention baseline, but its relative efficiency against other sparse-attention or cache-eviction methods is not established.

Overall assessment. The experiments convincingly demonstrate that LM-Infinite solves the catastrophic length generalization failure for language modeling and generation—the NLL curves in Figures 3 and 4, the perplexity results in Table 1, and the generation quality in Table 3 all show stable performance across lengths where vanilla models fail completely. The ablation in Figure 5 provides strong causal evidence for the three-factor diagnosis. The downstream task results are promising but preliminary: Passkey Retrieval shows large gains, Qasper shows marginal improvement, and no other long-context benchmarks are evaluated. The 200M-token result is a striking demonstration but based on a single sequence from a single model. The efficiency gains are well-characterized for the full-attention baseline but not compared against other efficient long-context methods. Overall, the paper provides strong evidence for its central thesis—that three computational distribution shifts cause length generalization failure, and that constraining attention can prevent them—but the breadth and depth of the downstream evaluation leave room for more comprehensive validation of the method's practical benefits beyond language modeling.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Not Accounted For, Yet It Dominates the Practical Budget

LM-Infinite's central premise is that it achieves length generalization without any parameter updates, making it deployable "off-the-shelf" on pretrained models. However, the method does require a non-trivial hyperparameter selection step that is not factored into any of the reported cost or efficiency numbers. Specifically, the optional top-k middle-token reintroduction mechanism—which is essential for the downstream task improvements that constitute a major claimed contribution of the paper—requires selecting k (the number of middle tokens to reintroduce), h (the layer threshold above which to apply it), and the attention distance for middle tokens.

These hyperparameters are selected "based on a held-out validation set of Passkey Retrieval" (Section 4, Appendix A.3). The paper does not specify the size of this validation set, the number of hyperparameter configurations evaluated, or the computational cost of this tuning procedure. For a practitioner deploying LM-Infinite on a new model family or a new task domain, this tuning would need to be repeated, and the cost could be substantial—each hyperparameter evaluation requires running the model on long sequences with full LM-Infinite attention, measuring downstream task accuracy, and iterating. The paper presents the final selected values (k = 5, h = 5, distance = L_pretrain/2) and applies them identically to Qasper without further tuning, but there is no guarantee that these hyperparameters transfer across tasks, models, or sequence length distributions. The h parameter in particular shows high sensitivity in the ablation (Table 7): accuracy drops from 0.94 at h = 4–6 to 0.46 at h = 24, suggesting that incorrect selection of h could degrade downstream performance substantially.

The paper's headline efficiency claims (2.7× decoding speedup, 7.5× memory saving; Appendix G) are measured excluding any hyperparameter search cost, and the zero-shot claim is qualified by this hidden tuning requirement. To the paper's credit, the language modeling and generation results do not depend on top-k (the default LM-Infinite without middle tokens suffices), and the n_starting parameter is shown to be robust across a wide range (1–100, Table 4). But the downstream task results—the 37.2-point Passkey Retrieval gain and the 1.2-point Qasper gain—do depend on the tuned top-k mechanism. A practitioner who skips the hyperparameter tuning and uses the paper's published values may or may not achieve the reported performance on their specific task and model.

Mitigation status. The paper does not attempt to reduce or amortize this tuning cost. There is no proposed method for predicting optimal top-k hyperparameters from model or task characteristics, no sensitivity analysis showing that the published values work across a broader range of tasks, and no measurement of the tuning budget. Future work on automated hyperparameter selection or task-agnostic defaults would address this gap.

---

The 200M-Token Result Is a Single-Sequence Proof-of-Concept, Not a Systematic Evaluation

The paper's most striking empirical claim—that LM-Infinite enables generalization "to up to 200M length inputs while retaining perplexity" (Abstract, Section 5.1, Figure 4)—is based on one constructed sequence from one model (Llama-2) on one dataset (ArXiv, sampled with replacement and concatenated). This is an existence proof, not a systematic evaluation of performance at extreme lengths.

The consequence is that we do not know: (a) whether the stable NLL observed in Figure 4 would replicate across multiple independently constructed 200M-token sequences, (b) how performance scales between 32K (the maximum length in the systematic evaluations of Figures 3 and Table 1) and 200M—the NLL curve in Figure 4 shows periodic spikes whose cause (document boundaries? concatenation artifacts? genuine degradation?) is unexplained, (c) whether the model would eventually degrade at some length beyond 200M (1B tokens? 10B?), making the "infinite" in the paper's title an aspirational claim rather than a demonstrated property, or (d) whether generation quality, not just perplexity, is maintained at such extreme lengths—the generation experiments in Section 5.4 only go up to 32K.

The single-sequence nature also means that no variance or confidence intervals are reported for the 200M result. The NLL at any given position in Figure 4 is a single data point (one prediction in one sequence). Real-world deployment on diverse long documents would encounter varying content distributions, and the stability of LM-Infinite across that variation is untested at extreme scales.

Mitigation status. The paper acknowledges no limitation here—the 200M result is presented as a straightforward demonstration of capability. The computational constraints are understandable (the paper notes the experiment "runs for 20 hours" on a single A100; Appendix A.1), and evaluating multiple such sequences would indeed be expensive. However, the gap between "works on one 200M sequence" and "generalizes to 200M length inputs" (the paper's phrasing) is meaningful. A more conservative framing, or at minimum a note that systematic evaluation at these scales remains future work, would strengthen the paper's credibility.

---

Downstream Task Evaluation Is Limited to Two Benchmarks on One Model, Without Ablations Disentangling the Contributions of Each Component

The downstream evaluation in Section 5.3 tests LM-Infinite on exactly two tasks—Passkey Retrieval and Qasper—using exclusively Llama-2-7b-chat. This is a narrow evaluation for a method that claims to "enhance Transformer LLMs' capabilities for modeling long contexts" broadly. Several important gaps exist:

Only one model tested on downstream tasks. While the language modeling experiments cover four model families (LLaMA, Llama-2, GPT-J, MPT), the downstream evaluation is restricted to a single instruction-tuned model. We cannot tell whether the Passkey Retrieval gain (37.2 points) or Qasper gain (1.2 points) would replicate on other model families, particularly MPT-7B which uses Alibi rather than RoPE.

No ablation of top-k on Qasper. The paper states that "top-k attention is necessary for achieving good performance" on Qasper (Section 5.3), but does not report Qasper accuracy without the top-k mechanism. This makes it impossible to determine whether the 1.2-point gain over the truncation baseline comes from: (a) the Λ-shaped mask and distance ceiling alone (the core LM-Infinite components), (b) the top-k middle-token mechanism specifically, or (c) the combination. If the gain disappears without top-k, then the downstream task benefit is attributable to the optional enhancement, not the base LM-Infinite method. If it remains, then the base method alone provides some downstream benefit. Neither case is established.

No long-context benchmarks beyond retrieval. The paper does not evaluate on any task requiring multi-hop reasoning across long contexts, summarization of long documents, long-form question answering requiring synthesis, or code understanding across large repositories—all of which are mentioned in the introduction as motivating applications. Recent long-context benchmarks (LongBench, LVEval, ∞Bench, etc.) that postdate this paper's initial preprint are understandably absent, but even the benchmarks available at the time of writing (e.g., SCROLLS, NarrativeQA, SummScreen) are not evaluated. The Passkey Retrieval task is deliberately simple—finding a single number in distractors—and tests whether the model can locate information, not whether it can reason over it. Qasper is a more realistic reading comprehension task, but the 1.2-point gain is small and its practical significance for a practitioner choosing between LM-Infinite and truncation is questionable.

The Qasper gain may not be statistically reliable. With 1.5K test questions, a 1.2 percentage-point absolute improvement corresponds to roughly 18 additional correct answers. Without confidence intervals or multiple evaluation runs, it is unclear whether this difference exceeds sampling noise. The paper does not report any statistical testing or variance estimates for the downstream evaluation.

Mitigation status. The authors do not claim broad downstream task generalizability—they present the two-task evaluation as initial evidence and note future work on "long reasoning, long-dialogue, retrieval-augmented generation, or long literature generation" (Section 6). However, the strong claims in the abstract ("It also improves performance on downstream tasks") are not proportionally supported by the breadth of evaluation. A practitioner considering LM-Infinite for a production long-context application would need to run their own evaluation to determine whether the method provides meaningful gains on their specific task.

---

The Method Cannot Recover Information That Is Simultaneously Far from Both the Start and the End of the Sequence (Without Top-k), Creating a Fundamental Blind Spot

The Λ-shaped attention mask creates a structural blind spot: tokens in the middle of a long sequence—more than n_starting from the beginning and more than L_pretrain from the current position—are completely invisible to the attention computation. The default LM-Infinite (without top-k) provides no mechanism for a token at position i to attend to a token at position j if j ≥ n_starting and i - j > L_pretrain.

This is acceptable for language modeling and generation, where token prediction is dominated by local context (the ending window) and broad document-level signals from the beginning (the starting span). The language modeling results (Figures 3, 4; Table 1) and generation results (Table 3) empirically confirm that the blind spot does not degrade these capabilities. However, for any task that requires reasoning about specific content located exclusively in the middle of a long document—without that content being near either the start or the end—the default LM-Infinite provides no way to access it.

The optional top-k mechanism partially mitigates this by allowing each attention head to independently select k middle tokens with the highest attention logits. However, this mitigation has significant limitations:

• It is content-based, not position-based. The top-k selection uses q^T k similarity to identify relevant middle tokens. If the task requires attending to a specific position (e.g., "what does paragraph 47 say?") rather than content-based retrieval, the top-k mechanism may fail to select the right tokens. The Passkey Retrieval task succeeds because the passkey is a distinctive number that stands out from the distractor text; a task requiring retrieval of an arbitrary token at a specified position would not benefit from content-based selection.

• It only selects k tokens. With k = 5, the model can attend to at most 5 middle tokens per head per layer. If a task requires integrating information from dozens or hundreds of middle tokens—for example, summarizing a long contract where important clauses are distributed throughout the middle—the top-k budget is insufficient. The authors show that increasing k to 200 decreases Passkey Retrieval accuracy from 0.81 to 0.73 (Table 5), so simply raising k is not a solution; it reintroduces the entropy dilution problem.

• It is applied only above layer h. Lower layers cannot access middle tokens at all, meaning that low-level features (syntax, local coherence) computed from middle tokens are unavailable to the model even with top-k enabled. The task must be solvable using only higher-layer semantic features of the selected middle tokens.

Mitigation status. The paper acknowledges implicitly that middle tokens are discarded by default and provides the top-k mechanism as an optional enhancement. However, the structural limitation of the Λ-shape—that there exists a class of tokens that are architecturally inaccessible—is not discussed as a fundamental constraint. Future work could explore dynamic window placement (moving the attention window based on task needs rather than fixing it at the most recent tokens), hierarchical attention (coarse-to-fine middle token access), or learned gating that adaptively adjusts the number of attended middle tokens per query.

---

The Method Assumes Relative Positional Encodings and Has Not Been Demonstrated on Absolute Positional Encoding Architectures

LM-Infinite is designed explicitly for Transformer LLMs "that use relative positional encoding" (Section 3, opening sentence). The distance ceiling component (capping d to L_pretrain in the function w(q, k, d)) is meaningful only when the positional signal is computed from relative distances. Models using absolute positional encodings—where each position has a learned or sinusoidal embedding vector that does not depend on distance—have a fundamentally different failure mode for length generalization (unseen absolute position embeddings rather than unseen distances), and LM-Infinite provides no mechanism to address it.

The paper evaluates on four model families, all of which use relative encodings (RoPE or Alibi). While this covers the dominant paradigm in open-source LLMs as of the paper's writing, it excludes several important model families:

• Earlier GPT variants (GPT-2, GPT-3) use learned absolute position embeddings.
• BERT-family encoder models use absolute position embeddings (though these are less relevant for long-context generation tasks).
• Some proprietary models may use absolute or hybrid encodings.
• Future architectures that depart from the RoPE/Alibi paradigm would not benefit from LM-Infinite without redesign.

The theoretical analysis in Section 3 (Theorem 1, Proposition 1) is also specific to relative positional attention. Theorem 1's proof relies on the assumption that the logit function w(·, ·, d) maps distances to logits; absolute positional encodings do not fit this framework because the positional signal depends on i and j independently, not on i - j. The three identified failure modes—logit explosion from unseen distances, entropy growth from unbounded attention context, and initial-token representational shifts—may or may not manifest in absolute encoding architectures, but the paper provides no analysis.

Mitigation status. The paper is transparent about this scope limitation in Section 3 ("Our discussion assumes Transformer-based LLMs that use relative positional encodings") and in the Limitations section ("The model is designed on relative positional encoding Transformer models, which is the mainstream backbone for most modern LLMs"). It does not claim applicability to absolute encoding architectures. However, a practitioner using a model with absolute encodings would find the paper's entire diagnostic framework and solution inapplicable, and the paper provides no guidance on how to adapt the approach. This is a scope constraint rather than a flaw, but it bounds the "highly flexible and can be used with most modern LLMs" claim to the relative-encoding subset of modern LLMs.

---

The Revision Chain Correct-to-Incorrect Reversion Problem Has No Architectural Solution

The paper identifies a significant practical issue with sequential revisions: "approximately 38% of correct answers get converted back to incorrect ones" during the revision chain (Section 6.1). This occurs because the revision model was trained only on sequences where all in-context answers are incorrect (followed by a correct target). At test time, when a revision in the chain happens to produce a correct answer, the model's next revision may "revise" it into an incorrect answer, because the model has never been trained to recognize when no revision is needed.

The paper mitigates this with a post-hoc selection mechanism: rather than always taking the final revision output, the system uses majority voting or verifier-based selection to pick the best answer from any point in the revision chain. This is an imperfect patch for several reasons:

• It wastes computation. If 38% of correct answers are reverted, the model is spending a substantial fraction of the revision budget undoing its own successes. The computational cost of generating revisions is paid regardless of whether they improve or degrade the answer.
• The selection mechanism itself is imperfect. Majority voting can select an incorrect answer if the model converges on a wrong consensus, and verifier-based selection depends on the verifier's quality, which is subject to the same over-optimization issues documented in Section 5.3.
• There is no mechanism to stop early. The model cannot signal "this answer is correct, stop revising." It generates a fixed-length chain and hope the selection mechanism picks the right point.

Evidence in the paper. The 38% reversion rate is reported in Section 6.1 but is not systematically characterized. It is unclear whether the reversion rate varies with difficulty (harder problems may be more prone to reversion because the model is less certain), with revision depth (later revisions may be more or less stable), or with the specific incorrect-correct pairing used in training. The paper's mitigation (chain-level selection) is evaluated only through the overall compute-optimal revision results (Figure 8), which do not isolate the contribution of the reversion problem versus other factors.

Mitigation status. The paper does not attempt an architectural solution. The training data construction—which deliberately excludes trajectories where the model should "do nothing" because the current answer is already correct—is the root cause, but fixing it would require a different data generation procedure (e.g., including some trajectories where the correct answer is repeated, training an explicit "stop revising" token). The paper does not discuss this possibility. The ReST-EM experiment (Appendix K, Figure 16) further demonstrates the fragility of revision training: attempting to optimize the revision model with RL-style training caused performance to degrade substantially with sequential revisions, suggesting the approach is sensitive to training methodology in ways not fully understood.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper fundamentally reframes the length generalization problem from a learning problem (the model hasn't been trained on long sequences and needs more data) to a computational distribution-shift problem (the model's internal quantities drift out of the ranges seen during pretraining, and constraining them suffices). This is not an incremental refinement—it changes what a solution must do. Prior to this work, the dominant paradigm was that length generalization requires either architectural modification during pretraining (Longformer, BigBird) or expensive fine-tuning on long sequences (MPT-7B-Storywriter, LongLLaMA). The paper's demonstration that an off-the-shelf 4K-pretrained model can process 200M tokens with stable perplexity using only an attention mask and a distance cap—no weight updates, no additional training tokens, no architectural changes—establishes a new lower bound on what is necessary for length generalization.

The diagnostic framework (the three factors in Section 3) resolves a prior tension in the literature that was hiding in plain sight. Relative positional encodings like RoPE and Alibi were theoretically motivated to enable arbitrary-length generalization—distance-based attention should, in principle, be invariant to absolute sequence length. Yet every model using these encodings failed catastrophically on inputs exceeding pretraining length. The field lacked a coherent explanation for why the invariant representation didn't produce invariant behavior. This paper supplies the missing mechanistic account: invariance of the representation function does not imply invariance of the numerical quantities flowing through the computation. The attention logits grow because the function class must distinguish more distances; the entropy grows because softmax must distribute weight over more tokens; the initial-token value vectors must be preserved because they span a distinct region of the feature space. None of these are failures of learning—they are failures of numerical stability and distribution matching that no amount of training data at a fixed length can prevent.

This reframing has immediate consequences for which research directions become more versus less attractive:

More attractive: constraint-based approaches to Transformer limitations. If length generalization is a distribution-shift problem, then other apparent limitations of Transformers—sensitivity to input order, difficulty with recursive structures, brittleness under domain shift—may also be addressable through constrained computation rather than additional training. The paper's methodology (diagnose the specific computational quantity that drifts → apply a minimal constraint to keep it in-range) provides a template for investigating other failure modes.

More attractive: attention manipulation as a first-class inference-time tool. The paper demonstrates that the attention pattern can be radically altered at test time without destroying model performance, as long as the alteration respects the three functional regions identified in Figure 2(b). This opens the door to task-specific attention policies—different masking patterns for retrieval versus summarization versus generation—applied to the same frozen model, analogous to how prompt engineering manipulates the input distribution without changing weights.

Less attractive: fine-tuning on long sequences as the default solution. The paper shows that MPT-7B-Storywriter (fine-tuned on 65K sequences) achieves only modestly better perplexity than LM-Infinite applied to the base 4K-pretrained MPT-7B (Table 1: 3.6 vs. 4.6 at 32K on ArXiv). This narrows the case for expensive long-context fine-tuning—the base model already possesses the necessary knowledge; it just couldn't access it stably. Fine-tuning remains valuable for pushing the performance frontier when the best possible quality is required, but for many applications, zero-shot constraint-based methods may capture most of the benefit at a fraction of the cost.

Less attractive: novel positional encoding schemes as the primary path forward. If the length generalization failure is not fundamentally about how positions are encoded but about how the encoded values interact with attention softmax at scale, then incremental improvements to RoPE or Alibi are unlikely to solve the core problem. The paper shows that both RoPE-based models (LLaMA, Llama-2, GPT-J) and Alibi-based models (MPT-7B) suffer from the same three failure modes despite their different encoding mechanics—suggesting the encoding scheme is not the bottleneck. Research attention would be better directed at the downstream numerical consequences of positional signals rather than at the signals themselves.

The paper also establishes an existence proof that unbounded length generalization is achievable without architectural change—the 200M-token result in Figure 4 demonstrates that the Transformer, as currently architected, is not fundamentally limited to its training length. This is important because it shifts the burden of proof: a new model claiming to handle long contexts must now justify why it requires architectural innovation or retraining rather than simply applying constraint-based attention at inference time.

Follow-Up Research This Work Enables

Direct measurement of the three failure modes under LM-Infinite across length scales. The paper diagnoses the three factors in the vanilla model (Figure 1) and ablates them by removing LM-Infinite components (Figure 5), but never directly plots attention logit magnitudes, attention entropies, or initial-token PCA projections under LM-Infinite on long sequences. A direct measurement study would plot the same metrics as Figure 1 but for the LM-Infinite model at 8K, 32K, 100K, and 1M lengths. The prediction is that max attention logits remain bounded at ~5–10 (the 4K pretraining range), attention entropy stabilizes at ~ln(4000) ≈ 8.3 nats regardless of total sequence length, and initial-token PCA regions remain accessible. Confirming these predictions would close the loop between the theoretical diagnosis and the empirical success, transforming inferential evidence into direct evidence. Deviations from these predictions at extreme lengths would reveal secondary failure modes not captured by the three-factor framework.

LM-Infinite on a state-of-the-art long-context benchmark suite with fine-grained position-dependent evaluation. The paper's downstream evaluation is limited to Passkey Retrieval and Qasper on a single model. Since the preprint's initial release, several comprehensive long-context benchmarks have emerged: LongBench (Bai et al., 2023), LVEval (Yuan et al., 2024), ∞Bench (Zhang et al., 2024b), and LongWanjuan (Lv et al., 2024). Evaluating LM-Infinite across a full benchmark suite would answer: (a) Does LM-Infinite's performance generalize beyond retrieval and QA to tasks requiring multi-hop reasoning, summarization, and code understanding? (b) How does performance vary with the position of critical information—do questions about content in the first 100 tokens, the middle, and the last 100 tokens show systematic differences? The middle region is LM-Infinite's structural blind spot (only accessible through the optional top-k mechanism); quantifying the performance drop for middle-located information would precisely characterize this limitation. (c) How does LM-Infinite compare head-to-head against recent fine-tuning-based context extension methods (YaRN, LongLoRA, Position Interpolation) on the same models and benchmarks, providing a FLOPs-matched comparison of the "constrain at inference" versus "fine-tune on long data" approaches?

Automated hyperparameter selection for the top-k mechanism across models and tasks. The paper's top-k hyperparameters (k = 5, h = 5, distance = L_pretrain/2) are tuned on a Passkey Retrieval validation set and applied unchanged to Qasper. The sensitivity analysis (Tables 5–7) shows that h is particularly brittle: accuracy drops from 0.94 to 0.46 when h goes from 5 to 24. A systematic study would evaluate whether these hyperparameters transfer across: (a) different tasks (is Passkey Retrieval-optimal h also optimal for Qasper, or does each task need separate tuning?), (b) different model families (do Llama-2, LLaMA, GPT-J, and MPT all benefit from the same h?), (c) different model scales (does the optimal h shift when moving from 7B to 13B or 70B parameters?). The study would produce either a set of robust defaults (if transfer is good) or a characterization of when retuning is necessary (if transfer is poor). A particularly valuable contribution would be a predictor of optimal h from model properties (number of layers, pretraining length, encoding scheme) that would eliminate the need for task-specific validation sets.

LM-Infinite as a pretraining strategy: train short, deploy long from the start. The paper applies LM-Infinite as a post-hoc inference-time modification to already-trained models. But if the Λ-shaped mask and distance ceiling keep the model's computations in-distribution regardless of sequence length, what happens if you pretrain with LM-Infinite applied? A controlled experiment would train two identical-architecture models from scratch on the same data: one with standard full causal attention on 4K segments, and one with LM-Infinite attention on 4K segments. At test time, both would be evaluated with LM-Infinite applied on long sequences. The hypothesis: the LM-Infinite pretrained model might perform better on long sequences because its training distribution exactly matches its inference distribution—there is no shift in attention patterns between training and testing. The counter-hypothesis: full attention during pretraining teaches the model to use the middle-token information that LM-Infinite discards, and this learned capability is important even if the middle tokens are later removed. The experiment would also characterize whether LM-Infinite pretraining changes the effective learning dynamics (different gradient flow through sparse attention, potentially faster training due to reduced computation per step).

Stress-testing the "infinite" claim: at what sequence length does LM-Infinite eventually fail, and why? The paper demonstrates stable NLL at 200M tokens for a single constructed sequence. A systematic stress test would construct sequences of exponentially increasing length (1M, 10M, 100M, 1B tokens) from multiple data sources, evaluate on multiple models, and track NLL, generation quality, and downstream task performance as a function of length. The goal is to identify: (a) whether there is a soft failure regime where performance degrades gradually rather than catastrophically, (b) whether the optional top-k mechanism's effectiveness degrades at extreme lengths (as the number of middle tokens grows, the fixed-k selection becomes sparser relative to the total middle region), (c) whether KV-cache management or numerical precision eventually becomes a bottleneck independent of the attention mechanism, and (d) whether different data distributions (random tokens, structured documents, repetitive text) produce different failure thresholds. Finding that LM-Infinite truly maintains stable performance indefinitely would validate the "infinite" claim; finding a failure point would characterize the method's true ceiling and motivate extensions.

Combining LM-Infinite with dynamic cache eviction for memory-constrained deployment. LM-Infinite reduces the attention context to n_starting + L_pretrain + k tokens, but for very long sequences, the KV cache for even those tokens grows linearly with sequence length (every new token must store its key and value for the starting and ending spans). Methods like H2O (Zhang et al., 2024d) or StreamingLLM (Xiao et al., 2024) propose KV-cache eviction policies that selectively discard cached entries. Combining LM-Infinite's structured sparsity (which determines which tokens to attend) with cache eviction (which determines how many of the attended tokens' KV entries to retain) could yield further memory savings for extreme-length deployment. A concrete experiment: apply LM-Infinite with L_pretrain = 4096, but only retain the most recent 1024 tokens' KV entries in the ending span (with an eviction policy for the remaining 3072), and measure the NLL vs. memory tradeoff curve. The prediction is that the Λ-shape's structured sparsity makes eviction more effective (because the model explicitly does not need middle tokens, so eviction can be aggressive in the ending span without losing information the model was going to use).

Practical Applications and Downstream Use Cases

Cost-efficient processing of long documents without fine-tuning infrastructure. An organization with a pretrained Llama-2-7B model and a collection of 50K-token scientific papers for question answering can deploy LM-Infinite immediately, without any GPU clusters for fine-tuning. The paper's efficiency numbers are directly applicable: at 32K input lengths, LM-Infinite uses 4.41 GB GPU memory per sequence versus 33.2 GB for the vanilla model (7.5× savings; Appendix G), and decodes 2.7× faster (2.70 vs. 7.34 seconds per token). This means a single A100 can process long documents that would otherwise require multiple GPUs or DeepSpeed ZeRO-3 just to fit in memory, substantially reducing infrastructure costs for long-context inference. The quality tradeoff is quantified in Table 1: at 32K on ArXiv, LM-Infinite achieves perplexity 4.2 on Llama-2 versus the vanilla model's NaN—the choice is not between "good" and "worse" but between "works" and "crashes."

Deployment of on-device long-context assistants with bounded memory. A mobile or edge deployment scenario—a local LLM processing a user's long conversation history or a lengthy document—faces hard memory constraints. LM-Infinite's 7.5× memory reduction at 32K (Appendix G) translates directly to longer usable context on memory-constrained devices. A device that could previously support a 4K-token context window with full attention can now support ~30K tokens with LM-Infinite. The method requires no model weight modification, meaning existing on-device optimized model files (quantized, compiled) can be used as-is with only an attention mask modification in the inference code. The n_starting parameter's robustness (1–100 all work; Table 4) means the implementation does not need per-device tuning.

Long-document summarization and information extraction pipelines. A batch processing pipeline that needs to summarize 100,000 research papers, each averaging 15K–40K tokens, faces a choice: truncate inputs (losing potentially critical information from the paper's introduction or methods), fine-tune a model on longer sequences (requiring training infrastructure and long-context training data), or apply LM-Infinite at inference time. LM-Infinite provides the third option with the generation quality demonstrated in Table 3: Llama-2 + LM-Infinite achieves BLEU of 23.9 and ROUGE-L of 28.8 at 8K on ArXiv, compared to 0.0 for the vanilla model. The optional top-k mechanism (k = 5, h = 5, fixed distance = L_pretrain/2) can be enabled when the task requires locating specific information in the middle of documents (e.g., extracting methodology details), providing Passkey Retrieval-level access (81.2% accuracy averaged across lengths; Table 2) to middle-document content that truncation would discard.

Enabling LLM-based code review on large pull requests and repositories. A code review system that must reason about a 20K-token diff in the context of a 100K-token repository faces the same length generalization challenge. LM-Infinite's Λ-shaped mask naturally maps to this use case: the "starting tokens" can preserve the repository's overall structure and the PR description; the "ending window" captures the immediate diff context. The distance ceiling ensures that cross-file references (tokens tens of thousands of positions apart) are processed with in-distribution positional signals. The 2.7× decoding speedup (Appendix G) directly reduces latency for interactive code review scenarios where the model must generate comments or suggestions in real time. While the paper does not evaluate on code tasks, the identical underlying architecture (decoder-only Transformer with RoPE) and the domain-agnostic nature of the three failure modes suggest transferability.