An Efficient Recipe for Long Context Extension via Middle-Focused Positional Encoding

Abstract: Recently, many methods have been developed to extend the context length of pre-trained LLMs, but they often require fine-tuning at the target length ($\gg4K$) and struggle to effectively utilize information from the middle part of the context. To address these issues, we propose $\textbf{C}$ontinuity-$\textbf{R}$elativity ind$\textbf{E}$xing with g$\textbf{A}$ussian $\textbf{M}$iddle ($\texttt{CREAM}$), which interpolates positional encodings by manipulating position indices. Apart from being simple, $\texttt{CREAM}$ is training-efficient: it only requires fine-tuning at the pre-trained context window (e.g., Llama 2-4K) and can extend LLMs to a much longer target context length (e.g., 256K). To ensure that the model focuses more on the information in the middle, we introduce a truncated Gaussian to encourage sampling from the middle part of the context during fine-tuning, thus alleviating the "Lost-in-the-Middle" problem faced by long-context LLMs. Experimental results show that $\texttt{CREAM}$ successfully extends LLMs to the target length for both Base and Chat versions of $\texttt{Llama2-7B}$ with "Never Miss A Beat". Our code is publicly available at https://github.com/bigai-nlco/cream.

• The paper introduces CREAM, achieving efficient long context extension by fine-tuning within the original context window.
• It utilizes truncated Gaussian sampling to enhance middle context representation, addressing the 'Lost-in-the-Middle' issue.
• Experimental results show CREAM outperforms baselines on benchmarks while maintaining low perplexity even at 256K tokens.

Efficient Context Window Extension in LLMs: The CREAM Approach

Introduction

The paper "Never Miss A Beat: An Efficient Recipe for Context Window Extension of LLMs with Consistent 'Middle' Enhancement" presents an innovative method named CREAM (Continuity-Relativity indExing with gAussian Middle) for extending the context windows of pre-trained LLMs efficiently. This approach addresses two critical challenges: computational overhead due to fine-tuning at target lengths and the degradation of performance when processing the middle sections of long contexts, commonly referred to as the "Lost-in-the-Middle" problem.

Main Contributions

CREAM leverages the strengths of positional encoding (PE) methods, which are known for their straightforward implementation and rapid adaptability, and introduces several key improvements:

1. Efficiency in Fine-tuning: CREAM requires fine-tuning only within the pre-trained context window (e.g., 4K tokens for Llama 2), yet it enables effective extension to much longer target context lengths.
2. Middle-focused Enhancement: By incorporating a truncated Gaussian distribution, CREAM prioritizes the sampling of positions from the middle part of the context during fine-tuning, significantly mitigating the "Lost-in-the-Middle" issue.
3. Superior Positional Indexing: CREAM makes strategic changes to positional indices to balance continuity and relativity, ensuring better long-range dependency learning and reducing computational complexity.

Methodological Insights

Context Division and Indexing Strategies

CREAM divides the pre-trained context into three segments: head, middle, and tail. The head and tail segments are fixed at small values, improving continuity, while the lengths for the middle segment indices are determined using truncated Gaussian sampling, emphasizing relativity and enhancing performance in "middle" positions. By ensuring the coverage of relative positions across all segments efficiently, CREAM optimizes both ends of the spectrum: short and long-range dependencies.

Truncated Gaussian Sampling

The middle segment positions are sampled using a truncated Gaussian function, which fosters better focus and learning of the middle contexts. This method ensures that the model dedicates more resources to understanding and retrieving content from the middle of the context, addressing a vulnerability seen in most PE-based extension methods.

Experimental Results

CREAM was evaluated through extensive experiments using the Llama 2-7B and Llama 2-7B-Chat models, with context sizes extended up to 256K tokens. The results demonstrate:

1. Performance on Long-Context Benchmarks: In the LongChat-Lines and “Lost-in-the-Middle” tasks, CREAM markedly outperformed baseline methods like PoSE and RandPos. At 32K tokens, CREAM-Linear outperformed PoSE-Linear by 21.2% in middle index retrieval tasks.
2. Instruction Tuning Efficiency: CREAM-Chat required only 100 steps of instruction-tuning to achieve strong results on Needle-in-a-Haystack and LongBench benchmarks, outperforming models such as LongChat-v1.5-7B-32k on average by 1.6%.
3. Perplexity Metrics: Across different evaluation datasets, CREAM displayed lower perplexity scores indicating better performance without sacrificing the language modeling capabilities of the base models. When extended to extremely long contexts (up to 256K), the perplexity increase was minimal, showcasing CREAM’s stability and effectiveness.

Implications and Future Directions

The proposed CREAM method illustrates a significant advancement in efficiently extending the context windows of LLMs without considerable computational overhead. Practically, this enables more effective deployment of LLMs in applications requiring long-term context understanding such as document summarization, question answering, and dialogue systems.

From a theoretical standpoint, CREAM’s balanced approach between continuity and relativity in positional encoding highlights a promising direction for future research. Potential areas for further exploration include:

• Alternative Positional Index Strategies: Testing other positional interpolation methods and their integration with CREAM’s Gaussian-based sampling.
• Application-Specific Fine-tuning: Investigating the optimal fine-tuning strategies for domain-specific LLM applications, ensuring that the middle-focused enhancement yields consistent improvements.
• Scalability: Extending this approach to even larger models and more diverse datasets to further verify its robustness and effectiveness in real-world scenarios.

In conclusion, CREAM demonstrates substantial improvements for context window extension in LLMs, providing an efficient and effective recipe for leveraging large-scale pre-trained models in long-content processing tasks, with minimal loss in performance and significant gains in middle-context understanding.