位置编码

位置编码 (Positional Encoding)

由于 Transformer 的 Self-Attention 机制在计算时不考虑 Input Token 的顺序（即具有 Permutation Invariant 排列不变性），为了让模型理解序列的顺序信息（如 “Tom chased Jerry” vs “Jerry chased Tom”），必须显式地向输入中注入位置信息。

位置编码主要分为三大类：绝对位置编码、相对位置编码 和 旋转位置编码 (RoPE)。

1. 绝对位置编码 (Absolute Positional Encoding)

最直观的方法：给每个位置索引 $t$ 分配一个固定的向量 $PE_t$，直接加到 Input Embedding 上。

1.1 正弦位置编码 (Sinusoidal Positional Encoding)

这是原始 Transformer (“Attention Is All You Need”) 中提出的方法。

公式： 对于位置 $pos$ 和维度 $i$（$d_{model}$ 为总维度）：

核心特性：

• 相对位置线性关系：对于任意偏移 $k$，是可以表示为 $PE_{pos}$ 的线性函数的。即 $PE_{pos+k} = T_k \cdot PE_{pos}$，其中 $T_k$ 是一个与 $pos$ 无关的旋转矩阵。这使得模型理论上可以学习到相对位置信息。
• 固定不需训练：减少了参数量。
• 外推性：理论上可以处理比训练时更长的序列，但实际效果有限。

PyTorch 实现：

import torch
import torch.nn as nn
import math

class SinusoidalPositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 5000):
        super().__init__()
        # 创建一个足够长的 PE 矩阵 [max_len, d_model]
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        
        # 注册为 buffer，不是参数，不更新梯度
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        # x shape: [batch_size, seq_len, d_model]
        # 截取当前序列长度对应的 PE 相加
        return x + self.pe[:, :x.size(1), :]

1.2 可学习位置编码 (Learnable Positional Embedding)

BERT, GPT-2 等早期模型使用。直接初始化一个 nn.Embedding(max_seq_len, d_model)，让模型自己学习每个位置的向量表示。

• 缺点：不能处理超过训练时 max_seq_len 的序列。

2. 相对位置编码 (Relative Positional Encoding, RPE)

核心思想：位置信息不应该加在 Input Embedding 上，而应该直接作用于 Attention Score 的计算中。我们要编码的是 Token $i$ 和 Token $j$ 之间的距离 $i-j$。

2.1 经典相对位置编码 (Shaw et al., 2018)

这是最早期的 RPE 方案。它为 Key 和 Value 引入了可学习的相对位置向量。 公式：

其中 $R^K$ 和 $R^V$ 是根据相对距离 $i-j$ 查表得到的向量。

2.2 简化相对位置偏置 (Relative Position Bias, e.g., T5)

T5 模型简化了这一过程，不再修改 $K$ 和 $V$ 向量，而是直接在计算出的 Attention Logits 上加一个可学习的标量偏置 (Scalar Bias)。

• 分桶策略 (Bucketing): T5 对相对距离进行了对数分桶处理（例如距离 50 和 51 可能共享同一个 Bias），以减少参数量并增强对未见长距离的泛化能力。
• 广泛应用: 这种方式因为计算高效，被 Swin Transformer 等许多后续模型采用。

3. 旋转位置编码 (Rotary Positional Embedding, RoPE)

目前最主流的位置编码方式（LLaMA, PaLM, ChatGLM, Qwen）。它巧妙地结合了绝对位置编码的实现便捷性和相对位置编码的数学性质。

3.1 核心思想

RoPE 不改变 Embedding 的模长，而是对 Query 和 Key 向量进行旋转。 通过绝对位置的旋转，使得两个向量的点积（Attention Score）仅依赖于它们的相对位置。

这里 $R_{n-m}$ 仅与相对距离 $n-m$ 有关。

3.2 数学推导 (Simplified)

在 2D 情况下，为了满足 $f(q, m) \cdot f(k, n) = g(q, k, m-n)$，RoPE 定义变换为复数域的旋转：

对应矩阵形式：

对于 $d$ 维向量，将其切分为 $d/2$ 个子空间，每两个维度为一组进行旋转。

3.3 PyTorch 实现

下面的代码演示了相邻配对的实现方式，更加简洁直观，适合理解原理：

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    """
    预计算复数形式的旋转角度 e^{i m theta}
    """
    # 计算词向量元素两两分组之后，每组对应的旋转角度theta_i
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device)
    # 外积生成所有位置的所有频率
    freqs = torch.outer(t, freqs).float()
    # 转换为复数形式: cos(x) + i*sin(x)
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # shape: [end, dim/2]
    return freqs_cis

def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):
    """
    应用 RoPE 旋转
    xq, xk: [batch, seq_len, n_head, head_dim]
    """
    # 1. Reshape 为复数形式: [..., head_dim] -> [..., head_dim/2, 2] -> complex
    # 这里假设是相邻配对 (Adjacent Pairs)
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    
    # 2. 调整 freqs_cis 形状以匹配 seq_len 并广播
    freqs_cis = freqs_cis[:xq.shape[1], :] # [seq_len, head_dim/2]
    
    # 广播维度: [1, seq_len, 1, head_dim/2]
    # 这样可以自动广播到 batch 和 n_head 维度
    freqs_cis = freqs_cis.view(1, xq.shape[1], 1, -1) 
    
    # 3. 旋转 (复数乘法)
    # (a+ib)(cos+isin) = ... 这一步 PyTorch 的复数乘法会自动完成
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    
    return xq_out.type_as(xq), xk_out.type_as(xk)

3.4 长文本外推 (Length Extrapolation)

当推理长度超过训练长度时，直接外推效果很差。RoPE 相关的改进主要为了解决这个问题：

1. 线性插值 (Linear Interpolation):

   • 由 Meta 在 LLaMA 中验证。
   • 原理：将推理时的长序列位置索引 $pos'$ 压缩回训练范围内。即令 $pos_{new} = pos / \kappa$，其中 $\kappa = L_{new} / L_{train}$。
   • 效果：简单有效，微调少量步数即可适配更长窗口。

2. NTK-Aware Scaling:

   • 原理：高频分量（低位）对位置信息敏感，不宜压缩；低频分量（高位）主要提供长距离信息，适合压缩。
   • 方法：通过修改 Base ($10000 \to 10000 \cdot \kappa^{d/(d-2)}$) 间接实现非线性插值。
   • 优点：无需微调也能在一定程度上外推。

3. YaRN (Yet another RoPE extension):

   • 在 NTK 的基础上，根据不同频段引入温度系数 $t$ 来平滑注意力分布（Entropy theory），是目前长文本效果最好的方案之一。

4. ALiBi (Attention with Linear Biases)

不使用 Embedding，直接修改 Attention 矩阵。

• 对 Attention Score 加上一个与距离成线性比例的惩罚项。
• $m$ 是每个 Head 特有的固定斜率（Geometric Sequence）。
• 优点：不需要 Position Embedding，外推性极强（State-of-the-art until RoPE scaling took over）。通常用于 BLOOM, MPT 等模型。

总结