In-Context Pretraining: Language Modeling Beyond Document Boundaries
Authors: Weijia Shi
Sewon Min
Maria Lomeli
Chunting Zhou
Margaret Li
Xi Victoria Lin
Noah A. Smith
Luke Zettlemoyer
Wen-tau Yih
Mike Lewis
Venue: ICLR 2024 spotlight (2024)
Keywords: Large Language Models
Primary area: representation learning for computer vision, audio, language, and other modalities

Abstract: Language models are currently trained to predict tokens given document prefixes, enabling them to zero shot long form generation and prompting-style tasks which can be reduced to document completion. We instead present IN-CONTEXT PRETRAINING, a new approach where language models are trained on a sequence of related documents, thereby explicitly encouraging them to read and reason across document boundaries. Our approach builds on the fact that current pipelines train by concatenating random sets of shorter documents to create longer context windows; this improves efficiency even though the prior documents provide no signal for predicting the next document. Given this fact, we can do IN-CONTEXT PRETRAINING by simply changing the document ordering so that each context contains related documents, and directly applying existing pretraining pipelines. However, this document sorting problem is challenging. There are billions of documents and we would like the sort to maximize contextual similarity for every document without repeating any data. To do this, we introduce approximate algorithms for finding related documents with efficient nearest neighbor search and constructing coherent batches with a graph cover algorithm. Our experiments show IN-CONTEXT PRETRAINING offers a scalable and simple approach to significantly enhance LM performance: we see notable improvements in tasks that require more complex contextual reasoning, including in-context learning (+8%), reading comprehension (+15%), faithfulness to previous contexts (+16%), long-context reasoning (+5%), and retrieval augmentation (+9%).

Published as a conference paper at ICLR 2024
IN-CONTEXT PRETRAINING: LANGUAGE MODELING BEYOND DOCUMENT BOUNDARIES
Weijia Shi
Sewon Min Maria Lomeli'
Chunting Zhou'
Margaret Li
Gergely Szilvasy*
Rich James'
Xi Victoria Lin
Noah A. 2m Luke Zettlemoyer Scott Yih!
Mike Lewis!
Meta AI University of Washington
Allen Institute for AI
swje419@ce wwshington.
ABSTRACT
Large language models (LMs) are currently trained to predict tokens given doc- ument prefixes, enabling them to directly perform long-form generation and prompting-style tasks which can be reduced to document completion. Existing pretraining pipelines train LMs by concatenating random sets of short documents create input contexts but the prior documents provide signal for predicting the next document. We instead present IN- NN-CONTEXT PRETRAINING where language models are pretrained on sequence of related documents, thereby explicitly encouraging them to read and reason across document boundaries. We can do IN-CONTEXT PRETRAINING by simply changing the document ordering that each context contains related documents, and directly applying existing pretraining pipelines. However, this document sorting problem is challenging. There are billions of documents and we would like the sort to maximize contextual similarity for every document without repeating any data. To do this, we intro- duce approximate algorithms for finding related documents with efficient nearest neighbor search and constructing coherent input contexts with graph traversal algorithm. Our experiments show IN-CONTEXT PRETRAINING offers simple and scalable approach to significantly enhance LMs' performance: we see notable improvements tasks that require more complex contextual reasoning, including in-context learning (+8%), reading comprehension (+ (+15%), faithfulness to previous contexts (+ (61),) long-context reasoning (+5%), and retrieval augmentation (+9%).
new approach
INTRODUCTION
Large language models (LMs) are trained to complete documents; each token is predicted given the context provided by the prefix of the document appears in. Such contexts can be widely varied, especially at pretraining scale, allowing models to excel diverse tasks such as instruction- following (Ouyang et al. 2022), conversational interfaces (OpenAl (023). reading comprehen- sion /Zhag al.) 2020), and in-context learning (Brown al. 2020). However, recent studies highlight that LMs sometimes struggle to understand more complex contexts: they can fail to follow instructions accurately (McKenzie al. [2233 Efrat & Levy 2020 Liu Liu| 2023), struggle with reasoning over conditioned documents (Liu et al. 20033] Shi et al. (023a), and exhibit high variance in in-context learning (Zhao al. 2021). 1j. In this paper, we present IN-CONTEXT PRETRAINING, new pretraining method that learns to predict tokens conditioned on sequence of related documents, explicitly enabling the model read and reason about much more varied and longer contexts that go beyond document boundaries.
Current LM training pipelines concatenate random sets of shorter documents to create longer con- text windows. However, the prior documents provide no signal for predicting the next document, incurring unnecessary computational overhead for tokens that do not require communication between them (de Vries] 2023). IN-CONTEXT PRETRAINING instead reorders the pretraining data by combin- ing several semantically related documents to create a coherent input context, thereby exposing LMs to long relevant contexts and providing pretraining signals beyond document boundaries. We illustrate this via an example in Figure when predicting the following tokens for the phrase "For 2022, FIFA the prize money at $42m," " previous document stating that the "World Cup never awarded

---

Published as a conference paper at ICLR 2024
Pretraining Documents
In-Context Pretraining
the highest for
World Cup
World Cup award
Language Model
2022,
Input
World Cup never swarded 510M before 2022
For 2022, FIFA set the price money at $42m,
Messi scored seven
Standard
the highest so for
Paris bisected
Language Model
Paris, France's capital
Input Contexts
Paris is bisected by the River Seine, which flows
For 2022, FIFA set the prize money st $42m,
Figure 1:
Overview of IN-CONTEXT PRETRAINING. Different from the standard pretraining strategy that place randomly shuffled documents in the input context, IN- IN-CONTEXT PRETRAINING places related documents in the same context, making models learn to reason across prior documents. For example, when predicting the following tokens for the phrase "For 2022, FIFA set the prize money at $42m," " LMs could reference prior documents stating "World Cup never awarded more than $10M before 2022" and learn infer that "the highest so far."
more than $10M before 2022" could be in the context, enabling the prediction of continuation like "the highest so far." As IN-CONTEXT PRETRAINING only changes document ordering and leaves all other aspects of LM pretraining untouched, it can be easily integrated into existing pretraining pipelines for large-scale LMs.
However, this document sorting problem is challenging. LMs are typically trained on billions of documents and we would like to sort them to maximize document similarity in the input context windows without repeating any data. We introduce two new approximate algorithms to tackle these challenges. We use retrieval model paired with an efficient search index to build a document graph that pairs each document with its nearest-neighhors based on its semantic similarity in the embeddings space. We also formulate document sorting as a travelling salesman problem, for which we develop an effective algorithm that maximizes similarity of documents to their context while also ensures that each document is included only once.
To evaluate the effectiveness of IN-CONTEXT PRETRAINING, we pretrain language models from 0.3 to billion parameters on 300 billion tokens from the CommonCraw) dataset (Wenzek et al. 2020). Across all model scales, IN IN-CONTEXT PRETRAINING LMs (ICLM) demonstrate strong language modeling and downstream task performance, outperforming LMs pretrained using the standard approach on the same corpus. We observe various improvements resulting from IN-CONTEXT PRETRAINING compared with existing LMs: (1) in-context learning with an average increase of 8% across 8 datasets; (2) reading comprehension, with an average of 15% improvement on reading comprehension tasks; (3) outputs that are more faithful to prior contexts (+16%)) (4) long context reasoning. showing a 5% boost; and (5) retrieval augmentation, leading to 9% gains when augmenting with external knowledge such as documents retrieved from Wikipedia. Our results demonstrate that, by simply altering order of the pretraining documents, IN-CONTEXT PreTrAININg offers scalable and simple approach to significantly enhance understanding and reasoning over their full contexts. Code are publicly released at atgiita withub.coom
2 IN-CONTEXT PRETRAINING
The standard practice in pretraining is to form input contexts by concatenating random documents until reaching the maximum context length. It then trains the LM using a language modeling objective on the input contexts. However, training LMs on randomly concatenated documents does not offer additional learning signals compared with training each document individually. In contrast, IN-CONTEXT PRETRAINING generates more coherent input contexts by concatenating semantically related documents together during pretraining. As depicted in Figure 2 IN- IN-CONTEXT PRETRAINING consists of two steps: it first finds related documents at scale (2.T) and then constructs input contexts using these related documents (.2). Successively,
the contexts formed with
2

---

Published as a conference paper at ICLR 2024
semantically related documents to pretrain LMs with language modeling objective. Since IN- CONTEXT PRETRAINING is identical to existing pretraining recipes for LMs, except for changing how input contexts are built, it can be easily integrated into existing pretraining pipelines for large- scale L.s.
2.1 FINDING RELATED DOCUMENTS AT SCALE: RETRIEVING NEIGHBOR DOCUMENTS
To find related documents at scale, we link documents within the pretraining corpus D using retrieval model. Specifically, for each document dj E DD a dense retrieval model is used retrieve the top-k most similar documents, represented as N (d4). The retrieval model uses approximate nearest neighbours search for efficient pairwise similarity comparison between any two documents, making it scalable for finding related documents in web-scale pretraining corpora.
Our retrieval process employs the contriever model (Izacard al., 2022 This model maps each document d; E to an embedding E(d,) by taking the mean pooling of the last hidden representation over the tokens in d,. The cosine similarity then used to determine the similarity between any two documents:
Retrieval.
s s(d,,d)) cos (d()), B(d,))
The retrieval model uses approximate nearest neighbour search with the faiss library (Johnson et al. 2019]Douzee et al.) 2024). We use product quantization (Jégou al. al.] (2011) to reduce the memory footprint and an IVF (inverted file) index structure to conduct efficient pairwise similarity search together with faiss big batch search. The OIVFBBS faiss framework is leveraged for this task, OIVFBBS refers to conducting offline search with queries of big batches with faiss inverted indexes.
Further details can be found in Appendix A.2 and in the OIVFBBS demo in the faiss github repository github. wcto///accbbooreeerrhhheeo http:aarrcc/aaaissseeeee///aaaaa wwttttttteee/eeeeeeeeeeeeetttttt http://///maammemmmmmeeeee/eee/e dww.s//oooooooooo..eoe /ow.n/eefffffif..ooo
During the retrieval process, when computing pairwise similarity among each document in the pretraining corpus, we found that the pretraining corpus contains many near duplicate documents. Hence, we further leverage the retrieval scores to eliminate near duplicate documents from the pretraining corpus. More details be found in Appendix A.1 In $4.2. show that this deduplication step is crucial for achieving good performance of language models.
2.2 CREATING INPUT CONTEXTS: DOCUMENT GRAPH TRAVERSAL
Given set of documents D = }d.} and nearest neighbours for each document ((d)). our goal is to sort the documents to create input contexts such that each of them consists list of related documents. Formally, we aim to form set of input contexts C₁ ... C. where each context C; = (d.,,..d)] C 'D m
and U C, D. Ideally, documents in C, are nearest neighbors of each others.
A straightforward approach to form C Cm is to directly place each document and its retrieved top- k documents together in the same input context (referred to as kNN), which has been used in some surrival-auugenter pretraining methods (Guu et al.] 2020 Levine et al.) 2022). This KNN approach maintains document similarity within each context but creates the data repeating prob- lem: some documents frequently appear as nearest neighbors of other documents, causing that differ- ent input contexts contain overlapping documents,
Algorithm 1 Maximum Traveling Salesman
Input: Document graph & (D,
(d,) returns nearest neighbors for d; min_deg( (D) returns a min-degree doc Output: A path P
1:
P P+
2:
while [D) d
3:
d, + min_dee((D)
4:
P.appere((d.)
5:
D.remove(d)
6:
while ((d)) n do
i.e. 3i # J. C. nC; # U. The data repeating problem exposes LMs to a less diverse set doc- uments given a fixed computational budget and could lead to overfitting of popular documents. In- stead, we aim to build set of contexts in way that each document is included only once, which can be cast graph traversal problem.
d, + -arg minggn sin(d, sim(d:,
d. dd
9:
P.oppend(d.)
10:
D.remove(d))
11:
end while
12: end while
13:
return P
3

---

Published as a conference paper at ICLR 2024
Pretraining Documents
Doc
Doc
..0.
Doc
Doc
Doc
Doc
0.9
0.6
0.2
Doc
Doc
Pretraining Input Contexts
8192 context window
0.7
Doc
Doc
Doc
Doc
0.8
Doc
Doc
Doc
Doc
Doc
Doc7
Step 1: Finding Related Does
Step 2: Creating Input Contexts
Figure 2: Illustration of IN-CONTEXT PRETRAINING IN-CONTEXT PRETRAINING first finds related documents at scale to create a document graph (. and then builds pretraining input contexts by traversing the document graph (22.2). Along the path, documents are concatenated into sequence and subsequently divided to form fixed-sized input contexts (c.g., 8192 token length).
Document graph traversal. To achieve our goal maximizing the chance that the related doc- uments are concatenated together, intuitive approach is to find single path that visits each document once and maximize the chance that related documents are visited sequentially. Then we subsequently segment the path into multiple input contexts. We formulate it the maximum traveling salesman problem (Flood 1956) that aims to find the maximum weight path that traverses all nodes exactly once. We represent each document as a node in the graph and use document similarity edge weight. We design an undirected weighted graph representing the documents, symbolized as (D, L). Here, represents the set of documents, while (d, d²) E is edge if d" E (d,)) or d, E (dd).. The weight of each edge corresponds to the document similarity (Equation
Solving large traveling salesman problems exactly is NP hard, but greedy algorithms are known to provide an efficient approximate solution. We adopt this approach, introducing modifications to better suit our context. Algorithm shoows the method construct a maximum weight path. We show path identified by our algorithm in Figure 2 Our algorithm starts by selecting a set-to-be-visitee document with minimum degree as the starting node (Doc 0). The algorithm then progressively extends the current path by navigating to its unvisited neighboring document with highest weight (Doc 9), adding the document node to the path. This process continues until the path reaches a node where all neighboring documents have been visited, which happens because our graph is not complete, and only contains edges between documents where one is within the other's k nearest neighbors. In this case, extend the graph with edge of weight 0 to random unvisited minimum degree document (Doc 1), and continue the above process. The motivation for starting at minimum degree documents is that they are most likely have all their neighbors visited first, and therefore be connected to dissimilar documents in the final path.
As final step, we traverse the documents along the path and concatenate them create fixed-sized input contexts suitable for pretraining. It is important to note that when forming the input training batches, we ensure the diversity among different input contexts within the same batch.
3 EXPERIMENTS
In this section, we describe details of our pretraining setup ($3.1) the baseline methods we use for comparison ($3.2) and experimental results (3.3)
3.1 PRETRAINING SETUP
Since IN-CONTEXT PRETRAINING leaves other details of model training unchanged, and only changes the document ordering so that each context contains related documents, we can directly integrate it into pretraining pipelines preprocessing step during batching. For experiment, we adopt the model architecture and pretraining objective of LLaMA (Touvron et al., 2023a b) and pretrain LMs from scratch.
4

---

Published as a conference paper at ICLR 2024
Pretraining Datasets. We use the English Commoncrawl dataset (Wenzek et al. 2020), the widely- used data source for pretraining LMs. Due to resource constraints, we randomly sample 235 million documents from this dataset, amounting 306 billion tokens in total. We use the same pretraining data for all models.
Model Details. We take the model architecture from LLaMA (Tuvron al. 2023a) and train models across various sizes: 0.3, 0.7, 1.5, and 7.0 billion parameters, all with an 8 8192-length context window. Following LLaMA, we employ the AdamW optimizer (Loshchllov & Hutter 2018 with parameters $1 = 0.9 and B2 0.95, and a cosine learning rate schedule. The 7B model is pretrained using 128 A A100 GPUs across 16 nodes with a batch size of 4 million tokens. It takes 9 days train the 7B model on our pretraining dataset. Due to the long context window of our models, we use flash attention (Dao et al. 2022) to reduce memory consumption during pretraining.
To perform the retrieval over our pretraining datasets, we employ the contriever model (Izacard et al. 2022) and encode the first 512 token of each document into an embedding. We then construct FAISS big batch search that is designed for conducting efficient similarity search with big batches of vectors (typically 50M-100M vectors per batch). Given each query document, we retrieve top 10 documents (k=10). We split the data in batches of 50M embeddings, the search step is conducted in each batch before merging the results using 8 GPUs per batch. The total search time 6 hours over 32 GPUs with average search time per batch is 4,738s. The document graph traversal phase requires 12 hours on a setup of 20 CPUs.
More details are provided in the Appendix A.2
3.2 BASELINES
We compare IN N--COTEXXT PRETRAINING with the following baselines: (1) Standard is the prior standard in pretraining that places randomly shuffled documents in the input contexts. This method is commonly adopted by existing models (Zhang et al. [022 Scao et al., 2022] Touvron et al. 2023a). (2) ANN (also referred to as retrieval-ane -ugmeetdd language model pretraining Guu et al. 2020 Levine et al. )) directly places each document and its retrieved top-k documents together in the same input context. Given the same number of training steps, &NN exposes LMs to a less diverse set of documents, since documents can repeat. For fair comparison, both standard and &NN methods are trained using the same pretraining data as IN-CONTEXT PRETRAINING and undergo an identical number of training steps, ensuring the same computation cost.
3.3 RESULTS
We perform evaluations on tasks that require understanding of contexts including language modeling
($ $.3.ID, D. in-context learning ($ 3.3.2). reading comprehension ($ 3.3.3) and open-book question
answering ($ 3.3.4). factuality ($ 3.3.5 and long context reasoning ($ 3.3.6).
3.3.1 LANGUAGE MODELING
Datasets & Metrics. We evaluate the language modeling perplexity of IN IN-CONTEXT PRETRAIN- ING and baselines the Wikipedia, Arxiv, and Books corpora. We follow the standard language modeling evaluation in concatenating randomly-ordered documents when computing perplexity.
Results. Figure 3 shows average perplexity across different model sizes. First, ANN does not improve over the standard LM, likely due the overfitting problem discussed in 2.2 ICLM, in contrast, outperforms both the standard LM and kNN on all three datasets, even when the evaluation documents are not sorted, The gains are consistent or larger as the size of the model scales. These improvements suggest that IN-CONTEXT PRETRAINING provides better pretraining signals, enabling LMs better hone their language modeling abilities.
3.3.2 IN-CONTEXT LEARNING FOR TEXT CLASSIFICATION
Datasets & Metrics. In-context learning requires to perform task without fine-tuning by condi- tioning on a few demonstration examples about the task. We evaluate the in-context learnig ability of (CLM using 32 demonstration examples. We use seven text classification datasets, including
5

---

Published as a conference paper at ICLR 2024
Standlard
KNN
ICLM
13.0
11.0
24.0
11.5
9.7
21.2
Perplexity
8.5
7.1
15.6
7.0
5.8
12.8
5.5
4.5
10.0
0.0
2.0
4.0
6.0
8.0
0.0
2.0
6.0
8.0
0.0
2.0
4.0
8.0
Porameters (Billion) (a) Wikipedia
Poromctors (Bi lion)
Parameters (Billion) (b) Books
Anxiv
Figure Language modeling perplexity (the lower better) on Wikipedia, Arxiv, and Books ICLM outperforms the baselines consistently across all model sizes.
3.3.1
Table 1:
In-context learning performance on seven classification datasets (33.22). We use 32 in- context examples for all datasets. ICLM outperforms baselines on all datasets.
Sentiment
Hate Speech
Topic Classification
Method
Average
Amazon
SST2
Yelp
Hate
Offensive
Agnews
Dhpedia
Standard
94.6
83.7
74.3
52.7
55.7
68.3
61.5
66.0
ANN
88.0
80.2
65.1
50.1
53.1
65.7
56.4
61.8
ICLM
96.5
93.2
77.4
60.6
57.3
76.0
63.2
71.3
sentiment analysis (SST-2 (Socher et al., 2013), Amazon and Yelp (Zhang et al. 2015a)), topic classificaiton (AGN (Zhang al. 20155) and Dbepdia (Lehmann al. 20155) and hate speech detection (Barbieri et al. (020). We use label words from Min et al. (2022) and report accuracy as the metric.
Results. As shown Table ICLM consistently demonstrates better performance across all text classification datasets, leading 8% gain average. This result suggests that ICLM is better at learning from demonstration examples. We later analyze the relationship between the number of demonstration examples and the performance of the in-context learning in $4.3
3.3.3 READING COMPREHENSION
Datasets & Metrics. Reading comprehension requires to answer the question based on the given paragraph. We consider the RACE reading comprehension benchmark (RACE-High and RACE- Middle) LLai et al. 2017). SQuAD (Rajpurkar et al. 2016), BoolQ (Clark et al. 2019). DROP (Dua et al. 2019), and HotpoiQA (Yang et al. 2018). We use 2-shot in-context learning for evaluation; we did not use more because some documents in reading comprehension tasks are very long. We for HotpotQA and SQuAD, and accuracy for other datasets that
report the exact match multi-choice tasks (RACE, BoolQ, DROP), following the standard in prior work.
Results. Table highlights that ICLM consistently surpasses both the standard and ANN baselines across all datasets with an average improvement of 14%. In particular, we observe significant gains on HotpotQA, which requires multi-hop understanding of multiple related documents. The performance gain on reading comprehension tasks demonstrates that IN-CONTEXT PRETRAINING improves LMs' ability of undestanding and reasoning the given context.
3.3.4 RETRIEVAL-AUGMENTATIE
Datasets Metrics. Retrieval-augmentation a method to retrieve a set passages from the external text corpus (e.g.. Wikipedia) and prepend it the input query in order to better handle input
queries that require factual knowledge (Lin et al.) 2023) Xu et al. 2023: Su et al. 2023). We conduct well-studied open-domain QA datasets: Natural Questions (NQ) (Kwiatkowski jet al. 2019) and TriviaQA (Joshi et al. 2017). ) For both datasets, we report exact match scores (EM) and evaluate the model performance in both closed-book and open-book settings. In the closed-book
evaluation

---

Published as a conference paper at ICLR 2024
Table 2: Reading comprehension results, using 2-shot in-context learning ( 3.3.3 forms baselines on all six datasets.
ICLM outper-
Method
RACE-High
RACE-Middle
BoolQ
SQuAD
HotpotQA
DROP
Average
Standard
39.5
53.3
68.9
26.3
10.5
27.2
37.6
KNN
36.2
51.4
65.3
23.5
14.4
25.1
36.0
ICLM
41.5
56.9
73.0
30.3
21.9
35.7
43.2
Table 3: Results on NQ and TQA (.3.4) with-
Table 4: Results on two datasets with knowl- edge conflicts, requiring better reasoning of the given context (3..5).
out retrieval (closed) and with retrieval (open).
NQ
TQA
Method
Closed
Open
Closed
Open
Method
NQ-Swap
MemoTrap
Standard
17.0
28.5
49.3
48.1
Standard
39.6
48.4
ANN
13.5
40.2
43.2
ANN
42.1
54.3
ICLM
17.0
32.2
48.0
51.6
ICLM
45.8
56.2
setting, we only provide the question to the model and the model has to answer the question based on its parametric knowledge. In the open-book setting, we follow [Shi et al. (2023c) in providing the model with the top-10 retrieved documents from Wikipedia as additional context to the question.
Results. Results are reported in Table In the closed-book setting, ICLM performs comparably or slightly worse than the standard baseline, likely because our model memorizes less. Nonetheless, in the open-book setting, ICLM significantly outperforms the standard baseline in the open-book setting (+9%), obtaining much better performance than the closed book setting. It is also worth noting that the training objective of &NN is exactly the same as the aerrivallauumenntatin achieves better performance, likely due to the overfitting problem of I:NN discussed in 2.2
but ICLM still
3.3.5 FACTUALITY
Datasets & Metrics. Prior work has found that language models generate text that is not factual nor faithful to the given context, especially when the context contradicts to knowledge the model has
acquired during pretraining (often called parametric knowledge (Longpre et al. [2021; Zhou et al.
2023 Shi et al. 2023b; Wang et al. 2023)). We evaluate LMs' ablities to follow instructions and contexts knowledge conflict datasets: NQ-Swap (Longpre et al., 2021) and MemoTTra (Luu Liu 2023). Both datasets contain instruction and contexts that are in conflict with the models parametrie knowledge. We report exact match score the metric.
Results. Table 4 shows that ICLM is better than the standard and KNN baselines on both datasets, implying that IN CONTEXX PRETRAINING improves LMs' ability to generate outputs that are faithful to prior contexts. Gains are larger than those in other datasets, likely because NQ-Swar and MemoTrap highlight the challenge in reasoning about the given context, which the previous LMs struggle with.
3.3.6 LONG CONTEXT REASONING
Datasets & Metrics. To evaluate the ability of long context reasoning, we compare ICLM with the standard and KNN baselines on the SCROLL benchmark (Shaham et 2022) that evaluates LMs' ability to synthesize information over long texts. Following the original paper setting, we finetune the pretrained LMs (standard, ANN, IN- N-CONTEXT PRETRAINING) on the training datasets of the scroll and evaluate them on the test datasets. We report F1 score for Narrative QA, Qasper and ContractNI.I datasets and report ROUGE-1 score for QMSum and GovReport datasets the SCROLL benchmark.
Results. Results in Table that ICLM outperforms the baselines by around 5%, suggesting that ICLM is better at long context reasoning. We hypothesize that the gains from ICLM may fade
7

---

Published as a conference paper at ICLR 2024
Table Performance on long context reasoning benchmarks from SCROLL (hahaam et 3.3.6 ICLM outperforms baselines on all five datasets.
2022
NarrativeQA
Qasper
ContractNLI
QMSum
GovReport
Method
Average
F1
ROUGE-1
Standard
16.5
34.2
78.6
8.2
32.5
ANN ICLM
16.8
34.1
79.5
24.3
6.6
32.3
17.1
36.7
80.7
26.8
9.1
34,1
Standard
ICLM
Standardd
ICLM
Stanclard
ICLM
2.5
43.0
35.0
2.2
29.0
Training Loss
2.0
38.5
23.0
1.7
36.4
'7.0
1.5
34.2
'1.0
1.2
32.0
5.0
100
Billion of takens
Billion of tokens
Billion of takens (c) NO (Open)
(a) Training Loss
(b) Race-High
Figure 4: Training loss and performance evolution reading comprehension during pretraining. After training on around 150 billion tokens, ICLM is consistently better than the standard LM reading comprehension and retrieval augmentation tasks.
Standard
75.0
Method Design
Choice
PPL
730
Document
Random
Average
67.
Relevance
Clustering
7.9
Links (final)
7.3
62.
630
Semantic
No dedup
8.3
Dedup
Dedup (final)
7.3
Kishot
Figure 6:
Performance with respect to the
Figure 5: Ablation study of our method design.
number of in-context examples (k).
some extent when the LMs are fine-tuned, which may explain the relatively small gains in this evaluation compared to our other experiments.
ANALYSIS
4.1 EVOLUTION OF PERFORMANCE DURING PRETRAINING
Throughout the pretraining process, we closely monitor both the training loss and the downstream task performance for the ICLM as well as the standard LM. Figure 4illustrates the trajectory of the training loss and the performance on the RACE reading comprehension tasks for the 7B models. The training loss for ICLM consistently remains lower than that of the standard L.M This suggests that, when predicting the next token, ICLM benefits from a richer set of relevant prior documents to refer to, while the standard LM has limited information to rely on, leading to higher loss. Figure 4 (b, c) shows that after training on around 150 billion tokens, ICLM is consistently better than the standard LM reading comprehension tasks. This performance gap remains consistent throughout the remainder of the pretraining phase. This suggests the scale of improvements by IN- EX-CONTEXT PRETRAIN AINING does not diminish and remains consistent training more tokens.
8

---

Published as a conference paper at ICLR 2024
4.2 ABLATION STUDY ON IN-CONTEXT PRETRAINING DESIGN
We perform analysis on two design choices of IN-CONTEXT PRETRAINING: choice of methods for finding retrieved documents and deduplication. Ablations are done with 1.5B models and evaluated with perplexity on Wikipedia. The results are presented in Figure 5
Document relevance. A key design of IN-CONTEXT PRETRAINING is grouping documents by their relevance. We consider three levels of relevance: random (the standard baseline discussed in (3.2), clustering, and our document linking method in IN- IN-CONTEXT PRETRAINING. Clustering follows the method from Abbas al. 2023) in clustering documents into 11k clusters based on their embeddings and sample documents from each cluster to form the training inputs. Documents grouped by clustering are sourced from the same clusters, indicating topical similarity but not necessarily close relation. In contrast, ICLM links documents nearest neighbors, indicating a higher degree of similarity. The relevance between documents increases from random, clustering to linking. We observe that the perplexity of the language model decreases the relevance increases.
Deduplication. We compare perplexity of the models trained with and without the semantic deduplication step. Removing the semantic deduplication step leads to a significant decrease in perplexity. When near duplicate documents are present in the same context, language models might merely copy from the prior document, leading to training instability.
4.3 DEMONSTRATION EXAMPLES SIZE FOR IN-CONTEXT LEARNING
We evaluate the 7B models trained with the standard method and IN- C--CTTEET PRETRAINING, using a varying number of demonstration examples on text classification tasks described in $3.3.2 As depicted in Figure 6 ICLM maintains consistent performance gains over the standard method, even as the number of demonstration examples grows. While the performance improves as the number of demonstration examples increases, it plateaus after 32 examples.
5 RELATED WORK
Data batching based on similarity Previous work employs batching lexically similar segments in the same training batches to construct high-quality positive pairs for training -aggeeel- language models. For instance, [Zhong et (2022) use BM25 and same documents to ensure the segments the same batch are similar to each other, while Min et al. (2023) group segments from the same documents in the same batch. Our method shares the same spirit with these methods except we maintain the relevance of documents in the same context window, yet context windows within batches are shuffled. Additionally, our focus is to apply the batching method to train the standard language models.
CONCLUSION
We introduce IN-CONTEXT PRETRAINING, a new pretraining method that learns generate text conditioned on a set of relevant documents, exposing LMs to relevant contexts and providing training signals beyond document boundaries. Our method is highly scalable and simple, and works with any pre-training pipeline by simply changing the document ordering during preprocessing. Our comprehensive evaluation demonstrates our method leads to significant improvements in wide variety of settings that highlight the ability to understand and reason over the given context, including in-context learning, reading comprehension, retrieval augmentation, and more. Future research may delve into the inherent connections between documents within specific corpus domains or using multilingual retriever to group related multilingual documents in the same context. For example, the code scripts within the same repository related. This insight paves the way for future exploration, where concatenating entire repositories into a unified whole could lead to the creation of meaningful long-context data sets.

---

Published as a conference paper at ICLR 2024
REFERENCES
Amro Abbas, Kushal Tirumala, Dániel Simig, Surya Ganguli, and Ari S Morcos. Semdedup: Data- efficient learning web-scale through semantic deduplication. arXiv preprint arXiv:2303.094: 009540, 2023.
Francesco Barbieri, Jose Camacho-Collados Luis Espinosa Anke, and Leonardo Neves. TweetEval: Unified benchmark and comparative evaluation for tweet classification. In Findings of the Asso- ciation for Computational Linguistics: EMNLP 2020, pp. 1644-1650, Online, November 2020. findings-emnlp.1444 URL
Association for Computational Linguistics. doi: 10.
https://aclanthologgy org.//ooooooo////e org/2020 wtt.e.......innnn. meadings-emnlp.. 148
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems. httpss//cccccccceeg..s.o neuruigggnnnurippp volume 33, pp. 1877-1901 Curran Associates, Inc., 2020. 64a Paper. pdf] cc/paper/2020/filee http:////2220011113336e666662022 (bbprrecceb6771bbbaccccccae
Christopher Clark, Kenton Lee, Ming-Wai Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 2924-29366 Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 18 N19-1300. https:/ss:///acll lathoologthoooooooooooo www..org//ioo wrg/o1gg/// 9- 1300
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO -wwaereesss. In Advances in Neural Information Processing Systems, 2022.
Harm de Vries. In the long (context) run, 2023. httpss//wwwharrmmeerrss...o context-Length/
Matthijs Douze, Alexandr Guzhva, Chengqi Deng, Jeff Johnson, Gergely Szilvasy, Pierre-Emmanuel Mazaré, Maria Lomeli, Lucas Hosseini, and Hervé Jégou. The faiss library, 2024.
Dheeru Dua, Yizhong Wang, Pradeep Dasigi. Gabriel Stanovsky, Sameer Singh, and Matt Gardner. DROP: reading comprehension benchmark requiring discrete reasoning over paragraphs. In Proceedings of the 2019 Conference of the North American Chupler of the Association for Com- putational Linguistics: Human Language Technologies, Volume (Long and Short Pupers), pp. 2368 2378, Minncapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N199 1246. URL http:////......eeeeee.eeeearrrrp 22500211112
Avia Efrat and Omer Levy. The turking test: Can language models understand instructions? ArXiv, abs/2010.11 1982, 2020.
Merrill M Flood. The salesmaa problem. Operations research, 4(1):61-75, 1956.
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.
Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Unsupervised dense information retrieval with contrastive learning. Transactions on Machine Learning Research, 2022. httphhtttt///toeeeetttt wtw.ee.eeeooommmmccccm
jKN7pXi7b@
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion- scale similarity search with gpus. 1EEE Transactions Big Data, 7(3):535-547, 2019.
10

---

Published as a conference paper at ICLR 2024
Mandar Joshi, Eunsol Choi, Daniel 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, Vancouver, Canada, July 2017. Association for Computational Linguistics. doi: 10.18653/ v1/P17-1147.
Hervé Jégou, Matthijs Douze, and Cordelia Schmid. Product quantization for nearest neighbor search. IEEE transactions pattern analysis and machine intelligence, 33:117-28, 01 2011. doi: 10.1100/TPAM122000.55.
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-Wei 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:452-466, 2019. doi: 10. 10.1162/tacl_ a_02776. URL
https://aclanntooogg/ oor.aannooooooyooo//11-----0
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. RACE: Large-scale ReAding comprehension dataset from examinations. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing, pp. 785-794, Copenhagen, Denmark, September https
2017. Association for Computational Linguisties. doi: 10. 10.18633/1//D17-1088
1082.
Jens Lehmann, Robert Isele, Max Jakob, Anja Jentzsch, Dimitris Kontokostas, Pablo N Mendes, Sebastian Hellmann, Mohamed Morsey, Patrick Van Kleef, Sören Auer, et al. Dbpedia-a large- scale, multilingual knowledge base extracted from wikipedia. Semantic web, 6(2):167-19)) 2015,
Yoav Levine, Noam Wies, Daniel Jannai, Dan Navon, Yedid Hoshen, and Amnon Shashua. The inductive bias of in-context learning: Rethinking pretraining example design. In International Conference on Learning Representations, 2022. URL hhttps://ppenrevieew wtt.w/iwwwweeeooummmiiiii InEaqbTJIRZ
Xi Victoria Lin, Xilun Chen, Mingda Chen, Weijia Shi, Maria Lomeli, Rich James, Pedro Rodriguez, Jacob Kahn, Gergely Szilvasy, Mike Lewis, Luke Zettlemoyer, and Scott Yih. Ra-dit: Retrieval- augmented dual instruction tuning, 2023.
Alisa Liu and Jiacheng Liu. The memotrap dataset. pww.e//lwbbmain///eee datann/aaaaaaaaaa release/README. md] 2023.
Nelson Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts, 2023.
Shayne Longpre, Kartik Perisetla, Anthony Chen, Nikhil Ramesh, Chris DuBois, and Sameer Singh. Entity-based knowledge conflicts in question answering. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. 7052-7063, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653/ v1/2021 v1/2021.emnlpe www.lp main.
URL https: /tttps://aclantteee
wttggg//2/2 emnlp-main.55 565
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2018
Ian R. McKenzie, Alexander Lyzhov, Michael Pieler, Alicia Parrish, Aaron Mueller, Ameya Prabhu, Euan McLean, Aaron Kirtland, Alexis Ross, Alisa Liu, Andrew Gritsevskiy, Daniel Wurgaft, Derik Kauffman, Gabriel Recchia, Jiacheng Liu, Joe Cavanagh, Max Weiss, Sicong Huang, The Floating Droid, Tom Tseng, Tomasz Korbak, Xudong Shen, Yuhui Zhang, Zhengping Zhou, Najoung Kim, Samuel R. Bowman, and Ethan Perez. Inverse scaling: When bigger isn't better, 2023.
Sewon Min, Mike Lewis, Luke Zettlemoyer, and Hannaneh Hajishirzi. MetalCL: Learning to learn in context. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 2791-2809, Seattle, United States, July 2022. Association for Computational Linguistics. doi: 18.53/ h/2022.naacl...... 201. URL htt.:: wtt..nttsss//acllttoooooooooo//p naw.eaanccccc..aaaa laaal-mainn20)

---

Published as a conference paper at ICLR 2024
Sewon Min, Weijia Shi, Mike Lewis, Xilun Chen, Wen-tau Yih, Hannaneh Hajishirzi, and Luke Zettlemoyer. Nonparametric masked language modeling. In Findings of the Associa- fion for Computational Linguistics: ACL 2023, pp. 2097-2118, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 00.18633/11/0222 htt.///..nnnnnnsaaa.llll. 132. URL
https///accaanthoooy..orrr////// htt.g//222...nnnngge..... Findings-acl. 132
OpenAl. Gpt-4 technical report, 2023.
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback. 2022.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pp. 2383-2392, 2016.
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. arXiv preprint ar arXiv:2211.05100 2022.
Uri Shaham, Elad Segal, Maor Ivgi. Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. SCROLLS: Standardized CompaRison over long language sequences. In Proceedings of the 2022 Conference on Empirical Methods Natural Lan- guage Processing, pp. 12007-12021, Abu Dhabi, United Arab Emirates, December 2022. Associa- tion for Computational Linguistics. URL, https:://claanthtt
wwwgg2222..om emnlp-mminn 823
Freda Shi, Xinyun Chen, Kanishka Misra, Nathan Scales, David Dohan, Ed H Chi, Nathanael Schärli, and Denny Zhou. Large language models can be easily distracted by irrelevant context. In International Conference on Machine Learning, pp. 3 31210-1112222 PMLR, 2023a.
Weijia Shi, Xiaochuang Han, Mike Lewis, Yulia Tsvetkov, Luke Zettlemoyer, and Scott Wen tau Yih. Trusting your evidence: Hallucinate less with context-aware decoding, 2023b.
Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Rich James, Mike Lewis, Luke Zettle- moyer, and Wen-tau Yih. Replug: Retrieval-augmented black-box language models. arXiv preprint arXiv:2301. /2652, 2023c.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pp. 1631-1642, Seattle, Washington, USA, October 2013. Association for Computational Linguistics. URL https:://www wwwweb.wwwww...c.w.ww.www......b org. wto.o///noooooo///111111---
Hongjin Su, Weijia Shi, Jungo Kasai, Yizhong Wang, Yushi Hu, Mari Ostendorf, Wen tau Yih, Noah A. Smith, Luke Zettlemoyer, and Tao Yu. One embedder, any task: Instruction-finetueed text embeddings, 2023.
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 23XX::2302.1977,,, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei. Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cris- tian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian. Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen
12

---

Published as a conference paper at ICLR 2024
Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023b.
Yike Wang, Shangbin Feng, Heng Wang, Weijia Shi, Vidhisha Balachandran, Tianxing He, and Yulia Tsvetkov. Resolving knowledge conflicts in large language models. arXiv preprint arXiv:2310.00935. 2023.
Guillaume Wenzek, Marie-Anne Lachaux, Alexis Conneau, Vishrav Chaudhary, Francisco Guzmán, Armand Joulin, and Édouard Grave. Cenet: Extracting high quality monolingual datasets from web crawl data. In Proceedings of the Twelfth Language Resources and Evaluation Conference, pp. 4003 4013-401,2200 2020.
Fangyuan Xu, Weijia Shi, and Eunsol Choi. Recomp: Improving retrieeal-augmented Ims with compression and selective augmentation, 2023.
Zhilin Yang. Peng Qi, Saizheng Zhang, Yoshua Bengio, William Cohen, Ruslan Salakhutdinov, and Christopher D. Manning. HotpoiQA: A dataset for diverse, explainable multi-hop question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp.
Brussels, Belgium, October-November 2018, Association for Compu- tational Linguisties. doi: 10.18653/v1/D1S 1 URL BtLLLtttss///allanooooyoooooo 181 1259
Michihiro Yasunaga, Armen Aghajanyan, Weijia Shi, Richard James, Jure Leskovec, Percy Liang, Mike Lewis, Luke Zettlemoyer, and Wen-tau Yih. Retrieval-angmmnntd multimodal language modeling. 2023.
Susan Zhang, Mona Diab, and Luke Zettlemoyer. Democratizing access to large-scale language models with opt-175b. Meta AI, 2022.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text classification. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), Advances Neural Information Processing Systems, volume 28. Curran Associates. Inc., 2015a. URL https: #556688B1773FF00066e6676828-20
Xiang Zhang, Junbo Jake Zhao, and Yann LeCun. Character-level convolutional networks for text classification. In NIPS, 2015b.
Zhuosheng Zhang, Hai Zhao, and Rui Wang. Machine reading comprehension: The role of contextu- alized language models and beyond. arXiv preprint 0rXiv:2005.06249 2020.
Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving
few-shot performance of language models. International Conference on Machine Learning, pp- 12697-12706. PMLR, 2021.
Zexuan Zhong, Tao Lei, and Danqi Chen. Training language models with memory augmentation. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 5657-5673, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 122.8663311102222.mmmp.... ma.... URL https:://ccaaanolooogyooooo///// hww.hhhooloyoooogg//go emnlp-main.33 382
Wenxuan Zhou, Sheng Zhang, Hoifung Poon, and Muhao Chen. Context-faithful prompting for large language models. ArXiv, abs/2303. 13155, 2023.
13

---

Published as a conference paper at ICLR 2024
A ADDITIONAL BACKGROUND
A.1 DEDUPLICATION
Corpora often have semantic duplicates: pairs of documents that are semantically related, yet not entirely identical. Previous studies (Yasunaga et al. 2023) show that retraining highly similar documents in the input contexts during training hurts multimodal models' performance. We observed similar behaviur: when near duplicate documents are present in the same context, language models might merely copy from the prior document, leading to training instability. Given that our retrieval approach inherently assesses pairwise document similarity, we can easily filter out near duplicate documents that have high cosine similarity with other documents, We find this deduplication step to be crucial for achieving performance of good language models ($4.
A.2 FAISS INDEX
We used product quantised inverted file (IVFPQ) FAISS index with a code size of 256 and the corresponding number of inverted lists 32768 with total size of 62 gigabytes. The index contains 235266464 (68-dimensional embeddings originally in float 32. The index was trained on sample of 1572864 embeddings and the train time was 423 s. Successively, the data is split in batches of 50M embeddings and for each index shard the corresponding batch of embeddings is added to the trained index, the average adding embeddings time per index shard is 628.4 s. Finally, approximate nearest neighbor search is conducted per each shard before aggregating all results using faiss big batch search. The nprobe used for conducting approximate search is 64, this means that 0.2% of the inverted lists probed during the nearest neighbors search.
14


Reviews:
[{"review_id": "v0yjaHNxLz", "reviewer": "Reviewer_8aRp", "summary": "This paper introduces In-Context Pretraining for Language Models (ICLM), a novel approach that connects relevant documents during the pretraining of LLM. In contrast to existing pretraining methods that mainly use random concatenation, ICLM employs a more targeted approach.\n\nGiven the substantial number of pretraining documents, this paper presents an approximate algorithm for efficiently identifying related documents.\n\nThe experimental results of ICLM surpass those of the traditional method and a KNN-based approach across various tasks, particularly excelling in long-form reasoning tasks such as Hotpot QA and Drop.", "questions": "1. Are you using a multilingual Contriver or an English Contriver? I'm thinking that if you are using the multilingual retrieval, the LLM (Large Language Model) may be better for machine translation or multilingual downstream tasks, which is a strong point of the proposed method.\n\n2. Do you want to discuss the code data and the proposed method? Maybe for the code data, each repository should be concatenated together instead of relying on retrieval. Is this a possible limitation of the proposed method?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This paper addresses an evident and critical issue in LLM pretraining. \nThe experiments clearly demonstrate the effectiveness of the proposed method.", "weaknesses": "1. Lack of comparison with existing works. Some previous studies have used hyperlinks or timestamps to group related documents together. Although these papers are discussed in the Related Work section, it would be beneficial to understand how the proposed method compares to these existing approaches by reporting the results.\n\n2. Lack of information about the computational and time costs of the Retrieval process and DOCUMENT GRAPH TRAVERSAL process. It would be helpful to demonstrate the relationship between the size of the pretraining data and the time cost of the proposed algorithm.", "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": "UJkxWMtUQh", "reviewer": "Reviewer_AJRu", "summary": "A naive approach to language modeling on large corpora often considers randomly concatenating documents within the corpora together during training, resulting in long contexts which may not be relevant to the current next-word prediction. In this work, the authors conjecture that this naive approach harms a models ability to learn to attend to long-range dependencies during both language modeling, as well as down-stream tasks. To address this, the authors propose In-context pretraining (ICP), which, prior to language model pre-training, groups and sorts document by their similarity such that a language model is trained on several related documents contiguously.\n\nTo perform this grouping, ICP first uses _contreiver_ (a similarity-based retriever model from prior work) to first obtain a set of nearest neighbors for each document. Then the authors propose a modified version of a greedy algorithm to the traveling salesman problem to obtain a path through all the documents such that no documents repeat but most documents will appear next to their nearest neighbors. The resulting path is then treated as the document order for input to the LM, which is split into several contiguous chunks which serve as the pre-training corpora.\n\nThe authors consider two baselines: naive pre-training, which randomly shuffles and concatenates documents together, and kNN pre-training, which groups documents by similarity, but can have high document duplication and overlap.\nThe authors demonstrate that LLMs trained with ICP significantly outperform standard LMs and kNN LMs on both language modeling, as well as on several downstream tasks including in-context learning, reading comprehension, and retrieval augmentation, highlighting the importance of long-range consistency and dependency in LM pre-training.", "questions": "What is the difference between the results presented in 4.2 for Document Relevance, and the results presented in Figure 3 (a)? It seems to me as though this ablation analysis was already performed in the main results section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The method, and proposed algorithm, are very simple to understand but significantly outperform the reasonable baselines across several downstream evaluations.\n- The paper is very well written and clear. The core hypothesis and motivation are both easy to understand and reasonable.\n- The paper considers models that go up to a rather large size, up to 7B parameter models which are pre-trained from scratch on 306B tokens. While this is not state-of-the-art in terms of model size, these are nevertheless very convincing experimental settings.\n- There is some ablation study showing that the deduplication strategy is very important for the observed results, and that the benefits of ICP arise consistently over training (are not explained by variance).", "weaknesses": "- The core methodological contribution of the paper is perhaps relatively small. The novelty of the work comes from the identification that data duplication in kNN-based groupings is a problem, and in section 2.2 which addresses the data duplication problem via a greedy algorithm. Most of the paper instead covers the (extensive) evaluation of the proposed method.\n- The analysis section feels quite weak, outside of the ablation study over data deduplication (I'm not sure if document relevance is a repeat result, see questions). I'm not sure how much section 4 really helps us understand why ICLM works better than Standard LM, which feels important given how simple the proposed method really is.", "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": "MvtpKXu9yj", "reviewer": "Reviewer_Gjsg", "summary": "In this study, the authors argue that the conventional method of sequence construction, which involves random sentence concatenation from disparate documents, fails to provide adequate training signals for language pre-training.\nThey introduce an alternative approach termed 'in-context pre-training,' which concatenates relevant documents identified via a retrieval model. Specifically, they first construct a document graph in which the documents are nodes and edges are valued by the similarity between two documents. Then they employ a greedy strategy to concatenate the documents, i.e., constructing the sequence for language pre-training.\nEmpirical evaluations across multiple downstream tasks demonstrate that in-context pre-training consistently surpasses the performance of the baseline, except on close-booked question answering.", "questions": "I would appreciate it if the authors could provide answers to the following questions:\n1. Could the authors offer additional information regarding the 'standard way' of constructing the sequence? To my understanding, it's common to create a sequence using texts from successive documents[3], wherein the texts are naturally semantically relevant to each other. Is the approach different in this case because the sequence length is so extensive that most documents cannot fill a single sequence? I would appreciate further clarification on this matter.\n2. Could you provide an estimate of the time required for the graph construction and traversal processes at your experimental scale?\n3. I am intrigued by the scalability of the proposed approach in a practical context. If we receive new raw text data, will we have to repeat the similarity computation, graph construction and graph traversal processes?\n\n[3]  Liu et al. RoBERTa: A Robustly Optimized BERT Pretraining Approach. 2019", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. Clear motivation and presentation: enhancing LM’s understanding of context is intriguing and makes a lot of sense to me. The authors effectively articulate the problem, motivation, and proposed solution, making it easy to follow and comprehend the paper.\n2. Strong empirical results and interesting findings: the comprehensive experiments provide compelling evidence that in-context pre-training (achieved by simply constructing sequence with relevant documents) yields improvements across diverse downstream tasks, with the exception of close-booked question answering.\n3. Reproducibility: the proposed approach is straightforward to reimplement, requiring only minor modifications to existing pre-training procedures for further exploration and validation. The traversal process is also clearly delineated.", "weaknesses": "This paper  could further be enhanced from the following perspectives:\n1. In Section 3.3.4, the authors attribute the inconsistent performance in a closed-book setting to the “ICLM memorizes less”. This claim could benefit from further elaboration. As we see better perplexity on language modeling, seems ICLM should memorizes more with the help of relevant documents.\n2. According to the presented results, ICLM outperforms baselines in in-context settings. Experiments in out-of-context setting could be added for us to better understand the proposed ICLM.\n3. The related work section could be improved by discussing a broader range of previous studies.  For example, the first paragraph of Section 5 only refers to three papers to summarize ‘pre-training with related documents’, however, there are other works with similar ideas including [1] use dictionary as context, [2] use co-occurrence of words across sentences as context for language pre-training, among others.\n\n[1] Yu et al. Dict-BERT: Enhancing Language Model Pre-training with Dictionary. ACL Findings 2022.\n\n[2] Wu et al. Taking Notes on the Fly Helps Language Pre-Training. ICLR 2021.", "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": "6dnnlZvWhZ", "reviewer": "Reviewer_MeqG", "summary": "The paper introduces in-context pretraining, a new method to pretrain large language models on a sequence of related documents, improving their ability to read and reason across document boundaries. Authors present an algorithm to find related documents at web-scale and construct coherent input contexts for pretraining LLMs. The paper also shows that in-context pretraining leads to significant improvements on various tasks that require complex contextual reasoning, such as in-context learning, reading comprehension, factuality, and long context reasoning etc.", "questions": "NA", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper introduces a novel method of in-context pretraining of LLMs on a sequence of related documents. This approach is innovative as it enhances the LMs’ ability to read and reason across document boundaries. The paper is well-structured and clear. It provides a detailed explanation of the in-context pretraining method, the document sorting problem, and the experimental results. The significance of this work is evident in its potential impact on various tasks that require complex contextual reasoning. The paper shows that in-context pretraining leads to significant improvements in these areas.", "weaknesses": "While the paper presents some experimental results, more evidence could strengthen its claims. Conducting additional experiments or providing more detailed analysis of the existing results could enhance the credibility of the findings. Meanwhile, case studies are also necessary for better understanding the improvements from the proposed method.", "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 (spotlight)