Paper: In-Context Pretraining: Language Modeling Beyond Document Boundaries (ICLR-2024)

Reviewer comment (Evaluation):
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.

Author response:
This is a good point, thanks. It is true the ICLM and memorize more if you consider the documents in the datastore part of the complete model. This also has other advantages often seen in retrieval augmented methods. We will rewrite to clarify this point.

Reviewer comment (Experiments):
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.

Author response:
| Dataset     | Standard | ICLM  |
|-------------|----------|-------|
| RACE-High   | 36.4     | **37.8**  |
| RACE-Middle | 48.5     | **50.2**  |
| BoolQ       | 62.1     | **63.3**  |
| Amazon      | 83.4     | **83.5**  |

These results demonstrate that ICLM performs better than the standard method in zero-shot settings. We will conduct more extensive evaluation in zero-shot settings and include the results in the final version of the paper.

Reviewer comment (Novelty):
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.

Author response:
We thank the reviewer for pointing us to related works. [1][2] improves LMs’ understanding of rare words, where [1] integrates definitions of these words, and [2] includes contextual vectors from previously encountered contexts mentioning the same words. **We have incorporated these papers into our related work section (highlighted in yellow).**

Reviewer comment (Experiments):
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.

Author response:
Yes, the approach is different in this case because the sequence length is so extensive that most documents cannot fill a single sequence. The standard pretraining method concatenate documents that are randomly ordered, as seen in [3]. [3] creates a sequence using **consecutive sentences within a document**. ​​However, upon reaching a document’s end, the method still requires switching to a new, randomly selected document.  [3] pretrains language models with short context windows (512 tokens long). The context window size is short enough that a single document could potentially occupy the entire window.

However, with the trend of extending context window sizes in language model, the standard pretraining method encounters issues. The average pretraining document spans about 1,000 tokens, while newer models like LLAMA employ a 4,000-token context window and ChatGPT takes up to 16,000 tokens. In such scenarios, the standard pretraining method, which concatenates random documents, poses significant challenges in maintaining contextual coherence and relevance. Through our evaluation of downstream tasks, even if the task of interest does not necessarily require a long context, our results show that our in-context pretraining is better because it enhances the ability to read and understand relevant context more effectively.

Reviewer comment (Evaluation):
Could you provide an estimate of the time required for the graph construction and traversal processes at your experimental scale?

Author response:
As detailed in Appendix A.2 and Section 3.1, we have outlined the specific time and resource requirements for our proposed algorithm. To elaborate:

* The retrieval index building process requires approximately 7 mins.
* The search step is conducted on 32 GPUs and takes around 6 hours for processing 300 billion tokens (235 million documents).
* The data sorting phase requires 12 hours on a setup of 20 CPUs.

**We have updated Section 3.1 with a few additional details, which are highlighted in yellow for easy identification.**

Reviewer comment (Reproducibility):
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?

Author response:
Regarding the scalability of our approach when incorporating new raw text data, our method is designed to be efficient and practical. There is no need to restart the entire process. Instead, we can simply conduct a similarity search for the new data and integrate these new documents into our existing document path, based on their similarity with the old data.

Reviewer comment (Novelty):
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.

Author response:
We respectfully disagree with the reviewer that our methodological contribution is relatively small. A key strength of our method is it is simple and effective. It requires only modifications to the document ordering during pretraining, making it directly applicable to existing pretraining data pipelines. Despite its simplicity, our approach has demonstrated substantial improvements in downstream task performance.

Reviewer comment (Evaluation):
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.

Author response:
In our ablation study for ICLM, we study different levels of document relevance from least to most relevant: random, clustering, and links built by ICLM. In the clustering scenario, documents grouped in the same context are sourced from the same clusters, indicating topical similarity but not necessarily close relation. In contrast, ICLM links documents as nearest neighbors, indicating a higher degree of similarity.

Different from Figure 3 (a), Section 4.2 provides an in-depth exploration of how varying degrees of document similarity – categorized into three levels: random, clustering (documents on the same topic), and links (closely related) – affect the performance of LMs. **We have updated Section 4.2 ablation study with a few additional details, which are highlighted in yellow for easy identification.**

Reviewer comment (Presentation):
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.

Author response:
In our ablation study for ICLM, we study different levels of document relevance from least to most relevant: random, clustering, and links built by ICLM. In the clustering scenario, documents grouped in the same context are sourced from the same clusters, indicating topical similarity but not necessarily close relation. In contrast, ICLM links documents as nearest neighbors, indicating a higher degree of similarity.

Different from Figure 3 (a), Section 4.2 provides an in-depth exploration of how varying degrees of document similarity – categorized into three levels: random, clustering (documents on the same topic), and links (closely related) – affect the performance of LMs. **We have updated Section 4.2 ablation study with a few additional details, which are highlighted in yellow for easy identification.**

Reviewer comment (Experiments):
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.

Author response:
We agree with the reviewer that using hyperlinks is a great idea. However, as we mention in the related work section, hyperlinks may not always be available across all domains. For instance, a part of the commomncrawl pretraining corpus includes short flash fictions from sources like 365tomorrows.com, where individual stories typically lack hyperlinks. Our proposed retrieval method can generalize to corpora with potentially any text in any domain.

Reviewer comment (Experiments):
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.

Author response:
As detailed in Appendix A.2 and Section 3.1, we have outlined the specific time and resource requirements for our proposed algorithm. To elaborate:

* The retrieval index building process requires approximately 7 mins.
* The search step is conducted on 32 GPUs and takes around 6 hours for processing 300 billion tokens (235 million documents).
* The document graph traversal phase requires 12 hours on a setup of 20 CPUs.

**We have updated Section 3.1 with a few additional details, which are highlighted in yellow for easy identification.**

Reviewer comment (Experiments):
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.

Author response:
We are using English Contriever. Thanks for your suggestion. We agree that a multilingual approach could be useful and we would consider it for future research. **We have updated Section 6, the Conclusion to discuss this idea as a future work (highlighted in yellow)**

Reviewer comment (Experiments):
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?

Author response:
Thank you for suggesting this idea. We agree that concatenating code data from the same repository could be an effective approach. **We have revised Section 6, the Conclusion, to include a discussion of this idea as a future work (highlighted in yellow).**