前言
由于attetnion运算的特性，Transformer本身不感知顺序，位置编码是弥补这一缺陷的关键。
一、绝对位置编码
绝对位置编码的方式是通过将每个位置映射到一个高维空间中，该编码采用了正弦和余弦函数的组合
周期性：正弦和余弦函数具有周期性，这使得编码能够容易地表示固定范围的位置信息。这个特点允许模型有能力处理序列的循环性质，例如在语言中，某些词可能在不同上下文中重复出现，但它们的相对位置仍然是重要的。
不同频率：对于不同的维度 ( i )，使用不同的频率来编码位置。低维度的编码（例如低 ( i ) 值）对应较大的周期，能够捕捉较长的依赖关系；而高维度编码对应较小的周期，能够捕捉短距离的依赖关系。因此，模型可以通过不同的维度考虑不同尺度的位置信息。
任意长度的序列：这种方法能够处理任意长度的输入序列，因为正弦和余弦函数是定义于所有实数的，可以为任意的位置提供唯一的编码。
尽管绝对位置编码看似只提供了位置的信息，但模型在训练过程中会学会捕捉相对位置信息。原因如下：
1、位置之间的差异：通过将位置编码加到输入向量中，模型可以学习到位置之间的相对关系。例如，给定位置 ( i ) 和位置 ( j )，它们的编码向量可以表示为：
PE(i)−PE(j)
这一差值可以在一定程度上反映这两个位置之间的相对距离。
2、向量性质：在高维空间中，向量之间的方向和距离能够也反映相对位置。例如，当两个词在序列中相隔一定距离时，它们的相应位置编码的差异会隐含这种相对关系。
import torch
import torch.nn as nn
import math
class PositionalEncoding(nn.Module):
"""
实现经典的基于正弦和余弦函数的绝对位置编码。
"""
def __init__(self, d_model, max_len=5000, dropout=0.1):
"""
Args:
d_model (int): 模型的维度（或词嵌入的维度）。
max_len (int): 预先计算编码的最大序列长度。
dropout (float): Dropout 的比例。
"""
super(PositionalEncoding, self).__init__()
self.dropout = nn.Dropout(p=dropout)
# 创建一个足够长的位置编码矩阵
# 形状为 (max_len, d_model)
pe = torch.zeros(max_len, d_model)
# 创建一个位置索引张量
# position.shape: (max_len, 1)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
# 计算除法项 1 / (10000^(2i/d_model))
# div_term.shape: (d_model/2)
# 这里的 log 是为了数值稳定性，等价于 1 / (10000^(2i/d_model))
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
# 使用广播机制计算正弦和余弦值
# 偶数索引使用 sin
pe[:, 0::2] = torch.sin(position * div_term)
# 奇数索引使用 cos
pe[:, 1::2] = torch.cos(position * div_term)
# 增加一个 batch 维度，使其能够与输入张量 (batch_size, seq_len, d_model) 直接相加
# pe.shape from (max_len, d_model) to (1, max_len, d_model)
pe = pe.unsqueeze(0)
# 将 pe 注册为 buffer。
# buffer 是模型的状态的一部分，但不是参数 (parameters)，因此不会被优化器更新。
# 它会随着模型一起移动（例如 .to(device)），并且会保存在 state_dict 中。
self.register_buffer('pe', pe)
def forward(self, x):
"""
将位置编码添加到输入张量中。
Args:
x (torch.Tensor): 输入张量，形状为 (batch_size, seq_len, d_model)。
Returns:
torch.Tensor: 添加了位置编码的输出张量，形状与输入相同。
"""
# x.shape: (batch_size, seq_len, d_model)
# self.pe.shape: (1, max_len, d_model)
# 截取所需长度的位置编码并与输入相加
# self.pe[:, :x.size(1), :] 的形状变为 (1, seq_len, d_model)，可以与 x 广播相加
x = x + self.pe[:, :x.size(1), :]
return self.dropout(x)
二、相对位置编码
相对位置编码（Relative Positional Encoding）是一种改善模型捕捉序列中词语相对位置关系的技术。与绝对位置编码（Absolute Positional Encoding）不同，相对位置编码侧重于表示词与词之间的相对距离，从而增强模型的学习能力。
但相对位置编码的实现和绝对位置编码的实现差距还是挺大的，相对位置编码以矩阵的形式呈现，计算每个词i 和 词j 的相对位置，如果序列长度过大，会导致整个位置编码十分庞大，计算成本巨大。
标准的相对位置编码方法（如论文 “Self-Attention with Relative Position Representations” 中提出的）并不会直接预计算所有位置对的编码。相反，它创建了一个可学习的相对位置嵌入——relative position embedding查找表。
实现：
import torch
import torch.nn as nn
import math
class RelativeAttentionScore(nn.Module):
def __init__(self, d_model, max_len=50):
super(RelativeAttentionScore, self).__init__()
# 假设 d_k = d_model，在多头注意力中 d_k = d_model / num_heads
self.d_k = d_model
# 定义一个最大相对距离。超过这个距离的位置被视为相同距离。
# 这有助于模型泛化到比训练时更长的序列。
self.max_relative_position = max_len
# 创建一个可学习的嵌入查找表，用于相对位置。
# 大小为 2 * max_len - 1，覆盖从 -(max_len-1) 到 (max_len-1) 的所有相对位置。
# +1 是因为我们需要一个位置来存储被裁剪的距离
self.relative_embeddings = nn.Embedding(2 * self.max_relative_position + 1, self.d_k)
def forward(self, queries, keys):
"""
Args:
queries (torch.Tensor): 查询张量，形状为 (batch_size, seq_len, d_model)
keys (torch.Tensor): 键张量，形状为 (batch_size, seq_len, d_model)
Returns:
torch.Tensor: 加上了相对位置偏置的注意力得分，形状为 (batch_size, seq_len, seq_len)
"""
batch_size, seq_len, _ = queries.shape
# 1. 计算基于内容的注意力得分
content_score = torch.matmul(queries, keys.transpose(-2, -1))
# 2. 计算基于位置的注意力得分
# a. 生成相对位置矩阵
# range_vec.shape: (seq_len)
range_vec = torch.arange(seq_len, device=queries.device)
# relative_matrix.shape: (seq_len, seq_len)
# 每一行表示当前位置与所有其他位置的相对距离
relative_matrix = range_vec[None, :] - range_vec[:, None]
print(relative_matrix)
# b. 裁剪相对距离并移动到非负索引
# 将距离裁剪到 [-max_relative_position, max_relative_position] 范围内
# print(self.max_relative_position)
clipped_relative_matrix = torch.clamp(
relative_matrix,
-self.max_relative_position,
self.max_relative_position
)
# 将索引平移到 [0, 2 * max_relative_position] 范围，以用于Embedding查找
positive_indices = clipped_relative_matrix + self.max_relative_position
# c. 查找相对位置嵌入
# pos_embeddings.shape: (seq_len, seq_len, d_k)
pos_embeddings = self.relative_embeddings(positive_indices)
# d. 计算位置得分
# 我们需要计算 query 和 相对位置嵌入 的点积。
# (b, i, d) * (i, j, d) -> (b, i, j)
# b=batch_size, i=query_pos, j=key_pos, d=d_k
# torch.einsum 是实现这种复杂矩阵乘法的优雅方式
position_score = torch.einsum('bid,ijd->bij', queries, pos_embeddings)
# 3. 将内容得分和位置得分相加
attention_scores = content_score + position_score
# （可选）应用缩放
scaled_attention_scores = attention_scores / math.sqrt(self.d_k)
return scaled_attention_scores
--- 相对位置矩阵示例 (seq_len=5) ---
原始相对位置矩阵:
tensor([[ 0, 1, 2, 3, 4],
[-1, 0, 1, 2, 3],
[-2, -1, 0, 1, 2],
[-3, -2, -1, 0, 1],
[-4, -3, -2, -1, 0]])
裁剪后的矩阵 (max_len=2):
tensor([[ 0, 1, 2, 2, 2],
[-1, 0, 1, 2, 2],
[-2, -1, 0, 1, 2],
[-2, -2, -1, 0, 1],
[-2, -2, -2, -1, 0]])
用于嵌入查找的非负索引:
tensor([[2, 3, 4, 4, 4],
[1, 2, 3, 4, 4],
[0, 1, 2, 3, 4],
[0, 0, 1, 2, 3],
[0, 0, 0, 1, 2]])
由于裁剪机制的存在，任何超出 [-max_relative_position, max_relative_position] 范围的相对距离都会被“压缩”到边界值上。这个机制其实很好的帮助了模型：
1）假设用最大长度为512的句子训练了一个模型，max_relative_position 也设为512。现在，想用这个模型去处理一个长度为1024的句子。如果没有裁剪：模型在处理这个长句子时，会遇到它从未见过的相对距离，比如 +600 或 -800。由于嵌入表中没有这些距离的位置，或者这些位置的嵌入向量从未被训练过，模型的表现会变得非常不稳定，甚至完全崩溃。
有了裁剪：模型在训练时已经学会了一个对于“非常远”的距离（比如+512或-512）的通用表示。当它在推理时遇到一个新的、更远的距离（如 +600）时，它会将其裁剪到 +512，然后使用那个它已经熟知的“非常远”的嵌入。这使得模型能够平滑地泛化到比训练时更长的序列，而不会因为遇到未知的距离而失败。
2）在自然语言中，词语之间的关系强度通常与距离有关，但这种关系不是无限延伸的。模型通过裁剪，学会了一个“局部注意力窗口”（在 [-50, 50] 范围内），并对这个窗口内的位置进行精细建模。对于窗口外的所有位置，它只学习一个统一的“远距离”表示。这是一种非常合理的归纳偏置（inductive bias）。
三、旋转位置编码
好的，旋转位置编码（Rotary Positional Encoding, RoPE）是目前大型语言模型（如 LLaMA, PaLM）中非常流行且效果出色的一种位置编码方案。它由苏建林在论文《RoFormer: Enhanced Transformer with Rotary Position Embedding》中提出。
与传统的加性位置编码（Absolute PE）或在注意力分数上增加偏置（Relative PE from Shaw et al.）不同，RoPE 的思想极为巧妙：它通过旋转查询（Query）和键（Key）向量来注入位置信息。
其核心思想为：
绝对位置决定初始角度：一个词在序列中的绝对位置 m 决定了它的查询向量 q 和键向量 k 需要旋转的角度 mθ。相对位置体现在角度差：当计算两个词（位置m的查询q和位置n的键k）的注意力时，它们旋转后的点积结果，神奇地只与它们的内容 (q, k) 和它们的相对位置 m-n 有关，而与它们的绝对位置 m 和 n 无关
数学原理为：
RoPE 的魔法在于复数运算。将 d 维的向量两两配对，看作 d/2 个复数。对一个位于位置 m 的向量 x（它可以是 q 或 k），其旋转操作可以表示为：
x’_m = x_m * e^(i * m * θ) ; x_m 是原始向量（被看作复数）。 i 是虚数单位。 m 是绝对位置。θ 是一个预设的、与维度相关的常数（类似于传统PE中的频率）
当计算旋转后的查询 q’_m 和键 k’_n 的点积（在复数域中是取共轭后相乘再取实部）时：
Re[ (q_m * e^(imθ)) * (k_n * e(i*nθ)) ]
= Re[ q_m * k_n^* * e^(imθ) * e^(-inθ) ]
= Re[ q_m * k_n^* * e^(i*(m-n)θ) ]
最终结果仅依赖于m-n
import torch
import torch.nn as nn
import math
class RotaryPositionalEncoding(nn.Module):
def __init__(self, d_model, max_len=512):
super().__init__()
self.d_model = d_model
# 计算旋转角度 theta
# theta_i = 10000^(-2(i-1)/d) for i in [1, 2, ..., d/2]
inv_freq = 1.0 / (10000 ** (torch.arange(0, d_model, 2).float() / d_model))
# 预先计算所有可能位置 m 的 m*theta
t = torch.arange(max_len, dtype=inv_freq.dtype)
freqs = torch.einsum('i,j->ij', t, inv_freq)
# freqs 包含了所有位置的 m*theta, 形状是 (max_len, d_model/2)
# 将其扩展为 (max_len, d_model) 以便应用
# emb.shape: (max_len, d_model)
emb = torch.cat((freqs, freqs), dim=-1)
# 注册为 buffer，这样它就不会被视为模型参数，但会随模型移动 (e.g., .to(device))
# self.cos_cached.shape: (1, 1, max_len, d_model)
# self.sin_cached.shape: (1, 1, max_len, d_model)
self.register_buffer("cos_cached", emb.cos()[None, None, :, :])
self.register_buffer("sin_cached", emb.sin()[None, None, :, :])
def forward(self, x):
# x.shape: (batch_size, num_heads, seq_len, head_dim)
# head_dim == self.d_model
seq_len = x.shape[-2]
# 获取预计算的 cos 和 sin 值
cos = self.cos_cached[:, :, :seq_len, ...]
sin = self.sin_cached[:, :, :seq_len, ...]
# 执行旋转
# 1. 将 x 分为两半
# x1.shape, x2.shape: (batch_size, num_heads, seq_len, head_dim/2)
x1 = x[..., 0::2] # 偶数维度
x2 = x[..., 1::2] # 奇数维度
# 2. 应用旋转公式
# x_rotated = (x1 + i*x2) * (cos + i*sin) = (x1*cos - x2*sin) + i*(x1*sin + x2*cos)
rotated_x1 = x1 * cos[..., 0::2] - x2 * sin[..., 0::2]
rotated_x2 = x1 * sin[..., 1::2] + x2 * cos[..., 1::2]
# 3. 将旋转后的两半合并
rotated_x = torch.cat([rotated_x1, rotated_x2], dim=-1)
return rotated_x
# --- 集成到多头注意力中 ---
class RoPEMultiHeadAttention(nn.Module):
def __init__(self, d_model, num_heads, max_len=512):
super().__init__()
assert d_model % num_heads == 0
self.d_model = d_model
self.num_heads = num_heads
self.head_dim = d_model // num_heads
self.q_proj = nn.Linear(d_model, d_model)
self.k_proj = nn.Linear(d_model, d_model)
self.v_proj = nn.Linear(d_model, d_model)
self.out_proj = nn.Linear(d_model, d_model)
self.rotary_encoder = RotaryPositionalEncoding(self.head_dim, max_len)
def forward(self, x, mask=None):
batch_size, seq_len, _ = x.shape
# 1. 线性投影
q = self.q_proj(x)
k = self.k_proj(x)
v = self.v_proj(x)
# 2. 改变形状以适应多头
# shape: (batch_size, num_heads, seq_len, head_dim)
q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
# 3. 对 Q 和 K 应用旋转位置编码
q = self.rotary_encoder(q)
k = self.rotary_encoder(k)
# 4. 计算注意力得分
# scores.shape: (batch_size, num_heads, seq_len, seq_len)
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
attention = torch.softmax(scores, dim=-1)
# 5. 应用注意力到 V
# context.shape: (batch_size, num_heads, seq_len, head_dim)
context = torch.matmul(attention, v)
# 6. 恢复形状并进行最终投影
context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
output = self.out_proj(context)
return output
RoPE 作为目前最受欢迎的位置编码，其优势如下：
良好的长度外推性：由于其相对位置的性质，RoPE 能够很好地泛化到比训练时更长的序列。
随距离增加而衰减的注意力：RoPE 的数学性质天然地使得随着相对距离的增加，注意力得分会有一个衰减的趋势，这符合语言直觉（离得越远的词关系越弱）。
高性能：它不引入额外的模型参数，并且计算非常高效，可以无缝集成到现有的自注意力框架中。