Source-linked AI summary

LLM Maybe LongLM: Self-Extend LLM Context Window Without Tuning

Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia-Yuan Chang, Huiyuan Chen, Xia Hu

arXiv:2401.01325v3cs.CLcs.AIcs.LG

TL;DR

LLMs struggle with contexts longer than their training window, raising the question of whether this reflects missing long-context ability or unseen positional information. SelfExtend remaps unseen relative positions to positions observed during pretraining and applies this during inference without tuning. Across reported long-context evaluations, it improves performance and matches or surpasses many fine-tuning-based methods, while large group sizes can degrade performance and computation remains a consideration.

  • Problem

    LLMs behave unpredictably on sequences longer than their pretraining context window because inference exposes them to unseen relative positions.

  • Method

    SelfExtend elicits long-context ability by mapping unseen relative positions to positions seen during pretraining, using inference-time changes without fine-tuning.

  • Results

    SelfExtend improves long-context performance and matches or surpasses many fine-tuning-based methods across synthetic and real-world tasks.

  • Takeaways & Limitations

    The results support the paper’s view that LLMs can handle long contexts without additional tuning when positional O.O.D. information is addressed.

  • Takeaways & Limitations

    Performance degrades with large group sizes, and naive implementations increase computation cost because SelfExtend performs extra attention across query-key pairs.

Abstract

from arXiv · show

It is well known that LLMs cannot generalize well to long contexts whose lengths are larger than the training sequence length. This poses challenges when employing LLMs for processing long input sequences during inference. In this work, we argue that LLMs themselves have inherent capabilities to handle long contexts without fine-tuning. To achieve this goal, we propose SelfExtend to extend the context window of LLMs by constructing bi-level attention information: the grouped attention and the neighbor attention. The grouped attention captures the dependencies among tokens that are far apart, while neighbor attention captures dependencies among adjacent tokens within a specified range. The two-level attentions are computed based on the original model's self-attention mechanism during inference. With minor code modification, our SelfExtend can effortlessly extend existing LLMs' context window without any fine-tuning. We conduct comprehensive experiments on multiple benchmarks and the results show that our SelfExtend can effectively extend existing LLMs' context window length. The code can be found at \url{https://github.com/datamllab/LongLM}.

1. Introduction

LLMs often behave unpredictably on inputs longer than their pretraining window because inference exposes them to unseen relative positions. SelfExtend remaps these positions to training-seen positions without fine-tuning, and performs competitively with fine-tuning-based methods on long-context tasks.

  • Pretraining on fixed-length sequences limits most LLMs’ context windows, causing unpredictable behavior and severe perplexity increases on longer inputs.
  • The positional O.O.D. issue arises when inference introduces relative distances absent during pretraining.For RoPE, relative position is represented by m −n between two tokens’ absolute positions.
  • SelfExtend uses floor division to map unseen large relative positions to positions encountered during pretraining, without fine-tuning.Its motivation is that approximate position and information order can preserve meaning over long texts.
  • SelfExtend operates at inference time as a plug-in method evaluated on language modeling, synthetic long-context, and real-world long-context tasks.
  • Across synthetic and real-world long-context tasks, SelfExtend matches or surpasses many existing fine-tuning-based models.

2. Preliminary

The preliminary section introduces positional encoding and RoPE, where positional information is incorporated into queries and keys so their inner product depends on token relative distance.

  • Transformers encode position using absolute embeddings or relative positional encodings.Absolute embeddings add a position vector to each token embedding before model processing.
  • RoPE incorporates positional information into query and key vectors so their inner product contains relative positional information.
  • RoPE defines query and key transformations using position-dependent complex rotations before computing attention interactions.
  • The real part of the query-key inner product depends on the tokens’ relative distance m −n through an abstract mapping function g.

3. SelfExtend

SelfExtend addresses positional out-of-distribution failures by combining grouped attention for distant tokens with normal attention for neighboring tokens, extending inference context without fine-tuning.

  • Grouped attention: SelfExtend maps unseen relative positions to positions observed during pretraining using a FLOOR operation applied before self-attention.This grouped attention preserves the original self-attention mechanism apart from position mapping.
  • Grouped attention: Grouped attention reduces positional out-of-distribution problems by mapping large positions into a smaller discrete set within the pretraining context range.The group size determines the base of the FLOOR operation.
  • Preliminary analysis: With grouped attention, Llama-2-7b-chat maintains relatively low and stable PPL on sequences exceeding its pretraining context window.PPL remains nearly unchanged within the enlarged context window, but explodes again beyond that window.
  • Preliminary analysis: Grouped attention slightly increases PPL relative to the original model, while preserving similar PPL behavior within the extended context window.For Llama-2, the reported group sizes extend the window to 8192, 16384, and 32768 tokens.
  • SelfExtend: SelfExtend combines grouped attention for distant tokens with standard attention for adjacent tokens within a specified neighbor window.The method merges the two attention parts before softmax while leaving other components unchanged.
  • SelfExtend: SelfExtend shifts grouped-attention positions before merging attention values so the transition between neighbor and grouped regions remains smooth.The example extends a pretraining window of 7 to 10 tokens using group size 2 and neighbor window 4.

4. Experiments

SelfExtend is evaluated across language modeling, synthetic and real-world long-context, short-context, and ablation settings. It maintains performance beyond pretraining windows, achieves strong retrieval and benchmark results without fine-tuning, and exhibits hyperparameter trade-offs.

  • Language modeling: SelfExtend maintains low PPL beyond the pretraining context window for Llama-2-7b-chat and Mistral, whereas PPL explodes without SelfExtend after the window limit.Mistral with sliding-window attention also maintains low PPL, motivating evaluation beyond perplexity alone.
  • Synthetic long-context tasks: 100% passkey retrieval accuracy is achieved across all tested depths and context lengths without fine-tuning.Mistral with sliding-window attention cannot retrieve passkeys outside its 4096-token sliding window.
  • Real-world long-context tasks: SelfExtend achieves the best performance on most LongBench datasets for several base models and superior performance on nearly all L-Eval datasets.The LongBench comparison includes models with and without SelfExtend, while most alternative context-extension methods require fine-tuning.
  • Real-world long-context tasks: Across two long-context benchmarks, SelfExtend achieves comparable or better performance than methods requiring further fine-tuning.The reported results are obtained through inference-time operation without additional training.
  • Short-context and ablation analyses: SelfExtend maintains short-context performance while improving long-context performance, but group size and neighbor-window size each involve performance trade-offs.Large groups coarsen position information, small groups rely on less-trained larger position embeddings, and larger neighbor windows sharpen local information while coarsening global information.

5. Conclusion and Discussion

SelfExtend elicits long-context capabilities without tuning by remapping unseen relative positions to positions encountered during pretraining. The method improves long-context performance, but computation, context-length, and evaluation limitations remain.

  • Conclusion and Discussion: SelfExtend maps unseen relative positions to positions seen during pretraining, eliciting long-context capabilities without tuning or further training.It uses a simple FLOOR operation for this remapping.
  • Conclusion and Discussion: SelfExtend increases computation cost under naive implementations because it performs extra attention across all query-key pairs.Blocked-kernel optimizations can make the computation linear, with a small marginal cost for long inputs.
  • Conclusion and Discussion: Large group sizes degrade performance, preventing SelfExtend from supporting indefinitely long contexts.The paper identifies this as a limitation of the method.
  • Conclusion and Discussion: SelfExtend processes the entire sequence, whereas prompt-compression methods can shorten inputs to reduce computation.This marks a practical boundary relative to methods that reduce input length.
  • Conclusion and Discussion: Long-context evaluation methodologies remain open research questions, and the lack of standard practices complicates experimental results.The paper presents this as an evaluation limitation rather than a property of SelfExtend itself.

Impact Statement

The contribution is positioned as improving access to long-context LLM applications while simplifying deployment. Its implementation is presented as an inference-stage pseudocode procedure using queries, keys, values, positional encoding, grouped attention, and masking.

  • Impact Statement: SelfExtend can improve LLM usability for document analysis, long-form question answering, and retrieval augmented generation.The impact statement identifies these as domains requiring lengthy input sequences.
  • Impact Statement: The method can simplify deployment by extending context windows without fine-tuning, making advanced language models more accessible.The supplied pseudocode describes the implementation as a lightweight inference procedure.
  • Impact Statement: The pseudocode computes normal attention from position-encoded queries and keys, then applies a causal mask.The normal-attention branch uses q, k, and v as queries, keys, and values.
  • Impact Statement: The grouped-attention branch uses floor-divided positions and a shift to construct grouped query-key attention.This branch encodes the paper’s position remapping design before applying a causal mask.
  • Impact Statement: A mask combines local neighbor attention with grouped attention across the sequence.The pseudocode constructs the grouped mask relative to the neighbor-window size.

B. Perplexity as a Metric for Long Context Capabilities

Perplexity can remain low when local context predicts most tokens, even if critical tokens needed for long-context understanding are missed. Therefore, lower PPL does not necessarily indicate better long-context ability.

  • Perplexity as a Metric for Long Context Capabilities: Perplexity can be comparable to or lower than SelfExtend’s while a local-window method still fails to genuinely handle long contexts.The cited example, Infinite, divides a long sequence into short local subsequences.
  • Perplexity as a Metric for Long Context Capabilities: A local attention window can produce low PPL by modeling most tokens accurately while missing a few critical tokens.Those critical tokens matter for understanding long contexts and answering questions.
  • Perplexity as a Metric for Long Context Capabilities: Inference PPL is a static point analogous to a location on the pretraining loss curve, so it cannot alone determine performance.The paper contrasts this with the informative decreasing trend observed during pretraining.
  • Perplexity as a Metric for Long Context Capabilities: Lower PPL does not necessarily equate to better performance in understanding long contexts.The paper retains low PPL as essential for a good model while rejecting it as a sufficient long-context measure.

D. Detailed Experimental Setting

The appendix presents the details of the paper’s experiments.

  • Detailed Experimental Setting: The appendix presents the details of the experiments in the paper.

D.1. Experimental Setting on Language Modeling Tasks

The language-modeling evaluation uses a restricted PG19 setup and examines models with different context-window configurations, including Mistral variants with and without sliding-window attention.

  • Experimental setting: PG19 evaluation uses the first sentence from each of 100 test books rather than the standard perplexity-testing setting.Perplexity is computed with a sliding window using S = 256.
  • Experimental setting: Figure 8 evaluates Phi-2 with SelfExtend against vanilla Phi-2 (Phi-2-2k) on GSM8K, Helleswag, MMLU, and ARC-c.The x-axis represents group size, while the y-axis reports the corresponding task metric.
  • Experimental setting: Mistral without sliding-window attention has an 8192-token context window that SelfExtend extends beyond 16k.With sliding-window attention, Mistral uses a default 4096-token window and can process sequences of unlimited length.

D.2. Experimental Setting on Passkey Retrieval Task

The passkey retrieval task tests whether models can locate a key sentence embedded in long repeated text, using a prompt-insensitive synthetic setup illustrated by an example figure.

  • Task design: Passkey retrieval places a sentence carrying a passkey among repeated random text and tests whether the model can locate it.Figure 9 provides an example of the task.
  • Task design: Unlike Needle in a Haystack, passkey retrieval is designed to be less sensitive to prompt wording.The passkey sentence differs markedly from the surrounding repeated random text.

D.3. Experimental Setting on Varying-Length Passkey Retrieval Task

The varying-length passkey experiments compare SelfExtend and other long-context Llama-based models while documenting the evaluated model configurations and MPT’s non-RoPE positional-encoding setup.

  • Compared models: The comparison includes Llama2-7b-chat with SelfExtend, LongLora-7b-16k, vicuna-1.5-7b-16k, Together AI’s Llama-2-7b-32k, and Yarn-Llama-2-7b-64k.LongLora uses its fully fine-tuned variant for this task, while vicuna and Together AI’s Llama-2-7b-32k were fine-tuned using position interpolation.
  • Compared models: The paper provides links to additional details for the LLMs used in the experiments.These links accompany the listed model configurations.
  • Compared models: Table 7 lists the LLMs used in the experiments.The table serves as the experiment-model inventory.
  • Positional encodings: SelfExtend is also tested with MPT-7b-chat, which uses Alibi rather than RoPE positional encodings, on PG19.The reported results indicate that SelfExtend works with non-RoPE positional encodings similarly to models using RoPE.
  • Positional encodings: For the MPT experiment, the vanilla model has a 2k-token context window, and Table 8 reports perplexity across sequence lengths.SelfExtend uses a neighbor window of 512 and a group size of 6.

G. Hyperparameyer Selection for SelfExtend

SelfExtend’s hyperparameters are generally robust within a suitable range, but the preferred settings depend on the model and should be selected empirically when necessary.

  • Empirical rule: SelfExtend is not overly sensitive to group size or neighbor-window selection when both are neither too large nor too small.The paper reports that heuristic default values often provide satisfactory performance.
  • Empirical rule: Excessively small group sizes can degrade performance because they require less-well-trained relative positions.The explanation links precise position information to reliance on relative positions outside the model’s well-trained range.
  • Empirical rule: Excessively large neighbor windows can degrade performance because they likewise require less-well-trained relative positions.Larger windows provide more neighbor information but can move attention toward poorly trained positional relationships.
  • Model dependence: The hyperparameter rule may not apply across all models; for example, Llama3 uses a much smaller neighbor window of approximately 100.The paper suggests using a simple representative task to identify suitable settings.
Loading 2401.01325v3…