Make Your LLM Fully Utilize the Context

Summary

This paper introduces INformation-INtensive (IN2) training, a data-driven approach to address the "lost-in-the-middle" problem in large language models (LLMs) with long context windows. The authors hypothesize that this issue stems from insufficient explicit supervision during long-context training. They create a synthetic long-context question-answering dataset that requires models to utilize information from various positions within long contexts. By applying IN2 training to the Mistral-7B model, they develop FILM-7B, which demonstrates significantly improved performance on probing tasks designed to assess long-context utilization across different context styles and retrieval patterns. FILM-7B also shows enhanced performance on real-world long-context tasks while maintaining comparable performance on short-context tasks. The authors provide detailed analyses of their training strategies, including the impact of sliding windows and position encoding adjustments.

Strengths

- The paper introduces IN2 training, an innovative data-driven solution to address the "lost-in-the-middle" problem in long-context language models. This tackles a critical issue in the field of long-context LLMs, potentially improving their ability to utilize information from extended contexts. - The authors develop a thorough evaluation framework called VAL Probing, which assesses long-context capabilities across various context styles (document, code, and structured-data) and retrieval patterns (forward, backward, and bi-directional). This provides a more rigorous and diverse testing suite compared to existing benchmarks like Needle-in-the-Haystack. - FILM-7B, the model developed using IN2 training, demonstrates significant improvements not only on the probing tasks but also on real-world long-context tasks. Importantly, it maintains performance on short-context tasks, showing that the improvements in long-context handling do not come at the expense of other capabilities. The model even achieves performance comparable to or better than GPT-4-Turbo on some tasks.

Weaknesses

- I do not see a major flaw in the paper. But one potential concern is: since the data is synthesized and the chunk is randomly injected into the document, would the model memorize such synthetic data but lose the generalizability on other long-context tasks beyond such scenarios? - The authors may add more details about their construction of the training data.

Questions

- Will the training data be released to the community? - I saw a related paper a few months ago that developed a multi-hop qa style needly-in-the-haystack benchmark. I tried to find the paper during the review process but unfortunately, I didn't find it. The basic idea is similar to the training data construction method described in the paper. If the authors also came across this paper, it would be great if you could test your model on the benchmark. If not, it is totally fine. I will revise my comment and add the paper if I can find it later.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes.

Summary

The paper argues that long-context models suffer from a “lost in the middle” phenomenon due to the infrequency of important data at any one position in the middle of the window during training. To correct for this, the authors propose IN2, a method of training on synthetic long-context QA data to artificially control for when important information is at each point in the context window during training. The model trained with this method, FILM-7B, demonstrates stronger performance than its base model at the 32k context scale.

Strengths

S1. The paper clearly outlines a problem (the lost in the middle phenomenon), a proposed root cause (the insufficient supervision during training for important info in the middle context), and a proposed solution (IN2 training). The work is well-motivated and the proposed root cause, while not empirically demonstrated, seems reasonable. S2. The method clearly improves the performance of the base model on both synthetic (probing) and realistic long-context tasks. It’s exciting that this improvement comes with the use of very little genuinely-long-context data, and I find it very interesting that increasing the RoPE base is helpful here. S3. The design of the three probing tasks to focus around directionality of information is well-justified, and I believe the choice of different domains for source text for each task is appropriate, given that no general claims are being made about relative performance between “forward” and “backward” retrieval. The authors also correctly identify that these probing tasks are not sufficient for evaluation and present results on a reasonable set of existing long-context datasets. S4. I appreciate the concern taken to avoid data contamination and the explicit description of the steps taken (in Appendix A).

Weaknesses

W1. My biggest concern is that I’m not really convinced that this improvement is due to IN2 rather than the general benefits to downstream tasks of finetuning on QA and instruction tuning data, maybe coupled with some additional benefit from changing the RoPE base that would hold independently of IN2 training’s data intervention. Compared to Mistral-7b-Instruct-v0.2, FILM is trained on ~1.75M additional tokens of data, including quality instruction tuning data and a large amount of QA-style tasks. Comparing this model to Mistral-7b-Instruct-v0.2 directly does not seem to be a fair comparison for evaluating IN2’s qualities as a data selection strategy for finetuning. The real comparison here would be comparing FILM to a Mistral model finetuned on the same quantity and ordering of data but with the IN2-generated synthetic question/answer pairs removed. I understand that this is very computationally costly, but even a checkpoint at 20% of the data mixture (in the same style as the RoPE ablations) would be very helpful to understand how much benefit the method is actually providing. Ideally you could also compare finetuning on non-IN2 instruction data with and without modifying the RoPE base. W2. The insight that increasing the RoPE base improves performance here is really cool, but I think the reasons why are underexplained and underexplored. The paper attributes this to “higher information intensity” in the context window, but this phrase is not defined in the text. I think it’s okay if the improvements from scaling theta are a purely empirical finding, but any claims about the cause of the improvement require much more justification. W3. Notation in 2.1. I think the level of notational formality used here for the synthetic data generation isn’t strictly necessary, but if it’s included, more specificity is necessary. In particular, is [r_j] restricted to never include s_i / [s_i]? What number of segments might be included in [s_i]? In equation 2, it may be helpful to remind the reader of what [r_i] is, as it is not referenced in the text explaining equation 2 before the equation. And when, on line 108, you write that the length balancing is achieved “by reject[ion?] sampling on [r_j]” — does this simply mean you reduce the number of segments you sample? I don’t quite understand how this is rejection sampling, especially given that all segments are the same length.

Questions

Q1. There are existing QA datasets for single-hop and multi-hop reasoning. Why generate data instead of using an existing corpus and mixing in other pretraining data as the filler material? Q2. The paper focuses on improving performance within the pretrained context length. Do you think IN2 training would be beneficial when finetuning to *extend* the context length as well? Q3. What is information intensity (line 257), and why would training with higher information intensity require a larger RoPE base? Q4. Can you provide any evidence from your experiments that your method's gains come from the synthetic QA data used rather than the benefits of additional (instruction tuning) data and training time? minor suggestions/typos: - line 205: `lm_eval` provides a preferred citation format; please cite them instead of using a footnote here. - line 554: “illustrates how do we segment” → “illustrates how we segment” - line 556: capitalized “We” in the middle of the sentence - line 127: was it 300 GPU days for all experiments or for training a single model? It seems from Table 3 that there were multiple models (at least partially) trained; it would be helpful to have numbers for total compute used and compute used on the final FILM-7b. - Needle in the Haystack is often attributed jointly to [Ivgi et al 2023](https://aclanthology.org/2023.tacl-1.17/) and [Liu et al 2023](https://arxiv.org/abs/2307.03172). - Table 1: why are the FiLM numbers for database bolded when GPT-4-Turbo is outperforming it? - It would be nice to reference what kinds of experiments are in the appendices somewhere in the main text, to make them more discoverable (e.g. I don’t think Appendix G is mentioned anywhere in the main body). - In the second critiquing needle in the haystack, should cite work on the reversal curse and other critiques of needle in the haystack (there have been several, including critique specifically of the single-needle case). Some of this work is already discussed in the related work, but should be cited here to show that there is community discussion about these critiques of needle in the haystack. - The claim that sliding window strategies greatly hurt long-context capability— I think this is a bit too broadly stated for the experimental results provided. - Line 70: “This demonstrates that training on synthesized long-context data can be generalized to real-world scenarios.” I think care should be taken here to emphasize that this is *additional* training on synthesized long-context data, after already training on real long-context documents at that same length (during Mistral’s pretraining).

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

The authors identify insufficient ablations as a limitation, which I agree with. However, the concern of the limitations section seems to be that the method might either 1) be brittle to the choice of hyperparameters here (which I agree is unlikely) or 2) perform much better with careful tuning. My concern is really 3) the method may not outperform existing baselines of normal instruction tuning (possibly plus RoPE base scaling).

Summary

In this paper, the authors propose two types of training to mitigate the lose-in-the-middle issue of Large Language Models (LLMs). The first one creates some QA pairs. They first pick 128-token segments with which they use LLM to generate QA pairs and mix these 128 tokens with other contexts to form long context QA pairs. The second one extends this approach with multiple segments. The authors also propose three VAL Probing tasks that ask LLMs to print the context of certain phrases. Results show that the model fine-tuned with proposed methods can significantly mitigate the lose-in-the-middle issue on challenging long context tasks and do well on Val Probing while keeping good performance on short context tasks.

Strengths

1. Lost-in-the-middle is an important and challenging tasks that need to be solved to fully utilize the capability of LLMs. 2. The proposed fine-tuning approaches demonstrate effective performance on many real-world problems, showing the potential usage of the proposed method on Large Language Models (Perhaps GPT-4). Also, it can also keep the good performance on short tasks, showing the fine-tuning won't hurt the unrelated tasks 3. The paper is well-organized and easy to follow. I like figure one which shows the core idea of whole paper.

Weaknesses

1. The comparison with natural dataset. What is the difference between fine-tuning on existing long context tasks, such as L-EVAL, LongBench, and SCOLLS. Why do you use synthetic data rather than these data? Do you think your synthetic data can work better? This comparison should be helpful in understanding the effectiveness of proposed approaches. 2. Many choices are not justified. What is the context you use to mix 128-token segments with gold answers? Are they the best choices compared with using similar contexts like LongBench does? Do you have any criteria for choosing segments with gold answers? 3. Intermediate results are not evaluated. What is the quality of the synthetic QA? Can we ensure the multihop question is valid? Is the split of two or more segments natural (not breaking the context)? Which fine-tuning approach helps more on the final results? Some quality evaluation and ablation should be important for the analysis of the proposed approaches. 4. Lacks detail analysis. For instance, the fine-tuning is only conducted on QA tasks. How could it improve the performance of Summarization? Will you include more content in the middle of the source as well? Overall, the strong improvement in small LLM demonstrates the effectiveness of the proposed method. Thus, I would assume that the proposed approaches can obtain positive answers to most of the questions above.

Questions

See weakness. Also, what is bold in table 2?

Rating

6

Confidence

4

Soundness

2

Presentation

3

Contribution

3

Limitations

They have limitations and broader impact sections. I think more discussion on synthetic data usage will be helpful.

Summary

This work tackles an essential problem in LLMs, called lost-in-the-middle, that causes the lack of information in the middle of lengthy input. As a data-driven solution, the authors propose a new training approach, INformation-INtensive (IN2). IN2 aims to enhance the use of context information in an unbiased position by shuffling the clue part of the data during training. When creating the shuffled data, the authors use GPT-4-Turbo to maintain its grammatical and semantic appropriateness. Furthermore, the authors propose a new model, FILM-7B, additionally trained Mistral-7B on a dataset created by their proposed IN2. To comprehensively evaluate FILM-7B's long-context handling ability, the author introduces VArious Long-context (VAL) Probing. VAL Probing results show that FILM-7B can outperform even GPT-4-Turbo by avoiding being trapped by lost-in-the-middle. Experimental results on long-context tasks from the LongBench collection and short-context tasks, MMLU, BoolQ, RACE-High, CommonsenseQA, ARC-Challenge, HellaSwag, GSM8K, and MATH show that FILM-7B outperforms Mistral-7B models.

Strengths

- This paper proposes a new training method, INformation-INtensive (IN2), that can make trained LLMs handle long-context with mitigating the influence of lost-in-the-middle. - By using IN2, the authors train Mistral-7B to make a new model, FILM-7B, than can handle long-context. - The authors provide a newly constructed evaluation task, VArious Long-context (VAL) Probing, that covers various long-context tasks. - The evaluation result on VAL Probing shows the effectiveness of IN2 based on the improved performance of FILM-7B. - The evaluation results on long-context and short-context tasks also show that FILM-7B outperforms FILM-7B.

Weaknesses

- From the viewpoint of Scaling Laws, comparing Mistral-7B and its further trained model FILM-7B is unfair.

Questions

Currently, the evaluation is conducted in a zero-shot setting. Is there any possibility that FILM-7B outperforms proprietary models like GPT-4-Turbo on a few-shot setting by using its ability to handle long contexts equally?

Rating

8

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

Even though IN2 mitigates the lost-in-the-middle problem, experimental results show that this problem is still left in some tasks.

The response meets my expectations

I appreciate your in-depth responses based on the actual evaluation results. The result addressed every concern raised by me. I am really thankful to the authors for doing such high-quality research.