模型学习笔记 · 第一章

位置编码:Transformer 怎样记住顺序?

从正弦编码到 RoPE,理解位置信息放在哪里,以及“能算更长”与“能用好更长”的区别。

Muxuan Liu · 模型学习笔记

为什么还要专门告诉模型顺序

“狗追猫”和“猫追狗”用到的字差不多,意思却不同。模型需要知道谁在前、谁在后。没有位置信息、也没有顺序相关掩码的完整自注意力,只会随输入排列相应地改变输出排列,并没有专门表示“第一个、第二个”的部件。

从工程角度看,问题是:把顺序放进计算的哪一步?可以给输入附上位置向量,也可以直接改变两个位置之间的交互方式。因果掩码同样会引入顺序约束,所以不能笼统地说所有 Transformer 都完全没有顺序。

绝对位置:给每个位置一个地址

绝对位置编码回答“我在第几位”。原始 Transformer 在进入网络之前,把正弦位置向量加到输入嵌入上;它的定义并不是分别往每一层的 Q、K、V 都加一次。另一种做法是学习位置表,每一行代表一个位置,随训练更新。

可学习位置表有预设范围。想扩展到更长序列,需要插值、继续训练等额外设计。每个位置一行,并不意味着网络完全学不到位置之间的关系,只是这种表没有显式写出相对距离。

正弦编码:几根转速不同的时钟指针

可以把正弦编码想成一组转速不同的指针:合在一起,它们的读数用于表示当前位置。每两个维度是一对 sin 和 cos,频率逐渐降低。下面公式中的 p 是位置,d 是向量维度,i 是维度对的编号。

频率因子是 10000^(−2i/d),负号不能漏掉。写成正指数,得到的就是另一种频率安排。公式能计算训练中没出现过的位置,但“能计算”不等于“在超长文本中一定表现好”,后者还取决于模型训练。

PE(p, 2i) = sin(p · 10000^(−2i/d)) PE(p, 2i+1) = cos(p · 10000^(−2i/d))

从绝对位置到相对位移

我试着用三角函数的和角公式理解位置之间的关系。只说“sin 和 cos 能表示相对位置”还不够清楚,因此我把其中一对分量单独拿出来。令 ω 为这一对分量的频率,u(p) = [sin(ωp), cos(ωp)]ᵀ。

sin(ω(p+Δ)) = sin(ωp) cos(ωΔ) + cos(ωp) sin(ωΔ) cos(ω(p+Δ)) = cos(ωp) cos(ωΔ) − sin(ωp) sin(ωΔ) u(p+Δ) = A(Δ) u(p) A(Δ) = [[ cos(ωΔ), sin(ωΔ)], [−sin(ωΔ), cos(ωΔ)]]

我关注的是:这个变换矩阵只取决于偏移 Δ,不取决于起点 p。也就是说,同样的位移,在不同起点上可以用同一个线性变换表达。这让我理解了为什么绝对位置向量也能为学习相对关系提供结构,而不是把“绝对”和“相对”看成完全无关的两类。

不过,这个性质只描述位置向量之间的关系。把位置向量加到 token embedding 上,再经过学习到的投影,并不能保证最终注意力分数只依赖相对距离。RoPE 则把旋转直接用于 query 和 key,使点积中的位置变换显式体现相对位移;其中仍然包含内容向量的作用。

相对位置:直接告诉模型相隔多远

有时“在左边三格”比“在第十七格”更重要。相对位置编码就把这种距离放进交互里。Shaw 等人的方法在 key/value 的计算中加入相对位置向量;如果把距离截断到包含端点的 [−k, k],共有 2k+1 个距离桶。

相对位置方法是一类设计,不全是在 softmax 前往分数矩阵里加一项。有人加入分数偏置,有人改变 key/value 的贡献,RoPE 则旋转 query 和 key。看具体公式,比只记分类名更有用。

RoPE:把位置变成旋转角度

RoPE 根据位置旋转 query 和 key 的成对坐标。两个旋转后的向量做点积时,会体现相对角度。R(m)ᵀR(n) = R(n−m) 这个关系说明:各自按照绝对位置旋转,比较时可以得到相对位置差。

RoPE 本身不会把稠密注意力的平方复杂度变成线性,也不会自动赋予无限上下文。频率设置、缩放策略、训练长度和注意力实现都会影响结果。它解决的是位置信息的表示问题。

q̂ₘ = R(m)qₘ, k̂ₙ = R(n)kₙ q̂ₘᵀk̂ₙ = qₘᵀ R(n−m) kₙ

用张量形状读懂代码

示例输入是 [batch, sequence, d_model],位置表保存为 [1, max_len, d_model],利用广播加到整个 batch。register_buffer 让位置表随模型保存、迁移设备,但不把它当成可训练参数;它并不是禁止所有相关运算参与梯度计算。

代码兼容奇数和偶数维度。max_len 是这个实现预先分配的表长,不是正弦公式的理论上限。做确定性核对时关闭 dropout;正常使用时,dropout 会遵循模块的训练或评估模式。

展开 PyTorch 示例
import math
import torch
from torch import nn

class SinusoidalPositionEncoding(nn.Module):
    """Input: [batch, sequence, d_model]."""
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        if d_model < 1 or max_len < 1:
            raise ValueError('d_model and max_len must be positive')
        pos = torch.arange(max_len, dtype=torch.float32)[:, None]
        freq = torch.exp(
            torch.arange(0, d_model, 2, dtype=torch.float32)
            * (-math.log(10000.0) / d_model)
        )
        pe = torch.zeros(max_len, d_model)
        pe[:, 0::2] = torch.sin(pos * freq)
        pe[:, 1::2] = torch.cos(pos * freq[:d_model // 2])
        self.register_buffer('pe', pe[None, :, :])
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        if x.ndim != 3 or x.shape[-1] != self.pe.shape[-1]:
            raise ValueError('Expected [batch, sequence, d_model]')
        if x.shape[1] > self.pe.shape[1]:
            raise ValueError('Sequence exceeds configured max_len')
        return self.dropout(x + self.pe[:, :x.shape[1]].to(dtype=x.dtype))