PaperlayerView exact revision on arXiv

Source-linked AI summary

Make Your LLM Fully Utilize the Context

Shengnan An, Zexiong Ma, Zeqi Lin, Nanning Zheng, Jian-Guang Lou

arXiv:2404.16811v2cs.CLcs.AI

TL;DR

Long-context LLMs can overlook information in the middle of their inputs, motivating explicit supervision for position-robust information use. The paper introduces IN2 training with synthesized long-context QA data and applies it to Mistral-7B as FILM-7B, improving probing and real-world long-context performance without compromising short-context performance.

  • Problem

    Long-context LLMs struggle to robustly use information throughout the context, particularly at middle positions, raising the question of how to fully utilize long contexts.

  • Method

    IN2 training uses synthesized 4K–32K-token contexts assembled from randomly placed ~128-token segments, with questions requiring fine-grained awareness of one segment or reasoning across multiple segments.

  • Results

    FILM-7B improves probing and real-world long-context performance while preserving short-context performance, including 23.5→26.9 F1 on NarrativeQA and 59.3→59.2 accuracy on MMLU.

  • Takeaways & Limitations

    Synthesized long-context training generalizes to real-world tasks and enables robust information retrieval across context positions, including performance comparable to or stronger than GPT-4-Turbo on probing tasks.

  • Takeaways & Limitations

    Needle-in-the-Haystack may overestimate long-context capabilities because its document-style context resembles natural-language corpora used during pre-training.

Abstract

from arXiv · show

While many contemporary large language models (LLMs) can process lengthy input, they still struggle to fully utilize information within the long context, known as the lost-in-the-middle challenge. We hypothesize that it stems from insufficient explicit supervision during the long-context training, which fails to emphasize that any position in a long context can hold crucial information. Based on this intuition, our study presents information-intensive (IN2) training, a purely data-driven solution to overcome lost-in-the-middle. Specifically, IN2 training leverages a synthesized long-context question-answer dataset, where the answer requires (1) fine-grained information awareness on a short segment (~128 tokens) within a synthesized long context (4K-32K tokens), and (2) the integration and reasoning of information from two or more short segments. Through applying this information-intensive training on Mistral-7B, we present FILM-7B (FILl-in-the-Middle). To thoroughly assess the ability of FILM-7B for utilizing long contexts, we design three probing tasks that encompass various context styles (document, code, and structured-data context) and information retrieval patterns (forward, backward, and bi-directional retrieval). The probing results demonstrate that FILM-7B can robustly retrieve information from different positions in its 32K context window. Beyond these probing tasks, FILM-7B significantly improves the performance on real-world long-context tasks (e.g., 23.5->26.9 F1 score on NarrativeQA), while maintaining a comparable performance on short-context tasks (e.g., 59.3->59.2 accuracy on MMLU). Github Link: https://github.com/microsoft/FILM.

1 Introduction

Long-context LLMs often overlook information in the middle of their context, motivating IN2 training and the FILM-7B model. FILM-7B improves long-context performance while preserving short-context capability.

  • Problem: Long-context LLMs struggle to robustly use all provided information, often overlooking content in the middle of long contexts.This lost-in-the-middle challenge can affect even simple probing tasks such as Needle-in-the-Haystack and passkey retrieval.
  • Motivation: General training may introduce position bias by emphasizing nearby tokens and placing influential system messages at the beginning of contexts.This motivates explicitly supervising information use throughout the context.
  • Approach: IN2 training uses synthesized long-context QA data requiring fine-grained awareness of one segment or integration and reasoning across multiple randomly placed segments.Contexts range from 4K to 32K tokens and are built from short segments of approximately 128 tokens.
  • Results: 23.5→26.9 F1 score on NarrativeQA, while MMLU accuracy remains comparable at 59.3→59.2.These examples represent improved real-world long-context performance alongside preserved short-context performance.

2 Information-Intensive Training

IN2 constructs synthetic long-context QA data from short segments to supervise both precise retrieval and multi-segment reasoning. The resulting training dataset mixes long-context, short-context, and general instruction data before fine-tuning Mistral-7B into FILM-7B.

  • Training Data Construction: IN2 constructs long-context QA examples in which answers require information from randomly placed short segments within each context.The dataset formalizes long contexts as Li with questions qi and answers ai requiring selected segments.
  • Training Data Construction: The construction pipeline generates QA pairs from a natural-language corpus, then combines necessary source text with randomly sampled segments into a shuffled long context.It supports both fine-grained information awareness and integration of information at different positions.
  • Fine-Grained Information Awareness: Fine-grained awareness treats a 128-token segment as the minimum information unit and prompts an LLM to create a question and answer specific to that segment.The segment and distractor segments are shuffled into the synthesized context.
  • Integration and Reasoning: Integration and reasoning examples require answers based on information from two or more 128-token segments that may be far apart after joint shuffling.This extends the training objective beyond retrieving isolated segment-level details.
  • Data Mixture: The dataset contains 1.1M fine-grained long-context examples, 300K integration-and-reasoning examples, 150K short-context QA examples, and 200K general instruction examples.The corresponding proportions are approximately 63%, 17%, 9%, and 11%.
  • Training Details: Training uses instruction tuning on Mistral-7B-Instruct-v0.2, with long contexts and questions as instructions and answer-token loss for model updates.The training process uses a global batch size of 128 and approximately 300 GPU days.

3 Long-Context Probing

The paper argues that Needle-in-the-Haystack can overestimate long-context utilization and introduces VAL Probing to evaluate information retrieval across context styles and retrieval directions. VAL Probing uses approximately 32K-token contexts and covers document, code, and structured-data tasks.

  • 3.1 Near-Perfect Performance on Needle-in-the-Haystack: Are We There Yet?: Near-perfect Needle-in-the-Haystack results may overestimate long-context capability because the task uses familiar document contexts and only forward retrieval.The authors therefore argue that this task may not adequately reflect general long-context information utilization.
  • 3.1 Near-Perfect Performance on Needle-in-the-Haystack: Are We There Yet?: Forward retrieval places the answer immediately after the query keyword, making the task an easily learned follow-up copying pattern.The paper connects this simplification to the induction-head mechanism.
  • 3.2 VAL Probing: VAL Probing evaluates three context styles and three retrieval patterns, including forward, backward, and bi-directional retrieval.The contexts contain approximately 32K tokens, and the tasks cover document, code, and structured-data settings.
  • 3.2 VAL Probing: Document Sentence Retrieval uses natural-language sentences and bi-directional retrieval, with word-level recall as its evaluation metric.The expected sentence contains words both before and after a supplied piece, and the sentences are sampled from arXiv abstracts.
  • 3.2 VAL Probing: Code Function Retrieval uses Python functions and backward retrieval to recover a function name from lines within its definition, measured by exact-match accuracy.The function name precedes the selected definition lines, and code is sampled from StarCoder.

4 Experiments and Analysis

Experiments evaluate FILM-7B on probing, real-world long-context, and short-context tasks, while analyzing comparisons, retrieval hardness, sliding windows, and RoPE settings. FILM-7B improves long-context utilization and preserves short-context performance, whereas sliding windows weaken long-distance retrieval.

  • Probing Results: FILM-7B consistently performs robustly across positions in all three probing tasks, unlike vanilla Mistral-7B, which loses information in the middle.The probing tasks compare document, code, and structured-data contexts with varied retrieval patterns.
  • Probing Results: FILM-7B achieves performance comparable to, or better than, GPT-4-Turbo on VAL Probing, with greater robustness in document and code tasks.The evaluation uses average score and min-max gap to measure overall utilization and positional robustness.
  • Probing Results: VAL Probing remains challenging for LongAlign and InternLM2 despite their near-perfect Needle-in-the-Haystack performance.Its document and database tasks vary retrieval patterns and context styles relative to Needle-in-the-Haystack.
  • Real-World Long-Context Tasks: FILM-7B significantly improves real-world long-context performance and reaches SOTA-level results among approximately 7B open-source models.The improvements generalize from synthesized training contexts to real-world scenarios.
  • Short-Context Tasks: FILM-7B maintains nearly comparable short-context performance to its backbone model, with only minor variances.Figure 5 compares the two models across short-context tasks.
  • Training Strategy Analysis: Performances drop dramatically when retrieval distance exceeds the 4K sliding-window size, whether the window is used during pre-training and IN2 training or only IN2 training.The result indicates that sliding windows substantially hurt long-context capability.
  • Training Strategy Analysis: Increasing the RoPE base θ from 1.0 × 10^6 toward 1.0 × 10^8 improves VAL Probing performance, motivating a 10-times-default setting for IN2 training.The paper attributes this exploration to increased information intensity despite unchanged context length.

5 Related Work

Related work develops long-context LLMs through data engineering and other directions, while evaluations span real-world benchmarks and concise probing tasks. VAL Probing is introduced as the first listed probing benchmark to explicitly incorporate varied retrieval patterns.

  • Long-context LLMs: Long-context LLM research primarily follows data engineering and other development directions for extended context windows.The passage introduces two broad directions, though the supplied text details only the first incompletely.
  • Long-context evaluations: Long-context evaluations include real-world benchmarks for QA, summarization, and language modeling, alongside probing tasks targeting context lengths and positions.The probing category includes Needle-in-the-Haystack, passkey retrieval, RULER, and VAL Probing.
  • Long-context evaluations: VAL Probing is presented as the first listed probing task to explicitly incorporate a variety of retrieval patterns.Its purpose is to provide a concise reflection of long-context utilization across context lengths and positions.

6 Conclusion

The conclusion presents IN2 training as a way to address lost-in-the-middle. Applied to an open-source model, it improves probing and real-world long-context results without compromising short-context performance.

  • 6 Conclusion: IN2 training is introduced to overcome the lost-in-the-middle problem.The method is applied to the open-source FILM-7B model.
  • 6 Conclusion: FILM-7B improves probing and real-world long-context performance while preserving short-context performance.This is the paper’s summarized outcome across its evaluation settings.

A Data Filtering Strategy

The IN2 dataset construction applies pre-filtering to reduce evaluation-data contamination. Raw segments overlapping evaluation examples by 10 grams are excluded from question-answer generation and random-segment use.

  • A Data Filtering Strategy: A 10-gram overlap with any evaluation example excludes a sampled segment from IN2 dataset construction.The filter covers probing, real-world, and short-context evaluation data.
  • A Data Filtering Strategy: Filtered segments are used neither to generate question-answer pairs nor as random segments [r_j].The strategy is applied while sampling raw texts for training-data construction.

B Training Examples for IN2 Training

The IN2 training examples test fine-grained awareness of information in short segments and integration of information across multiple segments.

  • Fine-Grained Information Awareness: Fine-grained examples require answering highly specific questions from information in a short context segment.Examples cover document and code contexts, including identifying people charged with crimes, kernel drivers, and an Arlington County Board plan.
  • Training Example Coverage: The training examples span document, code, and structured information while varying the number and placement of relevant segments.The supplied examples include news-style passages, kernel documentation, legal text, and product descriptions.
  • Integration and Reasoning of Information: Integration examples require combining information from at least two context segments to answer the question.The examples include legal reasoning across segments and questions about benefits derived from multiple pieces of information.

C Performance on Needle-in-the-Haystack

Figure 7 evaluates FILM-7B on Needle-in-the-Haystack within its 32K context window. The reported performance is near-perfect across that window.

  • Needle-in-the-Haystack: Near-perfect performance is reported for FILM-7B on Needle-in-the-Haystack within its 32K context window.Figure 7 presents the model’s performance on this probing task.

D Prompts For Data Generation and Training

The data-generation prompts produce specific question-answer pairs from context, including questions that require linking information across multiple pieces.

  • Prompt For Equation 1: The first prompt asks for one highly specific, clear, and context-answerable question with its answer.It disallows generic questions and phrases such as “based on the provided context.”
  • Prompt For Equation 2: The second prompt requires questions whose answers use at least two context pieces and multiple logical connections or inferences.Its template formats the context as numbered pieces before the generated question.
  • Training Template: The IN2 training template defines examples as long-context question-answer triples in which answers depend on randomly placed short segments.This template operationalizes the goal that any position in a long context can contain crucial information.
Loading 2404.16811v2…