From Artificial Needles to Real Haystacks: Improving Retrieval Capabilities in LLMs by Finetuning on Synthetic Data
Authors: Zheyang Xiong
Vasilis Papageorgiou
Kangwook Lee
Dimitris Papailiopoulos
Venue: ICLR 2025 Poster (2025)
Keywords: Synthetic Data
LLM finetuning
Long Context
Retrieval
Primary area: foundation or frontier models, including LLMs

Abstract: Recent studies have shown that Large Language Models (LLMs) struggle to accurately retrieve information and maintain reasoning capabilities when processing long-context inputs. To address these limitations, we propose a finetuning approach utilizing a carefully designed synthetic dataset comprising numerical key-value retrieval tasks. Our experiments on models like GPT-3.5 Turbo and Mistral 7B demonstrate that finetuning LLMs on this dataset significantly improves LLMs' information retrieval and reasoning capabilities in longer-context settings. We present an analysis of the finetuned models, illustrating the transfer of skills from synthetic to real task evaluations (e.g., $10.5\%$ improvement on $20$ documents MDQA at position $10$ for GPT-3.5 Turbo). We also find that finetuned LLMs' performance on general benchmarks remains almost constant while LLMs finetuned on other baseline long-context augmentation data can encourage hallucination (e.g., on TriviaQA, Mistral 7B finetuned on our synthetic data cause no performance drop while other baseline data can cause a drop that ranges from $2.33\%$ to $6.19\%$). Our study highlights the potential of finetuning on synthetic data for improving the performance of LLMs on longer-context tasks.

Published as a conference paper at ICLR 2025
FROM ARTIFICIAL NEEDLES TO REAL HAYSTACKS: IM- PROVING RETRIEVAL CAPABILITIES IN LLMs BY FINE- TUNING ON SYNTHETIC DATA
Zheyang Xiong", Vasilis Papageorgiou", Kangwook Lee",, Dimitris Papailiopoulos sime
"University of Wisconsin-Madisen m Microsoff Research
ABSTRACT
Recent studies have shown that Large Language Models (LLMs) struggle accu- rately retrieve information and maintain reasoning capabilities when processing long-context inputs. To address these limitations, we propose finetuning approach utilizing carefully designed synthetic dataset comprising numerical key-value retrieval tasks. Our experiments on models like GPT-3.5 Turbo and Mistral 7B demonstrate that finetuning LLMs this dataset significantly improves LLMs' in- formation retrieval and reasoning capabilities in longer- context settings. We present analysis of the finetuned models, illustrating the transfer of skills from synthetic to real task evaluations (e.g., 10.5% improvement on 20 documents MDQA at position 10 for GPT-3.5 Turbo). We also find that finetuned LLMs' performance on general benchmarks remains almost constant while LLMs finetuned on other baseline long-context augmentation data can encourage hallucination (e.g., on TriviaQA, Mistral 7B finetuned on our synthetic data cause no performance drop while other baseline data can cause drop that ranges from 2.33% 6.19%). Our study highlights the potential of finetuning synthetic data for improving the performance LLMs on longer-context tasks.
INTRODUCTION
Recent studies have revealed that Large Language Models (LLMM) struggle to accurately retrieve information and maintain reasoning capabilities when processing longer context inputs when retrieval is required across different parts of their context (Liu et al. 2023 Levy et al. 2024). These limitations hinder their performance on tasks that involve processing and reasoning over extensive textual information, such as summarization or question answering over long passages.
To address these challenges, we propose a novel approach that involves finetuning LLMs on carefully designed fully numerical synthetic algorithmic dataset containing key-value dictionary retrieval tasks (i.e., see Figure [ for an example of such a task). We conduct extensive experiments on popular LLMs, including GPT-3.5 Turbo (OpenAT 2023) and Mistral 7B (Jiang et al.) 2023). and find that our method improves their performance on both information retrieval and long-context reasoning.
Specifically, our approach mitigates the "oost-i--themmidlle" phenomenon identified by Liu al. (2023) and significantly improves performance on the FLenQA benchmark (Levy et al. (2224) that measures LLMs' long-context reasoning capability. Interestingly, we observe that finetuning on proposed dataset often yields more significant improvement compared to finetuning corresponding benchmark's 's data. In addition, it results in only a slight degradation on popular benchmarks such MMLU (Hendrycks et al. 202T) and HellaSwag (Zellers et al. 2019). indicating that the overall capabilities of the models remain largely unaffected. Finally, another advantage of our proposed dataset is that it contains no factual information: as it was recently discovered by Gekhman et al. 2024), finetuning on previously unseen knowledge may encourage hallucinations.
the
Email: <zheyangtcs. sc. sc.edu>. Correspondence: <dim tri criapai 1,10>.

---

Published as a conference paper at ICLR 2025
Simple dictionary vayuv retrieval
Do task using the list of dictionaries below.
Dictionary [1] {122: 765, 4548: 1475, 4818: 4782)
Dictionary [2] {526: 290, 9205: 9318, 9278: 1565)
Dictionary [32] (2931: : 8364, 196: 1464, 812: 5363]
Dictionary [85] (344: 1579, 116: 617, 330: 411)
Above a 1ist of dictionaries such that each key and value is an integer. Report the value of key 2931 and the dictionary it is in.
Desired answer: The value of key 2931 is 8364 and it is in Dictionary [32].
Figure 1: An example prompt with desired answer of simple dictionary key-value retrieval task.
Thus, finetuning our key-value dataset improves LLMs' retrieval and reasoning without suffering from such unwanted characteristics.
Our findings highlight the potential of finetuning synthetic data as promising approach to enhancing the performance of LLMs on real downstream tasks. Our paper is organized as follows: in Section 2 we describe the format of the proposed dataset, and its variations that provide (or not) an answer template to the model, in Section we present our experimental results, Section 4 we discuss the main limitations and possible future directions of our work, and in Section 5 we discuss our main conclusions.
1.1 RELATED WORK
Long Context LLMs. Recent works have observed LLMs' limited retrieval and reasoning ca- pabilities in the long-context setting. Liu et al. (2023) discovered a positional bias when LLMs retrieve information from long contexts. In particular, the authors found out that the retrieval accuracy drops when the desired information lies in the middle of the context. Kamradt (2023) conducted the "needle-in-a aaee-in----h-ttte hh-aahhatttck" experiment by placing random fact (the "needle") in long input context (the "haystack"") and observed that LLMs struggle spot the needle as the input context length grows. To mitigate this behavior, Yu (2024) and An et al. (2024) finetuned LLMs on long-context augmentation data consisting long-context question-anse wnwwerinn tasks to enhance LLMs' long-context capabilities. Tang et al. (2023) shuffled the prompt and marginalized the prompt order biases in the long- cong-context setting and Zhang et al.] (2024) re-scaled the indices in positional encoding. Levy et al. (2024) introduced benchmark, FLennQ, by extending input samples with varying lengths and types of padding, discovering LLMs' significant degradation reasoning ability at context lengths much shorter than the maximum limit.
There are also other relevant works on long-context LLMs (Junging et al., 2023 Mohtashami &
Jaggi 2023 Chen al.) 2023b Bai et al. 2023 An et 2023). Xu al. (2023) showed that Retrieval Augmented Generation (RAG) can be as accurate as full finetuning on longer context windows. Chen et al. (2023a) extended the LLM's predetermined context limit by treating it as an interactive agent who processes the input through iterative prompting. Jin et al. (2024) extended LLM's context window by remapping the unseen relative positions during inference. Zhu al. (2024) introduced "LONGEMBED". embedding models' context window up to 32,768 tokens, leveraging Rotary Position Encoding (RoPE) in processing long contexts. Fu et al. (2024) proposed a data engineering recipe for scaling LLMs to 128k context lengths through lightweight continual pretraining on a balanced mixture of Sength-upsampler data. Peysakkovvch & Lerer (2023) proposed "attention sorting." method that improves long context models by iteratively sorting documents based on attention and generating responses with the re-ordered context.
benchmark and suite of training- free strategies to extend
2

---

Published as a conference paper at ICLR 2025
Multi-subkee dictionary kay-v retrieval
Do task using the list of dictionaries below.
Dictionary [1] { (141, 986, 163): 2528, (726, 947, 349, 820): 4130)
Dictionary [2] (555, 710, 424): 5756, (623, 141, 997) 1633, (957, 634, 969): 7871}
Dictionary [6] {(645, 417, 847): 6409, (141, 623, 616) 5617]
Dictionary [49] {(710, 105, 141, 799) 5369, (623, 210, 477): 8971, (899, 126, 999) 4409]
Above is a 1ist of dictionaries such that each key is a tuple of integers and each value is an integer. Report the key that contains the integers 616, 141, 623 (not necessarily in order), its value, and the dictionary it is in.
Desired answer: The key that contains the integers 616, 141, 623 is (141, 623, 616). Its value is 5617 and it is Dictionary [6].
Figure 2: An example prompt with desired answer of multi- dictionary key-value retrieval task. Here (141, 623, 616) is the gold key. Note that 141 and 623 in the gold key are also subkeys of other keys.
Data-centric AI. In recent years, the field of data-centric AI has emerged, which focuses improving the quality and efficiency of AI systems through data-oriented approaches rather than model-centric techniques (Sener & Savarese 2018 Ghorbani & Zoul 2019 Zha et al. 2023 Albalak et al. 2024). Gadre et a (2024) and Mazumder et al. (2024) proposed benchmarks that fix model training code, where the goal is to design better datasets to achieve better performance. [ee al.
(023) and Zhou et al. 2024) studied the data format in training transformers to learn arithmetic tasks.
LLM Benchmarks and Evals. Much research has been recently conducted towards the design of meaningful benchmarks that probe the capabilities of LLMs. Benchmarks such as GLUE (Wang et al. 2018). SuperGLUE (Wang et al.] (219) test whether a model has general language understanding capabilities. MMLU (Hendrycks et al. 2021) aims to measure the models' accuracy across wide variety of tasks that span STEM, humanities, social sciences, and more, while GSM8k (Cobbe et al., 2021) tests capabilities on school math. In HellaS (Zellers et al.) (219)) models are presented with an event description and must select the most likely follow-up sentence from set of carefully selected choices, while HumanEval (Chen et al.) 2021) measures their ability to generate code given docstrings. TriviaQA (Joshi et al., (2017) is a reading comprehension benchmark and NQ-Open (Lee) al., 2019 Kwiatkowski et 20199a) is open domain question-answering benchmark where the question-answer pairs are collected from a diverse set of fields.
SYNTHETIC DATASET OF RETRIEVAL TASKS
In this section, introduce the dataset on which we finetune the models. The dataset consists of two synthetic retrieval tasks: 1) simple dictionary key-value retrieval and 2) multi-subkey dictionary key-value retrieval.
Simple dictionary key-value retricval. this task, we provide the model with a list of dictionaries of integer keys and values, and ask it to retrieve the value of a specified key (denoted as the gold key). Figure T shows an example of this task and the detailed algorithm is shown in Algorithm
Multi-subkey dictionary key-value retrieval. For models that can already tackle the first task (e.g., for the first task GPT 3.5 Turbo achieves around 0.99 accuracy irrespective of the position of gold key), we design a harder version of the key-value retrieval task where each key is a tuple of subkeys. Other keys can share some but not all of the subkeys of the gold key. We increase the difficulty of this
3

---

Published as a conference paper at ICLR 2025
Simple dictionary key-value retrieval (with an answer template)
Do task using the list of dictionaries below.
Dictionary [1] {122: 765, 4548: 1475, 4818: 4782)
Dictionary [2] {526: 290, 9205: 9318, 9278: 1565)
Dictionary [32] (2931: : 8364, 196: 1464, 812: 5363]
Dictionary [85] (344: 1579, 116: 617, 330: 411)
Above a list dictionaries such that each key and value an integer. Report the value of key 2931 and the dictionary it is in. Answer in the following template: The value of key 2931 is <fill vall-en-vllue> and it is in Dictionary
[<fill-aaddictiooaryyaaam
Desired answer: The value of key 2931 is 8364 and it is in Dictionary [32]
The prompt of the simple dictionary key-value retrieval task is provided with an answer template.
Figure 3:
Instruction
Instruction
Repors the value of key 2931 and the dictionary it 1a in.
Report the value of key 2931 and the dictionary it is in. Anawer in the following template: The value key is <fill-in-valse>
Target Answer
it is in Dictionary [<fill-in-dicttonnayye
value of key 2931 is 8364 and i: in Dictionsry [32].
Target Answer
The value of key 2931 is B364 and it is Dictionary [32]-
Figure 4: Token-level loss on the target answer when provided with (right) and without (left) an answer template, where red indicates high and green low loss.
task by randomizing the order of subkeys the prompt so that the order is not necessarily the same as that of the gold key. Figure] shows example of this task and the detailed algorithm is shown in Algorithm
Prompt with answer template. Note that with the prompt in Figure I slightly different answers like "8364 the value key 2931 in dictionary 32" and "Dictionary [32] has the key 2931 with value 8364" are also correct. Therefore, since the model is finetuned on the entire answer, during supervised finetuning, it also learns the format of our provided answer besides learning to retrieve the desired value. In order to make the model only focus retrieving the correct value without being affected by the format of the answer, we provide the model with an answer template with which want the model to answer. Figure 3 shows Figure 4 we visualize the token-level loss on the target answer, where red indicates high and green low loss. an answer template is provided, the loss on the formatting part is small. This lets the model to focus on the important part and learn the right skill rather than how to answer the question.
example of prompt with answer template. In
3 EXPERIMENTS AND RESULTS
Our goal is to investigate whether finetuning LLMs (in particular, GPT-3.5 Turbo and Mistral 7B on our proposed synthetic numerical retrieval tasks improves their long context capabilities on
gpt 3opt-3.5-trtbe. 3.5-turbo- 106 and Mi stral 1 -7H- In-tnttruct--
4

---

Published as a conference paper at ICLR 2025
20 Documents MDQA (-4k tokens)
20 Documents MDQA (-4k tokens)
0.92
0.86
capoeatee templaaa)
esgovatus (Wle tamplatu)
0.90
0.84
Accuracy
Accuracy
0.82
0.86
gin 610
on kopsaue retiéva zomp
0,80
0.84
on MDGA
hepoceii retaac (wil templeri
0.78
0.82
0.76
0.80
0.74
0.72
10
15
20
10
20
Position of the gold document
Position of the gold document
(a) GPT-3.5 Turbo and the finetuned versions.
(b) Mistral 7B and the finetuned versions.
Figure 5: Performance of GPT-3.5 Turbo, Mistral 7B and their corresponding finetuned versions the MDQA task.
natural language tasks: multi-document question answering (MDQA) (Liu et al. 2023) and flexible length question answering (FLenQA) (Levy et al. 2024).
3.1 STAGE 1: FINETUNING LLMS ON SYNTIETIC RETRIEVAL TASKS
For Mistral 7B, our dataset consists of 350 samples of simple dictionary key-value retrieval tasks. Each task has 85 dictionaries and each dictionary has 3 to 4 keys, so each prompt has roughly 4K tokens. We finetune the model only the answer part (masking out the instruction part) for 2 epochs. More implementation details are in A.I Figure II shows Mistral 7B's s performance on simple dictionary key-value retrieval task before and after finetuning.
Since GPT-3.5 Turbo already performs well on simple dictionary key-value retrieval task, we finetune it multi-subkey dictionary key-value retrieval tasks. The dataset consists of 150 samples and each sample has 49 dictionaries. We finetune the model for epochs using OpenAl's API.
3.2 STAGE 2: EVALUATIONS ON LONG CONTEXT RETRIEVAL AND REASONING TASKS
3.2.1
MCLTI-DOCUMENT QUESTION ANSWERING (MDQA)
We test models' capabilities of retrieving important information long context setting. In MDQA, provide the model with i documents and prompt it to answer question such that only 1 of k documents (denoted as the gold document) contains the answer and the other k documents (denoted distractors) are completely irrelevant the question. We the setting
context with
20 documents (around 4K tokens) and place gold document at positions {1, 5, 10, 15, 20] 2 For each position, we test the model on 200 task samples and measure the accuracy using the maximum subspan exact match as in (Liu et al. 2023).
Finetuning LLMs synthetic key-value retrieval tasks enhances their perfor- mance on practical retrieval tasks, demonstrating effective transfer of learned capabilities.
Finding 1:
The result 20 documents MDQA is shown in Figure 5 where x-axis the position of gold document. In Figure 5a for the original GPT-3.5 Turbo model, there is U-shaped performance curve, indicating that the performance is highest if the important information is at the beginning or at the end of the input context, with the model struggling retrieve the answer if the important information is in the middle. Finetuning the models on synthetic retrieval tasks flattens the U-shaped curve and information is much more accurately retrieved over all positions across the input context. In Figure 5b) the original Mistral 7B model has primacy bias in the sense that it can more accurately Forr example, gold document placed at position it is the first document in the context.
5

---

Published as a conference paper at ICLR 2025
retrieve information that is at the beginning of the input context. Finetuning the models proposed data manages to improve the accuracy across all the positions in the input context. In addition, when the finetuning dataset contains template, Mistral seems to mitigate this primacy bias, showcasing a more uniform accuracy across all the positions in the input context.
Finding 2: Synthetic data is better than MDQA data even if the goal is to perform better MDQA task.
As a comparison, we also finetune the models on the MDQA dataset itself for roughly the same number of training tokens and see how finetuned models perform. Since the MDQA dataset only provides the ground truth answers in one or two words, we prompt GPT-3.5 Turbo with correct answers and let it form a complete sentence as the target answer. As shown in Figure 5a GPT-3.5 Turbo finetuned on our synthetic data perform better than the one finetuned on MDQA. In Figure 5b we can see that despite training on MDQA tasks, Mistral 7B still struggles to perform well on MDQA, with a significant performance drop when gold document is at the beginning of the prompt. These findings underscore the effectiveness of our synthetic data generation method, which enhances performance on specific datasets like MDQA, even surpassing direct finetuning on the target dataset.
3.2.2
FLEXIBLE LENGTH QUESTION ANSWERING (FLENQA)
1...5-tturbo-110 on FLenQA (cot)
Mistral-7B on FLeeeAnAAAccc (cot)
0.850
G87-3 5-tuubo---
0.80
Mozul-T feotureed (wfo tomp
CPT 3.5 findtuned (wl complate)
Martal-1d fawtured template)
0.825
GPT findtuned
0.75
Mistal-78
0.800
Accuracy
Accuracy
0.70
0.775
0.65
0.750
0.725
0.60
0.700
0.55
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
Context size
Context size
(a) GPT-3.5 Turbo and the finetuned versions.
(b) Mistral 7B and the finetuned versions.
Figure 6: Performance of GPT-3.5 Turbo, Mistral 7B and their corresponding finetuned versions on the FLenQA task, using chain-of-thought prompting.
We also test models' long context reasoning capabilities. FLenQA is a dataset comprising reasoning tasks with varying length that ranges from 250 tokens 3000 tokens. Each task consists of a context and a "True" or "False" question that can be answered by two key sentences from the context. We test chain-of-thougle 2000 task samples. For of-in-of-thougll prompting, we ask the model to produce the result step by step and derive the answer ("True" or "False") at the end, and in the non chain-of-thouge prompting we ask the model to directly answer "Truc" or "False".
(Wei et al., 2022) and non tfuthughh prompting, each with a total of
Finetuning LLMs on synthetic key-value retrieval tasks improves LLMs' long- context reasoning capabilities, even if explicit thain-of-thought reasoning is not allowed.
Finding 3:
In Figure 0 and we present our results on the FLenQA dataset. The x-axes represent the number of tokens in the context, while the y-axes represent the accuracy of the response. Figure shows results where chain-of-thought prompting is employed. In Figure 6a] we notice that although the model suffers from a performance drop if finetuned on data without answer template, finetuming GPT-3.5 Turbo on data with answer template significantly improves model's of-nn-of-thought reasoning capability. In Figure 6b we can also see that finetuning Mistral 7B on data with answer template improves models -hain-of-thougle capability. We hypothesize that the reason for this is that the

---

Published as a conference paper at ICLR 2025
GPT-3. www.5tttubo-11000 on FAenQA (no cot)
Mistral-7B on FLenQA (no cot)
0.90
GPY-3 S-trbb--ob
finotureed
GPT-3 5-tttb/bb--11100 finetuned twd tempiste!
GPT 3.5 turbo-1106 finetuned (wio temolate)
0.80
0.85
Accuracy
0.80
Accuracy
0.75
0.75
0.70
0.70
0.65
0.65
0.60
0.60
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
Context size
Context size
(a) GPT-3.5 Turbo and the finetuned versions.
(b) Mistral and the finetuned versions.
Figure 7: Performance of GPT-3.5 Turbo, Mistral 7B and their corresponding finetuned models on the FLenQA task without employing Shain-of-thought prompting.
finetuned models utilize their improved retrieval capabilities to capture relevant information more accurately, which helps them deduce the answer.
Figure pressntss results where models are required to directly answer with "True" or "False" without providing explicit reasoning. The results show notable improvement in performance for finetuned models. This improvement is significant because it demonstrates that, even if explicit reasoming (that related to retrieval capability) is not allowed, finetuning on our proposed synthetic tasks enhances the models' internal reasoning capabilities.
Finding 4: LLMs finetuned on synthetic tasks with answer templates are better.
From Figure 5 and we can observe that models finetuned synthetic key-value retrieval tasks with answer templates perform better on MDQA and FLenQA than that on without answer templates. This verifies our hypothesis that having an answer template helps the model learn the right skill more efficiently. This highlights key advantage of synthetic data: allows for greater control over the model's 's output format. Unlike real-world tasks where developing answer templates can be challenging, synthetic tasks allow for easy implementation of structured response formats, facilitating skill learning.
3.3 STAGE 3: EVALUATION OF FINETUNED MODELS' GENERAL CAPABILITIES
Finding:: 5:
Finetuning LLMs on synthetic key-value retrieval tasks does not hurt models' general capabilities.
One possible drawback of our approach is that finetuning on the proposed artificial tasks would severely harm the general purpose capabilities of the tested models. In order to assess this concern, we tested the original and finetuned versions of GPT-3.5 Turbo and Mistral 7B on some general purpose benchmarks. Note that for our assessments we used the codebases of Gao al. (2023) and Fu et al. (003)).
The results can be seen in Table I In particular, we consider five widely used benchmarks: MMLU (Hendrycks et al.) 2021 HellaSwag (Zellers et al. 2211)), GSM8k (Cobbe et al. 2021). TriviaQA (Joshi et al., 2017) and NQ-Open (Kwiatkowski et al.) 2019b). What we can observe is that all the finetuning strategies result in no significant degradation the general purpose benchmarks mentioned above.
Due to computational constraints, did not evaluate GPT-3.5 Turbo on all henchmarks, and for MMLU 20% of the full dataset.
7

---

Published as a conference paper at ICLR 2025
Table Model's performance evaluated on general ability benchmarks. All numbers are reported in percentage. Here "w/" and "w/o" denote the models that are finetuned on the the synthetic tasks that were described in Section 2
MODEL
MMLU
HellaSwag
GSM8K
Triviaqa
NQ-Open
Mistral-7B
53,42
56.31
34,65
47.63
11.61 11.98 (+0.377)
Mistral-7B ft (witemplate)
53.44 (+0.02)
(-0.09) 56.30 (-0.01)
(-0.3))
47.74 (+0.11))
Mistral-7B it (w/o template)
53.42 (-0.00)
34.14 (-0.51)
47.62 (-0.01)
11.40 (-0.21)
GPT-3.5-turbe
68.07
72.33
GPT-3.5-turbo ft (w/template)
67.75 (-0.32)
71.65 (00088)
GPT-3,5-turbo ft (w/ template)
68.16 (+0.09)
75.06 (-2.73)
3.4 STAGE 4: COMPARISONS WITH OTHER BASELINES
We also consider three additional long-context augmentation datasets as baselines: MultidocQA Ya 2024). IN2 (An al., 2024), and Needle-in-ae in-a-haystack (Kamradt) (023). MultidocQA is a dataset of multiple documents question and answering where the model needs to paraphrase the document before answering. IN2 is long-context question answering dataset where the answer can deduced from one multiple parts of the context. Needle-in-ae in-a-haystack is widely used long-context test set where the model is prompted to identify some key information (the needle) within long context (the haystack). We finetune Mistral 7B on these baselines, using roughly the same number of training tokens and report their performance on MDQA, FLenQA, and general purpose benchmarks.
20 Documents MDQA (-4k tokens)
Baselines on FlenQA (cot)
0.95
Mizz -8-iratuci-0011
MuseooGh
0.8
0.90
Nend--nstaytta
Accuracy
0.85
Accuracy
0.7
0.80
0.6
Mizu- 7B insuuss v6.1
Mu ACooDA
0.75
IN2
0.5
Nesds-na- srndd-nn-taayttce
Ours yel lumplats)
10 Position of the gold dictionary
15
20
500
1000
1500
2000
2500
3000
Context size
(a) MDQA
(b) FLenQA with chain-of- prompting
Baselines FlenQA (no cot)
eocti
0.80
IN2
Curs tsl lemplow)
0.75
Accuracy
0.70
0.65
0.60
500
1000
1500
2000
3000
Context size
(c) FLenQA without chain-of-though prompting
Figure 8: Performance of finetuned Mistral 7B on (a) MDQA, (b) FLenQA with chain-of of thoughtt prompting, and (c) FLenQA without tf-houht prompting.
8

---

Published as a conference paper at ICLR 2025
Table Mistral 7B and finetuned versions' performance evaluated on general ability benchmarks. All numbers are reported in percentage.
Finetuning dataset
MMLU
HellaSwag
GSM8K
Triviaqu
NQ-Open
Original Mistral-715
53.42
56.31
34.65
47.63
11.61
Ours (w/template)
53.44 (+0.00)
(-0.09)
34.34 (-0.31)
47.74 47.74++0.11)
11.98 (+0.37)
MultidocQA (Yu 2024
53.19 (-0.22)
56.27 (-0.04)
33.28 (-1.36)
(-2.43))2.43)
8.69 (-2.91)
IN2 (An al.
53.49 (+0.07)
56.44 (+0.13)
34.98 (+0.32)
45.44 (-2.19)
(-1.8)
Ncedle-in-a-hayee Kamradt 2025
52.83 (-0.59)
56.22 (-0.09)
33.79 (-0.86)
41.30
4.88 (46.73)
MDQA Liu et ai. 2023
52.94 (-0.47)
56.23 (-0.07
(-0.07)
44.77 (-2.85)
7.64 (-3.96)
Finding 6: Synthetic data do not encourage hallucinations that other baselines may yield.
From Figure X and Table 2 we can see that while some baselines outperform our proposed data on either MDQA FLenQA, they all have more significant degradation the general benchmarks test, especially on TriviaQA and NQ-Open. One possible reason that all other baselines contain factual information. Gekhman al. (2024) shows that finetuning on factual information encourages hallucinations, something that we verify observing the significant degradation on TriviaQA and NQ-Open, which are knowledge-basse basee benchmarks. In contrast, our proposed dataset is purely synthetic. comprising of key-value pairs, and as a result, does not encourage ballucinations. We also highlight another benefit our synthetic data: since it does not contain any factual information. it will not have the problem of containing potential outdated information that further encourages hallucinations, from which other long-context augmentation datasets may suffer.
3.5 STAGE 5: EVALUATION ON LONGER-CONTEXT SETTING
We also test the longer-context setting. We finetune Mistral-Pa Instruct-ve .2 on simple key-value retrieval task with maximum context length of 24K and test MDQA. We observe clear improvement over the original model as shown in Figure 9 and does not observe significant degradation in general capability, shown in Table 3
120 Documents MDQA (-24k tokens)
0.90
+ Mistral 7b indrucoo0222 templane)
Accuracy
0.86
0.84
0.82
0.80
100
120
Position of the gold document
Figure 9; Performance of finetuned Mistral-78 - v0.. .2 on 120 documents MDQA.
Table 3: The performance of Disstal-7B-Insttuuttvv... and finetuned version (on simple dictionary key-value retrieval with context length 24K) on general capability benchmarks. All numbers are reported in percentage.
Finetuning dataset
MMLU
HellaSwag
GSMSK
TriviaQA
NQ-Open
Mistral-BB IB-nttrucct--
58.74
65.98
41.85
32.65
14.46
Ours (w/ template)
58.42 (-0.32)
65.74 (-0.24)
41.72 (-0.13)))
33.02 (+0.37)
15.03 (+0.57)

---

Published as a conference paper at ICLR 2025
20 Documents MDQA with relevant distractors (-4k tokens)
0.60
Mistrr-1 (B-indraat--)) on key- resteval (wlo templast))
tempate)
0.55
Accuracy
0.50
0.45
0.40
0.35
2
10
15
Position of the gold document
Figure 10: Mistral 7B and the finetuned versions MDQA with relevant distractors. The finetuned variants do not show significant improvement over the original model.
4 LIMITATIONS AND FUTURE WORK
Our dataset does have limitation. MDQA benchmark also has another version where distractors relevant distractors, meaning that they are documents retrieved by a retrieval system (based on the relevance score) that do not contain the answer. Models finetuned on our dataset will not improve in this setting, as is shown in Figure 10 A possible future work of this study is to add our synthetic retrieval dataset as a small part of larger instruction finetuning dataset and see the difference between models finetuned with and without synthetic retrieval data and observe how they perform differently on long context retrieval and reasoning tasks.
CONCLUSION
In this work, we introduce novel finetuning approach that leverages carefully designed synthetic datasets enhance the information retrieval and reasoning capabilities of LLMs in real downstream tasks. Our study demonstrates that finetuning on our proposed synthetic data significantly improves the performance of the tested models on tasks like MDQA and FLenQA, mitigating the "lost-in- the-middle" behavior that was observed in Liu et al. (2023). On the other hand, we find that after finctuning, the models' performance on general benchmarks remains almost constant, something that indicates that their overall capabilities are mostly unaffected. We also find that compared to other long-conted augmentation datasets that contain factual information, our purely artificial data does encourage hallucinations. Moreover, will not have the problem of containing potential outdated information. Thus, we believe that our study demonstrates the potential of finetuning LLMs on carefully crafted synthetic datasets enhance their capabilities downstream tasks. We hope that our findings will inspire further research into the development of effective synthetic datasets.
ACKNOWLEDDEEMEN
This work supported by NSF CAREER Award CCF2339978, Amazon Research Award, grant from FuriosaAI, and ONR Grant No. N00014-21-1-2880 and No. N00014-23-1-2848.
REFERENCES
Alon Albalak, Yanai Elazar, Sang Michael Xie, Shayne Longpre, Nathan Lambert, Xinyi Wang, Niklas Muennighoff, Bairu Hou, Liangming Pan, Haewon Jeong, et al. A survey on data selection for language models. arXiv preprint arXiv:2402.1682. 2024.
10

---

Published as a conference paper at ICLR 2025
Chenxin An, Shansan Gong, Ming Zhong, Mukai Li, Jun Zhang, Lingpeng Kong, and Xipeng Qiu. L-eval: Instituting standardized evaluation for long context language models. arXiv preprint arXhv:2307.11088, 2023.
Shengnan An, Zexiong Ma, Zeqi Lin, Nanning Zheng, and Jian- Saan-Guang Lou. Make your Ilm fully utilize the context. arXiv preprint arXiv:2404.16811 2024.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. 14508, 2023.
Howard Chen, Ramakanth Pasunuru, Jason Weston, and Asli Celikyilmaz. Walking down the memory maze: Beyond context limit through interactive reading. arXiv preprint arXiv:2310.05029, 2023a.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained code. arXiv preprint arXiv:2107.0333 133.3444, 2021.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2.006. 15595, 2023b.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reichiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint arXiv:21 10. 14168, 2021.
Yao Fu, Litu Ou. Mingyu Chen, Yuhao Wan, Hao Peng, and Tushar Khot. Chain-of-thought hub: A continuous effort to measure large language models' reasoning performance. arXiv preprint arXiv:2305. /7306, 2023.
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. Data engincering for scaling language models to 128k context. arXiv preprint 2402:2240..10171, 2024.
Samir Yitzhak Gadre, Gabriel Ilharco, Alex Fang. Jonathan Hayase, Georgios Smyrnis, Thao Nguyen, Ryan Marten, Mitchell Wortsman, Dhruba Ghosh, Jieyu Zhang, et al. Datacomp: In search of the next generation of multimodal datasets. Advances in Neural Information Processing Systems, 36, 2024.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang. Kevin Wang, and Andy Zou. framework for few-shot language model evaluation, 12 2023. URL https:ss: /ttp:////eeooooooo
Zorik Gekhman, Gal Yona, Roce Aharoni, Matan Eyal, Amir Feder, Roi Reichart, and Jonathan Herzig. Does fine-tuning I1ms on new knowledge encourage hallucinations? arXiv preprint 00X5::2405.05904 2024.
Amirata Ghorbani and James Zou. Data shapley: Equitable valuation of data for machine learning. In International Conference on Machine Learning, pp. 2242-2251, 2019.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Mcasuring massive multitask language understanding. In International Confer- ence Learning Representations, 2021. URL https::/ wttp://wwwweeeeoummemeeeaea
d7иB I
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. 2023.
Mistral 7b. arXiv preprint
Hongye Jin, Xiaotian Han, Jingfeng Yang. Zhimeng Jiang, Zirui Liu, Chia- Yuan Chang, Huiyuan Chen, and Xia Hu. Lim maybe longlm: Selfextend Ilm context window without tuning. In Forty-first International Conference on Machine Learning, 2024.
11

---

Published as a conference paper at ICLR 2025
Mandar Joshi, Eunsol Choi, Daniel S Weld, and Luke Zettlemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1601- 1611, 2017.
He Junqing, Pan Kunhao, Dong Xiaoqun, Song Zhuoyang, Liu Yibo, Liang Yuxin, Wang Hao, Sun Qianguo, Zhang Songxin, Xie Zejian, et al. Never lost in the middle: Improving large language models via attention strengthening question answering. arXiv preprint wwx.v:2311.09/9000 2023.
G Kamradt. Needle in haystack pressure testing 1lms. hnttps://githusa www.a///..g..a.././aaaaa 2023.
SeuTeettNNedleIIAfaaay
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wer Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453 466, 2019a. doi: 10.1162/tacl -0000 URL a 0__00766
https: http://...../1..11//////.///////
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Matthew Kelcey, Jacob Devlin, Kenton Lee, Kristina N. Toutanova, Llion Jones, Ming-Wei Wei Chang, Andrew Dai, Jakob Uszkoreit, Quoe Le, and Slav Petrov. Natural questions: benchmark for question answering research. Transactions of the Association of Computational Linguistics, 2019b.
Kenton Lee, Ming-Woo Beig-WWiide Chang, and Kristina Toutanova. Latent retrieval for weakly supervised open domain question answering. In Proceedings of the 57th Annual Meeting of the Association for Com- putational Linguistics, pp. 6086 6096, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10. 11t55///1/1111/111 URL https: www ac acwweb. org/ anthology/ P19- P19-1612 1612
Nayoung Lee, Kartik Sreenivasan, Jason D Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. arXiv preprint arXiv:2307.0338/1 2023.
Mosh Levy, Alon Jacoby, and Yoav Goldberg. Same task, more tokens: the impact of input length on the reasoning performance of large language models. arXiv preprint arXiv:2402. 14848, 2024.
Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. arXiv preprint arXiv:2307.03112. 2023.
Mark Mazumder, Colby Banbury, Xiaozhe Yao, Bojan Karlaš, William Gaviria Rojas, Sudnya Diamos, Greg Diamos, Lynn He, Alicia Parrish, Hannah Rose Kirk, et al. Dataperf: Benchmarks for data-centric development. Advances in Neural Information Processing Systems, 36, 2024.
Amirkeivan Mohtashami and Martin Jaggi. Landmark attention: Random-access infinite context length for transformers. arXiv preprint arXiv:2305. 2023.
OpenAl. Chatgpt, 2023. URL 29.
htw.:///...cooooooo//o//a/tt//// Accessed: 2024-03-
Alexander Peysakhovich and Adam Lerer. Attention sorting combats recency bias in long context language models. arXiv preprint arXiv:2310.01427, 2023.
Ozan Sener and Silvio Savarese. Active learning for convolutional neural networks: core-set approach. In International Conference on Learning Representationss 2018. URL https:// openreview net/ http://ouummiiiiiooo
RM
Raphael Tang, Xinyu Zhang, Xueguang Ma, Jimmy Lin, and Ferhan Ture. Found in the middle: Permutation conssscnn improves listwise ranking in large language models. arXiv preprint arXiv:2310.077/2. 2023.
12

---

Published as a conference paper at ICLR 2025
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In Tal Linzen, Grzegorz Chrupala, and Afra Alishahi (eds.), Proceedings of the 2018 EMNLP Workshop Black- boxNLP: Analyzing and Interpreting Neural Networks for NLP,, pp. 353-355, Brussels, Belgium, November 2018. Association for Computational Linguisties. doi: 10. 1/ URL
httts::/aclaanthl..gg org. W18 5446
Alex Wang, Yada Pruksachatkun, Nikita Nangia, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. Superglue: stickier benchmark for general-purpose language understanding systems. Advances in neural information processing systems, 32, 2019.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of -hainghf- prompting elicits reasoning in large language models. Advances in 2022.
neural information processing systems,
Peng Xu, Wei Ping, Xianchao Wu, Lawrence McAfee, Chen Zhu, Zihan Liu, Sandeep Subramanian, Evelina Bakhturina, Mohammad Shoeybi, and Bryan Catanzaro. Retrieval meets long context large language models. In The Twelfth International Conference on Learning Representations. 2023.
Yijiong Yu. Training with "paraphrasing the original text" improves long-context performance, 2024.
Rowan Zellers, Ani Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can machine really finish your sentence? In Anna Korhonen, David Traum, and Lluís Mårquez (eds.), Proceedings of the 57th Annual Meering of the Association for Computational Linguistics, pp. 4791-4800, Florence, Italy, July 2019, Association for Computational Linguistics. doi: 1/119-14227
1472
Daochen Zha, Zaid Pervaiz Bhat, Kwei-Herng Lai, Fan Yang, and Xia Hu. Data-centric ai: Perspec- tives and challenges. Proceedings of the 2023 SIAM International Conference Data Mining (SDM), pp. 945-948. SIAM, 2023.
Zhenyu Zhang. Runjin Chen, Shiwei Liu, Zhewei Yao, Olatunji Ruwase, Beidi Chen, Xiaoxia Wu, and Zhangyang Wang. Found in the middle: How language models use long contexts better via plug-and-play positional encoding. arXiv preprint 1xXiv:2403.0470 2024.
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Trans- formers can achieve length generalization but not robustly. arXiv preprint arXiv:2402.03371, 2024.
Dawei Zhu, Liang Wang, Nan Yang. Yifan Song, Wenhao Wu, Furu Wei, and Sujian Li. Longembed: Extending embedding models for long context retrieval. arXiv preprint 2441:2204.129966 2024.
13

---

Published as a conference paper at ICLR 2025
A
EXPERIMENT DETAILS
A.1
FINETUNING MISTRAL 7B AND GPT 3.5 TURBO
Simple dictionary key-value retrieval (-4k tokens)
1.0
0.9
Accuracy
0.8
Mietral-7B Mestral-7BInnstuad-dd
keyvalue retrival template)
0.7
key-value relrivvl (wio temptale)
0.6
0.5
60
80
Position of the gold dictionary
Figure 11: Mistral 7B and the finetuned versions on simple dictionary key-value retrieval.
For Mistral 7B, we choose simple dictionary key-value retrieval as the task to finetune on. We use two prompting strategies to prepare the dataset: with and without an answer template as described in Section 2 For each prompting strategy generate 3 different datasets using the same configuration but with different seeds. Each dataset consists of 350 simple dictionary key-value retrieval tasks (roughly 4K tokens in each task). Each task has 85 dictionaries and each dictionary has 3 to 4 keys. Each key and value is an integer of 3 to 4 digits (in particular, we choose 4min = Tmin = 3, lmax Tmaz 4). We finetune Mistral 7B on all attention layers and use global batch size of 16 and finetune the model for epochs on each dataset with learning rate 5 x 10 10-6, For evaluation results, average across 3 runs, each with different training data and seed.
For GPT-3.55 Turbo, we choose multi-subkey key-value retrieval as the task to finetune on (in particular, choose num_dict 49, min Timin 3, Imax Tmaz 4, n_keys 3, n_common 2.Pshare = 0.5). For each prompting strategy, we generate 2 different datasets. Each dataset consists of 150 multi-subkey key-value retrieval tasks (roughly 4K tokens in each task). Each task has 49 dictionaries. We finetune GPT-3.5 Turbo for 2 epochs on each dataset using OpenAl API. For evaluation results, average across runs.
A.2 EVALUATION DETAILS
We use Im-eval-harees (ao et al. 2023) for evaluation and use the default configuration. In particular, for GSM8K we use 5-ohot prompting and for other tasks we use 0 shot prompting. An exception is for Mistral-Pee Pro-Instad Instruct-ut -vo .2, where use shot prompting for NQ-Open as the model cannot well answer in the desired format in 0 shot prompting setting.
14

---

Published as a conference paper at ICLR 2025
B DETAILS ON GENERATING RETRIEVAL TASKS
In this section we provide the pseudocodes on generating retrieval tasks introduced in the paper: (1) simple dictionary key-value retrieval and (2) multi-subkey dictionary key-value retrieval.
B. SIMPLE DICTIONARY VALUV RETRIEVAL.
Algorithm 1:
Gen key val
Input: min and max number of digits of key / value Emins Emazi gold key go gold Output: key and val where key different from gold.key
val <- randint(rmin) Emaz)
while True do
key randint( Tminn Fmaz)
if key = gol gold.key then return key, val
Algorithm 2: Simple dictionary key-value retrieval
Input: Number of dictionaries num_di ct;; min and max length of each dictionary Imin- Imazi range of all keys/ values (Imin Iman)
Output: A list of dictionaries dicts, the position of gold dictionary gold.pos. gold key çold.key and gold value gold_val.
Initialize gold foold_dico dtct as an empty dictionary gold did.dicttleer.. <- randint(Iminn (mminy)mar) gold.pos +- randint num_dict)
gold.key f randint Eminn: Imaæ)
gold_val randint (andint(Emmnn:. Imar)
Add (gold.key, gold_val) key-vluee pair to gold for i = 1, ... gowd.dict.len-
do
key, val + Gen Gen_key_noe voll_val(rrr)) min. r max, gold.key)
Add (key, val) key-value pair gold
Shuffle the order of gold..ceee gold.dict. u Initialize dicts to an empty array of dictionaries
for 1 num_dict
do
13
Initialize dict as an empty dictionary dict len randint((min. Umax) for k= 1,... dio...ddit_il dict_len do
15
16
key,val, + Gen Gen.key_ve -eekyyyyaaEmiinn Emans gold.key)
17
Add (key, val) key-value pair to dict
18
Append dict to dicts
19
Insert gold dict into dicts position gold pos
20
return di ct.
15

---

Published as a conference paper at ICLR 2025
B.2 MULTI-SUBKEY KEY-VALUE RETRIEVAL
Algorithm 3:
Gen multikey val
Input: range for all keys / values: (Tmins Tmaz) gold multi-key: gold.key.tuple number of keys in each multi-key: n_keys, keys from gold.key.tupie that can be shared with the output etypl: common conmon.subkey, probability of key sharing: Output: key-tuple and corresponding val
Pahare
assert 1en (common_subkey) n_keys
val <- randint((minn Emaz)
while True do
key; < randint((mmmn Tmaz), =i= 1,,, n_kyys key-tuple = (key1, key..... .eeey....,,,y..ee..
fori Ien (common subkey) do
L With probability Pshare replace key; with common common.subkey,
Shuffle the elements of key tuple.
if key_tuple and go holl_key-tuple share at most len (common_subkey) keys then return key-tuple, val
10
Algorithm 4: Multi-subkey dictionary retrieval
Input: Number of dictionaries: num.dict, min and max length of each dictionary: mmin: Lmazs range of each key / value: (Emins Imaz), number of keys in each multikey: n_keys, max number keys to share among key.typly syple'ss s: n_common, probability of key
sharing between keys: Pshare.
Output: A list of dictionaries dicts, the position of gold dictionary gold gold.pos, gold kel.kky.ttt. suppee and gold value gold val.
multi-key gold bolld.ey....
Assert n_common
n_keys.
Initialize goldddict
an empty dictionary
gold dict len +- randint(4min 4max)
gold.pos -- randint(1, nlm.iiittmmmmmt
gold.key, randint( mmin, Emaz), Vi 1,, keys
gold_key_tupa = (gold_
gold keys)
gold.val randint (Emin; Iman)
Choose n.common random keys from gold gold.key.tuple,
Add (golddkey goldival) key-value pair to golddict 10 for =eri=1..... gold dict len - do
key tuple, val
Genn.ullikey_vl((rTa... val(Tmin, marr; gol selldkkeytuplle n_kys,, Pshare) Gen_multikey.
12
Add (key-tuple, val) multikey-value pair to gold
Shuffle the order of gold.dict. Initialize dicts to an empty list. do
14
for j,...... num_dict
16
Initialize dict as an empty dictionary dict_len randint( Imax) for 1,.., dict do key_tuple, val Gen len__mltikee__all vaa(( Umaz: gold_key)
18
20
Add (key.tuple, val) multikey-value pair to dict
21
Append dict to dicts
22 Insert gold_dict into di s position gold.pos
23 return di sts
16


Reviews:
[{"review_id": "mqWYmll1kz", "reviewer": "Reviewer_G5rG", "summary": "The work proposes a set of synthetic tasks based dictionary-based key-value retrieval and uses it extend the context lengths of Mistral-7B-0.1 and GPT-3.5 turbo models. Finetuning on these datasets improves performance on MDQA and FLenQA benchmarks while discouraging hallucinations. Performance on general evaluation benchmarks like TruthfulQA, GSM8k is shown to be retained.", "questions": "<see Weaknesses>", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Finetuning on randomly sampled key-value datasets is able to fix the “lost-in-the-middle” phenomenon for GPT-3.5 while retaining original knowledge.\n- Finding that finetuning with answer templates performs better.\n- Work on improving LLMs through synthetic datasets is crucial for the community.", "weaknesses": "1) The proposed synthetic dataset targets only retrieval tasks, like MDQA, while other important applications of long-context, such as in-context learning and RAG, involve \\emph{understanding} the context as a whole. Hence, I have doubts about the scope of applications of this work. Concurrent works [1, 2] argue that finetuning/improving only retrieval capabilities does not capture all long-context applications.\n\n2) In line 263, it is claimed that \"..proposed data mitigates this primacy bias..\", but the curve in Fig 5(b) still shows a descent, although with a higher accuracy.\n\n3) About line 270, \"..we also finetune the models on the MDQA dataset ..\", from what I have understood, the difference between performance upon finetuning on the proposed dataset vs MDQA, can arise from two factors:\n  (a) Randomization of the position of the gold document.\n  (b) Complexity due to requirement of english understanding.\nSince point (a) is easy to track, further study is needed to understand the differences between the two. Point (b), for example, can be explored by including the <Mistral-7B ft with MDQA> row in Table 1.\n\n4) In Section 3.4, a possible reason MultiDocQA & IN2 outperform the proposed dataset is that the baselines require extracting information from multiple contexts. Can you solve this drop by adding a task that needs to retrieve the values of multiple keys? Such a task also discourages the hallucinations that MultiDocQA & IN2 seem to give rise to.\n\n5) Regarding line 214, It is unclear why the finetuning and evaluation is fixed within one sliding window. Since some works can improve the maximum context length of Mistral (32k), evaluation at longer context lengths is needed. Section 3.5 is the right step towards this, but quality evaluations still need to be included.\n\n6) Although manipulation of positional embeddings to extend context lengths without training is orthogonal to this work, Works like [3, 4] improve the context lengths of Mistral while maintaining performance on retrieval-like benchmarks. Either an explanation in related work or a comparison with baselines is needed.\n\nOverall, I feel further studies and evaluations are required to improve the work. The work claims multiple interesting findings that require more comprehensive analysis and explorations. Happy to discuss more during rebuttal.\n\n[1] Is It Really Long Context if All You Need Is Retrieval? Towards Genuinely Difficult Long Context NLP. CoRR abs/2407.00402 (2024)\n\n[2] Michelangelo: Long Context Evaluations Beyond Haystacks via Latent Structure Queries. CoRR abs/2409.12640 (2024)\n\n[3] LongEmbed: Extending Embedding Models for Long Context Retrieval. CoRR abs/2404.12096 (2024)\n\n[4] LLM Maybe LongLM: SelfExtend LLM Context Window Without Tuning. ICML 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1q9IVK5irh", "reviewer": "Reviewer_Vib9", "summary": "The paper targets the challenge of LLM accurately retrieving information and maintaining reasoning capabilities when processing long-context inputs. To this end, the author designs a synthetic dataset specifically for fine-tuning recent popular pretrained models. An interesting advantage of this dataset is it doesn't contain any factual information. Their evaluations on long context retrival and reasoning tasks show clear improvements. They also compare with three additional long-context augmentation datasets as baselines. Results show that their synthetic dataset achieves comparable performance in long-context retrieval and reasoning without causing the significant degradation on general benchmarks observed with other baselines.", "questions": "(1). Is there any reason the dataset is built with 3 or 4 digits? I think 4 digits will count as 2 tokens in the GPT tokenizer, while 3 digits is just 1.\n(2) Can GSM8K be considered a reasoning task? I think Triviaqa can also be considered a retrieval and reasoning task. Why can't those datasets improve performance?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and easy to follow. \n\nThe problem addressed in this paper is pretty meaningful.\n\nThe experiments are comprehensive and effectively support the main points.\n\nThe idea is straightforward but produces good results.", "weaknesses": "I have some concerns about the experiments. \n\n(1) For stage 1, why does the author fine-tune Mistral 7B for 2 epochs but fine-tune GPT-3.5 Turbo for 3 epoches? Is the performance sensitive to the number of training epochs? \n\n(2) For stage 1, does it need any training strategy such as early stopping? Will a longer training hurt the performance on general benchmarks? Also, what will happen if we train with a larger key-value retrieval dataset? Will the performance on long context retrival and reasoning tasks further improve, and might it negatively affect general benchmark performance?\n\n(3) Will the performance further improve if also fine-tune Mistral 7B with Multi-subkey dictionary key-value retrieval dataset?\n\n(4) When fine-tuning with MDQA, does it train with 20 documents and place a gold document in some positions? If so, does this mean that the model is fine-tuned with too few samples? \n\n(5) The author presents fine-tuning with MDQA as a baseline in Section 3.2.1 but does not provide similar results for FLENQA in Section 3.2.2. Is there a specific reason for this? I’m curious whether the conclusions would remain consistent.\n\n(6). How many tokens are used for training for the other baselines in stage 4? I think it's a little bit hard to say if this method beats the other baseline because the baselines actually shows good performance on long context retrival and reasoning tasks especially on FlenQA(cot), the performance is also ok on some datasets in the general benchmarks. \n\nA minor suggestion: I think it would be helpful to include the average accuracy/gap across all datasets in Table 2, as it would clarify the overall average degradation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "d6YzKKLSYY", "reviewer": "Reviewer_19vD", "summary": "The paper proposes synthetic data generating for the tasks requiring  understanding of long context. Specifically, it generates a key-value retrieval task where an LLM is fine-tuned to find out a dictionary with the specified key value. To make the task more difficult,  the key becomes a tuple of integers and the order of integers are randomly shuffled. After fine-tuning the LLM, LLM is evaluated on MDQA and FlenQA dataset. The proposed synthetic data generation leads to performance improvement of  GPT-3.5-Turbo and Mistral.", "questions": "- Did you full fine-tune the Mistral or use parameter efficient fine-tuning (peft) such as lora?\n\n- What is the rationale for choosing the number of samples in key-value retrieval tasks?\n\n- In Figure 5b, Mistral-v0.1 is used, while Mistral-v0.2 is used in Figure 7b. Is there any reason why you use different version of Mistral?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed method is simple and effective.\n\n- It does not degrade the performance on other benchmark datasets such as MMLU, HellaSwag, GSM8K, Triviaqa, NQ-open.\n\n- The paper is well written.", "weaknesses": "- Details of how key-value retrieval task is generated are missing, which seems to be critical.\n\n- More extensive experiments need for validating the proposed method. For example, I highly recommend to evaluate the method on benchmark datasets such as RULER [1] and Long Bench [2] with different backbone LLMs such as Llama and Gemma.\n\n\n\n## References\n[1] Hsieh, Cheng-Ping, et al. \"RULER: What's the Real Context Size of Your Long-Context Language Models?.\" arXiv preprint arXiv:2404.06654 (2024).\n\n[2] Bai, Yushi, et al. \"Longbench: A bilingual, multitask benchmark for long context understanding.\" arXiv preprint arXiv:2308.14508 (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)