LLM：位置编码详解与实现

---

前言

由于attetnion运算的特性，Transformer本身不感知顺序，位置编码是弥补这一缺陷的关键。

---

一、绝对位置编码

绝对位置编码的方式是通过将每个位置映射到一个高维空间中，该编码采用了正弦和余弦函数的组合
周期性：正弦和余弦函数具有周期性，这使得编码能够容易地表示固定范围的位置信息。这个特点允许模型有能力处理序列的循环性质，例如在语言中，某些词可能在不同上下文中重复出现，但它们的相对位置仍然是重要的。

不同频率：对于不同的维度 ( i )，使用不同的频率来编码位置。低维度的编码（例如低 ( i ) 值）对应较大的周期，能够捕捉较长的依赖关系；而高维度编码对应较小的周期，能够捕捉短距离的依赖关系。因此，模型可以通过不同的维度考虑不同尺度的位置信息。

任意长度的序列：这种方法能够处理任意长度的输入序列，因为正弦和余弦函数是定义于所有实数的，可以为任意的位置提供唯一的编码。

尽管绝对位置编码看似只提供了位置的信息，但模型在训练过程中会学会捕捉相对位置信息。原因如下：1、位置之间的差异：通过将位置编码加到输入向量中，模型可以学习到位置之间的相对关系。例如，给定位置 ( i ) 和位置 ( j )，它们的编码向量可以表示为：
PE(i)−PE(j)
这一差值可以在一定程度上反映这两个位置之间的相对距离。
2、向量性质：在高维空间中，向量之间的方向和距离能够也反映相对位置。例如，当两个词在序列中相隔一定距离时，它们的相应位置编码的差异会隐含这种相对关系。

import torch
import torch.nn as nn
import mathclass 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))# 使用广播机制计算正弦和余弦值# 偶数索引使用 sinpe[:, 0::2] = torch.sin(position * div_term)# 奇数索引使用 cospe[:, 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 mathclass RelativeAttentionScore(nn.Module):def __init__(self, d_model, max_len=50):super(RelativeAttentionScore, self).__init__()# 假设 d_k = d_model，在多头注意力中 d_k = d_model / num_headsself.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 mathclass 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*thetat = 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_modelseq_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 == 0self.d_model = d_modelself.num_heads = num_headsself.head_dim = d_model // num_headsself.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 的数学性质天然地使得随着相对距离的增加，注意力得分会有一个衰减的趋势，这符合语言直觉（离得越远的词关系越弱）。
高性能：它不引入额外的模型参数，并且计算非常高效，可以无缝集成到现有的自注意力框架中。