[论文解读] In-context Pretraining: Language Modeling Beyond Document Boundaries

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

Topic ModelingComputer Science被引用 3

一句话总结

本文提出了一种可扩展的上下文预训练方法（In-Context Pretraining），通过将语义相关的文档分组重组预训练数据，使语言模型能够学习跨文档边界的推理能力。该方法利用近似最近邻搜索与基于图的排序算法，在不修改现有预训练流程的前提下，显著提升了少样本提示学习（+8%）、阅读理解（+15%）和生成忠实度（+16%）的性能。

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%).

研究动机与目标

• 为解决标准预训练中随机拼接文档导致缺乏跨文档信号的问题，通过重新排序训练数据，将语义相关的文档分组。
• 通过简单且可扩展的预训练数据重排方法，使语言模型能够学习超越单个文档边界的上下文推理能力。
• 开发高效的近似算法，实现大规模文档检索与无重复的连贯上下文构建。
• 评估在语义连贯的多文档上下文中进行训练是否能提升需要复杂推理的下游任务表现。
• 证明该方法可在不改变现有预训练流程的前提下，增强语言建模与推理能力。

提出的方法

• 利用检索模型与嵌入空间中的高效近似最近邻搜索，对预训练数据进行重排，将语义相似的文档分组。
• 通过图遍历算法构建连贯输入上下文，将文档排序问题视为旅行商问题，以最大化相似性并避免重复。
• 通过仅在数据预处理阶段改变文档顺序，保留所有其他训练组件，从而复用现有预训练流水线。
• 使用语义检索模型基于嵌入相似性识别最近邻文档，实现可扩展的相关文档序列构建。
• 在重排后的、语义连贯的输入上应用相同的预训练目标（掩码语言建模），使模型能够学习跨文档推理。
• 通过优化索引与近似搜索技术，将该方法扩展至数十亿文档，同时保持效率与连贯性。

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

实验结果

研究问题

• RQ1 在语义相关的文档序列上进行训练，是否能提升语言模型在文档边界间的推理能力？
• RQ2 如何在最大化上下文相似性并避免数据重复的前提下，高效地将文档重排扩展至数十亿规模？
• RQ3 上下文预训练是否能在需要复杂上下文理解的下游任务中带来可测量的性能提升？
• RQ4 在少样本提示学习与阅读理解等任务中，上下文预训练相较于标准预训练的性能表现如何？
• RQ5 该方法能否在不进行架构修改的前提下，无缝集成到现有预训练流水线中？

主要发现

• 上下文预训练在8个基准数据集上的少样本提示学习性能平均提升8%。
• 在8项下游阅读理解任务中，阅读理解准确率平均提升15%。
• 模型输出的忠实度较之前文上下文提高16%，显著减少幻觉与不一致性。
• 长上下文推理任务性能提升5%，表明模型更有效地利用了扩展的上下文窗口。
• 在检索增强生成任务中，当结合来自维基百科的外部知识时，性能提升9%。
• 该方法在不修改预训练目标或模型架构的前提下，仅通过重排训练数据即实现性能增益。

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