Extending LLMs' Context Window with 100 Samples

📄 arXiv: 2401.07004v1 📥 PDF

作者: Yikai Zhang, Junlong Li, Pengfei Liu

分类: cs.CL

发布日期: 2024-01-13

🔗 代码/项目: GITHUB

💡 一句话要点

提出新方法扩展LLMs上下文窗口以解决长输入问题

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

关键词: 大型语言模型 上下文窗口 旋转位置嵌入 注意力机制 长文本处理 模型优化

📋 核心要点

  1. 现有方法在扩展LLMs的上下文窗口时资源消耗大，且缺乏有效的比较实验。
  2. 本文提出通过调整RoPE的基频和缩放注意力对数来扩展上下文窗口，旨在提高效率和稳定性。
  3. 实验结果显示，所提方法在多种任务中表现优异，能够显著提升LLMs的上下文处理能力。

📝 摘要（中文）

大型语言模型（LLMs）在其预训练的上下文窗口之外的外推能力有限，限制了其在长输入下游任务中的应用。近期研究通过修改旋转位置嵌入（RoPE）来扩展LLMs的上下文窗口，但现有方法如位置插值（PI）和YaRN资源消耗大且缺乏比较实验。本文提出了一种新方法，通过调整RoPE的基频和缩放注意力对数，帮助LLMs高效适应更大的上下文窗口。实验表明，该方法在不同上下文窗口大小的任务中表现出色，能将LLaMA-2-7B-Chat的上下文窗口扩展至16,384，仅需100个样本和6个训练步骤，展示了卓越的效率。

🔬 方法详解

问题定义：本文旨在解决大型语言模型在长输入任务中的上下文窗口限制问题。现有方法如位置插值和YaRN在资源消耗和适用性上存在不足，导致无法有效扩展上下文窗口。

核心思路：论文提出了一种新颖的RoPE扩展方法，通过调整RoPE的基频和缩放注意力对数，帮助LLMs更高效地适应更大的上下文窗口。这种设计旨在保持注意力熵的稳定性，从而提高模型的性能。

技术框架：整体方法包括两个主要模块：一是调整RoPE的基频，二是缩放注意力对数。这两个模块共同作用，使得模型能够在较少的样本和训练步骤下，扩展上下文窗口。

关键创新：最重要的技术创新在于结合了RoPE的基频调整与注意力对数缩放，形成了一种新的上下文窗口扩展机制。这与现有方法的根本区别在于其高效性和低资源消耗。

关键设计：在参数设置上，本文通过实验确定了RoPE的最佳基频和注意力对数的缩放比例。此外，损失函数和网络结构的设计也经过优化，以确保模型在扩展上下文窗口时的稳定性和性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，所提方法在不同上下文窗口大小的任务中表现优异，能够将LLaMA-2-7B-Chat的上下文窗口扩展至16,384，仅需100个样本和6个训练步骤，相较于传统方法显著提高了效率和稳定性。

🎯 应用场景

该研究的潜在应用领域包括对话系统、文本生成和信息检索等需要处理长文本的任务。通过有效扩展LLMs的上下文窗口，能够提升这些系统的理解和生成能力，进而提高用户体验和系统的实用性。未来，该方法可能在更多复杂的自然语言处理任务中发挥重要作用。

📄 摘要（原文）

Large Language Models (LLMs) are known to have limited extrapolation ability beyond their pre-trained context window, constraining their application in downstream tasks with lengthy inputs. Recent studies have sought to extend LLMs' context window by modifying rotary position embedding (RoPE), a popular position encoding method adopted by well-known LLMs such as LLaMA, PaLM, and GPT-NeoX. However, prior works like Position Interpolation (PI) and YaRN are resource-intensive and lack comparative experiments to assess their applicability. In this work, we identify the inherent need for LLMs' attention entropy (i.e. the information entropy of attention scores) to maintain stability and introduce a novel extension to RoPE which combines adjusting RoPE's base frequency and scaling the attention logits to help LLMs efficiently adapt to a larger context window. We validate the superiority of our method in both fine-tuning performance and robustness across different context window sizes on various context-demanding tasks. Notably, our method extends the context window of LLaMA-2-7B-Chat to 16,384 with only 100 samples and 6 training steps, showcasing extraordinary efficiency. Finally, we also explore how data compositions and training curricula affect context window extension for specific downstream tasks, suggesting fine-tuning LLMs with lengthy conversations as a good starting point. We release our code and SFT data at https://github.com/GAIR-NLP/Entropy-ABF.