Effective Length Extrapolation via Dimension-Wise Positional Embeddings Manipulation

本文提出DPE，一种无需训练的长文本外推方法，通过检测RoPE不同维度组的有效相对距离并识别关键维度，有选择地调整这些关键维度的位置索引，显著扩展了LLM的上下文窗口并提升了长文本任务性能。

Large Language Model, Long Context, Representation Learning, Efficiency, Pre-training

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

Fudan University, The University of Hong Kong, Northeastern University

Generated by gemini-2.5-flash-preview-04-17

Background Problem

大型语言模型（LLMs）在处理超出预训练长度的文本时性能会显著下降。现有的长文本扩展方法通常需要昂贵的持续训练，或者采用对所有维度进行统一修改的无训练方法，这些方法未能充分考虑RoPE不同维度对长文本处理的差异化贡献，导致性能受限。RoPE在长文本外推中存在明显的分布外（OOD）问题，这与未见过的相对位置有关。虽然有方法尝试通过缩放频率或统一调整相对位置来缓解，但忽略了RoPE不同维度/频率子空间的差异性。

Method

• 核心思想: 针对RoPE在长文本外推中的OOD问题，发现不同维度对长文本的有效性不同，且存在关键维度。DPE通过维度层面的精细化操作，而非统一修改，来调整位置编码。RoPE中查询向量 $\mathbf{q}^m$ 和键向量 $\mathbf{k}^n$ 的点积可以分解为不同维度对的贡献之和： $\mathbf{q}\_m^{\top}\mathbf{k}\_n = \sum\_{j=0}^{d/2-1} \left( \mathbf{q}\_{\left[2j,2j+1\right]}\mathbf{R}(\theta\_{j\prime}n - m)\mathbf{k}\_{\left[2j,2j+1\right]}\right)$ 这表明每个维度对的贡献依赖于相对距离 $n-m$ 和频率 $\theta_j$。
• 具体步骤:
  1. 检测不同维度的有效相对距离: 将RoPE的维度分组，通过实验（如NIAH任务）检测每个维度组在不同最大相对距离下的性能表现，确定其“有效相对距离”。实验中，对于某个维度组，其位置索引 $P_l(m,n)$ 按如下方式缩放以限制最大相对距离 $t$： $\mathbf{P}\_l(m,n) = \begin{cases} \left\lfloor \frac{(n-m-w)t}{L} \right\rfloor + w, & \text{if } n-m > w, \\ n-m, & \text{otherwise.} \end{cases}$ 发现不同维度组的最佳有效距离不同，特别是中间频率维度对外推能力较弱。
  2. 识别关键维度: 使用2-norm Attention Contribution指标量化每个维度对注意力机制的贡献，识别出对上下文扩展起关键作用的维度。对于每个注意力头 $h$，关键维度 $\mathbf{D}_h$ 定义为： $\mathbf{D}\_h = \{ j | \text{top-}k(||\mathbf{q}\_\*^{(j)}|| ||\mathbf{k}\_\*^{(j)}||) \}$ 实验表明，仅调整这些关键维度即可恢复甚至超越调整所有维度的性能。
  3. 缩放位置索引: 对于识别出的关键维度 $j \in g_i$，根据其所属维度组 $g_i$ 检测到的有效相对距离 $e_i$，计算相应的缩放因子 $s_i = \lfloor (N-w)/(e_i-w) \rfloor$ (其中 $N$ 是目标长度，如128k)。然后，在这些关键维度的位置矩阵中，对超出局部窗口 $w$ 的相对距离 $n-m$ 应用缩放： $\mathbf{p}\_{\rangle}(m,n) = \begin{cases} \left\lfloor \frac{n-m-w}{s\_i} \right\rfloor + w, & \text{if } n-m > w, \\ n-m, & \text{otherwise.} \end{cases}$ 这确保了每个关键维度都能在其最有效的范围内工作，同时避免了OOD位置索引的影响。
• 实现细节: 将头维度分为8组，设置局部窗口w=1k，选择前48个关键维度进行操作。与FlashAttention 2集成，额外开销可忽略不计。

Experiment

• 模型: Llama3-8B-Instruct (8k), Mistral-7B-Instruct-v0.2 (32k) 用于外推测试；Llama-3.1-8B-Instruct (128k), Qwen-2.5-7B (128k), Llama-3.1-70B-Instruct (128k) 用于训练长度内性能提升测试。
• 基线: NTK-Dynamic, YaRN (修改频率); ReRoPE, Self-Extend, DCA (修改位置矩阵)。
• 数据集/任务: Needle-in-a-Haystack (NIAH，多针设置), RULER (包含NIAH, Variable Tracing, Aggregation, QA等13个任务), InfiniteBench。测试长度主要为128k。
• 结果:
  • NIAH: DPE在六个模型上的平均准确率达到90.25%，远高于原始RoPE (36.08%) 和最佳基线Self-Extend (85.29%)。
  • RULER: DPE在外推任务上显著优于所有基线（Llama3-8B提升8点，Mistral-v0.2-7B提升2点）。在训练长度内性能提升显著（Llama-3.1-8B提升8点，Qwen-2.5-7B提升40点）。Llama3.1-70B+DPE在RULER上的平均得分 (86.39) 甚至超过了GPT-4-128K (81.2)。
  • InfiniteBench: Llama3.1-8B+DPE和Qwen2.5-7B+DPE的性能与Claude2和Kimi-chat等商业模型相当。
• 消融实验: 验证了选择最有效的相对距离和识别关键维度的重要性。选择非最优有效距离或非关键维度（如贡献最小的维度）会导致性能显著下降。
• 效率: DPE与FlashAttention 2集成后，推理速度和峰值内存使用与原始模型相比，额外开销可忽略不计。

Further Thoughts

这项工作通过深入分析RoPE不同维度的特性，提出了一个无需训练且效果显著的长文本外推方法。其核心洞察在于“维度差异性”和“关键维度”的存在。这不仅为RoPE的改进提供了新思路，也提示我们在处理其他模型组件（如Attention Head、FFN层等）时，或许也应该考虑其内部维度或子组件的差异化贡献和特性，而非一概而论。例如，Attention Head的差异性已被广泛研究，但FFN内部维度的作用差异是否也值得探索？此外，DPE的“检测有效长度”步骤依赖于下游任务表现，这是否意味着其检测结果具有一定的任务依赖性？未来的工作或许可以探索更通用的维度特性分析方法，或者研究如何将这种维度层面的精细化调整应用于其他位置编码方式或模型架构。