深度学习进阶(十九)相对位置编码 RPE

深度学习进阶(十九)相对位置编码 RPE 深度学习进阶十九相对位置编码 RPE在 Transformer 模型中位置编码是至关重要的组成部分。传统的绝对位置编码Absolute Position Encoding, APE将每个位置映射为一个固定的向量但这种方式在处理长序列、可变长度序列以及捕捉相对位置关系时存在局限。相对位置编码Relative Position Encoding, RPE应运而生它通过建模 token 之间的相对距离而非绝对位置显著提升了模型的泛化能力和对序列结构的理解。本文将从实战角度出发通过代码示例深入解析 RPE 的原理、实现与优势。## 绝对位置编码的局限在 Transformer 中绝对位置编码通过添加正弦波或学习向量来标记每个 token 的位置。例如对于序列长度 ( n )位置 ( i ) 的编码为[PE_{(i,2k)} \sin(i / 10000^{2k/d}), \quad PE_{(i,2k1)} \cos(i / 10000^{2k/d})]这种编码的缺点在于- 无法利用序列内部的相对距离如“单词 A 距离单词 B 3 个位置”。- 对序列长度变化敏感长序列的编码可能超出训练时的最大长度。- 在注意力计算中绝对位置信息与内容信息混合可能导致模型难以捕捉局部依赖。相对位置编码通过直接建模 token 之间的相对偏移 ( i - j ) 来解决这些问题广泛应用于 BERT、T5、Transformer-XL 等模型。## RPE 的核心思想相对位置编码的核心是在注意力机制中引入一个偏置项该偏置项取决于查询Query和键Key之间的相对距离。具体来说标准注意力计算为[\text{Attention}(Q, K, V) \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V]在 RPE 中我们修改为[\text{Attention}(Q, K, V) \text{softmax}\left(\frac{QK^T \text{RPE_bias}}{\sqrt{d_k}}\right)V]其中 ( \text{RPE_bias}_{i,j} ) 是位置 ( i ) 和 ( j ) 的相对距离 ( i - j ) 的函数。常见的实现方式包括-可学习的相对位置偏置为每个可能的相对距离分配一个可训练的参数。-基于正弦波的距离嵌入将相对距离映射为固定或可学习的向量。## 实战实现一个带 RPE 的 Transformer 注意力层下面我们用 PyTorch 实现一个简单的相对位置编码注意力模块。我们将使用可学习的相对位置偏置限制最大相对距离为 ( k )。### 代码示例 1基础 RPE 注意力实现pythonimport torchimport torch.nn as nnimport torch.nn.functional as Fclass RelativePositionAttention(nn.Module): 带可学习相对位置偏置的注意力层 def __init__(self, d_model, n_heads, max_relative_position16): super().__init__() self.d_model d_model self.n_heads n_heads self.d_k d_model // n_heads # 每个头的维度 self.max_relative_position max_relative_position # 定义 Q, K, V 的线性变换 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) # 可学习的相对位置偏置形状为 (2*max_relative_position1, n_heads) # 索引映射相对距离 d 映射到 idx d max_relative_position self.relative_position_bias nn.Parameter( torch.randn(2 * max_relative_position 1, n_heads) ) def forward(self, x, maskNone): x: (batch_size, seq_len, d_model) mask: (batch_size, seq_len) 或 None batch_size, seq_len, _ x.size() # 线性变换并拆分为多头 q self.W_q(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) k self.W_k(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) v self.W_v(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2) # 计算标准注意力分数 (batch, n_heads, seq_len, seq_len) scores torch.matmul(q, k.transpose(-2, -1)) / (self.d_k ** 0.5) # 生成相对位置索引矩阵 (seq_len, seq_len) # 对于位置 i 和 j相对距离 d i - j # 我们限制 d 在 [-max_relative_position, max_relative_position] 之间 d torch.arange(seq_len).unsqueeze(0) - torch.arange(seq_len).unsqueeze(1) # (seq_len, seq_len) d_clamped torch.clamp(d, -self.max_relative_position, self.max_relative_position) # 截断 # 将 [-max, max] 映射到 [0, 2*max] d_index d_clamped self.max_relative_position # (seq_len, seq_len) # 获取相对位置偏置形状为 (seq_len, seq_len, n_heads) bias self.relative_position_bias[d_index] # (seq_len, seq_len, n_heads) # 调整维度以匹配 scores: (1, n_heads, seq_len, seq_len) bias bias.permute(2, 0, 1).unsqueeze(0) # (1, n_heads, seq_len, seq_len) # 添加偏置 scores scores bias # 应用 mask如果有 if mask is not None: # mask 形状 (batch, seq_len)补 0 表示有效1 表示无效 mask_expanded mask.unsqueeze(1).unsqueeze(2) # (batch, 1, 1, seq_len) scores scores.masked_fill(mask_expanded 0, float(-inf)) # Softmax 和加权求和 attn_weights F.softmax(scores, dim-1) context torch.matmul(attn_weights, v) # (batch, n_heads, seq_len, d_k) # 拼接多头并输出 context context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) output self.W_o(context) return output# 测试代码if __name__ __main__: # 创建模型实例 d_model 64 n_heads 4 max_rel_pos 8 model RelativePositionAttention(d_model, n_heads, max_rel_pos) # 模拟输入batch_size2, seq_len10 x torch.randn(2, 10, d_model) mask torch.ones(2, 10, dtypetorch.bool) # 全部有效 output model(x, mask) print(输出形状:, output.shape) # 应为 (2, 10, 64)代码说明-relative_position_bias是一个可学习参数大小为(2*max_relative_position1, n_heads)每个头有自己的偏置向量。- 索引映射d_index d max_relative_position将相对距离可能为负转换为非负索引。- 在注意力分数计算中scores加上bias从而让模型学习到不同相对距离的重要性。## 实战与绝对位置编码的对比实验为了直观展示 RPE 的效果我们设计一个简单的序列分类任务判断两个 token 是否相邻。我们将训练一个带 APE 的 Transformer 和一个带 RPE 的 Transformer并比较它们的性能。### 代码示例 2对比实验pythonimport torchimport torch.optim as optimfrom torch.utils.data import DataLoader, TensorDataset# 生成合成数据判断序列中两个指定位置的 token 是否相邻def generate_data(num_samples1000, seq_len8): 生成数据每个样本包含一个序列和两个位置索引标签表示这两个位置是否相邻。 data [] labels [] for _ in range(num_samples): seq torch.randint(0, 10, (seq_len,)) # 随机 token pos1 torch.randint(0, seq_len, (1,)).item() pos2 torch.randint(0, seq_len, (1,)).item() label 1 if abs(pos1 - pos2) 1 else 0 # 是否相邻 data.append((seq, pos1, pos2)) labels.append(label) return data, labels# 定义带 APE 的 Transformer 分类器class AbsolutePositionTransformer(nn.Module): def __init__(self, vocab_size11, d_model32, n_heads4, seq_len8): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) self.position_encoding nn.Parameter(torch.randn(1, seq_len, d_model)) # 绝对位置 self.attention nn.MultiheadAttention(d_model, n_heads, batch_firstTrue) self.fc nn.Linear(d_model, 2) # 二分类 def forward(self, seq, pos1, pos2): # seq: (batch, seq_len) x self.embedding(seq) self.position_encoding # 添加绝对位置 attn_out, _ self.attention(x, x, x) # 自注意力 # 取两个位置的输出 h1 attn_out[torch.arange(attn_out.size(0)), pos1] h2 attn_out[torch.arange(attn_out.size(0)), pos2] out self.fc(h1 h2) # 加和后分类 return out# 定义带 RPE 的 Transformer 分类器class RelativePositionTransformer(nn.Module): def __init__(self, vocab_size11, d_model32, n_heads4, max_rel_pos8): super().__init__() self.embedding nn.Embedding(vocab_size, d_model) self.attention RelativePositionAttention(d_model, n_heads, max_rel_pos) self.fc nn.Linear(d_model, 2) def forward(self, seq, pos1, pos2): x self.embedding(seq) # 无绝对位置编码 attn_out self.attention(x) # 使用 RPE 注意力 h1 attn_out[torch.arange(attn_out.size(0)), pos1] h2 attn_out[torch.arange(attn_out.size(0)), pos2] out self.fc(h1 h2) return out# 训练函数def train_model(model, dataloader, epochs10, lr0.001): criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lrlr) model.train() for epoch in range(epochs): total_loss 0.0 for batch in dataloader: seq, pos1, pos2, labels batch optimizer.zero_grad() output model(seq, pos1, pos2) loss criterion(output, labels) loss.backward() optimizer.step() total_loss loss.item() print(fEpoch {epoch1}, Loss: {total_loss/len(dataloader):.4f})# 测试准确率def evaluate(model, dataloader): model.eval() correct 0 total 0 with torch.no_grad(): for batch in dataloader: seq, pos1, pos2, labels batch output model(seq, pos1, pos2) _, predicted torch.max(output, 1) total labels.size(0) correct (predicted labels).sum().item() return correct / total# 主程序if __name__ __main__: # 生成数据 data, labels generate_data(2000, seq_len8) seqs torch.stack([d[0] for d in data]) # (2000, 8) pos1s torch.tensor([d[1] for d in data]) pos2s torch.tensor([d[2] for d in data]) labels torch.tensor(labels) dataset TensorDataset(seqs, pos1s, pos2s, labels) dataloader DataLoader(dataset, batch_size32, shuffleTrue) # 训练 APE 模型 print(训练绝对位置编码模型...) ape_model AbsolutePositionTransformer() train_model(ape_model, dataloader, epochs5) ape_acc evaluate(ape_model, dataloader) print(fAPE 模型准确率: {ape_acc:.4f}) # 训练 RPE 模型 print(训练相对位置编码模型...) rpe_model RelativePositionTransformer() train_model(rpe_model, dataloader, epochs5) rpe_acc evaluate(rpe_model, dataloader) print(fRPE 模型准确率: {rpe_acc:.4f})实验结果分析- 在“判断相邻位置”这类依赖相对距离的任务中RPE 模型通常收敛更快且准确率更高。- 绝对位置编码模型有时会过拟合于绝对位置如位置 3 和 4 相邻而 RPE 能更好地泛化到任意序列长度。## RPE 的变体与演进除了上述实现RPE 还有其他形式1.T5 风格在注意力分数上添加可学习的偏置但偏置仅取决于相对距离的桶bucket而非精确距离。2.Transformer-XL将相对位置信息融入键K和查询Q的计算中使用正弦波编码。3.RoPE旋转位置编码通过旋转矩阵对 Q 和 K 进行变换隐式编码相对位置无需额外偏置。## 总结相对位置编码 RPE 通过建模 token 间的相对距离克服了绝对位置编码在长序列和相对关系捕捉上的不足。本文从原理出发通过两个可运行的代码示例展示了 RPE 在注意力机制中的实现及其在序列分类任务中的优势。实战中RPE 尤其适用于需要理解局部结构或序列长度变化的任务如文本生成、时间序列预测。掌握 RPE 不仅有助于优化现有模型也是理解现代 Transformer 架构如 GPT-4、Llama的关键一步。建议读者进一步尝试将 RPE 集成到自己的 Transformer 模型中并对比不同变体的效果。