Effective Length Extrapolation via Dimension-Wise Positional Embeddings Manipulation

📄 arXiv: 2504.18857v1 📥 PDF

作者: Yi Lu, Wanxu Zhao, Xin Zhou, Chenxin An, Chenglong Wang, Shuo Li, Yuming Yang, Jun Zhao, Tao Ji, Tao Gui, Qi Zhang, Xuanjing Huang

分类: cs.CL, cs.AI

发布日期: 2025-04-26

💡 一句话要点

提出维度感知位置编码操控(DPE)，无需训练即可有效外推LLM上下文长度。

🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本建模 上下文扩展 位置编码 RoPE 大型语言模型

📋 核心要点

  1. 现有长文本扩展方法需要大量计算资源来训练大规模模型，成本高昂，限制了其应用。
  2. DPE通过分析RoPE不同维度，自适应调整位置编码，无需额外训练即可扩展LLM的上下文窗口。
  3. 实验表明，DPE显著优于现有方法，并能提升模型在训练长度内的性能，甚至超越商业模型。

📝 摘要（中文）

大型语言模型(LLM)在处理和生成超过预训练长度的连贯上下文时常常遇到困难。最近在长上下文扩展方面的进展显著扩大了LLM的上下文窗口，但需要昂贵的开销来训练具有更长上下文的大规模模型。本文提出维度感知位置编码操控(DPE)，这是一个无需训练的框架，通过深入研究RoPE的不同隐藏维度来外推LLM的上下文窗口。DPE不是平等地操控所有维度，而是检测每个维度的有效长度，并找到用于上下文扩展的关键维度。我们重用来自预训练模型的原始位置索引及其嵌入，并将关键维度的位置索引操控到其最有效长度。通过这种方式，DPE以最小的修改调整预训练模型，同时确保每个维度达到其外推的最佳状态。DPE显著超越了YaRN和Self-Extend等知名基线。DPE使Llama3-8k 8B能够支持128k tokens的上下文窗口而无需持续训练，并与Flash Attention 2无缝集成。除了其令人印象深刻的外推能力外，DPE还显著提高了模型在训练长度内的性能，例如Llama3.1 70B，在流行的长上下文基准RULER上提高了超过18个点。与商业模型相比，带有DPE的Llama 3.1 70B甚至实现了比GPT-4-128K更好的性能。

🔬 方法详解

问题定义：现有方法在扩展LLM的上下文窗口时，通常需要重新训练模型，这带来了巨大的计算成本和时间开销。此外，简单地扩展所有维度可能并非最优，因为不同维度对于捕捉长距离依赖关系的重要性可能不同。因此，如何高效且有效地扩展LLM的上下文窗口，同时避免昂贵的重新训练，是本文要解决的核心问题。

核心思路：DPE的核心思路是，并非所有位置编码维度都同等重要，某些维度可能更擅长捕捉长距离依赖关系。因此，DPE通过分析RoPE的不同维度，识别出对于上下文扩展至关重要的“关键维度”，并针对这些维度进行选择性地操控，使其达到最有效的长度。这种维度感知的处理方式，避免了对所有维度进行统一处理的低效性。

技术框架：DPE框架主要包含以下步骤：1) 维度重要性评估：分析RoPE的各个维度，确定哪些维度对于长上下文建模至关重要。2) 有效长度检测：针对每个关键维度，确定其最有效的长度，即在该长度下，模型能够最好地捕捉长距离依赖关系。3) 位置索引操控：对于关键维度，将其位置索引操控到其最有效的长度，从而扩展模型的上下文窗口。4) 模型集成：将操控后的位置编码集成到预训练模型中，即可实现长上下文扩展。

关键创新：DPE的关键创新在于其维度感知的处理方式。与以往方法不同，DPE不是平等地对待所有位置编码维度，而是通过分析和选择性地操控关键维度，实现了更高效和有效的上下文扩展。这种方法避免了对所有维度进行统一处理的低效性，并能够更好地利用预训练模型的信息。

关键设计：DPE的关键设计包括：1) 维度重要性评估方法：论文可能采用某种指标来衡量每个维度对于长上下文建模的重要性，例如信息熵、梯度等。2) 有效长度检测方法：论文可能通过实验或理论分析，确定每个关键维度的最有效长度。3) 位置索引操控策略：论文需要设计一种策略，将关键维度的位置索引操控到其最有效的长度，同时保持模型的稳定性和性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DPE在长上下文基准测试中取得了显著的性能提升。例如，DPE使Llama3-8k 8B能够支持128k tokens的上下文窗口而无需持续训练，并与Flash Attention 2无缝集成。在RULER基准测试中，带有DPE的Llama3.1 70B比基线模型提高了超过18个点，甚至优于GPT-4-128K。

🎯 应用场景

DPE具有广泛的应用前景，可用于各种需要处理长文本的场景，如长文档摘要、机器翻译、问答系统、代码生成等。通过扩展LLM的上下文窗口，DPE可以提高模型在这些任务中的性能，并使其能够处理更复杂的输入。此外，DPE无需重新训练的特性，使其能够快速部署到现有模型中，降低了应用成本。

📄 摘要（原文）

Large Language Models (LLMs) often struggle to process and generate coherent context when the number of input tokens exceeds the pre-trained length. Recent advancements in long-context extension have significantly expanded the context window of LLMs but require expensive overhead to train the large-scale models with longer context. In this work, we propose Dimension-Wise Positional Embeddings Manipulation (DPE), a training-free framework to extrapolate the context window of LLMs by diving into RoPE's different hidden dimensions. Instead of manipulating all dimensions equally, DPE detects the effective length for every dimension and finds the key dimensions for context extension. We reuse the original position indices with their embeddings from the pre-trained model and manipulate the key dimensions' position indices to their most effective lengths. In this way, DPE adjusts the pre-trained models with minimal modifications while ensuring that each dimension reaches its optimal state for extrapolation. DPE significantly surpasses well-known baselines such as YaRN and Self-Extend. DPE enables Llama3-8k 8B to support context windows of 128k tokens without continual training and integrates seamlessly with Flash Attention 2. In addition to its impressive extrapolation capability, DPE also dramatically improves the models' performance within training length, such as Llama3.1 70B, by over 18 points on popular long-context benchmarks RULER. When compared with commercial models, Llama 3.1 70B with DPE even achieves better performance than GPT-4-128K.