通过层特定位置嵌入缩放缓解Transformer中的位置偏差

# 通过层特定位置嵌入缩放缓解 Transformer 中的位置偏差  
来源：https://arxiv.org/html/2606.27705  

Changze Lv¹,³, Zhenghua Wang¹†,¹,³, Yiran Ding¹†,², Yixin Wu¹,³, Tianlong Li¹,³, Zhibo Xu¹,³, Muling Wu¹,³, Tianyuan Shi¹,³, Shizheng Li¹,³, Qi Qian¹,³, Xuanjing Huang¹,³, Xiaoqing Zheng¹,³  

¹复旦大学，²西湖大学，³上海市智能信息处理重点实验室  
{zhenghuawang23, czlv24}@m.fudan.edu.cn  
{yiran.ding}@hdu.edu.cn  
{xjhuang, zhengxq}@fudan.edu.cn  

###### 摘要  

大型语言模型（LLMs）仍然面临"中间丢失"问题，即位于长上下文输入中间的关键信息往往被低估或丢失。尽管现有方法尝试通过组合多尺度旋转位置编码（RoPE）来解决这一问题，但它们通常存在高延迟或依赖于次优的手工设计缩放策略。为了克服这些局限性，我们提出了一种层特定位置嵌入缩放（LPES）方法，该方法为每一层分配不同的缩放因子。LPES 在不微调模型参数或增加推理延迟的情况下实现了更均衡的注意力分布。我们采用了一种专门设计的遗传算法，通过引入贝塞尔曲线来显著缩小搜索空间，从而高效地为每一层选择最优缩放因子。大量实验表明，LPES 有效缓解了位置注意力偏差，并在多个长上下文基准上带来持续改进，在键值检索数据集上最高可获得 11.2% 的准确率提升。  

---

# 通过层特定位置嵌入缩放缓解 Transformer 中的位置偏差  

Changze Lv††同等贡献。¹,³, Zhenghua Wang¹†,¹,³, Yiran Ding¹†,², Yixin Wu¹,³, Tianlong Li¹,³, Zhibo Xu¹,³, Muling Wu¹,³, Tianyuan Shi¹,³, Shizheng Li¹,³, Qi Qian¹,³, Xuanjing Huang¹,³, Xiaoqing Zheng††通讯作者。¹,³  

¹复旦大学，²西湖大学，³上海市智能信息处理重点实验室  
{zhenghuawang23, czlv24}@m.fudan.edu.cn  
{yiran.ding}@hdu.edu.cn  
{xjhuang, zhengxq}@fudan.edu.cn  

## 1 引言  

使大型语言模型（LLMs）能够处理长输入，对于支持长文本摘要（Feng 等，2021 (https://arxiv.org/html/2606.27705#bib.bib11)；Zhang 等，2021 (https://arxiv.org/html/2606.27705#bib.bib14)）、代码生成（Zheng 等，2023 (https://arxiv.org/html/2606.27705#bib.bib12)；Liu 等，2024a (https://arxiv.org/html/2606.27705#bib.bib13)）以及长上下文问答（Li 等，2024 (https://arxiv.org/html/2606.27705#bib.bib10)）等复杂任务至关重要。旋转位置编码（RoPE）（Su 等，2021 (https://arxiv.org/html/2606.27705#bib.bib42)）被广泛用于基于 Transformer 的 LLMs，旨在编码输入 token 之间的相对距离，从而更有效地处理长上下文输入。然而，随着上下文长度的增加，基于 RoPE 的 LLMs 仍然存在位置偏差。这一问题的一个典型表现是众所周知的"中间丢失"现象（Liu 等，2024c (https://arxiv.org/html/2606.27705#bib.bib15)），即模型倾向于过度关注输入开头和结尾附近的 token，而相对忽略位于中间的信息。  

为了解决位置偏差问题，已有几种方法提出通过组合多个具有不同基座或缩放因子的 RoPE（Chen 等，2023b (https://arxiv.org/html/2606.27705#bib.bib19)；Zhang 等，2024 (https://arxiv.org/html/2606.27705#bib.bib52)；Lin 等，2024 (https://arxiv.org/html/2606.27705#bib.bib20)）。Chen 等（2023b (https://arxiv.org/html/2606.27705#bib.bib19)）观察到，不同基座的 RoPE 会在特定位置产生注意力低谷，从而削弱模型捕捉相应内容的能力。为了缓解这一问题，他们提出了一种名为 Attention Buckets 的方法，通过模型并行组合多个不同基座的 RoPE，以获得更均衡的注意力分布。类似地，Lin 等（2024 (https://arxiv.org/html/2606.27705#bib.bib20)）提出了 MoICE 方法，为每个注意力头分配多个 RoPE 基座，并通过加权和聚合输出。然而，这些方法高度依赖手工规则来确定缩放因子或基值，并且在推理过程中需要多次前向传播——每次针对特定基座或缩放因子进行一次，随后对结果进行集成。尽管某些操作可以并行化，但这一过程不可避免会增加推理时间和计算成本。  

参考说明  
图 1：所提出的 LPES 与两种代表性现有方法的比较。(a) Attention Buckets 通过模型并行组合多个不同基座的 RoPE。(b) MoICE 为每个注意力头分配多个基座。与这些需在推理中进行多次前向传播的现有方法不同，我们的 LPES (c) 通过单次前向传播实现更优性能，显著减少推理时间。  

在整个模型上变化 RoPE 基座可以看作是模型级集成，而为单个注意力头应用多个基座则对应于模块级集成（图 1 (https://arxiv.org/html/2606.27705#S1.F1)）。模型级集成需要多次模型推理，产生大量计算开销，而模块级缩放由于粒度较细，搜索空间大，限制了自动搜索算法的适用性。为了平衡效率和灵活性，我们在层级别应用多个缩放后的 RoPE，通过单次前向传播获得具有竞争力甚至更优的性能，从而避免了相关的推理开销。  

为每一层选择合适的缩放因子仍然是一个非平凡的问题。设 \( L \) 表示基于 Transformer 的网络中的层数，\( M \) 表示缩放因子的可能取值数目；组合总数为 \( M^L \)，这使得穷举搜索在计算上变得不可行。确定最优缩放因子本质上是一个组合优化问题，因此难以通过基于梯度的方法求解。为了克服这一困难，我们利用贝塞尔曲线，该曲线通过一组少量的离散控制点，在层深度和缩放因子之间定义一个平滑、连续的映射。设 \( C \) 表示控制点数量，搜索空间被缩减为 \( (M \times L)^C \)。除了减少搜索空间，我们还发现基于曲线的平滑缩放保留了层间表示结构，并作为一种有益的归纳偏置。我们进一步开发了一种曲线约束的遗传算法来解决这一组合优化问题。通过将搜索空间限制在贝塞尔曲线上，我们能够高效地优化层特定缩放因子，通常只需在四块 H100 GPU 上使用几百个示例（例如 200 个实例）即可在 3 到 4 小时内完成。  

在长文本任务中，我们的方法不引入额外的推理延迟，同时提供优于现有方法的性能。本研究做出以下贡献：  

- • 我们提出了一种层特定位置嵌入缩放方法，称为 LPES，它有效地缓解了位置偏差，且不增加推理延迟。LPES 实现了显著的加速，比 MoICE（Lin 等，2024 (https://arxiv.org/html/2606.27705#bib.bib20)）快 2.42 倍，比 Ms-PoE（Zhang 等，2024 (https://arxiv.org/html/2606.27705#bib.bib52)）快 1.45 倍，同时提高了模型处理长上下文任务的能力。  
- • 我们引入了一种高效的遗传搜索算法，其中搜索空间受贝塞尔曲线约束，仅使用少量示例即可快速优化层特定缩放因子。  
- • 在多个基准数据集上的大量实验表明，我们的方法在保留模型通用能力的同时，无需昂贵的微调即可产生更均衡的注意力分布，使其广泛适用于不同模型和任务。  

## 2 相关工作  

参考说明  
图 2：所提出的层特定位置嵌入缩放（LPES）方法的说明。左图：贝塞尔曲线可以表示多种形状。中图：通过我们的搜索算法找到的一条优化后的贝塞尔曲线，该曲线使用一组有限的离散控制点定义了一条平滑、连续的曲线。右图：缩放因子与优化后的贝塞尔曲线之间的关系，以及它们在基于 Transformer 网络的注意力机制中的应用。  

Chen 等（2023b (https://arxiv.org/html/2606.27705#bib.bib19)）观察到，不同基座的 RoPE 会在特定位置产生注意力低谷，称为"注意力波"，从而削弱模型捕捉相关内容的能力。为了解决这一问题，他们的"Attention Buckets"方法通过模型并行推理集成多个 RoPE 基座，以获得更均匀的注意力分布。Zhang 等（2024 (https://arxiv.org/html/2606.27705#bib.bib52)）认为注意力中的长期衰减可能导致位置偏差，并提出了 Ms-PoE，该方法根据注意力头对位置信息的相对敏感性为其分配不同的缩放因子。MoICE（Lin 等，2024 (https://arxiv.org/html/2606.27705#bib.bib20)）在 Chen 等（2023b (https://arxiv.org/html/2606.27705#bib.bib19)）的工作基础上，采用梯度下降来学习组合不同基座结果的权重，在单个注意力头级别进行。然而，这些方法的一个主要局限性是计算成本高且推理延迟大。具体来说，Attention Buckets 需要多次前向传播，而 Ms-PoE 和 MoICE 都需要重复的注意力计算来整合多尺度 RoPE 信息。它们还依赖启发式或手工规则来选择基座或缩放因子。相比之下，我们的方法通过单次前向传播实现更优性能，并提出了一个自动搜索算法，仅使用几百个示例即可有效确定最优缩放因子。  

## 3 方法  

### 3.1 问题定义  

在本研究中，我们关注 RoPE，其定义如下：  

\[
\langle f(\mathbf{q}, i), f(\mathbf{k}, j) \rangle = \mathbf{q}^{\mathrm{T}} R(i-j) \mathbf{k} \tag{1}
\]

其中 \( f(\mathbf{x}, i) \) 表示在位置 \( i \) 对查询 \( \mathbf{q} \) 应用的位置相关旋转，\( f(\mathbf{k}, j) \) 表示在位置 \( j \) 经过 RoPE 旋转后的键。符号 \( \langle \cdot, \cdot \rangle \) 表示两个位置感知向量之间的内积，\( R(\Delta) \) 是对应于相对偏移 \( \Delta = i - j \) 的旋转。该方程表明内积仅取决于向量 \( \mathbf{q} \)、\( \mathbf{k} \) 以及它们之间的相对距离。  

Chen 等（2023a (https://arxiv.org/html/2606.27705#bib.bib40)）表明，可以通过对位置索引应用缩放因子 \( s \) 来扩展上下文窗口，如下所示：  

\[
f'(x, i) = f(\mathbf{x}, i/s) \tag{2}
\]

我们进一步证明，缩放因子可以缓解长期衰减并诱导出多样的注意力模式（附录 A (https://arxiv.org/html/2606.27705#A1)）。因此，我们的目标是为每一层搜索一个唯一的缩放因子 \( s \)，以组合来自多个缩放 RoPE 的信息，减轻长期衰减和注意力波效应，从而减少位置偏差。  

我们使用贝塞尔曲线对层深度和缩放因子进行建模，这通过少量控制点确定所有层缩放因子，大幅缩小了搜索空间。详细分析见附录 B (https://arxiv.org/html/2606.27705#A2)。此外，在第 4.2 节 (https://arxiv.org/html/2606.27705#S4.SS2) 中，我们表明基于曲线的平滑连续建模保留了层间表示结构。暴力搜索表明，平滑缩放自然成为一种高性能配置，突出了跨层连续性作为一种有益的归纳偏置。  

如图 2 (https://arxiv.org/html/2606.27705#S2.F2) 所示，贝塞尔曲线可以看作是一条连接二维平面中所有缩放因子的平滑曲线。为所有层选择缩放因子的问题可以转化为搜索一条合适的贝塞尔曲线。幸运的是，贝塞尔曲线仅使用少量离散控制点即可模拟多种形状，这显著减小了搜索空间。  

一个 \( d \) 次贝塞尔曲线，具有 \( d+1 \) 个控制点，定义如下（Mortenson，1999 (https://arxiv.org/html/2606.27705#bib.bib26)）：  

\[
B(t) = \sum_{k=0}^{d} b_k^d(t) P_k, \quad 0 \leq t \leq 1. \tag{3}
\]

其中 \( t \) 是控制曲线上点位置的参数坐标，\( P_k \) 是曲线的控制点，\( b_k^d \) 是 Bernstein 基多项式，定义为：  

\[
b_k^d(t) = \frac{d!}{k!(d-k)!} t^k (1-t)^{d-k}, \quad k = 0, \dots, d. \tag{4}
\]

一旦贝塞尔曲线确定，第 \( h \) 层的缩放因子 \( s_h \) 可以通过下式计算：  

\[
s_h = \text{proj}_y \left[ B(t(x_h)) \right] \tag{5}
\]

其中符号 \( \text{proj}_y[\cdot] \) 表示提取二维点的 \( y \) 坐标的操作。函数 \( t(\cdot) \) 将 \( x_h \) 映射到对应的参数 \( t \)（见附录 D (https://arxiv.org/html/2606.27705#A4)），其中 \( x_h \) 表示层 \( h \) 在由控制点最小值和最大值定义的均匀间隔 \( x \) 坐标中的位置。\( x_h \) 的值可以通过下式计算：  

\[
x_h = P_0^x + \frac{P_d^x - P_0^x}{L-1} \cdot h, \quad h = 0, \dots, L-1. \tag{6}
\]

其中 \( L \) 表示网络中的层数，\( P_t^x \) 是贝塞尔曲线第 \( t \) 个控制点的 \( x \) 坐标。  

给定一个训练数据集 \( \mathcal{D} = \{(x_i, y_i)\}_{i=1}^N \)，包含 \( N \) 个示例，其中 \( x_i \) 是大型语言模型的输入，\( y_i \) 是相应的真实输出，我们的目标是最大化以下函数：  

\[
\mathcal{L}_{\mathcal{D}}(\boldsymbol{\theta}) = \frac{1}{N} \sum_{i=1}^N \mathbb{I} \{ \text{LLM}(x_i, \boldsymbol{\theta}) \simeq y_i \} \tag{7}
\]

其中 \( \boldsymbol{\theta} = (P_0, \dots, P_d) \) 表示定义一条 \( d \) 次贝塞尔曲线的控制点集合（每个控制点 \( P_k \) 是一个二维点），\( \text{LLM}(x_i, \boldsymbol{\theta}) \) 表示语言模型给定输入 \( x_i \) 的输出，其中所有缩放因子根据方程 (5) (https://arxiv.org/html/2606.27705#S3.E5) 由 \( \boldsymbol{\theta} \) 指定的贝塞尔曲线确定，\( \mathbb{I}\{\cdot\} \) 是一个指示函数，输出为 0 或 1。  

我们构建了训练数据集，使得包含生成正确答案所需信息的文本出现在输入中的不同位置，从而鼓励模型更均匀地分配注意力。