Long-Context Language Modeling with Parallel Context Encoding

📄 arXiv: 2402.16617v2 📥 PDF

作者: Howard Yen, Tianyu Gao, Danqi Chen

分类: cs.CL

发布日期: 2024-02-26 (更新: 2024-06-11)

备注: ACL 2024. Code, models, and data are available at https://github.com/princeton-nlp/CEPE

💡 一句话要点

提出CEPE框架以扩展长上下文语言模型的处理能力

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长上下文处理 语言模型 并行编码 上下文扩展 自然语言处理 信息检索 模型优化

📋 核心要点

  1. 现有的长上下文语言模型在处理长输入时面临计算成本高和位置编码泛化能力不足的挑战。
  2. 本文提出的CEPE框架通过小型编码器逐块处理长输入，利用交叉注意力扩展上下文窗口。
  3. CEPE在处理8K-token文档时，将LLAMA-2的上下文窗口扩展至128K tokens，提升了10倍的吞吐量，且内存使用仅为1/6。

📝 摘要（中文）

扩展大型语言模型（LLMs）以处理更长输入对于多种应用至关重要。然而，变压器的计算成本高和位置编码的泛化能力有限，限制了上下文窗口的大小。本文提出了上下文扩展与并行编码（CEPE）框架，可以应用于任何现有的解码器-only LLMs，以扩展其上下文窗口。CEPE利用小型编码器逐块处理长输入，使得冻结的解码器能够通过交叉注意力利用额外的上下文。CEPE在语言建模和上下文学习中表现出色，并在检索增强应用中优于现有的长上下文模型。我们还介绍了一个CEPE变体，可以仅使用未标记数据扩展指令调优模型的上下文窗口，并展示了其在LLAMA-2-CHAT上的有效性。

🔬 方法详解

问题定义：本文旨在解决大型语言模型在处理长输入时的上下文窗口限制，现有方法在计算效率和位置编码泛化能力上存在不足。

核心思路：CEPE框架通过引入小型编码器逐块处理长输入，使得解码器能够在不改变的情况下利用更多上下文信息，从而扩展上下文窗口。

技术框架：CEPE的整体架构包括一个小型编码器和一个冻结的解码器，编码器负责处理长输入并生成上下文，解码器通过交叉注意力机制整合这些上下文信息。

关键创新：CEPE的主要创新在于其并行编码的设计，使得长输入的处理更加高效，且能够在现有模型上无缝集成，显著提升了上下文处理能力。

关键设计：在CEPE中，使用了特定的参数设置以优化编码器的性能，损失函数设计上也考虑了上下文信息的有效利用，确保了模型在长上下文任务中的表现。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

CEPE在实验中将LLAMA-2的上下文窗口从8K扩展至128K tokens，提升了10倍的处理吞吐量，并且内存使用仅为1/6。该框架在语言建模和上下文学习任务中表现优异，尤其在检索增强应用中超越了现有长上下文模型。

🎯 应用场景

CEPE框架具有广泛的应用潜力，特别是在需要处理长文本的自然语言处理任务中，如文档摘要、长篇对话生成和信息检索等。其高效的上下文处理能力将推动相关领域的研究和应用发展，提升用户体验和模型性能。

📄 摘要（原文）

Extending large language models (LLMs) to process longer inputs is crucial for a wide range of applications. However, the substantial computational cost of transformers and limited generalization of positional encoding restrict the size of their context window. We introduce Context Expansion with Parallel Encoding (CEPE), a framework that can be applied to any existing decoder-only LLMs to extend their context window. CEPE employs a small encoder to process long inputs chunk by chunk, enabling the frozen decoder to utilize additional contexts via cross-attention. CEPE is efficient, generalizable, and versatile: trained with 8K-token documents, it extends the context window of LLAMA-2 to 128K tokens, offering 10x the throughput with only 1/6 of the memory. CEPE yields strong performance on language modeling and in-context learning. CEPE also excels in retrieval-augmented applications, while existing long-context models degenerate with retrieved contexts. We further introduce a CEPE variant that can extend the context window of instruction-tuned models using only unlabeled data, and showcase its effectiveness on LLAMA-2-CHAT, leading to a strong instruction-following model that can leverage very long contexts on downstream tasks.