点击“AladdinEdu,你的AI学习实践工作坊”,注册即送-H卡级别算力,沉浸式云原生集成开发环境,80G大显存多卡并行,按量弹性计费,教育用户更享超低价。
1. 引言:当自注意力忘记了顺序
Transformer架构的核心是自注意力机制。其计算公式为:
[
\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left( \frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{d_k}} \right) \mathbf{V}
]
仔细观察这一公式,你会发现一个容易被忽略的事实:它对输入序列的排列是等变的。如果将输入序列中的两个词交换位置,自注意力输出的对应位置也会相应交换,但每个位置的输出值本身不会因位置改变而变化。换句话说,纯自注意力不包含任何关于token顺序的信息。
对于语言而言,顺序至关重要。“狗咬人”与“人咬狗”含义截然相反;“我喜欢你”与“你喜欢我”情感关系完全不同。因此,必须设法将位置信息注入到模型中。这一注入方式,即为位置编码。
从Transformer原论文的正弦位置编码,到BERT的可学习绝对位置嵌入,再到相对位置编码、旋转位置编码(RoPE)以及最新的ALiBi,位置编码的设计经历了深刻的演进。不同的位置编码方案,不仅影响模型对序列长度的泛化能力(长度外推),也关乎训练稳定性与计算效率。在LLaMA、GPT-NeoX、PaLM等大模型中,RoPE已成为事实标准;而在某些高效架构中,ALiBi以其简洁有效崭露头角。
本文将系统性地解析位置编码的技术脉络。我们将从绝对位置编码出发,剖析其“为每个位置赋予唯一向量”的思想及局限性;继而深入相对位置编码,探讨如何在注意力分数中注入相对距离偏置;最后重点阐释旋转位置编码的数学之美——如何用复数旋转将相对位置信息隐式编码进内积。文章将结合数学推导、代码实现与模型案例,为你揭开位置编码的神秘面纱。
2. 绝对位置编码:为每个位置赋予身份
绝对位置编码的核心思想是:为序列中的每一个绝对位置索引 (pos) 生成一个唯一的向量表示 (\mathbf{p}_{pos}),并将其与词嵌入相加,作为模型的输入。
2.1 正弦位置编码
Transformer原论文采用了正弦位置编码。对于位置 (pos) 和维度索引 (i)((0 \le i < d_{\text{model}})):
[
PE_{(pos, 2i)} = \sin\left( \frac{pos}{10000^{2i / d_{\text{model}}}} \right)
]
[
PE_{(pos, 2i+1)} = \cos\left( \frac{pos}{10000^{2i / d_{\text{model}}}} \right)
]
设计直觉:
- 不同频率的正余弦函数沿位置轴变化,使每个位置产生唯一编码。
- 波长范围从 (2\pi) 到约 (10000 \cdot 2\pi),覆盖了短程和长程位置变化。
- 对于固定偏移量 (k),(PE_{pos+k}) 可以表示为 (PE_{pos}) 的线性函数(因正余弦的和角公式),这一性质理论上有利于模型学习相对位置信息。
优缺点:
- 优点:无需训练参数,可处理任意长度序列(理论上可外推)。
- 缺点:外推至远超训练长度的序列时性能下降明显;位置信息与词嵌入简单相加,交互方式固定。
2.2 可学习位置嵌入
BERT、GPT等模型则采用可学习的位置嵌入:定义一个嵌入矩阵 (\mathbf{P} \in \mathbb{R}^{L_{\max} \times d}),其中 (L_{\max}) 为最大序列长度。位置 (pos) 的嵌入为 (\mathbf{P}_{pos}),与词嵌入相加。
优势:位置表示完全由数据驱动学习,灵活性高,在训练长度范围内通常效果优于正弦编码。
劣势:无法处理超过 (L_{\max}) 的序列,遇到更长输入需截断或插值;参数随序列长度线性增长。
2.3 绝对位置编码的共性问题
无论是正弦还是可学习,绝对位置编码都存在两个根本缺陷:
- 绝对位置依赖:模型学习的是“第5个位置”的表示,而非“当前词与上文相距多远”。当测试序列长度与训练分布不一致时,绝对位置编码的外推能力有限。
- 加法注入的刚性:位置信息与词义信息以简单相加混合,模型无法根据上下文动态调整位置与内容的交互权重。
这些局限性催生了相对位置编码的设计。
3. 相对位置编码:关注距离而非绝对坐标
相对位置编码的核心思想是:不直接表示绝对位置,而是在计算注意力分数时,注入两个token之间的相对距离信息。
3.1 原始Transformer中的相对位置偏置
在原始的Transformer中,位置信息已经通过正弦编码加到了输入中,注意力计算时位置信息通过非线性变换间接交互。而显式的相对位置编码则直接将相对距离作为一个可学习的偏置项加入注意力分数。
基本形式:设 (i) 和 (j) 为两个位置,它们的相对距离为 (i-j)。定义一个可学习的嵌入表 (\mathbf{B} \in \mathbb{R}^{(2L_{\max}-1) \times \text{num_heads}})(或每个头独立),在计算注意力分数时添加:
[
\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left( \frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{d_k}} + \mathbf{B}_{i-j} \right) \mathbf{V}
]
这里 (\mathbf{B}_{i-j}) 是一个标量,取决于相对距离。为控制参数量,通常设定一个最大相对距离 (K),超过该距离的偏置共享同一个值(裁剪)。
3.2 Transformer-XL与T5的相对位置编码
Transformer-XL提出了一种更精细的相对位置编码方案。它将注意力分数分解为内容-内容、内容-位置、位置-内容、位置-位置四项,其中位置项使用相对位置嵌入。具体地:
[
\mathbf{A}{ij}^{\text{rel}} = \underbrace{\mathbf{E}{x_i}^\top \mathbf{W}q^\top \mathbf{W}k \mathbf{E}{x_j}}{(a)} + \underbrace{\mathbf{E}{x_i}^\top \mathbf{W}q^\top \mathbf{W}k \mathbf{R}{i-j}}{(b)} + \underbrace{\mathbf{u}^\top \mathbf{W}k \mathbf{E}{x_j}}{©} + \underbrace{\mathbf{v}^\top \mathbf{W}k \mathbf{R}{i-j}}_{(d)}
]
其中 (\mathbf{R}_{i-j}) 是可学习的相对位置嵌入,(\mathbf{u}, \mathbf{v}) 为可训练向量。
T5采用了一种简化的相对位置偏置:每个注意力头有一个可学习的标量偏置 (b_{\text{rel}}),仅依赖于相对距离,且进行分桶处理以限制参数。
3.3 ALiBi:线性偏置的简洁之美
ALiBi(Attention with Linear Biases)是一种极简的相对位置编码:不添加任何位置嵌入,仅在注意力分数中加上一个与相对距离成正比的负偏置。
[
\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left( \frac{\mathbf{Q} \mathbf{K}^\top}{\sqrt{d_k}} - m \cdot |i-j| \right) \mathbf{V}
]
其中 (m) 是头特定的斜率(如第 (h) 个头的斜率为 (2^{-8h/H}))。这个简单的线性偏置强迫模型更关注邻近的token,对远处token施加更大惩罚。
ALiBi的优势:
- 零参数:无需学习任何位置嵌入。
- 极强外推能力:由于偏置随距离线性增长,模型在训练时未见过的长序列上也能合理分配注意力,外推性能显著优于RoPE和绝对位置编码。
ALiBi已被用于BLOOM等模型,展示了无位置嵌入架构的潜力。
3.4 相对位置编码的小结
相对位置编码通过直接建模相对距离,部分解决了绝对位置编码的长度外推问题。但许多实现(如T5偏置、Transformer-XL)仍受限于最大相对距离的裁剪,外推时对超长距离的泛化依然受限。ALiBi的线性偏置在理论上可无限外推,但牺牲了对复杂位置模式的细粒度建模能力。
4. 旋转位置编码:用复数旋转编码相对位置
旋转位置编码(Rotary Position Embedding, RoPE)由Su等人在2021年提出,现已成为LLaMA、GPT-NeoX、PaLM等主流大模型的标准配置。其核心思想极其优雅:通过复数域中的旋转变换,将token的绝对位置信息转化为query和key的内积中的相对位置信息。
4.1 数学动机:内积只依赖相对位置
设词嵌入经线性投影得到查询向量 (\mathbf{q}_m)(位置 (m))和键向量 (\mathbf{k}_n)(位置 (n))。RoPE的目标是找到一个变换 (f),使得:
[
\langle f(\mathbf{q}_m, m), f(\mathbf{k}_n, n) \rangle = g(\mathbf{q}_m, \mathbf{k}_n, m-n)
]
即,变换后的内积仅依赖于向量本身和相对位置 (m-n),而不依赖于绝对位置 (m) 和 (n)。
4.2 二维情形下的旋转
首先考虑二维向量 (\mathbf{q} \in \mathbb{R}^2)。将其视为复平面上的点 (q e^{i\theta})。RoPE将位置 (m) 编码为一个旋转角度 (m\theta),通过复数乘法(旋转矩阵)施加于向量:
[
f(\mathbf{q}, m) = \begin{pmatrix} \cos(m\theta) & -\sin(m\theta) \ \sin(m\theta) & \cos(m\theta) \end{pmatrix} \begin{pmatrix} q_0 \ q_1 \end{pmatrix}
]
此时,两个旋转后向量的内积为:
[
f(\mathbf{q}, m)^\top f(\mathbf{k}, n) = \mathbf{q}^\top \mathbf{R}_{m-n} \mathbf{k}
]
其中 (\mathbf{R}_{m-n}) 是角度为 ((m-n)\theta) 的旋转矩阵。内积仅依赖于相对位置 (m-n)。
4.3 推广至高维
对于 (d) 维向量,RoPE将维度分成 (d/2) 对,每对作为一个二维子空间独立旋转。不同子空间采用不同的旋转频率:
[
\theta_i = 10000^{-2i/d}, \quad i = 0, 1, \dots, d/2-1
]
对于查询向量 (\mathbf{q}),位置 (m) 的RoPE变换为:
[
\begin{pmatrix} q_{2i} \ q_{2i+1} \end{pmatrix} \leftarrow
\begin{pmatrix} \cos(m\theta_i) & -\sin(m\theta_i) \ \sin(m\theta_i) & \cos(m\theta_i) \end{pmatrix}
\begin{pmatrix} q_{2i} \ q_{2i+1} \end{pmatrix}
]
键向量 (\mathbf{k}) 同理。
4.4 RoPE的实现与计算效率
在实际代码中,RoPE不显式构造旋转矩阵,而是利用复数运算高效实现:
defrotate_half(x):x1,x2=x.chunk(2,dim=-1)returntorch.cat([-x2,x1],dim=-1)defapply_rotary_pos_emb(q,k,cos,sin):# cos, sin: (seq_len, dim) 预计算好的频率q_embed=(q*cos)+(rotate_half(q)*sin)k_embed=(k*cos)+(rotate_half(k)*sin)returnq_embed,k_embed其中cos和sin根据位置和频率预计算,形状为(max_seq_len, dim)。
4.5 RoPE的优势
- 相对位置的内化:内积自然依赖于相对距离,模型对相对位置的感知更加直接。
- 良好的外推能力:通过调整推理时的频率缩放(如NTK-aware插值),RoPE可有效外推到训练长度数倍的序列。
- 与自注意力无缝融合:RoPE仅修改了(\mathbf{Q})和(\mathbf{K})的计算方式,不改变Transformer其他部分,易于集成。
- 保持向量模长:旋转是正交变换,不改变向量范数,有助于训练稳定性。
4.6 RoPE的变体与扩展
- 线性RoPE:调整频率基数为更大值(如500,000),使高频分量衰减更慢,改善长距离外推。
- NTK-aware插值:在外推时,对高频维度进行非线性频率缩放,避免直接插值导致的性能下降。
- YaRN:结合NTK插值与温度缩放,进一步提升极长上下文下的外推能力。
5. 位置编码的对比与选型
5.1 核心特性对比表
| 特性 | 绝对位置(正弦) | 绝对位置(可学习) | 相对位置偏置 | RoPE | ALiBi |
|---|---|---|---|---|---|
| 参数学习 | 无 | 有 | 有 | 无 | 无 |
| 长度外推能力 | 较弱 | 差(硬截断) | 中等(受裁剪限制) | 良好(配合插值) | 优秀 |
| 对相对距离建模 | 间接 | 间接 | 直接 | 直接(内积中) | 直接(线性偏置) |
| 计算复杂度 | 低 | 低 | 中 | 中 | 极低 |
| 代表模型 | Transformer原论文 | BERT, GPT, GPT-2 | T5, Transformer-XL | LLaMA, GPT-NeoX, PaLM | BLOOM |
5.2 选型建议
- NLU任务、训练长度固定:BERT式的可学习绝对位置嵌入足够,简单有效。
- 需要生成长文本、重视外推:RoPE是当前大模型首选,配合插值技术可外推至32k甚至100k上下文。
- 极致外推需求、算力受限:ALiBi几乎零成本且外推能力最强,适合长文档建模。
- 追求可解释性与理论优雅:RoPE通过复数旋转将相对位置编码内嵌,数学形式优美。
6. 代码实践:实现与对比
6.1 正弦绝对位置编码
importtorchimportmathdefsinusoidal_position_encoding(seq_len,d_model):pe=torch.zeros(seq_len,d_model)position=torch.arange(0,seq_len).unsqueeze(1)div_term=torch.exp(torch.arange(0,d_model,2)*-(math.log(10000.0)/d_model))pe[:,0::2]=torch.sin(position*div_term)pe[:,1::2]=torch.cos(position*div_term)returnpe6.2 RoPE实现(简化版)
classRotaryPositionEmbedding(torch.nn.Module):def__init__(self,dim,max_seq_len=2048,base=10000):super().__init__()inv_freq=1.0/(base**(torch.arange(0,dim,2).float()/dim))t=torch.arange(max_seq_len).float()freqs=torch.einsum('i,j->ij',t,inv_freq)emb=torch.cat((freqs,freqs),dim=-1)self.register_buffer('cos',emb.cos())self.register_buffer('sin',emb.sin())defforward(self,q,k,seq_len):cos=self.cos[:seq_len,:].unsqueeze(0).unsqueeze(0)# (1, 1, seq_len, dim)sin=self.sin[:seq_len,:].unsqueeze(0).unsqueeze(0)q_embed=(q*cos)+(rotate_half(q)*sin)k_embed=(k*cos)+(rotate_half(k)*sin)returnq_embed,k_embed6.3 ALiBi偏置生成
defget_alibi_slopes(num_heads):# 生成指数递减的斜率returntorch.tensor([2**(-8*i/num_heads)foriinrange(1,num_heads+1)])defapply_alibi(attn_scores,alibi_slopes,seq_len):# attn_scores: (batch, heads, seq_len, seq_len)distance=torch.arange(seq_len).unsqueeze(0)-torch.arange(seq_len).unsqueeze(1)distance=distance.abs().unsqueeze(0).unsqueeze(0)# (1, 1, seq_len, seq_len)alibi_bias=-alibi_slopes.view(-1,1,1)*distancereturnattn_scores+alibi_bias7. 前沿探索:无位置编码与混合方案
7.1 无显式位置编码的可能性
近期研究表明,通过特殊的注意力掩码或递归架构,模型可隐式习得位置信息。例如,Hyena和Mamba等状态空间模型不使用显式位置编码,在长序列任务上表现出色。这为未来彻底摆脱位置编码提供了新思路。
7.2 长度外推的持续挑战
尽管RoPE配合插值技术已将上下文窗口扩展至百万token级别,但真正有效的长程依赖建模仍面临注意力稀疏化、内存瓶颈等问题。位置编码是必要条件,但远非充分条件。
8. 结语:位置,序列的灵魂
位置编码的设计演进,是Transformer架构中最能体现数学之美的篇章之一。从绝对位置的直白赋予,到相对位置的偏置注入,再到旋转编码的复数优雅,每一步都凝结着研究者对“如何让注意力感知顺序”这一根本问题的深刻洞察。
在LLaMA、GPT-4等巨人身躯中,RoPE以其卓越的外推性能和理论完备性稳坐王座;ALiBi则以极致简洁在某些场景下异军突起。理解这些位置编码的原理与差异,不仅是读懂当今大模型技术栈的必要基础,更是未来设计新一代序列模型时的灵感源泉。
当你下次调用model.generate时,或许可以多想一想:那些流畅生成的文本背后,是无数复数旋转在默默编织着位置的丝线,让词语在时间的长河中找到属于自己的坐标。