[Paper Review] In-context Pretraining: Language Modeling Beyond Document Boundaries

Weijia Shi, Sewon Min+10 more|arXiv (Cornell University)|Oct 16, 2023

Topic ModelingComputer Science3 citations

TL;DR

This paper introduces In-Context Pretraining, a scalable method that reorders pretraining data to group semantically related documents together, enabling language models to learn reasoning across document boundaries. By using approximate nearest neighbor search and graph-based sorting, the approach boosts performance on in-context learning (+8%), reading comprehension (+15%), and faithfulness (+16%) without modifying existing pretraining pipelines.

ABSTRACT

Large language models (LMs) are currently trained to predict tokens given document 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 to create input contexts but the prior documents provide no signal for predicting the next document. We instead present In-Context Pretraining, a new approach where language models are pretrained on a 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 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 input contexts with a graph traversal algorithm. Our experiments show In-Context Pretraining offers a simple and scalable approach to significantly enhance LMs'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%).

Motivation & Objective

• To address the limitation of standard pretraining, where randomly concatenated documents provide no cross-document signal, by reordering training data to group semantically related documents.
• To enable language models to learn contextual reasoning beyond individual document boundaries through a simple, scalable reordering of pretraining data.
• To develop efficient approximate algorithms for large-scale document retrieval and coherent context construction without data repetition.
• To evaluate whether training on semantically coherent, multi-document contexts improves downstream tasks requiring complex reasoning.
• To demonstrate that this method enhances language modeling and reasoning capabilities without altering existing pretraining pipelines.

Proposed method

• Reordering pretraining data by grouping semantically similar documents using a retrieval model and efficient approximate nearest neighbor search in embedding space.
• Constructing coherent input contexts via a graph traversal algorithm that treats document sorting as a traveling salesman problem to maximize similarity and avoid repetition.
• Leveraging pre-existing pretraining pipelines by only changing document ordering during data preprocessing, preserving all other training components.
• Using a semantic retrieval model to identify nearest-neighbor documents based on embedding similarity, enabling scalable construction of related document sequences.
• Applying the same pretraining objective (masked language modeling) on the reordered, contextually coherent inputs to train models to reason across documents.
• Scaling the method to billions of documents using optimized indexing and approximate search to maintain efficiency and coherence.

Figure 1: Overview of In-Context Pretraining . Different from the standard pretraining strategy that place randomly shuffled documents in the input context, In-Context Pretraining places related documents in the same context, making models learn to reason across prior documents. For example, when pr

Experimental results

Research questions

• RQ1 Can training language models on sequences of semantically related documents improve their ability to reason across document boundaries?
• RQ2 How can we efficiently scale document reordering to billions of documents while maximizing contextual similarity and avoiding data repetition?
• RQ3 Does In-Context Pretraining lead to measurable improvements in downstream tasks requiring complex contextual understanding?
• RQ4 How does the performance of In-Context Pretraining compare to standard pretraining on tasks like in-context learning and reading comprehension?
• RQ5 Can this method be integrated into existing pretraining pipelines without architectural changes?

Key findings

• In-Context Pretraining improves in-context learning performance by an average of 8% across 8 benchmark datasets.
• Reading comprehension accuracy improves by 15% on average across 8 downstream reading comprehension tasks.
• Model outputs are 16% more faithful to prior context, reducing hallucination and inconsistency.
• Long-context reasoning tasks see a 5% performance gain, indicating better utilization of extended context windows.
• Retrieval-augmented generation tasks benefit from a 9% improvement when using external knowledge from Wikipedia.
• The method achieves these gains without modifying the pretraining objective or architecture, only by reordering training data.

Figure 2: Illustration of In-Context Pretraining . In-Context Pretraining first finds related documents at scale to create a document graph (§ 2.1 ) and then builds pretraining input contexts by traversing the document graph (§ 2.2 ). Along the path, documents are concatenated into a sequence and su