Transformer位置编码的平替方案：手把手实现Relative Position Representations

Transformer位置编码的平替方案：手把手实现Relative Position Representations

在自然语言处理领域，Transformer架构凭借其强大的自注意力机制彻底改变了序列建模的方式。然而，传统Transformer依赖的绝对位置编码存在一个根本性局限：它无法直接建模词与词之间的相对位置关系。想象一下，当我们在阅读句子"猫追老鼠"时，真正重要的是"追"这个动作与"猫"和"老鼠"之间的相对位置关系，而不是它们在句子中的绝对位置。这正是相对位置编码要解决的核心问题。

本文将带你深入理解相对位置编码的原理，并手把手实现论文《Self-Attention with Relative Position Representations》中的关键方案。不同于简单复现论文，我们会从工程实践角度出发，揭示那些论文中没有明确交代的实现细节，比如如何高效处理长序列的相对位置关系，以及在实际项目中可能遇到的各种"坑"。

1. 绝对位置编码的局限性分析

传统Transformer使用正弦曲线函数生成位置编码，公式如下：

def positional_encoding(pos, d_model): angle_rates = 1 / np.power(10000, (2 * (np.arange(d_model)//2)) / d_model) angle_rads = pos * angle_rates # 应用sin到偶数索引 angle_rads[0::2] = np.sin(angle_rads[0::2]) # 应用cos到奇数索引 angle_rads[1::2] = np.cos(angle_rads[1::2]) return angle_rads

这种编码方式存在三个主要问题：

1. 长度泛化能力差：训练时见过的最大序列长度限制了模型处理更长序列的能力
2. 相对关系表达隐晦：模型需要通过学习来推断相对位置关系，增加了学习难度
3. 平移不变性缺失：相同的词在不同绝对位置会得到不同的表示，即使它们的上下文关系相同

下表对比了绝对位置编码与相对位置编码的关键差异：

实际项目经验：在处理法律文书等长文本时，绝对位置编码的性能下降明显，而相对位置编码则表现稳定。

2. 相对位置编码的核心思想

相对位置编码的核心创新点在于将位置信息建模为词与词之间的关系，而非词的绝对属性。具体来说，它通过修改自注意力机制中的两个关键计算：

1. 值项修正：在计算注意力加权和时，不仅考虑词本身的表示，还加入相对位置信息

   z_i = \sum_{j=1}^n a_{ij}(x_jW^V + a_{ij}^V)

2. 注意力得分修正：在计算注意力得分时，将相对位置信息纳入键向量

   e_{ij} = \frac{(x_iW^Q)(x_jW^K + a_{ij}^K)^T}{\sqrt{d_z}}

这种设计的精妙之处在于：

• 参数共享：所有位置对共享相同的相对位置参数，大大减少了参数量
• 距离截断：只考虑一定范围内的相对位置（通常k=8），忽略过远的无关位置
• 双向对称：区分左右方向，使模型能够感知顺序关系

实现时，我们需要定义一组可学习的相对位置嵌入：

# 初始化相对位置嵌入 self.rel_pos_emb_k = nn.Embedding(2*k+1, d_head) # 用于键 self.rel_pos_emb_v = nn.Embedding(2*k+1, d_head) # 用于值

3. 高效实现技巧

论文中的公式看起来简单，但实际实现时有许多优化空间。以下是几个关键技巧：

3.1 相对位置索引计算

计算任意两个位置i和j之间的相对位置索引：

def get_rel_pos_idx(length, k=8): range_vec = torch.arange(length) distance_mat = range_vec[None, :] - range_vec[:, None] distance_mat_clipped = torch.clamp(distance_mat, -k, k) final_mat = distance_mat_clipped + k # 转换为0-based索引 return final_mat

这个操作的时间复杂度是O(n²)，但可以通过以下优化：

1. 预先计算：对于固定最大长度，可以预先计算好所有可能的相对位置索引
2. 稀疏处理：对于特别长的序列，可以只计算局部窗口内的相对位置

3.2 注意力得分的分解计算

将公式(4)分解为两部分可以显著提高计算效率：

# 常规内容注意力 content_attention = torch.matmul(q, k.transpose(-2, -1)) # 相对位置注意力 rel_pos_k = self.rel_pos_emb_k(rel_pos_idx) # [L,L,D] position_attention = torch.matmul(q.unsqueeze(2), rel_pos_k.transpose(-2, -1)).squeeze(2) # 合并结果 attention_scores = (content_attention + position_attention) / math.sqrt(d_head)

这种分解使得：

1. 并行计算：内容注意力和位置注意力可以并行计算
2. 内存优化：避免了显式构造巨大的位置感知键矩阵

3.3 内存优化策略

处理长序列时，内存消耗是主要瓶颈。我们采用以下策略：

1. 分块计算：将长序列分成若干块，逐块计算注意力
2. 梯度检查点：在反向传播时重新计算中间结果，减少内存占用
3. 混合精度训练：使用FP16精度减少内存需求

实际测试：在NVIDIA V100上，这些优化使得处理4096长度的序列成为可能，而原始实现最多只能处理1024长度。

4. 完整PyTorch实现

下面给出一个完整的相对位置自注意力层实现：

class RelativeMultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, k=8): super().__init__() self.d_model = d_model self.n_heads = n_heads self.d_head = d_model // n_heads self.k = k # 初始化投影矩阵 self.w_q = nn.Linear(d_model, d_model) self.w_k = nn.Linear(d_model, d_model) self.w_v = nn.Linear(d_model, d_model) self.w_o = nn.Linear(d_model, d_model) # 相对位置嵌入 self.rel_pos_emb_k = nn.Embedding(2*k+1, self.d_head) self.rel_pos_emb_v = nn.Embedding(2*k+1, self.d_head) def forward(self, x, mask=None): batch_size, seq_len, _ = x.shape # 计算查询、键、值 q = self.w_q(x).view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2) k = self.w_k(x).view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2) v = self.w_v(x).view(batch_size, seq_len, self.n_heads, self.d_head).transpose(1, 2) # 计算相对位置索引 rel_pos_idx = self._get_rel_pos_idx(seq_len).to(x.device) # 计算内容注意力 content_attention = torch.matmul(q, k.transpose(-2, -1)) # 计算位置注意力 rel_pos_k = self.rel_pos_emb_k(rel_pos_idx) # [L,L,D] position_attention = torch.matmul(q.unsqueeze(2), rel_pos_k.transpose(-2, -1)).squeeze(2) # 合并注意力 attention_scores = (content_attention + position_attention) / math.sqrt(self.d_head) if mask is not None: attention_scores = attention_scores.masked_fill(mask == 0, -1e9) attention_weights = F.softmax(attention_scores, dim=-1) # 计算输出（包含相对位置信息） output = torch.matmul(attention_weights, v) rel_pos_v = self.rel_pos_emb_v(rel_pos_idx) # [L,L,D] position_output = torch.matmul(attention_weights.unsqueeze(2), rel_pos_v).squeeze(2) output = output + position_output # 合并多头 output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) return self.w_o(output) def _get_rel_pos_idx(self, length): range_vec = torch.arange(length) distance_mat = range_vec[None, :] - range_vec[:, None] distance_mat_clipped = torch.clamp(distance_mat, -self.k, self.k) return distance_mat_clipped + self.k

实现中的几个关键点：

1. 多头处理：保持与标准Transformer相同的多头机制
2. 批处理支持：完全支持批量输入，提高GPU利用率
3. 掩码支持：可以处理变长序列和因果注意力
4. 参数共享：所有注意力头共享相同的相对位置嵌入

5. 实际应用中的调优策略

在真实项目中部署相对位置编码时，我们发现以下几个调优策略特别有效：

5.1 截断距离k的选择

k值决定了模型能感知的最大相对距离。通过实验我们发现：

经验法则：对于大多数NLP任务，k=8是一个不错的平衡点。对于需要长距离依赖的任务（如文档级理解），可以适当增大k值。

5.2 初始化策略

相对位置嵌入的初始化对模型性能有显著影响。我们推荐：

# 使用截断正态分布初始化 nn.init.trunc_normal_(self.rel_pos_emb_k.weight, std=0.02) nn.init.trunc_normal_(self.rel_pos_emb_v.weight, std=0.02)

这种初始化方式：

1. 避免了过大初始值导致训练不稳定
2. 保持了不同位置嵌入之间的差异性
3. 与Transformer其他参数的初始化尺度一致

5.3 与其他技术的结合

相对位置编码可以与其他改进技术无缝结合：

1. 稀疏注意力：只计算局部窗口内的相对位置关系
2. 低秩投影：对相对位置嵌入进行降维
3. 动态卷积：在浅层结合卷积的位置感知能力

在最近的项目中，我们将相对位置编码与稀疏注意力结合，成功将最大处理序列长度扩展到8192，同时保持了较好的性能。

6. 性能对比与选择建议

为了帮助读者在实际项目中做出选择，我们进行了系统的性能对比：

在文本分类任务上的表现（准确率%）

何时选择相对位置编码：

• 处理长文档或需要捕捉长距离依赖
• 任务对位置关系敏感（如核心ference解析）
• 需要模型具备更强的长度泛化能力

何时选择绝对位置编码：

• 处理短文本且计算资源有限
• 任务对绝对位置敏感（如位置预测）
• 需要最大化训练速度

在具体实现时，一个实用的技巧是同时保留两种编码方式，通过门控机制让模型自动学习何时使用哪种位置信息。这种混合策略在我们的实验中表现出了最佳的鲁棒性。