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. Smith, Luke Zettlemoyer, Scott Yih, Mike Lewis
分类: cs.CL, cs.AI, cs.LG
发布日期: 2023-10-16 (更新: 2024-06-24)
💡 一句话要点
提出上下文预训练以解决文档边界推理问题
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 上下文预训练 语言模型 文档推理 信息检索 自然语言处理 复杂推理 文本生成
📋 核心要点
- 现有的语言模型预训练方法无法有效利用文档间的上下文信息，导致推理能力不足。
- 本文提出上下文预训练，通过对相关文档序列进行预训练，增强模型跨文档推理能力。
- 实验显示，采用上下文预训练后，模型在复杂推理任务上性能显著提升，如阅读理解提升15%。
📝 摘要（中文）
大型语言模型（LMs）目前通过预测文档前缀下的标记进行训练，从而直接执行长文本生成和提示式任务。然而，现有的预训练流程通过随机拼接短文档来创建输入上下文，导致前文档无法为下一个文档的预测提供信号。本文提出了一种新的方法——上下文预训练，旨在通过对相关文档序列进行预训练，明确鼓励模型跨文档边界进行阅读和推理。我们通过改变文档排序，使每个上下文包含相关文档，从而实现上下文预训练。实验结果表明，该方法显著提升了模型在复杂上下文推理任务中的表现。
🔬 方法详解
问题定义：本文解决的问题是如何有效利用文档间的上下文信息进行语言模型的预训练。现有方法通过随机拼接短文档，导致模型无法学习到文档间的相关性，从而影响推理能力。
核心思路：论文的核心思路是通过上下文预训练，使语言模型在预训练阶段就能学习到相关文档的上下文信息，从而提高其在推理任务中的表现。通过改变文档的排序，使得每个上下文包含相关文档，模型能够更好地理解和推理。
技术框架：整体架构包括文档排序模块和预训练模块。首先，通过近邻搜索算法找到相关文档，然后利用图遍历算法构建连贯的输入上下文，最后将这些上下文输入到现有的预训练管道中进行训练。
关键创新：最重要的技术创新点在于提出了上下文预训练的概念，并通过文档排序和相关文档选择的算法实现了这一目标。这与现有方法的本质区别在于，现有方法未能考虑文档间的关系。
关键设计：在参数设置上，采用了高效的近邻搜索算法来寻找相关文档，并使用图遍历算法来确保输入上下文的连贯性。此外，损失函数和网络结构与现有预训练模型保持一致，以便于直接应用。
🖼️ 关键图片
📊 实验亮点
实验结果表明，采用上下文预训练后，模型在多个复杂推理任务上表现显著提升，包括上下文学习提升8%、阅读理解提升15%、对先前上下文的忠实度提升16%、长上下文推理提升5%以及检索增强提升9%。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理中的文本生成、问答系统和信息检索等。通过提升语言模型的推理能力，能够在更复杂的任务中提供更准确的结果，具有重要的实际价值和广泛的应用前景。
📄 摘要（原文）
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%).