Writing in the Margins: Better Inference Pattern for Long Context Retrieval (2408.14906v1)

Abstract: In this paper, we introduce Writing in the Margins (WiM), a new inference pattern for LLMs designed to optimize the handling of long input sequences in retrieval-oriented tasks. This approach leverages the chunked prefill of the key-value cache to perform segment-wise inference, which enables efficient processing of extensive contexts along with the generation and classification of intermediate information ("margins") that guide the model towards specific tasks. This method increases computational overhead marginally while significantly enhancing the performance of off-the-shelf models without the need for fine-tuning. Specifically, we observe that WiM provides an average enhancement of 7.5% in accuracy for reasoning skills (HotpotQA, MultiHop-RAG) and more than a 30.0% increase in the F1-score for aggregation tasks (CWE). Additionally, we show how the proposed pattern fits into an interactive retrieval design that provides end-users with ongoing updates about the progress of context processing, and pinpoints the integration of relevant information into the final response. We release our implementation of WiM using Hugging Face Transformers library at https://github.com/writer/writing-in-the-margins.

• The paper introduces a novel inference pattern (WiM) that uses segment-wise KV cache prefill and margin generation to enhance long-context retrieval.
• The method achieves approximately 7.5% improvement in multi-hop reasoning accuracy and a 30% F1-score increase in aggregation tasks.
• The approach offers an interactive retrieval design with real-time progress updates and early exit capabilities to optimize computational efficiency.

Writing in the Margins: Better Inference Pattern for Long Context Retrieval

The paper, "Writing in the Margins: Better Inference Pattern for Long Context Retrieval," introduces a novel inference pattern for optimizing the handling of extensive input sequences in retrieval-oriented tasks for LLMs. The proposed method, Writing in the Margins (WiM), leverages the chunked prefill of the key-value (KV) cache to perform segment-wise inference, which significantly enhances the processing of long contexts. In particular, WiM guides models through a series of intermediate steps, generating and classifying auxiliary information termed "margins," leading to better task-specific performance.

Context and Motivation

As LLMs process larger input sequences, their performance tends to degrade, primarily due to fixed context windows and inherent limitations in the attention mechanisms. This degradation is especially pronounced in retrieval tasks where relevant information is embedded within large volumes of text. Previous attempts to extend the context window have included sparse attention, length extrapolation, and context compression, as well as prompt engineering techniques like Chain of Thought (CoT).

The WiM approach aims to bridge the gap between efficient transformer architectures and new prompting strategies by introducing a novel KV cache-aware reasoning pattern. This method retains intermediate outputs at each segment step, improving the model's final prediction capabilities with minimal extra computational cost.

Key Contributions

1. Novel Inference Pattern:
   • The WiM pattern processes long context window tasks by prefilling the KV cache segment-wise and generating intermediate extractive summaries, or "margins," at each step.
   • Margins are reintegrated at the end of computations to enhance the model's final output.
   • This approach approximates a 7.5% improvement in accuracy for multi-hop reasoning tasks such as HotpotQA and MultiHop-RAG and a significant 30% increase in the F1-score for aggregation tasks like CWE, compared to traditional inference methods.
2. Interactive Retrieval Design:
   • WiM is implemented within an interactive retrieval setup, providing end-users with ongoing updates about the processing progress. This feature supports real-time transparency and reduces latency.
   • Users can exit early if the streamed margins already address the query, optimizing the computational load.
3. Implementation and Accessibility:
   • The implementation of WiM using the Hugging Face Transformers library is shared, allowing other researchers and developers to integrate and build upon this work.

Detailed Methodology

Chunked Prefill

The inference process starts with dividing the context into fixed-size segments to prefill the KV cache. This chunked prefill method reduces memory complexity from $O(L^2)$ to $O(LK)$, where $L$ is the prompt length and $K$ is the chunk size. The adjusted attention mask ensures each new chunk attends to all previous chunks while maintaining causality.

Writing in the Margins

For a given prompt $P$ composed of context $C$ and instruction $I$: $P = C + I$ The context $C$ is divided into $N$ segments: $C = c_1 + c_2 + \cdots + c_N$

Each chunk $c_k$ is processed with prefilled past key values $\text{pkv}_{[1..k-1]}$. At each step, an additional extractive instruction $I_A$ is appended to generate intermediate outputs or margins $M_i$, which are retained as plain text for subsequent steps.

Experimental Setup

Datasets

The performance of WiM was evaluated using curated datasets linked to three main skills: Multi-Hop Reasoning (HotpotQA, MultiHop-RAG), Needle Retrieval/Single-Hop Reasoning (SQuAD), and Aggregation (CWE). Each dataset contained long context examples with lengths ranging from 13k to 64k tokens.

Long Context Window LLMs

Seven off-the-shelf LLMs supporting up to 128k tokens were assessed: Phi-3-small-128k-instruct, Qwen2-7B-Instruct, Meta-Llama-3.1-8B-Instruct, Phi-3-medium-128k-Instruct, Palmyra-4-Chat-128K, Meta-Llama-3.1-70B-Instruct, and Qwen2-72B-Instruct. All models were tested under identical conditions using a 0-shot prompt configuration.

Results

The results demonstrated the robustness and capability of WiM in enhancing task-specific performance across all model and dataset combinations. WiM consistently outperformed both the conventional Long Context LLM (LLM) and the Retrieval-Augmented Generation (RAG). The improvement was especially apparent in complex multi-hop reasoning and aggregation tasks.

Implications and Future Directions

WiM shows substantial potential in improving the efficiency and effectiveness of LLMs in handling long-context retrieval tasks without requiring model fine-tuning. The segment-wise processing and integrated margin notes provide not only enhanced model performance but also greater transparency and user engagement.

Future research directions include refining KV cache management for even greater computational efficiency, exploring the applicability of WiM to other transformer models performing segmented context windows, and leveraging innovative techniques like PagedAttention for optimizing memory usage.

Conclusion

By introducing the WiM inference pattern, this paper pushes the boundaries of how LLMs handle extensive input sequences in retrieval tasks. The promising results and practical implications underscore the potential of this approach to reshape long-context LLM implementations and applications in real-world scenarios.