Quest: Query-centric Data Synthesis Approach for Long-context Scaling of Large Language Model
Authors: Chaochen Gao
Xing W
Qi Fu
Songlin Hu
Venue: ICLR 2025 Poster (2025)
Keywords: longcontext
pre-training
scaling
Primary area: foundation or frontier models, including LLMs

Abstract: Recent advancements in large language models (LLMs) have highlighted the importance of extending context lengths for handling complex tasks. While traditional methods for training on long contexts often use filtered long documents, these approaches lead to domain imbalances, limiting model performance. To address this, techniques like random document concatenation (Standard) and similarity-based methods (KNN, ICLM) have been developed. However, they either sacrifice semantic coherence or diversity. To balance both aspects, we introduce Quest, a query-centric data synthesis method aggregating semantically relevant yet diverse documents. Quest uses a generative model to predict potential queries for each document, grouping documents with similar queries and keywords. Extensive experiments demonstrate Quest's superior performance on long-context tasks, achieving remarkable results with context lengths of up to 1M tokens and confirming its scalability across various model sizes.

Published as a conference paper at ICLR 2025
QUEST: QUERY-CENTRIC DATA SYNTHESIS AP- PROACH FOR LONG-CONTEXT SCALING OF LARGE LANGUAGE MODEL
Chaochen Gao!?, Xing W Qi Fu', Songlin Hul
IIstitute of Information Engineering, Chinese Academy of Sciences School of Cyber Security, University of Chinese Academy of Sciences "Xiaohongshu Inc
[gaochaochen, wwwpeaaooceenwuxig..hhee www.ng @iie.ac.en fuqi @ xiaohongshu.com
ABSTRACT
Recent advancements in large language models (LLMs) have highlighted the im- portance extending context lengths for handling complex tasks. While tradi- tional methods for training on long contexts often filtered long documents, these approaches lead to domain imbalances, limiting model performance. To address this, techniques like random document concatenation (Standard) and similarity-basee methods (KNN, ICLM) have been developed. However, they either sacrifice semantic coherence or diversity. To balance both aspects, in- troduce Quest, query-centric data synthesis method aggregating semantically relevant yet diverse documents. Quest uses generative model to predict po- tential queries for each document, grouping documents with similar queries and keywords. Extensive experiments demonstrate Quest's superior performance on long-context tasks, achieving remarkable results with context lengths of up 1M tokens and confirming its scalability across various model sizes.
Quest-LLaMA-3-BB 1-M-1---------M Acc: 1.0
90.0
60.0
70.0
6
60.0
Score
20.0
10.0
25k
150k
275k
400k
525k
650k
775k
900x
1000k
Token Limit
Figure I: The Keedle-in-a-Haystane task evaluates a model's ability to retrieve specific information (the needle) from large collection of documents (the haystack). Following Long (Zhang al.) 2024a) and LWM (Liu et al., 2024). where the x-axis represents the document length and the y-axis indicates the position of the "needle" within the document, ranging from 25K IM tokens. To the best of our knowledge, Quest is the first base model (without instruction tuning) to achieve 100% accuracy with a 1M context length.
INTRODUCTION
Large Language Models (LLMs) are typically pre-trained using fixed context lengths. Recent ad- vancements, however, have highlighted the importance of extending the context lengths. For in- stance, the LLaMA series has progressively increased its context lengths from 2k tokens in LLaMA to 4k in ILaMA2 and 8k in LLaMA3 (ouuroo al. 2023a b Meta 2024a). LLMSs equipped with

---

Published as a conference paper at ICLR 2025
Longbench Results and Similarity
LongbookOA Results and Similarity
0.8
0.8
Model
128 Model
0.90 Model
Cosine-siin si---si (Nom-perfoo (Nor--peformance
performancee
1.40 Model
1.40 Model
0.7
32k CosSin
0.7
128k CosSim
6.6588
Cosine-sim (Norm-performmal
0.
0.6097
0.6
0.6
0.5
0.4965
0.5
0.4054
0.4
0.3536
0.3445
0.3
0.3
Standard
Quest
ICLM
KNN
Standard
Quest
ICLM
KNN
Method
Method
Figure 2: The cosine similarity of aggregated documents and the corresponding performance. The dotted lines indicate the performance of the models, with all results normalized to align within the specified similarity range. High similarity means the semantic correlation is strong, and low similarity indicates good context diversity. Quest balances the semantic correlation and context diversity, resulting in the best performance.
longer context lengths excel handling complex tasks (Caciularu al., 2023 Bairi et al. 2023 Mazumder & Liu 2022). especially those with long input, e.g., document summarization. The goal of long-context modeling is to improve a model's s ability to capture long-range dependencies by training on extended contexts. To handle long contexts-suchs
128k tokens common approach is to continue training LLMs on long- long-context data (Roziere et al. 2023 Xiong et al. 2023 Fu et al.) 2024 Previous works (Xiong et al. 2023 Fu et 2024) select such data by filtering long docu- ments from the training set that fit the target context length. However, those documents often come from few specific domains like Books3 Arxiv, leading skewed distribution, which impacts model performance after continued training (Jung & van der Plas 2024 Cai et al., 2023).
Previous studies synthesise long-context data by concatenating shorter documents achieve bal- anced domain distribution. Those methods can be classified into two categories: the Standard method (Roziere et al., 2023 Ouyang et al. [202]] Le Scao et al., 2023 Touvron et al. 2023a) and similarity-based methods like KNN (Guu et al. Levine et al. (022]) and ICLM (Shi et al, 2023). The Standard method randomly concatenates short documents to meet specified target length. While this approach ensures diversity within the context, the weak semantic relationship between concatenated documents hinders learning long-range dependencies. In contrast, similarity- based methods aggregate semantically similar documents, e.g., by concatenating document with its top k most similar counterparts from the corpus. However, similarity- -based methods overempha- size semantic correlation. They are prone to falling into narrow context (high redundancy) due to concatenating similar or even repeated documents. Figure 2 compares semantic correlation lev- els of different methods with their performance prioritizing context diversity at the expense of semantic correlation (Standard) or overemphasiz- ing semantic correlation while sacrificing context diversity (KNN and ICLM) leads to suboptimal performance. Therefore, both context diversity and semantic correlation are crucial for effectively modeling long texts, highlighting the need for a method to balance both aspects.
long-context tasks. The results show that either
To achieve the balance effectively, this paper proposes Quest, a query-centric data synthesis ap- proach that simultaneously ensures semantic correlation and context diversity within the long- context data. Our inspiration stems from the observation that similar queries can aggregate se- mantic relevant but low-redundancy documents via search engines (Malli et al., 2021 Babenko & Lempitsky 2014; Kaushik et al.) 2004). A straightforward way is to collect enormous queries and cluster articles related to each query, mimicking the functionality of search engines. However, collecting massive queries is very time-consuming, and it is hard to guarantee diversity. To over- come this, we predict potential queries using a generative model for each document in the training dataset. Specifically, Quest begins by employing a lightweight query-prediction model (Raffel et al. 2020 Noguuera et al.) 2019 Ww et al.) (2022) to predict varied potential queries for each document. Documents sharing the same query grouped relevant, simulating an inverse search process.
2

---

Published as a conference paper at ICLR 2025
Data Synthesis
doc1
How does the apple taste?
Keyword Extraction
What the origin of the sofa?
apple
doc1
doc3
Query Generation
doc3
How long is the shell life an apple?
&
sofa
doc2
doc4
Index Building
What the materials used for sofa?
Training Process
QDS (Topic-related)
Standard (Random)
Iteratio1 [apple related docs) doc1
doc3
Iteration1 (random sample)
doc2
doc3
Iteration2 (sofa related docs):
doc2
doc4
Iteration2 (random sample)
doc1
doc4
Overview of Query-centric data synthesis (Quest) method. Unlike the Standard pre- training strategy that randomly shuffled documents in the input context, Quest places relevant docu- ments in the same context.
Figure 3:
To cluster similar queries, we extract more coarse-grained and high-level keywords from queries. Thus, the same keywords further index documents associated with similar queries. Finally, Quest randomly samples from documents indexed by the same keywords and concatenates the selected documents build long-context data.
Through extensive experiments, we show that Quest significantly outperforms other data synthesis methods multiple long-context benchmarks with context lengths ranging from 32k 128k. Fig- ssows that applying the Quest method to 1M context length achieves impressive performance on the widely used Needle-in-a-Haysttel task. Additionally, we further investigate the scaling laws of synthesized long-com long-context data across various model scales and confirm the predictability of the Quest method, making a reliable solution for advancing long-context models.
Our contributions are summarized as follows:
1. We propose a query-centric data synthesis method to alleviate long-context data scarcity and uneven domain distribution.
2. Extensive experiments 32k and 128k context lengths show that method outperforms existing approaches.
3. We provide a detailed analysis of key design elements in Quest and offer valuable insights for future research in the long-context area.
4. We investigate the scaling law of synthesized long-contex. data and confirm the predictability of our method.
RELATED WORK
Long-Context Language Models The success of LLMs has sparked interest in enabling them to process longer texts. Some works adapt methods for longer texts without additional training by modifying position encoding. For example, Han et al. (2023) and Xiao et (2023) adjust the attention matrix to generate long contexts, while Jin et al. (202)) compresses position encoding into the pre-trained position range. Other works involve continued training for better performance. Xiong et al. (2023) demonstrates that long-co long-contett capabilities can be acquired by continually pre- training from short-context models. Chen et al. distribution of position encoding, and Yen et al. (2024a) proposes context expansion with parallel encoding. Advances have also been made using RoPE (Su et al. 2021), enabling LLMs to handle longer positions. PoSE (Zhu et al. 2023) employs skip-win position indices, allowing position encoding to adapt to different lengths. However, those approaches often overlook the scarcity and uneven distribution of long text data during continued training, relying filtering long documents
(002bb) uses position interpolation to change the

---

Published as a conference paper at ICLR 2025
from existing corpora (Xiong al. 2023 Fu et al. 2024 Liu et al. 2024) randomly splicing
short documents to reach a longer length (Roziere et al., 2023; Chen et al. 2023c Tworkowski
et al. 2024 Chen et al. 2023a Li al. 2023).
Acquiring effective long-comm context data for training challenging. Some previous serrivaal-augmenee pre-training works (Guu al., 2020) Levine et al.) 2021) can synthesize long-contex data. Guu et al. (2020) clusters semantically similar texts within same context window, while Levine al. (2021) shows that incorporating seman- tically related but non-adjacent sentences within the same pre-training example enhances sentence representations. Shi et al. (2023) uses a traveling salesman algorithm to address document redun- dancy in the KNN method.
Data Synthesis and Augmentation for Long-Context
Scaling Laws For a broad spectrum of factors x, scaling laws (Kaplan et al. 2020 Henighan al., 2020; Hoffmann et al. 2022) indicate that their impact on the loss L of pre-trained model follows a power law relationship. Here, x may represent model sizes, quantities of training data, or training steps, with parameters be determined. Previous research (Alabdulmohsin et al. 2022 OpenAl 2023 [Bi et al. 2024, Su et al., 2024) Xiong et al. (024) highlights the impressive pre- dictive power scaling laws. Notably, fitting this relationship to a set of smaller models, training datasets, or computational resources enables precise extrapolation to predict the test loss for much larger cases across several orders magnitude. This capability allows practitioners to estimate performance of pre-trained larger language model without incurring the substantial cost of com- pleting extensive training runs. However, the scaling law for synthesized long-context data remains unexplored despite its importance for long-context modeling. Therefore, using Quest method, we investigate the scaling laws of synthesized long-context data across various model sizes and confirm the predictability of the Quest method.
METHOD
This section details our proposed Query-centric data synthesis (Quest) method. Algorithm presents
the method for synthesizing long-context data. Given dataset with diverse documents D {d; }
our goal is to effectively aggregate relevant but low-redundancy documents for synthesizing training texts with context length of L. An overview of our approach is illustrated in Figure 3 Quest
mainly includes steps.
Algorithm Query-centric Data Synthesis (Quest) Method
Require: Dataset D {di}. Context length L, Split ratio Ensure: Training texts with context length
1: Initialize lists and K
2: for each d, E do
3: Q- U doc2query(d:)
4: end for
5: for each q, E Q do
6:
K; + {k E Rake(q;) A score(k) 2 3.0}
7: K K+ K (random(K,,) K;)
8: end for
E D} 10: Sort / by size and split: I, (i T rank(i) Srx [[], I = I
9: It [(k,,d,)
for each training step do
12: Sample 1 E I, Ii (oversample I2)
13: T concat mancat(sampl((yy, [T] 2
14: Train with T
15:
end for
1. Query Prediction: We utilize the open-source doc2query model (Nogueira et al. 2019) to predict queries } for each document {d,). We segment texts that exceed the context length limit of the doc2query model into parts and generate a query for each segment. Conse- quently, for document (d,). ) list of queries Qi = (g] ........00 q,' is predicted. Appendix C.2

---

Published as a conference paper at ICLR 2025
demonstrates that higher query quality leads to better model performance, while Appendix C.3 shows that query prediction can effectively improve the quality of keywords.
2. Keyword Extraction: We extract keywords from each query {q:} with an efficient tool, Rake] For texts with multiple queries, Rake generates several lists of keywords K; (R. (k....,,k,)) }. To ensure the quality of extracted keywords, we adopt two filtering strategies. First, we filter out keywords with Rake score below 3.0. Second, we remove frequent but non-informative keywords such as "following sentence" or "best way" (see Appendix AA.3 for details). Then, randomly select one of the remaining keywords to serve as the representa- tive keyword for the document. Appendix C.4 presents ablation studies on different methods for selecting keywords, demonstrating that randomly selecting keywords improve both keyword diversity and model performance.
3. Building Keyword-based Inverted Index: We then build a keyword-basee based inverted index after we map each document to its representative keyword. Documents with identical representative keyword are indexed together and treated as topically similar ones.
4. Indexes Split: We found that the number of documents associated with different keywords varies significantly. To address that imbalance and achieve a more balanced data distribution, we implement oversampling for documents with less frequent keywords. Specifically, we rank the keywords in ascending order based on the number of documents indexed by each keyword and divide the sorted keywords into two sets. The top-ranked split ratio% of the keywords are denoted as the short-index set T., while the remainder is denoted as the long-index set Tj. Appendix A.4 shows the impact of split_ratio%.
5. Training Process: We perform sampling without replacement from the documents within sampled keyword and concatenate the selected documents to the target context length L for training. We oversample the short-index set ensure that the number of tokens participating in training is evenly distributed between the keywords in both I, and It. Oversampling re- distributes the sampling probability to prioritize I,, ensuring it gets larger share of the total samples. We provide the mathematical formula for oversampling: Assuming I. contains training samples, It contains ny training samples, p is oversampling probability, and the total number of training samples to be used is N. The number of samples drawn from Ia can be calculated as:
n3 + (.-(
1
+D)
NN
(1)
The number of samples drawn from I can be calculated as:
I = N
(2)
EXPERIMENTS
In this section, first introduce the experimental settings (Section 4.1) Then we provide detailed
description of our baseline methods (Section 4.2 and the experimental results (Section 4.3 and Section
4.1 EXPERIMENTAL SETUP
We conduct continued training on Pythia (Biderman et al.) 2023) models of different scales, specifi- cally 1.4B, 6.9B, and I2B. Pythia a series of models trained on the Pile (Gao et al. (020)) dataset, explicitly designed for research. Experiments conducted with Pythia offer good reproducibilityy We use identical training sets for all methods to ensure strictly fair comparison. The only variation between the methods is how to rearrange documents into long-context data.
We apply the Quest method on Pythia's pre-training data, i.c., the Pile dataset, which does not lead to domain transfer issues. Specifically, we extract 30B tokens of keyword-indexed documents from the 300B tokens of the original Pile dataset. Documents indexed by an identical keyword are randomly concatenated to form long-context data that reach the training context length.
5

---

Published as a conference paper at ICLR 2025
Table 1: Comparison of Longhench results across methods. "Avg." represents the average
multiple test sets, The proposed Quest consistently outperforms baseline methods across various
model sizes in the 32K context length setting. Detailed results can be found in Appendix [
Train& Test
Model size
Method
Avg.
Sgl.
Multi.
Sum.
Few.
Syn.
Code.
Standard
20.94
19.24
17.46
20.65
26.75
2.04
36.41
KNN
19.97
17.26
13.01
22.97
24.16
2.33
39.22
32k
1.4B
ICLM
19.82
20.01
14.71
21.95
23.09
1.94
35.31
Quest
22.06
17.97
17.98
21.91
28.06
2.33
42,25
Standard
22.48
18.07
16.83
22.33
30.23
3.86
40.91
KNN
21.65
18.5
13.64
22.56
28.18
3.76
41.88
32k
6.9B
ICLM
20.86
17.82
15.34
22.35
25,86
1.21
41.15
Quest
23.23
19,21
14.13
22.45
30.14
2.96
50.55
Standard
24.85
22.18
21.94
22.30
32.05
3.78
43.73
KNN
22.95
20.55
20.48
23.51
29.19
2.47
37.44
32k
12B
ICLM
24.07
22.67
23.29
23.41
30.99
1.5
37.09
Quest
25.24
22.34
21.08
23.74
31.91
3.22
46.8
Table 2: Comparison of Longbook QA results across methods. The proposed Quest outperforms baseline methods across various model sizes in the 128 context length setting.
Train& Test
Model size
Method
Longbook QA
Standard
9.94
KNN
10.36
128k
1.4B
ICLM
10.70
Quest
11.30
Standard
14,47
KNN
13.38
128k
6.9B
ICLM
14.92
Quest
17.95
Standard
17.81
KNN
16.42
128k
12B
ICLM
18.44
Quest
18.92
During training, we use the open-source framework GPT-NeoX with batch size of 4M tokens for all settings. The AdamW optimizer (Loshchilov Hutter 2017) with B1 0.9 and 3g 0.95 parameters and cosine learning rate schedule is employed. To optimize memory and performance, We use Flash Attention2 (Dao) 2023) and ZeRO (Rajbhandari al. (020). The learning rates are 5e-5 for the 1.4B model, 4e for the 6.9B model, and 2e for the T2B model. For more details, refer Appendix
4.2 BASELINES METHODS
We compare the proposed Quest method with the existing remarkable data synthesis methods:
1. Standard Method shuffles and concatenates documents randomly in the input context and
has been the mainstream practice in pre-training (Ouyang al. 2022 Le Scao et al. 2023 Touvron et al. 2023a).
2. KNN (Retrieval-augnee augmented Language Model Pre-training) (Guu et al. 2020 Levine al.) 202T) places each document along with the top most similar retrieved documents in the same input context.
3. ICLM (Shi et al., 2023) Method is recently proposed method that utilizes traveling sales- man algorithm to alleviate the document redundancy problem in the KNN method by ranking similarities and determining the optimal training path.
https://githhb..ee

---

Published as a conference paper at ICLR 2025
Table 3: Comparison of short text performance across methods. Overall, Quest shows almost
degradation in short-text performance on average.
Model
Avg
Win
PIQA
LogiQA
LAMBADA
Hella
ARC-E
ARC-C
Pythia 0.4830
0.5746
0.7095
0.2120
0.6163
0.4042
0.6048
0.2594
+ Standard
0,4802
0.5675
0.6975
0.2227
0.6507
0.3943
0.5821
0.2466
+ KNN
0.4769
0.5651
0.7089
0.2028
0.6480
0.3946
0.5737
0.2449
ICI.M
0.4816
0.5785
0.7024
0.2120
0.6546
0.3941
0.5753
0.2543
Quest
0.4831
0.5691
0.7024
0.2304
0.6472
0.3961
0.5770
0,2594
To implement KNN, we utilize a product quantized inverted file (IVFPQ) FAISS index with code size of 32. For ICLM, we follow the GitHub repository to synthesize long-context data.
4.3 EVALUATION RESULTS
We evaluate four methods, including Quest and three bascline methods, with evaluation lengths ranging from 32k to 128k. To comprehensively compare Quest with baseline methods, the datasets from different evaluation tasks are divided into categories: long-text benchmark and short-text benchmark.
1. Long-text Benchmark: For 32k context length, we adopt the widely-used Longbench (Bai et al. 2023) Benchmark, testing six task types: Single-document QA (Sgl.). Multi-document QA (Multi.), Summarization (Sum.), Few-shot learning (Few.), Synthetic (Syn.). and Code completion (Code.), 17 datasets. For 128k context length, following Fu al. (2024), focus on the widely-used Longbook QA task (Zhang et al. 2024b), on which the pre-trained models perform reasonably well without instruction tuning.
2. Short-text Benchmark: To assess the effectiveness maintainence of long-text models
short-text tasks, we select widely-used short-text datasets: WinoGrande (Sakaguchi et al. 2021), PIQA (Bisk et al. 2020). Logiqa (Liu et al.; 2020), Lambada (OpenAl) (Paperno al. 2016), HellaSwag (Zellers et al. 2019). ARC-Easy, and ARC-Challenge (Clark et al. 2018).
Quest demonstrates better performance on average. Table compares showing that Quest outperforms other methods on average, especially on the Code dataset. We attribute it Levenshtein distance, more stringent and discriminative metric than score or ROUGE, as it accounts for character sequence. To further assess the efficacy of the Quest method in extended long-context settings, we extend the context length to 128k and evaluate the trained models on the Longbook QA task. Table 2 shows that Quest consistently outperforms other methods across various model sizes.
comparess the Longbench results,
The Standard method performs well in 32k tasks but shows the poorest performance the 128k task. Table T shows that the Standard method demonstrates good performance in 32k tasks and superior KNN and ICLM, especially in few-shot tasks. However, in Table data synthesis methods outperform the Standard method, proving the importance of long-context modeling. This improvement can be attributed the abundance of 32k-length documents, whereas 128k-length documents have a far lower count and uneven domain distributions.
Quest retains good performance on short text. To verify how well Quest maintains model per- formance on short text tasks, we evaluate it on 7 commonly reported tasks, as shown in Table Compared with the base model, the model trained with Quest long-context data shows almost no degradation in short-text performance average.
4.4 APPLYING QUEST ON THE SFATE-OF-THF-AA MODEL..
To further verify the effectiveness of Quest, we experiment with the current state-of-the (SOTA)
open-source model, i.e., LLaMA3 (Meta 2024a). Following Fu al. (2024), evaluate the Quest-
LLaMA3-8B model on the widely used N Needle-in- task and Longbook QA task. There

---

Published as a conference paper at ICLR 2025
Table 4: Comparisons with the state-of-the-art long-context pre-trained models the Longbook
QA task. Quest-LLaMA-3-8B- --------- achieves the best performance among open-source models,
surpassing -LaMA-3.1-8--D- and only inferior to the remarkable GPT-4-Turbo- 1 O results
from Fu et al. (2024) results evaluated by us.
Method
Model size
Test Len
Longbook QA
GPT-4-Turbo-128N1
37.40
LLaMA-3-8B (Meta 2024a
8B
13.87
LongLoRA Chen
202
2023ch
7B
24.30
LongLoRA Chen al.
2023c
13B
24.60
YaRN Mistral PPng cf al.
2023
7B
26.30
Yi-9B-200K (AT al.
21124
9B
30,35
128k
I.AMAA2-7B--UUK (Fu et al.
2024 >
7B
27.40
DeepSeek-V2-Line VV2LLite (DeepSeck-A) 2024
16B
21.56
Qwen2.5-77 [Team] 2024
16.57
LLaMA-3.1-8B--e (Dubey et al.
2024
8B
30.11
Quest-LLaMA-3-8B- (28k(ours)
8B
32.39
Table 5: Performance comparison of using only the existing long documents the pretraining corpus versus -uest-synthesized long-context data. "Avg." represents the average over multiple test sets. Quest-synthesized long-come nongcconextt data consistently outperforms the existing long documents on Longbench.
Method
Avg.
Sgl.
Multi.
Sum.
Few.
Syn.
Code.
Long document
21.11
19.77
15.15
22.11
24.81
2.46
41.84
Quest
22.06
17.97
17.98
21.91
28.06
2.33
42.25
approaches evaluate the Needle-in-a-Haystack task: retrieving text sentence (Fu al.
2024) for the 128k context length and retrieving numeric string (Zhang et al.) 2024a Liu et al. 2024 for the IM context length. As shown in Figure 4 our Quest-LLaMA3-81 achieves 97% accuracy on the Needle-in-at Lasell-in-a-Hayyee Haystackk task (retrieving a text sentence), significantly surpassing the previous highest accuracy of 88% (Fu et al. 2024). In addition, we increased the length 1M, and Figure [ggure[shhve shows that Quest achieved 100% accuracy within 1M context length (retrieving numeric string). Table shows that Quest-LLaMA achieves the highest score among open-source models, further narrowing the gap with GPT-4 Turbo under 1 128K length setting. These results demonstrate that Quest has strong robustness and scalability when dealing with ultra-long context data.
5 ANALYSIS
This section provides an in-depth analysis of Quest. Due to the high computational cost of LLM experiments, our ablation studies are conducted with 32k context length and 1.4B model size, unless otherwise noted.
5.1 QUEST'S ADVANTAGE GRADUALLY EXPANDS WITH TRAINING PROGRESS
This section studies the performance trends of the training progress using data synthesized by the Quest method. As shown in Figure 7 on the Longbench benchmark, the Quest method consistently outperforms other data synthesis methods during the whole training process in terms of superior performance. Additionally, the training progress using the Quest method saturates significantly later. In contrast, other methods generally reach performance saturation within the first 40% of the training progress. Such distinct advantages further demonstrate that Quest's long-context dataset is more diverse and beneficial for long-context training.
5.2 QUEST BALANCES DOCUMENT SIMILARITY FOR SUPERIOR PERFORMANCE
To explore how document similarity within the same context influences performance, randomly sample contexts derived from different methods and calculate the similarity among documents grouped within each context. As shown in Figure 2 model performance initially increases and
8

---

Published as a conference paper at ICLR 2025
DespSeek-V2 Lit- Acc: 0.25
Yi-9B-200K Acc: 0.32
Qwen2.5-17 5-B7 Acc: 0.54
09.0
50.3
44.
22
Dopth Porsont
11.
LLaMA-3-8B 10--A-3-B-BBss--NT-- Acc: 0.18
Meta- Mott--LLaMM33.1-80 Acc: 0.75
Quest-LLend Mate-LLLMA----B---- Aoc:
Score
29.0
78.0
57.0
44.0
250
21.0
6.0
17%
432
1200
175 051 114: 1256 432 1254
Figure 4: Performance comparison on the Needle-in-a-HHasttac task for collection base models (without instruction tuning). 3osst-LLMM----B--BBB exhibits strong performance, significantly outperforming other open-source models of similar or larger sizes. Unlike Figure D task difficulty increased within the 128k context length by relrieving a sentence ruther than random numeric string. "Acc" denotes the percentage of model responses rated as fully accurate (scoring 10 GPT- 4's evaluation) out of all responses generated.
Document Distribution in 32k Context
Document Distribution in 128k Context
KNN
Standard
4
7.5
5.0
Quest
0.0
Standaro
-5.0
-4
-7.5
ICLM
Figure 5: t-SNE visualization of aggregated documents from different methods. The proposed Quest maintains balanced distribution across varying context lengths. See Appendix D.2 for more exam- ples.
then decreases as similarity rises. It indicates that highly similar or irrelevant document aggrega- tions can lead to performance degradation in long-contex, modeling. When document similarity excessively high, different documents convey identical information, leading the LLM to repeat previously encountered text for predictions. Conversely, if documents are entirely unrelated, it hin- ders the model's 's long-range pattern recognition. Quest achieved the best performance by balancing semantic correlation and context diversity.
We use t-SNE (van der Maaten Hinton 2008) to visualize the aggregated documents within the same context. Figure 5 (left) shows that the Standard method forms the most dispersed clusters due to random aggregation, while KNN and ICLM yield tighter clusters in the 32k-context-length setting. The proposed Quest exhibits moderate clustering, ensuring semantic consistency and text diversity within the synthesizing long context. Figure55 illustrates document distribution for various methods under the 128k context length. Quest continues aggregate relevant, low- redundancy documents.
5.3 QUEST- QUFSTT-YYNTHESIZE LONG DOCUMENTS OUTPERFORM EXISTING ONES
Some long documents have already reached the target context length in the pre-training corpus. Therefore, compare the performance of using Quest-synthesized long-context data with using

---

Published as a conference paper at ICLR 2025
Change Training
Performance Change with Training
3.4
22
awo 5w144,
Volidetion (324)
3.2
mocel sow
20
aud B.1.12, y-2.43
teg Performence
3.0
awQ.12, f-1 y=2.30
128 model sioe
1.6
2.8
14
Standard
2.6
KNN
2.4
10
Quest
0
6
0
20
60
80
100
Training Tokens (Billions)
Training Progress (%)
Figure 6: Scaling law of synthesized long- context data under different model sizes. The Quest approach scales well and predictably.
Figure 7: Performance trends during the train- ing progress using data synthesis methods. The Quest consistently outperforms other methods.
only the existing long documents in the pre-training corpus for training. The amount of data in both the synthetic data and the long documents is identical. Table shows that using Quest-synthesized long-context data achieves better results on Longbench. Existing long documents perform worse because they only exist in a few domains, resulting in skewed data distribution, as illustrated in Appendix C.5 The Quest method, on the other hand, can cover nearly all domains, resulting in more diverse synthesized long-context data and better performance in evaluation tasks. We also attempt further comparisons with a context length of 128k. However, long documents exceeding 128k the Pile dataset are rare and inadequate to support a fair comparison experiment. As the target context length increases, the scarcity problem becomes more pronounced, further highlighting the necessity of developing effective long-context synthesis methods.
SCALING LAW OF SYNTHESIZED LONG-CONTEXS DATA
To explore the scaling law of synthesized long-context data, we vary the amount of training data for different model sizes (1.4B, 6.9B, and 12B) under the 32k context length setting. Formally, we formulate the scaling law of the validation loss by studying different model sizes N and dataset sizes D:
L(D) a cxex(-33) +
This formula applics to each model size, where fa, 3, are variables to be learned. our ex- periments, each model is trained separately datasets of different sizes: 250 million, 500 million, billion, billion, and billion tokens. Then, we fit a curve for each model size, showing the relationship between the data scaling and the validation loss the end of each training, shown in Figure 6
We validated the learned scaling law on billion data size by comparing the relative error between each model's final validation loss and its predicted value. The relative errors were 0.5% for the 1.4B model, -0.5% for the 6.9B model, and 0.4% for the 12B model, demonstrating the scalability and accuracy of Quest's 's data synthesis approach with minimal deviation.
7 CONCLUSION
In this paper, we introduce Quest, a novel method for synthesizing balanced long-context data by grouping and concatenating relevant but low-redundant documents associated with similar queries. The proposed Quest ensures semantic correlation and context diversity in long-context data to im- prove the long-context modeling capability of pre-trained models. Extensive experiments demon- strate that Quest outperforms existing approaches across various long-contex: and short-context benchmarks, proving be an effective and reliable solution for advancing long-context models.
ACKNOWLEDGEMENT
This work was supported by the National Natural Science Foundation of China (No. U24A20335).
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Iige Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint (eXXxx2333308770 2023.
01. AI, :, Alex Young, Bei Chen, Chao Li, Chengen Huang. Ge Zhang, Guanwei Zhang, Heng Li, Jiangcheng Zhu, Jianqun Chen, Jing Chang. Kaidong Yu, Peng Liu, Qiang Liu, Shawn Yue, Senbin Yang, Shiming Yang. Tao Yu. Wen Xie, Wenhao Huang, Xiaohui Hu, Xiaoyi Ren, Xinyao Niu, Pengcheng Nie, Yuchi Xu, Yudong Liu, Yue Wang. Yuxuan Cai, Zhenyu Gu, Zhiyuan Liu, and Zonghong Dai. Yi: Open foundation models by Ol. 2024.
Ibrahim M Alabdulmobsin, Behnam Neyshabur, and Xiaohua Zhai. Revisiting neural scaling laws in language and vision. Advances in Neural Information Processing Systems, 35:22300-2233 2022.
Artem Babenko and Victor Lempitsky. The inverted multi-index. IEEE transactions pattern analysis and machine intelligence, 377)):247- 2014.
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.
Ramakrishna Bairi, Athary Sonwane, Aditya Kanade, VageeshD C, Arun Iyer, Suresh Parthasarathy, Sriram Rajamani, B. Ashok, and Shashank Shet. Codeplan: Repository-level coding using Ilms and planning. Sep 2023.
Yoshua Bengio and Yann LeCun. Scaling learning algorithms towards AI. In Large Scale Kernel Machines. MIT Press, 2007.
Xiao Bi, Deli Chen, Guanting Chen, Shanhuang Chen, Damai Dai, Chengqi Deng, Honghui Ding, Kai Dong, Qiushi Du, Zhe Fu, et al. Deepseek Ilm: Scaling open-source language models with longtermism. arXiv preprint MxXiv:2401.02954 2024.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O' O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical com- monsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Avi Caciularu, MatthewF. Peters, Jacob Goldberger, Ido Dagan, and Arman Cohan. Peek across: Improving multi-document modeling via cross-document question-answering May 2023.
Xunxin Cai, Meng Xiao, Zhiyuan Ning, and Yuanchun Zhou. Resolving the imbalance issue in hier- archical disciplinary topic inference via Ilm-based data augmentation. In 2023 IEEE International Conference on Data Mining Workshops (ICDMW), pp. 1424-1429. IEEE, 2023.
Guanzheng Chen, Xin Li, Zaiqiao Meng, Shangsong Liang, and Lidong Bing. Clex: Continuous length extrapolation for large language models. arXiv preprint arXiv:2310. 10..20000,664,, 2023a.
Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2306. 15595, 2023b.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-contex: large language models. arXiv preprint arXiv:2309, 12307, 2023c.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. 2018
arXiv preprint arXiv:

---

Published as a conference paper at ICLR 2025
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08661, 2023.
DeepSeek-AI. Deepseck-v2: A strong, economical, and efficient mixture-of-expprts language
model, 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, al. The llama 3 herd of models. arXiv preprint arXiv:2407.2 1783, 2024.
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models 128k context. arXiv preprint arXiv:2402.1017.. 2024.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, al. The pile: An 800gb dataset of diverse text for language modeling, arXiv preprint arXiv:2101.00027 2020.
Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1. MIT Press, 2016.
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training. In International conference on machine learning, pp. 3929-3938. PMLR, 2020.
Chi Han, Qifan Wang, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Simple on-the-fly length generalization for large language models. arXiv preprint arXiv:2308. 16137, 2023.
Tom Henighan, Jared Kaplan, Mor Katz, Mark Chen, Christopher Hesse, Jacob Jackson, Heewoo Jun, Tom B Brown, Prafulla Dhariwal, Scott Gray, et al. Scaling laws for autoregressive generative modeling. arXiv preprint (tt::2000.14001.1 2020.
Geoffrey E. Hinton, Simon Osindero, and Yee whye Teh. A fast learning algorithm for deep belief nets. Neural Computation, 18:1527-1554, 2006.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia- Yuan Chang, Huiyuan Chen, and Xia Hu. Lim maybe longlm: Self-extend I1m context window without tuning, 2024.
Vincent Jung and Lonneke van der Plas. Understanding the effects of language-specific class imbal- ance multilingual fine-tuning. arXiv preprint arXiv:2402. / 13016, 2024.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arxiv preprint rxXi:2001108861, 2020.
Raghav Kaushik, Rajasekar Krishnamurthy, Jeffrey F Naughton, and Raghu Ramakrishnan. On the integration of structure indexes and inverted lists. In Proceedings of the 2004 ACM SIGMOD international conference on Management of data, pp. 779-790, 2004.
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilié, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, et al. Bloom: A 176b- parameter open-access multilingual language model. 2023.
Yoav Levine, Noam Wies, Daniel Jannai, Dan Navon, Yedid Hoshen, and Amnon Shashua. The inductive bias of in-context learning: Rethinking pretraining example design. arXiv preprint arXiv:2/ 10.04541, 2021.
Shanda Li, Chong You, Guru Guruganesh, Joshua Ainslie, Santiago Ontanon, Manzil Zaheer, Sumit Sanghai, Yiming Yang, Sanjiv Kumar, and Srinadh Bhojanapalli. Functional interpolation for relative positions improves long context transformers. arXiv preprint arXiv:2310.044118 2023.
12

---

Published as a conference paper at ICLR 2025
Hao Liu, Wilson Yan, Matei Zaharia, and Pieter Abbeel. World model on million-length video and language with ringattention. arXiv preprint arXiv:2402.08226, 2024.
Jian Liu, Leyang Cui, Hanmeng Liu, Dandan Huang, Yile Wang, and Yue Zhang. Logiqa: A challenge dataset for machine reading comprehension with logical reasoning. arXiv preprint arXiv:2007.08124, 2020.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 17X1v:/711.05/01/ 2017.
Antonio Mallia, Omar Khattab, Torsten Suel, and Nicola Tonellotto. Learning passage impacts for inverted indexes. In Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 1723-1727, 2021.
Sahisnu Mazumder and Bing Liu. Lifelong and continual learning dialogue systems. Nov 2022.
Meta. Introducing meta llama 3: The most capable openly available 11m to date, 2024a. URL https://llama meta. com hammiilma3//
AI Meta. Llama 3.2: Revolutionizing edge and vision with open, customizable models. Meta AI, 2024b.
Rodrigo Nogueira, Wei Yang. Jimmy Lin, and Kyunghyun Cho. Document expansion by query prediction. arXiv preprint arXiv:1904 08377,,,,, 2019.
OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to fol- low instructions with human feedback. Advances in neural information processing systems, 35: 27730-27744, 2022.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring broad discourse context. arXiv preprint
2016.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint arXiv:2309.0000 2023.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimizations toward training trillion parameter models. In SC20: International Conference for High Perfor- mance Computing, Networking. Storage and Analysis, pp. 1-16. IEEE, 2020.
Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi. Jingyu Liu, Tal Remez, Jérémy Rapin, al. Code llama: Open foundation models for code. arXiv preprint arXiv:2308. 12950, 2023.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generaliza- tion. arXiv preprint Xrviv: 1911118731, 2019.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver- sarial winograd schema challenge at scale. Communications of the ACM, (49))99-106, 2021.
Weijia Shi, Sewon Min, Maria Lomeli, Chunting Zhou, Margaret Li, Victoria Lin, Noah A Smith, Luke Zettlemoyer, Scott Yih, and Mike Lewis. In-context pretraining: Language modeling be- yond document boundaries. arXiv preprint arXiv:2310. 10638, 2023.
Daria Soboleva, Faisal AI- Khateeb, Robert Myers, Jacob R Steeves. Joel Hestness, and Nolan Dey. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama, 2023. URL https: 7/hugg ng a ce. co wtt.w//ttt.../eeetee http///aa.attteeeeerrrssss
imP ama 627B
13

---

Published as a conference paper at ICLR 2025
Hui Su, Zhi Tian, Xiaoyu Shen, and Xunliang Cai. Unraveling the mystery of scaling laws: Part arXiv preprint arXiv:2403.0656. 2024.
Jianlin Su. Rectified rotary position embeddings. https//gitib..om//bjjooceeee 2023.
Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer
with rotary position embedding. 2021.
Qwen Team. Qwen2.5: A party of foundation models, September 2024. Utthttss:///wwnlln
github. wt/blo/g///aaaee htwpll/l/ggww..////
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint 2itiii:330..1977./ 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.0922 092888 2023b.
Szymon Tworkowski, Konrad Staniszewski, Mikotaj Pacek, Yuhuai Wu, Henryk Michalewski, and Piotr Milos. Focused transformer: Contrastive training for context scaling. Advances in Neural Information Processing Systems, 36, 2024.
Laurens van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of Ma- chine Learning Research, ((86):2579-266 2008. URL http: //mp::://mmm.../// fmlr. wwtpp/pprpppppppppppppp
wandermaaten html
Lei Wang, Chen Ma, Xueyang Feng, Zeyu Zhang, Hao Yang, Jingsen Zhang, Zhiyuan Chen, Jiakai Tang, Xu Chen, Yankai Lin, et al. A survey large language model based autonomous agents. Frontiers of Computer Science, 18(6):1-26, 2024.
Xing Wu, Guangyuan Ma. Wanhui Qian, Zijia Lin, and Songlin Hu. Query-as-context pre-training for dense passage retrieval. arXiv preprint axXiv:22/2.09598, 2022.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming
language models with attention sinks. arXiv preprint 2XXiv:330.1745:: 2023.
Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy S Liang, Quoe V Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. Advances in Neural Information Processing Systems, 36, 2024.
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint 23X.:23399.10339 2023.
Yizhe Xiong, Xiansheng Chen, Xin Ye, Hui Chen, Zijia Lin, Haoran Lian, Jianwei Niu, and Guiguang Ding. Temporal scaling law for large language models. arXiv preprint arXiv:2404.17788 2024.
Michihiro Yasunaga, Armen Aghajanyan, Weijia Shi, Rich James, Jure Leskovec, Percy Liang, Mike Lewis, Luke Zettlemoyer and Wen-tau Yih. Retrieval-augmeener multimodal language modeling. arXiv preprint 2rXiv:2211.1255 2022.
Howard Yen, Tianyu Gao, and Danqi Chen. Long-context language modeling with parallel context encoding. arXiv preprint arXiv:2402. 16617, 2024a.
Howard Yen, Tianyu Gao, Minmin Hou, Ke Ding, Daniel Fleischer, Peter Izsak, Moshe Wasserblat, and Danqi Chen. Helmet: How to evaluate oughly, 2024b. htttttss:: //tpai/////xvvva/// org/ abs 2410.02699
language models effectively and thor-
Haofei Yu, Yue Zhang, Wei Bi, et al. Trams: Training-free memory selection for long-range lan- guage modeling. arXiv preprint arXiv:2310. 15494,, 2023.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a ma- chine really finish your sentence? arXiv preprint arXiv: /905.0783000 2019.
14

---

Published as a conference paper at ICLR 2025
Peiyuan Zhang. Kaichen Zhang, Bo Li, Guangtao Zeng, Jingkang Yang, Yuanhan Zhang, Ziyue Wang, Haoran Tan, Chunyuan Li, and Ziwei Liu. Long context transfer from language to vision. arXiv preprint 1XX1:22406.665222 2024a. URL httpp:://aa.....rg/////o http://.ab///44.....
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, et al. oobench: Extending long context evaluation beyond 100k tokens. arXiv preprint arXiv:2402.13711 2024b.
Dawei Zhu, Nan Yang, Liang Wang, Yifan Song, Wenhao Wu, Furu Wei, and Sujian Li. Pose: Efficient context window extension of Ilms via positional skip-wise training. arXiv preprint arXiv:2309. 10400, 2023.
A IMPLEMENTATION DETAILS OF QUEST
A.1 MODEL. CONFIGURATION
We present the model configuration in Table 6 For the other baselines, we only altered the training dataset, keeping the default model configuration unchanged.
A.2 TIME CONSUMPTION
During the dataset construction, query generation utilized 16 H800 GPUs and took 24 hours to process the 30 billion document tokens in The Pile dataset. After that, 128 CPUs were employed for distributed processing: it took 0.5 hours to generate the keywords and hours to construct the inverted index based on those keywords. Since the process is highly parallelizable, scaling it up would not pose significant challenges.
32K
128K
model size
1.4B
6.9B
12B
1.4B
6.9B
12B
rotary-pet
0.25
rotary-emb-base
100000
5000000
B1
0.9
Ba
0.95
eps
1e-8
Ir
56
4e
-5
5
5e-5
-5
-5
precision
6
Zero stage
1
gradient-clipping
1.0
weight-decay
0.1
Ir-decay-style
cosine
train-iters
1000
warmup-iters
200
seq-length
32768
131072
GPU-type
H800
GPU-numbers
16
32
32
32
32
32
training-time
6.3h
14h
20.6h
9.5h
30h
39h
Table 6: Model Training Configuration
best way
get rid main differences
had idea valid way better way
good way
following sentence
sentences
mean
passage mean
following data correct way following passage
good idea
ways
sentence mean
next word
part
current state fallowing equation
Table 7: Stop Keywords
15

---

Published as a conference paper at ICLR 2025
Quest Split
Single dbc Improvement (%6
- Quest
1.5
Ingroveniont
0.8
0.4
0.5
0.0
0.0
20
60
Split
Split
Figure 8: The performance changing trend with split_ratio increasing, For detailed results, please
refer to Appendix B
A.3 FILTERING OF KEYWORDS
We used Rake for keyword extraction and found many free-fry but meaningless keywords. Therefore, we maintained a list of stop words and performed a keyword extraction on the generated query avoid selecting those stop words. Some of the stop words listed in Table Moreover, to enhance the quality of keyword extraction, we applied a post-processing step to clean the keywords generated by the Rake algorithm. That involves removing punctuation and filtering out keywords that are either less than four characters in length or have score below 3. Such a cleaning process tends ensure that the extracted keywords are both meaningful and relevant.
A.4 IMPACT OF SPLIT RATIO
This section examines the impact of split_ratio, which controls the proportion of oversampled keyword-based inverted indexes. Figure shows that performance initially improves but then de- clines as split_ratio increases. The best results occur when oversampled indexes comprise 10-30%. It suggests that moderate oversampling of indexes with fewer documents benefits long context mod- eling, highlighting the importance balanced data distribution for long-context tasks.
Table 8: Performance of different methods across various Longbench subtasks.
Model Size Method
Few-shol Learning
Synthetic Tasks
Code Completion
triviaqa
sumsum
passuge
passage retrieval
repobench-p
Standard
32.29
19.99
24.75
29.95
2.47
1.62
35.66
39.16
1.4
ICLM
31.75
17.91
16.08
26.60
1.00
2.88
34.56
36.07
Quest
40.75
18.68
19.17
33.65
0.96
3.71
40.89
43.60
38.75
22.38
23.06
4.75
KNN
35.83
24.29
16.54
36.07
2.69
4.83
42.86
40.89
6.9B
38.08
Quest
38.50
21.59
24.22
36,26
2.41
3.50
49.98
51.16
Standard
39.25
26.87
26.97
35.12
3.12
4.44
42,03
45.43
12B
ICLM
41.08
23.05
22.24
37.60
1.75
1.25
34.65
39.52
Quest
26.60
Table Performance of different methods across various Longbench subtasks.
Single-Doe QA
Multi-Doc QA
Summarization qusum
Model Size Method narrativeqa
qasper multihellqqacc
hotpotqa
Zwikimqa
musique
gov .report
multi
Sundard 13.55
14.8
29.
21.71
22.96
7.7
23.28
14.53
24.13
KNN
13.29
26.24
26.29
15.31
27.32
1.48
ICLM
14,66
15.79
29.59
16.56
20.13
7.44
25. 28
14.33
26.25
Quest
12
12.77
29.14
21.63
24.7
7.62
25.53
14.35
25.86
Stundardl
KNN
16.10
10.41
29.00
19.66
17.30
3.97
36.70
17.06
23.91
6.9B
ICLM
14.62
28.19
19.48
6.11
26.08
16.41
24.56
Quest
17.77
8.63
31.23
19.46
17.60
5,33
26.69
16.11
24.56
Standard
20.32
13.85
32.36
28.79
22.09
14.94
24.57
18.36
23.78
ICLM
20.33
16.84
30.84
31.92
23,83
14,11
28.01
18.97
23.24
Quest
19.12
14.17
33.72
27.53
21.76
13.94
28.22
19.33
23.68
16

---

Published as a conference paper at ICLR 2025
Table 10: Performance Change with different split_ratio values across various Longbench subtasks.
Split Ratio (%)
Few-shot
Learning
Synthetic Tasks
Completion
free
triviaga
samsum
Bq
.ountt
passage retrieval
lec
repobench-p
0
33.46
21.94
22,50
29.30
1.48
2.50
32.29
37.06
30
36.75
18.06
27.06
27.85
0.89
3,50
33.74
38.32
36.17
22.02
24.29
31.60
4.69
32.28
37,00
70
33.83
19.17
25.58
27.13
1.81
3.67
29.78
35.66
90
32.33
21.23
26.34
23.98
2.41
2.60
30.81
34.52
33.46
21.94
22.50
29.30
1.48
2.50
32.29
37.06
Table 11: Performance Change with different splil. ralio values across various Longbench subtasks.
Single-Doc
Mulli-Doc
Summarization qmium
Split Ratio (%)
narrativega
qusper
mullifieldga.er
hotpolqa
2wikimqa
musique
goy_report
mulli_news
0
12.41
13.13
27.87
20.05
16.65
8.05
24.57
15.50
22,6-1
10
13.79
16.22
27.15
18.38
18.74
9.11
26.72
15.01
22.69
30
14.04
17.19
27.77
20.39
20.70
7.87
26.20
14.94
23.99
70
11.32
14.34
28.17
17.85
22.21
9.20
25.18
14.89
24.35
90
13.18
18.30
29.68
19.71
21.14
7.14
24.06
100
12.41
13.13
27.87
20.05
16.65
8.05
24.57
15.50
22.64
Table 12: Performance comparison across different methods on HELMET.
Method
Avg.
Recall
RAG
ICL
Re-rank
LongQA
Standard
48.08
62.33
58.67
71.24
19.18
28.99
KNN
46.31
64.24
56.00
60.28
18.77
32.27
ICLM
46.62
64.04
54.48
72.36
14.04
28.17
Quest
50.97
69.13
57.47
72.08
22.35
33.82
Table 13: Performance comparison across different query replacement ratios.
Replacement Ratio
Avg.
Recall
RAG
ICL
Re-rank
LongQA
0%
50.97
69.13
57.47
72.08
22.35
33.82
20%
50.33
67.35
58.07
71.72
22.55
31.98
50%
49.74
63.96
58.19
73.32
20.93
32.31
100% (Standard)
48.08
62.33
58.67
71.24
19.18
28.99
B 32K LONGBENCH RESULTS
We report the performance of Longbench on 17 English subtasks. Table 8 and Table 9 are the detailed results of Table Table 10 and Table 11 are the detailed results of Figure
C ABLATION STUDY IN QUEST
In this section, we conduct evaluations on the long-context benchmark HELMET (Appendix C.1 performing ablations query quality (Appendix C.2), keyword quality (Appendix C.3) and key- word sampling strategies (Appendix C.4
C.I RESULTS ON HELMET BENCHMARK
We evaluate the Quest method on HELMET benchmark (Yen et al.) 2024b). HELMET provides more reliable and consistent rankings for frontier long-context models. It spans evaluation lengths of 8k, 16k, 32k, 64k, and 128k across five task types and 17 subtasks. including Recall (4 subtasks), RAG (4 subtasks), ICL (5 subtasks). Re-rank (1 subtask), and LongQA (3 subtasks). Table 12 represents the average performance across lengths from 8k 128k. The comprehensive evaluation on HELMET shows that Quest achieves non-trive trivial improvements across multiple datasets, with a many 17 subtasks.
notable +2.89% average increase in performance
17

---

Published as a conference paper at ICLR 2025
C.2 QUERY QUALITY AND MODEL PERFORMANCE SHOW POSITIVE CORRELATION
We generate queries using large language models (LLMs). which intuitively could produce higher- quality queries. However, the computational cost prohibitively high and beyond our resource capacity. To explore query quality generation within a manageable range, we simulate the creation of low-quality queries by randomly replacing keywords in the original queries with a certain proba- bility. Table [3] shows that model performance gradually declines as the replacement ratio increases. This observation suggests that improving the quality of queries can further enhance model perfor- mance.
C.3 THE KEYWORDS FROM QUERIES ARE HIGH-QUALITY
We evaluate various strategies for keyword extraction using GPT-4o. Three distinct strategies are designed for generating keywords:
1. Keywords from Queries: Extracting keywords from the generated queries.
2. Keywords from Documents: Extracting keywords directly from the original documents.
3. Keywords from Summaries: Generating summaries from the original documents using a model, then extracting keywords from the summaries. We utilize the recently introduced Llama-3.2-IB MMeta 2024b) model to generate summaries.
To assess the quality of the extracted keywords, conduct evaluations using human annotations and GPT-4o on 100 samples. Scores ranged from 1 to 5. Two PhD students are hired for annotation, alongside GPT-4o,, all following the scoring criteria below:
Table 14: Performance comparison across different keyword extraction strategies.
Strategy
Avg Score
Score 2 3
Score 2 4
Score > 5
Keywords from Queries
3.574
74.3%
63.0%
36.4%
Keywords from Summaries
3.009
61.8%
48.7%
14.5%
Keywords from Documents
2.634
49.3%
31.2%
6.0%
Table 15: Comparison of entropy scores between keyword selection strategies. Strategy Highest RAKE Score
Entropy Score
15.31
Random Sampling
18.12
Table 16: Performance comparison of keyword selection strategies.
Method
Avg.
Recall
RAG
ICL
Re-rank
LongQA
Highest RAKE Score
49.51
66.76
56.27
72.72
22.32
29.47
Random Sampling
50.97
69.13
57.47
72.08
22.35
33.82
The keyword has little to no relevance to the article, may have been extracted incor- rectly, is significantly off-topic.
Score 1:
Score 2: The keyword has some relevance the article but is of low importance or overly generalized. It does not represent core theme or focus of the article.
Score 3: The keyword is clearly connected to the article but is not one of its primary concepts or focal points.
Score 4: The keyword is strongly related to the article and represents an important concept or secondary theme.
Score 5: The keyword perfectly reflects the article's central theme, main concepts, or key discussion points.
18

---

Published as a conference paper at ICLR 2025
Native Long-Form Data Distribution
Combined Data Distribution
ArXiv
Gibub
OpenWeNTeu2
StackEschange
BonkCupaa2
Gutenbery PG-19
PhilPapers
Backgrounds
Books
HackerNews
Pik-CC
Ubuntu IRC
DM Mathematics
NIH ExPorter
PobMed Abstracts
Wikipedix
Enros Emails
OpenSubeiles
PobMed Central
YoutubeSubiiiles
Freelaw
Figure 9: (Left) Distribution of long documents exceeding 128k tokens in the Pile dataset. (Right) Distribution of >= 128k tokens long-context data synthesized by Quest. By employing Quest, we successfully increase the domain diversity and balance of the long-context training set.
Table 17: Performance comparison across KNN strategies and Quest.
Methods
Cosine-Sim
Avg.
Recall
RAG
ICL
Re-rank
LongQA
KNN (TopK)
60.97
46.31
64.24
56.00
60.28
18.77
32.27
KNN
48.53
47.96
65.71
56.97
68.24
16.54
32.33
KNN (Random Sampling)
35.51
46.78
64.23
55.24
70.96
12.23
31.22
KNN (Reverse Order)
35.49
46.18
62.19
54.28
66.28
17.20
30.93
Quest
46.54
50.97
69.13
57.47
72.08
22.35
33.82
The Pearson correlation coefficient between human annotations and GPT-4o scores is calculated as 0.7986, indicating strong alignment between the two evaluation methods.
Subsequently, GPT-4o is used to score 1,000 samples. Table 14 shows that extracting keywords from queries significantly improves the quality of the extracted keywords, and these keywords achieve the highest score. The keyword from Generated summaries is of inferior quality. We attribute this to the broader semantic space of summaries compared to queries, which interfere with effective keyword extraction.
C.4 RANDOM SAMPLING ENHANCES KEYWORDS DIVERSITY AND IMPROVES MODEL PERFORMANCE
Through entropy analysis, we evaluate the diversity of two strategies: selecting keywords with the highest RAKE score and random sampling (for scores greater than 3). Table [5 shows that random sampling led to greater diversity. Furthermore, we construct a training datasett based on the highest RAKE score, and the results are presented in Table 16 These results demonstrate that random sampling enhances diversity and improves model performance.
C.5 THE DISTRIBUTION COMPARISON OF LONG CONTEXT DOCUMENTS
Figure shows the distribution of long-context data before and after the application of Quest. While the distribution of loog-coonte sources initially highly uneven, the implementation of Quest has significantly increased the percentage of data in domains such as ArXiv, FreeLaw, Open WebText2, Pile-CC, and PhilPapers, where there previously minimal native long-context data.
19

---

Published as a conference paper at ICLR 2025
D ANALYSIS WITH DOCUMENT SIMILARITY
D.1 CHANGING THE SIMILARITY OF KNN-AGGREGATED DOCUMENTS CANNOT ACHIEVE QUEST'S PERFORMANCE
As we mentioned in Section 5.2 Quest achieves significant performance improvements by aggregat- ing relevant, low-redundancy documents. This raises the question of whether selecting moderately similar documents instead of the most similar ones in KNN-based method can achieve comparable results. To perform such similarity in a KNN-based method, we designed four strategies:
1. TopK: The normal KNN selects and concatenates documents with top-ranking similarity.
2. Mid-Ranking: Documents with mid-ranking similarity are selected and concatenated.
3. Random Sampling: Documents are concatenated by random sampling.
4. Reverse Order: Documents concatenated in reverse order of similarity.
As shown in Table 17 for the KNN method, reducing the similarity of the selected documents results in a slight performance improvement, followed by a decline. Furthermore, these methods consistently perform worse than the Quest method. Additionally, while the "Mid-Ranking" ap- proach exhibits similarity extent close to the Quest method, its performance remains far inferior. It indicates that document aggregation has a more significant impact overall performance than the extent of similarity, highlighting it as the critical factor driving Quest's ' superior results.
D.2 EXAMPLES OF SIMILARITY VISUALIZATIONSS.
We present more visualization results. Figure 10 shows the t-SNE visualization documents within 32k context, while Figure TT illustrates documents within 128k context. The Standard method's random concatenation of documents results an overly dispersed distribution, disrupting docu- ment relationships and leading to poorer performance. In 32k context, ICLM causes excessive clustering due to shorter similarity-path lengths, mirroring a KNN-like distribution and impairing performance on the Longbench benchmark. However, the 128k context allows ICLM form longer similarity paths, dispersing document distribution and enhancing performance on the LongbookQA benchmark.
Notably, Quest maintains an evenly dispersed document distribution in both contexts, leading to its superior performance both benchmarks. The findings above indicate that overly dispersed or concentrated document semantics can harm model performance, while Quest improves performance by clustering query-related documents, ensuring relevance and avoiding redundancy.
20

---

Published as a conference paper at ICLR 2025
Distribution
Document Distribution
Standard
4
KNN
ICLM
Quest
-2
Standard
KNN
6
-10
Quest
-15
-10
-5
10
Distribution
Document Distribution
8
Standard
6
KNN
10
4
ICLM
Quest
Standard
-4
KNN
-6
-10
-8
15
Figure 10: Visualizing Documents Comprising 32k Context.
Document Distribution in 128k Context
Dacument Distribution in 128k Context
Standard
KNN
ICLM
Quest
-2.5
Standard
KNN
-5.0
-10
ICLM
-7.
-15
-5
-8
-6
Document Distribution 128k Context
Document Distribution 128k Context
Standard
KNN
7.5
5.0
Quest
0.0
Standard
-2.
-10
ICLM
Quest
-7.5
-15
-10
-5
10
-10
-5
0
10
15
Figure 11: Visualizing Comprising a 128k Context.
21


Reviews:
[{"review_id": "aBGrcYocr2", "reviewer": "Reviewer_ki5C", "summary": "This paper proposes a method to generate synthetic long documents for model training. It aims to solve the problem of standard method that may concatenate unrelated short document to form a long document, and the similarity-based methods that concatenate too similar documents leading to the lack of diversity. The proposed method - Quest - selects the documents to be concatenated by the common keywords they would have in the generated queries to which they may answer. This is intended to concatenate documents that may be on a common topic (keywords), without being too similar.\nThe experiments on several datasets show that the synthetic data generated by the proposed method can lead to some improvements in the test tasks (on average).", "questions": "Have you tried simpler methods to select documents to be concatenated, e.g. based on controlled similarity measure, so that the documents are within some range of similarity? Would this achieve the same goal as the proposed method?", "limitations": "", "rating": 3, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The problem observed in the previous methods to generate synthetic data is interesting: the documents to be concatenated should be similar to some extent, but not too similar.\n2. The experiments show that the method can result in improvement on some tasks.", "weaknesses": "1. The paper proposes a quite complex algorithm using doc2query, keyword indexing, then sampling through keywords. The paper does not motivate why these steps are necessary. One could imagine that much simpler methods would be able to lead to synthetic data that are to some extent similar, as is required by the authors. For example, if it is observed that the method based on KNN lead to concatenating too similar documents, it would be possible to select documents to be concatenated that are similar to some extent, but not the most similar. Could this simpler method lead to similar results?\n2. The paper tends to over-claim the advantage of the proposed method. For example, it is said \"Table 1 compares the Longbench results, showing that Quest consistently outperforms other methods across model sizes, ...\". Looking at Table 1, one can see that the proposed method outperforms the others on average, but underperform them on several datasets. The advantage of the method is not so clear. This over-claim appear in other observations as well (e.g. Table 3 shows that Pythia obtains the best performance on 4 datasets, while Quest only on 2). Globally, if one consider all the datasets, the demonstration that the proposed method is better is very weak.\n3. Some of the figures are unclear. Fig. 1 is not well explained. It intends to show the perfect performance of Quest on a specific task. However, the figure does not provide very useful information in addition to say it is perfect. One would also like to see some other methods in comparison, and to understand better the task itself. In Fig. 2, the dotted lines are not explained. One can guess later that they correspond to the performance of the model in some task.\n4. Fig. 5 compares the results using the documents of some length in the original dataset, and the synthetic documents. It shows that the latter perform slightly better. Do these data have equivalent sizes? If the synthetic data are much more than the real subset of data, the observation would not be surprising. However, it would be difficult to conclude that the synthetic data are better than the real data, only they are more.", "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": "mkexCzV1Cm", "reviewer": "Reviewer_9HBG", "summary": "This paper introduces QUEST, a long-context scaling method for LLM pretraining. The proposed method focuses on data sythesis and introduces a simple algorithm QUEST, in which queries are first generated by documents, followed by key words generation of the queries. Then, the keywords are used as inverted index to group similar documents, where the similar documents are arranged within the same training batch to improve long-context modeling by cross-document reasoning. The authors perform experiment on continual pretraining and demonstrate the effectiveness of the proposed QUEST on both long- and short-context benchmarks. The authors also provide additional analysis on the training dynamics and scaling properties of long-context modeling using QUEST.", "questions": "1. In fig3, Llama 3 8B 128k has accuracy of 0.97, so why even longer context of 1M as in fig 1 would result in 1 accuracy?\n2. Why the performance in appendix from table 8 to table 10 are in consistent, I find it hard to inteprete these results and how do they differ from the tasks that are present in the main text?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The authors propose a simple and scalable method to group documents for long-context pretraining\n2. The proposed method seem to work and achieve performance gains on the selected short- and long-context benchmarks\n3. The authors also provide many interesting observations and analysis of QUEST and long-context modeling, which may be helpful for LLM pretraining", "weaknesses": "1. Although the proposed method is simple, it largely follows previous approaches but with a different indexing and sampling method. In addition, no ablation is provided so we are not sure which part of the QUEST is the most helpful (query-based indexing or sampling etc.).\n2. The performance of QUEST is quite similar to ICLM on long-context benchmarks. Interestingly, in appendix, it seems QUEST often underperforms compared to ICLM or even KNN. THis raises questions on the real performance of the proposed method.\n3. The authors provide many interesting observations and analysis, yet these do not necessarily demonstrate the validity of QUEST. For instance, the examples in 5.2 show that most baseline methods also have good clustering quality.", "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": "KmAvkOfwUU", "reviewer": "Reviewer_W5Zm", "summary": "This paper introduces Quest, a query-centric data synthesis approach addressing data scarcity and domain imbalance issues in long-context modeling of LLMs. Quest achieves a balance between semantic correlation and context diversity by predicting potential queries for each document and aggregating semantically relevant but low-redundancy documents based on query similarity and keywords. Experiments demonstrate Quest's superior performance on long-context tasks ranging from 32k to 1M tokens.", "questions": "Please see the weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The motivation is clear. The authors identifie two critical challenges in long-context training data: scarcity and domain distribution imbalance, then provide analysis of limitations in existing methods (Standard, KNN, ICLM).\n2. The authors propose a query-centric data synthesis approach to alleviate the above problem, includes: query-based document aggregation moduls, keyword extraction and filtering modules, split ratio mechanism.\n3. Extensive evaluation across various model scales (1.4B-12B) and context lengths (32k-1M) supports the method's effectiveness.", "weaknesses": "1. The author used a query generation model to generate different queries for different documents, which would make the quality of query generation affect the final performance of the model. The author should discuss in more detail the impact of query quality or different query generation methods on model performance.\n2. The model approach is more like an integration of existing tools, without any contribution to methodological or theoretical innovation. For example, the query generation and selection process directly employ the off-the-shelf tools (doc2query, RAKE) without significant enhancement.", "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": "ISITvD4nXt", "reviewer": "Reviewer_YyBr", "summary": "This paper introduces Quest, an approach for grouping documents to train long-context LLMs. Existing approaches for document grouping either use random grouping or grouping by similarity. However, the former leads to aggregated documents that contain unrelated pieces, while the latter yields aggregated documents with repeated information coming from similar documents. Quest then seeks to balance the coherence of the aggregated documents and the diversity of their content.\n\nThe Quest algorithm first associates each document to a query, which is used to extract a representative keyword associated to the document. Documents are then indexed based on their respective keyword, in the same spirit as an inverted index. During training, a keyword is sampled and an aggregated document is formed by randomly sampling documents associated to this keyword up to the desired context size. The keyword sampling procedure also takes into account the rarity of the keyword (i.e., the number of documents associated to a keyword) by oversampling rarer keywords, to have a more balanced domain representation. Quest is validated through extensive experiments on long-context tasks and is also shown to preserve the short-context capabilities of the original pre-trained model.", "questions": "- The first step of the Quest algorithm is to associate documents with keywords. This is done by first formulating a query from the document, before extracting keywords from this query. It's not very intuitive why this document-to-query strategy is favored over applying any standard summarization technique to the document. A summary may retain more of the document's information than a single query, and thus might lead to higher-quality keywords. The paragraph lines 93-105 attempts to make some links with search engines, but the motivation could benefit from being further clarified and strengthened.\n- At the end of the keyword extraction step, the keyword that is eventually retained for a document is randomly selected among the different keywords associated to the document. This strategy seems a bit unintuitive and risky, as it could lead to the selection of a keyword that is only marginally relevant to the document. Why isn't the Rake score used here to decide which keyword to retain? Or even simply retain the keyword with the largest Rake score?\n- Was there any evaluation conducted to validate the quality of the keywords identified for the documents? This could for example be done by asking some human annotators to provide keywords for a few documents and compare these against the keywords obtained through Quest, or instead ask them to rate the relevance of the Quest keywords. This seems like an important sanity check to conduct given the importance of proper keyword-document matching in the proposed approach.\n- How is the oversampling of keywords from the short-index set $I_s$ done in practice? How are the probability of sampling a keyword from $I_s$ and that of sampling from $I_l$ defined, exactly? Adding the mathematical formula for this would be helpful for the reader's understanding.\n- What is the motivation for splitting $I$ into only two sets $I_s$ and $I_l$ instead of having a more fine-grained approach? For example, one could define an adaptive probability of sampling a keyword based on the number of tokens in its associated documents. Such an approach would enable a more tailored treatment of the different keywords to compensate for their rarity.\n- Figure 3 contains several typos: \"query genaration\" -> \"query generation\", \"random smaple\" -> \"random sample\"\n- In the results on Longbench in Table 1, it seems that the superior average performance of Quest over baselines is essentially due to its significantly better results on the Code Completion task. The good performance of Quest on this specific task is attributed to the use of Levenshtein distance as the metric. It is however not entirely clear why using Quest would specifically boost this metric, so more detailed explanations would be appreciated.\n- It is surprising that Quest achieves 100% accuracy on the needle-in-a-haystack task for a context of 1M tokens, but gets 97% accuracy for a context of 128k tokens (which should intuitively be easier). Are there any intuitive explanations for this phenomenon?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper is overall well-written, clear and easy to follow.\n- The experiments are comprehensive and convincing.", "weaknesses": "- Quest is defined as a simple heuristic that is not backed by any theoretical guarantees or motivations.\n- Several design choices in the approach lack justifications or intuitive explanations, and thus feel arbitrary (see specific points in the \"Questions\" section).", "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)