3. 位置编码

概述

Transformer 的核心是注意力机制，但标准的注意力机制是排列不变（permutation invariant）的——打乱输入 token 的顺序，输出不变。为了让模型感知 token 的位置信息，需要显式地注入位置信号。

本章将介绍位置编码的演进历程：

绝对位置编码 → 相对位置编码 → RoPE 旋转位置编码 → Partial RoPE (Qwen3.5 采用)

3.1 绝对位置编码（Sinusoidal）

原理

原始 Transformer（Vaswani et al., 2017）使用正弦/余弦函数生成位置编码：

$$P{E}_{(pos,2i)}=\sin \left(\frac{pos}{{10000}^{2i/d_{model}}}\right)$$$$P{E}_{(pos,2i+1)}=\cos \left(\frac{pos}{{10000}^{2i/d_{model}}}\right)$$

其中：

• $pos$：token 在序列中的位置（0, 1, 2, ...）
• $i$：维度索引（0, 1, ..., $d_{model}/2-1$）
• $d_{model}$：模型隐藏维度

直觉理解：

• 不同维度使用不同频率的正弦/余弦波
• 低维度（$i$ 小）→ 高频率 → 变化快 → 编码细粒度位置
• 高维度（$i$ 大）→ 低频率 → 变化慢 → 编码粗粒度位置
• 类似于二进制编码：最低位变化最快，最高位变化最慢

代码实现

代码文件

import torch
import torch.nn as nn
import math


class AbsolutePositionEmbedding(nn.Module):
    """
    Transformer 经典的正弦余弦绝对位置编码
    核心原理：通过不同频率的正弦/余弦函数对不同位置、不同维度进行编码
    """

    def __init__(self, d_model, max_seq_len):
        super().__init__()
        # 1. 生成位置索引: (max_seq_len, 1)
        pos = torch.arange(max_seq_len).unsqueeze(1)

        # 2. 生成维度索引（偶数维度）: (1, d_model//2)
        dim_pos = torch.arange(0, d_model, 2).unsqueeze(0)

        # 3. 计算频率衰减项: 10000^(-2i/d_model)
        div_term = torch.exp(-math.log(10000) * dim_pos.float() / d_model)

        # 4. 计算位置-角度: (max_seq_len, d_model//2)
        angle = pos * div_term

        # 5. 初始化位置编码矩阵: (max_seq_len, d_model)
        pe = torch.zeros(max_seq_len, d_model)

        # 6. 偶数维度用 sin，奇数维度用 cos
        pe[:, 0::2] = torch.sin(angle)
        pe[:, 1::2] = torch.cos(angle)

        # 7. 注册为 buffer（不参与梯度更新）
        self.register_buffer("pe", pe)

    def forward(self, x):
        # x: (batch_size, seq_len, d_model)
        return self.pe[:x.size(1)]  # (seq_len, d_model)

使用方式

直接将绝对位置编码使用相加的方式在embedding后的向量上。在如何融合位置编码信息后面的方式不一样，需要留意。

# 位置编码加到 token embedding 上
token_emb = embed_tokens(input_ids)           # (batch, seq, d_model)
pos_emb = position_embedding(token_emb)        # (seq, d_model)
hidden_states = token_emb + pos_emb            # 广播相加

绝对位置编码的局限性

1. 位置信息与内容解耦：位置编码是固定的，不随输入内容变化
2. 泛化长度有限：超过 max_seq_len 的位置没有定义
3. 相对位置信息隐式：虽然正弦编码理论上能表示相对位置（通过线性变换），但模型需要学习这种关系

3.2 相对位置编码（Relative Position Encoding）

原理

相对位置编码不关心 token 的绝对位置，而是关心两个 token 之间的相对距离。核心思想是在注意力计算中加入相对位置偏置。

Shaw et al. (2018) 的方法：

1. 定义最大相对距离 $k$（如 50），超出范围的截断
2. 为每个相对距离学习两个 Embedding：$R_k$（用于 Key）和 $R_v$（用于 Value）
3. 在注意力计算中，将相对位置信息融入得分和输出

代码实现

代码文件

class RelativePositionEncoding(nn.Module):
    def __init__(self, d_model, max_relative_position=50):
        super().__init__()
        self.d_model = d_model
        self.max_relative_position = max_relative_position
        vocab_size = 2 * max_relative_position + 1  # [-50, 50] → 101 个值
        self.relative_position_k = nn.Embedding(vocab_size, d_model)
        self.relative_position_v = nn.Embedding(vocab_size, d_model)

    def get_relative_positions(self, seq_len):
        """计算相对位置矩阵"""
        range_vec = torch.arange(seq_len)
        # distance_mat[i][j] = j - i
        distance_mat = range_vec.unsqueeze(0) - range_vec.unsqueeze(1)
        # 截断到 [-max, max]
        distance_mat = torch.clamp(distance_mat, -self.max_relative_position, self.max_relative_position)
        # 映射到 [0, 2*max]
        return distance_mat + self.max_relative_position

    def forward(self, seq_len):
        relative_positions = self.get_relative_positions(seq_len)
        relative_positions_k_emb = self.relative_position_k(relative_positions)  # (seq, seq, d_model)
        relative_positions_v_emb = self.relative_position_v(relative_positions)  # (seq, seq, d_model)
        return relative_positions_k_emb, relative_positions_v_emb

融入注意力计算

标准注意力得分

Scores = Q · K^T / √d_k

相对位置注意力得分

Scores = Q · K^T / √d_k + Q · R_k^T / √d_k = 内容-内容得分 + 内容-位置得分

最终输出（扩展版）

Output = softmax(scores) · V + softmax(scores) · R_v= 常规注意力输出 + 相对位置值输出

相对位置编码的局限性

1. 显存开销大：需要存储 $(seq\mathrm{\_}len,seq\mathrm{\_}len,d\mathrm{\_}model)$ 的相对位置 Embedding
2. 计算开销大：注意力计算多了两次矩阵乘法
3. 不支持高效推理：相对位置矩阵依赖序列长度

3.3 RoPE：旋转位置编码（Qwen3.5 采用）

原理

RoPE（Rotary Position Embedding）由 Su et al. (2021) 提出，是目前最流行的位置编码方式。核心思想：通过旋转变换来编码位置信息，使得内积天然包含相对位置信息。

数学推导：

对于位置 $m$ 的 query 和位置 $n$ 的 key，RoPE 编码后：

$$⟨f_q(x_m,m),f_k(x_n,n)⟩=g(x_m,x_n,m-n)$$

即：编码后的内积只依赖相对位置 $m-n$，而不是绝对位置 $m$ 和 $n$。

具体实现：将 $d$ 维向量分成 $d/2$ 组，每组 2 维，进行 2D 旋转：

$$f(x,m)=\left(\begin{array}{c}{x}_0\\ x_1\\ x_2\\ x_3\\ ⋮\end{array}\right)\odot \left(\begin{array}{c}\cos (m{\theta }_0)\\ \cos (m{\theta }_0)\\ \cos (m{\theta }_1)\\ \cos (m{\theta }_1)\\ ⋮\end{array}\right)+\left(\begin{array}{c}-x_1\\ x_0\\ -x_3\\ x_2\\ ⋮\end{array}\right)\odot \left(\begin{array}{c}\sin (m{\theta }_0)\\ \sin (m{\theta }_0)\\ \sin (m{\theta }_1)\\ \sin (m{\theta }_1)\\ ⋮\end{array}\right)$$

其中 ${\theta }_i={10000}^{-2i/d}$ 是不同维度的频率。

直觉理解：

• 每对维度 $(x_{2i},x_{2i+1})$ 看作一个 2D 向量
• 位置 $m$ 对应旋转角度 $m\cdot {\theta }_i$
• 不同维度对有不同的旋转频率
• Q 和 K 的内积 = 原始内积 × 旋转角度差的 cos → 天然包含相对位置信息

代码实现

basic 版本

代码文件

class RotaryEmbedding(nn.Module):
    def __init__(self, hidden_size, num_heads, base=10000, max_seq_len=512):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads
        self.base = base
        self.max_seq_len = max_seq_len

        # 缓存 cos 和 sin
        self.register_buffer("cos_pos_cache", torch.empty(max_seq_len, self.head_dim))
        self.register_buffer("sin_pos_cache", torch.empty(max_seq_len, self.head_dim))
        self._compute_pos_emb()

    def _compute_pos_emb(self):
        """计算旋转位置编码缓存"""
        # 1. 计算频率: theta_i = 1 / (base^(2i/d))
        theta = 1.0 / (self.base ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))

        # 2. 生成位置序列: [0, 1, ..., max_seq_len-1]
        positions = torch.arange(self.max_seq_len).float()

        # 3. 计算位置-角度乘积: (max_seq_len, head_dim//2)
        pos_theta = positions.unsqueeze(1) * theta.unsqueeze(0)

        # 4. 扩展为完整的 head_dim（奇偶维度复用同一组 cos/sin）
        cos_pos = torch.cos(pos_theta).repeat_interleave(2, dim=-1)
        sin_pos = torch.sin(pos_theta).repeat_interleave(2, dim=-1)

        self.cos_pos_cache.copy_(cos_pos)
        self.sin_pos_cache.copy_(sin_pos)

    def forward(self, q):
        """
        对 Query 进行旋转位置编码
        q: (batch_size, num_heads, seq_len, head_dim)
        """
        batch_size, num_heads, seq_len, head_dim = q.size()

        # 取缓存
        cos_pos = self.cos_pos_cache[:seq_len].unsqueeze(0).unsqueeze(0)
        sin_pos = self.sin_pos_cache[:seq_len].unsqueeze(0).unsqueeze(0)

        # 构造旋转后的 q2：偶数位取负，奇数位保留
        q2 = torch.stack([-q[..., 1::2], q[..., ::2]], dim=-1).reshape(q.shape)

        # RoPE 核心公式: q' = q * cos + q2 * sin
        q_rot = q * cos_pos + q2 * sin_pos
        return q_rot

miniqwen 版本（与 Qwen3.5 对齐）

代码文件

class RotaryEmbedding(nn.Module):
    """RoPE 旋转位置编码"""

    def __init__(self, dim: int, max_position_embeddings: int = 262144, theta: float = 10000000.0):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.theta = theta

        # 计算频率: 1 / (theta ^ (2i / dim))
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Args:
            position_ids: (batch_size, seq_len)
        Returns:
            cos: (batch_size, seq_len, dim)
            sin: (batch_size, seq_len, dim)
        """
        # (batch_size, seq_len, dim/2)
        freqs = position_ids.unsqueeze(-1).float() * self.inv_freq.unsqueeze(0).unsqueeze(0)
        # (batch_size, seq_len, dim)
        emb = torch.cat([freqs, freqs], dim=-1)
        return emb.cos(), emb.sin()


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """旋转隐藏维度的后半部分: [-x2, x1]"""
    x1 = x[..., :x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2:]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    unsqueeze_dim: int = 1,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    对 Q 和 K 应用 RoPE

    Args:
        q: (batch, heads, seq_len, head_dim)
        k: (batch, heads, seq_len, head_dim)
        cos: (batch, seq_len, rotary_dim)
        sin: (batch, seq_len, rotary_dim)
    """
    cos = cos.unsqueeze(unsqueeze_dim)  # (batch, 1, seq_len, rotary_dim)
    sin = sin.unsqueeze(unsqueeze_dim)

    # 对齐官方：把 cos/sin 转到 q/k 的 dtype，保证旋转在模型精度（如 bf16）下进行，
    # 否则 bf16 模型下 q_rot * cos(fp32) 会把 attention 抬高到 fp32，产生逐层累积差异
    cos = cos.to(q.dtype)
    sin = sin.to(k.dtype)

    # 分离需要旋转的部分和不需要旋转的部分（Partial RoPE）
    rotary_dim = cos.shape[-1]
    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]

    # 应用旋转: q' = q_rot * cos + rotate_half(q_rot) * sin，后半部分直通
    q_embed = torch.cat([q_rot * cos + rotate_half(q_rot) * sin, q_pass], dim=-1)
    k_embed = torch.cat([k_rot * cos + rotate_half(k_rot) * sin, k_pass], dim=-1)

    return q_embed, k_embed

两种实现的对比：

特点	basic 版本	miniqwen 版本
输入	直接接收 q	接收 position_ids，返回 cos/sin
应用方式	在类内部完成	外部调用 apply_rotary_pos_emb
旋转方式	[-q[..., 1::2], q[..., ::2]]	rotate_half(x)
等价性	✅ 逻辑等价	✅ 逻辑等价

RoPE 的优势

1. 天然包含相对位置信息：内积只依赖相对位置差
2. 无需额外参数：频率是固定的，不需要学习
3. 支持外推：通过 NTK-aware 缩放等技术可以扩展到更长序列
4. 高效实现：只需要逐元素乘法和加法

3.4 Partial RoPE（Qwen3.5 的创新）

原理

标准 RoPE 对 head_dim 的所有维度都应用旋转。但 Qwen3.5 使用 Partial RoPE，只对 head_dim 的前 25%（64 维）应用 RoPE，剩余 192 维直通。

head_dim = 256
rotary_dim = 64 (= 256 × 0.25)

Q = [q_rot_0, q_rot_1, ..., q_rot_63 | q_pass_64, q_pass_65, ..., q_pass_255]
     ↑ 前 64 维应用 RoPE                        ↑ 后 192 维直通

K = [k_rot_0, k_rot_1, ..., k_rot_63 | k_pass_64, k_pass_65, ..., k_pass_255]
     ↑ 前 64 维应用 RoPE                        ↑ 后 192 维直通

实现方式

def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)

    # 分离需要旋转的部分和不需要旋转的部分
    rotary_dim = cos.shape[-1]
    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]

    # 只对前 rotary_dim 维应用旋转
    q_embed = torch.cat([q_rot * cos + rotate_half(q_rot) * sin, q_pass], dim=-1)
    k_embed = torch.cat([k_rot * cos + rotate_half(k_rot) * sin, k_pass], dim=-1)

    return q_embed, k_embed

为什么使用 Partial RoPE？

1. 位置信息 vs 语义信息的权衡：

   • 位置编码帮助模型理解 token 间的距离关系
   • 但并非所有维度都需要位置信息，有些维度更适合编码语义信息
   • Partial RoPE 让模型在不同维度上分别学习位置和语义

2. 计算效率：更少的维度需要旋转，减少计算量

3. 实验效果好：Qwen3.5 使用 partial_rotary_factor=0.25，即只有 25% 的维度做 RoPE

Qwen3.5 的配置

# config.py
head_dim: int = 256
partial_rotary_factor: float = 0.25
rotary_dim = int(head_dim * partial_rotary_factor)  # 64
rope_theta: float = 10000000.0        # 官方 0.8B 用 1e7
max_position_embeddings: int = 262144

3.5 位置编码对比总结

位置编码	位置信息	额外参数	计算开销	长度泛化	代表模型
绝对位置（Sinusoidal）	绝对	无	低	差	原始 Transformer
绝对位置（Learned）	绝对	有	低	差	BERT, GPT-2
相对位置	相对	有	高	中	T5
RoPE	相对（旋转）	无	低	好	Qwen3.5, LLaMA
Partial RoPE	部分维度编码	无	更低	好	Qwen3.5

Qwen3.5 的选择：

• 使用 RoPE，官方 Qwen3.5-0.8B 的 rope_theta = 10000000（1e7，比经典 1e4 大得多，支持更长上下文的外推）
• partial_rotary_factor = 0.25，只对 head_dim 的前 64 维做 RoPE
• 支持最大 max_position_embeddings = 262144（256K）

---

上一章：2. 归一化层 | 下一章：4. 注意力机制