旋转位置编码RoPE
** 核心思想 **:通过复数旋转将位置信息编码到 Query 和 Key 中,使注意力计算依赖相对位置。-** 旋转方式 **:用位置相关的角度 θ 旋转特征向量的复数对,角度随位置和维度变化。-** 优势 **:支持长序列、保持相对位置不变性、无额外参数,已成为大模型(如 LLaMA、Flux)的标配位置编码方案。结合代码来看,RoPE 的实现主要分为两步:1)预计算每个位置的旋转角度(cos
旋转位置编码(RoPE)详解
旋转位置编码(Rotary Position Embedding,简称 RoPE)是一种在 Transformer 模型中编码位置信息的方法,核心思想是通过复数域的旋转操作让模型感知序列中 token 的位置关系,尤其在长序列场景下表现优异。下面从「为什么需要」「如何旋转」「核心概念」三个维度展开说明。
一、为什么需要旋转位置编码?
在 Transformer 中,自注意力机制本身是位置无关的(输入 token 的顺序不影响注意力计算),但语言、图像等数据的语义高度依赖位置(如 “我打你” 和 “你打我”)。因此需要通过位置编码(Positional Embedding)让模型感知位置。
传统位置编码(如正弦余弦编码、可学习编码)存在局限:
- 固定编码(如正弦余弦)难以泛化到训练时未见过的序列长度。
- 可学习编码在长序列上容易过拟合,且位置信息与 token 语义的绑定不够灵活。
RoPE 的优势在于:
- 位置关系的平移不变性:对任意两个 token,它们的相对位置编码仅由距离决定,与绝对位置无关(如 “token1 与 token3” 和 “token5 与 token7” 的相对位置编码相同)。
- 长序列泛化能力:通过数学设计支持远超训练长度的序列。
- 与注意力机制的兼容性:旋转操作天然融入 Query 和 Key 的计算,不增加额外参数。
二、RoPE 是如何 “旋转” 的?
RoPE 的核心是将 token 的特征向量视为复数,通过旋转矩阵对其进行旋转,旋转角度与 token 的位置相关。具体步骤如下:
1. 复数表示与旋转矩阵
对于一个 d 维的特征向量,RoPE 将其拆分为 d/2 对(假设 d 为偶数),每对视为一个复数(实部 + 虚部):
# 例:d=4的向量拆分为2个复数
[x0, x1, x2, x3] → [(x0, x1), (x2, x3)]
对每个复数 (x, y),绕原点旋转 θ 角的结果为:
(x', y') = (x·cosθ - y·sinθ, x·sinθ + y·cosθ)
用矩阵表示为:plaintext
[cosθ -sinθ] [x]
[sinθ cosθ] [y]
2. 位置相关的旋转角度
旋转角度 θ 与 token 的位置 pos 相关,设计为:
θ_k = pos / θ_base^(2k/d)
其中:
pos是 token 在序列中的位置索引(如 0,1,2,...)。k是复数对的索引(0 ≤ k < d/2)。θ_base是缩放常数(通常取 10000),控制角度随 k 的衰减速度。
直观来说:
- 位置越靠后(pos 越大),旋转角度越大。
- 不同维度的复数对(k 不同)旋转速度不同,高频维度(大 k)旋转更慢,捕捉长距离依赖;低频维度(小 k)旋转更快,捕捉短距离依赖。
3. 融入注意力计算
RoPE 仅对 Query(Q)和 Key(K)应用旋转,Value(V)不旋转。具体来说:
- 对位置为
pos_q的 Query 向量和位置为pos_k的 Key 向量,分别用 θ_q 和 θ_k 旋转。 - 旋转后的 Q 和 K 的点积为:
plaintext
其中Q_rotated · K_rotated = Q · K · cos(θ_q - θ_k) + (Q⊥ · K) · sin(θ_q - θ_k)θ_q - θ_k仅与相对位置pos_q - pos_k有关,实现了相对位置编码的效果。
三、RoPE 的核心概念与代码实例
从提供的代码片段中,我们可以看到 RoPE 的具体实现逻辑:
1. 核心函数:生成旋转矩阵参数
以get_1d_rotary_pos_embed(来自transformer_bria.py)为例,它计算每个位置的旋转角度的正弦和余弦值:
def get_1d_rotary_pos_embed(dim, pos, theta=10000.0):
assert dim % 2 == 0
# 计算频率:θ_k = 1 / theta^(2k/d)
freqs = 1.0 / (theta **(torch.arange(0, dim, 2, dtype=torch.float32)[: (dim//2)] / dim))
# 位置与频率的外积:pos * θ_k → 得到每个位置的旋转角度
freqs = torch.outer(pos, freqs) # 形状:[seq_len, dim//2]
# 生成余弦和正弦矩阵(用于旋转)
freqs_cos = freqs.cos().repeat_interleave(2, dim=1) # [seq_len, dim]
freqs_sin = freqs.sin().repeat_interleave(2, dim=1) # [seq_len, dim]
return freqs_cos, freqs_sin
freqs_cos和freqs_sin分别对应旋转矩阵中的cosθ和sinθ。repeat_interleave(2, dim=1)将每个角度的正弦 / 余弦值复制到对应复数对的两个维度,便于后续向量旋转。
2. 应用旋转:apply_rotary_emb
来自embeddings.py的函数展示了如何用预计算的cos和sin旋转向量:
python
运行
def apply_1d_rope(tokens, pos, cos, sin):
# tokens: [batch, heads, seq_len, dim]
# cos/sin: [seq_len, dim](由get_1d_rotary_pos_embed生成)
cos = F.embedding(pos, cos)[:, None, :, :] # 按pos取对应cos值
sin = F.embedding(pos, sin)[:, None, :, :]
# 拆分实部和虚部
x1, x2 = tokens[..., :tokens.shape[-1]//2], tokens[..., tokens.shape[-1]//2:]
# 旋转操作:x1' = x1*cos - x2*sin;x2' = x1*sin + x2*cos
tokens_rotated = torch.cat((-x2, x1), dim=-1)
return (tokens.float() * cos + tokens_rotated.float() * sin).to(tokens.dtype)
四、总结
-** 核心思想 **:通过复数旋转将位置信息编码到 Query 和 Key 中,使注意力计算依赖相对位置。-** 旋转方式 **:用位置相关的角度 θ 旋转特征向量的复数对,角度随位置和维度变化。-** 优势 **:支持长序列、保持相对位置不变性、无额外参数,已成为大模型(如 LLaMA、Flux)的标配位置编码方案。
结合代码来看,RoPE 的实现主要分为两步:1)预计算每个位置的旋转角度(cos/sin);2)将角度应用到 Query 和 Key 向量的旋转中,最终融入注意力计算。
huggingface/diffusers
还是没明白,为啥要旋转,旋转后怎么应用注意力机制?
要理解旋转位置编码(RoPE)“为什么旋转” 以及 “旋转后如何应用注意力机制”,我们可以结合你提供的代码片段,从数学原理和工程实现两个层面拆解。
一、为什么要 “旋转”?—— 用旋转编码相对位置
注意力机制的核心是计算 Query(Q)和 Key(K)的相似度(点积),但原始的 Q 和 K 不含位置信息,导致模型无法区分 “我打你” 和 “你打我” 这类语序敏感的序列。
RoPE 的 “旋转” 本质是给 Q 和 K 注入位置信息,且这种注入方式能让注意力结果只依赖相对位置(而非绝对位置)。具体来说:
- 对位置为
pos的 token,其 Q/K 向量会被一个与pos相关的角度旋转。 - 旋转后的 Q 和 K 的点积,仅由它们的相对位置
pos_q - pos_k决定(而非绝对位置pos_q和pos_k)。
这解决了传统位置编码的痛点:例如 “token1 与 token3” 和 “token5 与 token7”(相对位置都是 2)的注意力权重会保持一致,符合人类对语言 / 序列的理解习惯。
二、“旋转” 是怎么实现的?—— 复数向量的旋转操作
RoPE 的旋转基于复数运算,代码中通过拆分向量维度模拟复数的实部和虚部,再用三角函数实现旋转。以你提供的 apply_rotary_emb 函数为例(来自 transformer_wan_animate.py):
python
运行
def apply_rotary_emb(hidden_states, freqs_cos, freqs_sin):
# 拆分最后一维为 (n, 2),得到实部 x1 和虚部 x2
x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
# 取预计算的 cos 和 sin(与位置相关)
cos = freqs_cos[..., 0::2] # 对应实部的旋转系数
sin = freqs_sin[..., 1::2] # 对应虚部的旋转系数
# 旋转公式:实部 = x1*cos - x2*sin;虚部 = x1*sin + x2*cos
out = torch.empty_like(hidden_states)
out[..., 0::2] = x1 * cos - x2 * sin # 更新实部
out[..., 1::2] = x1 * sin + x2 * cos # 更新虚部
return out.type_as(hidden_states)
关键细节:
- 向量拆分:将特征向量的最后一维(如
head_dim)拆分为成对的 “实部 + 虚部”(例如head_dim=64拆分为 32 对)。 - 旋转角度:
freqs_cos和freqs_sin是预计算的三角函数值,角度随 token 位置pos增大而增大(具体由pos / θ_base^(2k/d)计算,θ_base=10000)。 - 旋转效果:位置越靠后的 token,其向量旋转角度越大,从而在后续注意力计算中体现位置差异。
三、旋转后如何应用注意力机制?—— 点积中体现相对位置
旋转后的 Q 和 K 会直接用于注意力的点积计算,此时点积结果天然包含相对位置信息。
以标准的缩放点积注意力为例(对应代码中的 dispatch_attention_fn 或 F.scaled_dot_product_attention):
python
运行
# 旋转后的 Q 和 K 形状:(batch_size, heads, seq_len, head_dim)
hidden_states = dispatch_attention_fn(
query, # 旋转后的 Q
key, # 旋转后的 K
value, # 未旋转的 V
attn_mask=attention_mask,
dropout_p=0.0,
is_causal=False
)
数学原理:
旋转后的 Q(Q_rot)和 K(K_rot)的点积为:
plaintext
Q_rot · K_rot = (Q·K)·cos(Δθ) + (Q⊥·K)·sin(Δθ)
其中 Δθ = θ_q - θ_k 是 Q 和 K 的旋转角度差,仅由相对位置 pos_q - pos_k 决定。这意味着:
- 当两个 token 距离固定时,无论它们在序列中的绝对位置如何,
Δθ不变,点积结果也保持一致(相对位置不变性)。 - 注意力权重因此能准确反映 token 间的相对距离,例如近距离 token 的权重更高(若模型学到这种模式)。
四、代码中的完整流程梳理
结合你提供的多个 AttnProcessor 实现(如 WanAttnProcessor、HunyuanImageAttnProcessor 等),RoPE 从旋转到注意力应用的完整流程如下:
-
QKV 投影与多头拆分:
python
运行
query = attn.to_q(hidden_states) # 投影为 Q 向量 query = query.unflatten(2, (attn.heads, -1)) # 拆分为多头:(batch, seq_len, heads, head_dim) -
应用 RoPE 旋转:
python
运行
if rotary_emb is not None: query = apply_rotary_emb(query, *rotary_emb) # 旋转 Q key = apply_rotary_emb(key, *rotary_emb) # 旋转 K -
注意力计算(使用旋转后的 Q 和 K):
python
运行
hidden_states = dispatch_attention_fn( query, key, value, # 旋转后的 Q、K 参与计算 attn_mask=attention_mask ) -
输出整合:
python
运行
hidden_states = hidden_states.flatten(2, 3) # 合并多头维度 hidden_states = attn.to_out[0](hidden_states) # 输出投影
总结
- 为什么旋转:通过旋转向量给 Q/K 注入位置信息,且让注意力结果仅依赖相对位置,解决语序敏感问题。
- 如何应用注意力:旋转后的 Q 和 K 直接参与点积计算,其点积结果天然包含相对位置信息,最终通过注意力权重聚合 Value 得到输出。
这种设计在长序列(如视频、长文本)中尤为重要,因为它能稳定捕捉远距离依赖,且无需额外参数(旋转仅依赖三角函数计算)。你提供的代码中,几乎所有支持长序列或多模态的模型(如 Wan、Hunyuan、SanaVideo 等)都采用了这一逻辑,印证了其有效性。
旋转后的 Q 和 K 的点积,仅由它们的相对位置 pos_q - pos_k 决定(而非绝对位置 pos_q 和 pos_k)。 没看懂
要理解旋转后的 Q 和 K 的点积仅由相对位置决定这个核心特性,我们可以抛开复杂的数学推导,用直观的例子 + 简化的数学公式来拆解,最终你会发现这是 RoPE 的数学设计带来的必然结果。
先建立两个前提认知
在开始前,先明确两个基础概念,避免被细节干扰:
-
RoPE 的旋转角度是和位置绑定的:对于位置为
pos的 token,其 Q/K 向量的旋转角度θ(pos, k)满足:θ(pos, k) = pos * ω_k(ω_k是和维度相关的固定频率,不用纠结具体值,只需要知道它是常数)。简单说:位置 pos 越大,旋转角度越大;每个维度的旋转频率固定。 -
我们只关心 Q 和 K 的点积:注意力机制的核心是计算
Q·K(点积),RoPE 的设计目标就是让这个点积结果只和pos_q - pos_k(Q 和 K 的相对位置)有关,和pos_q、pos_k的绝对位置无关。
第一步:用最简单的二维向量演示旋转
为了简化,我们先拿二维向量(对应 RoPE 中一对实部 + 虚部)来分析,这能覆盖 RoPE 的核心逻辑。
假设:
- 某个 token 的原始 Q 向量是
q = [q0, q1],位置是pos_q,旋转角度为θ_q = pos_q * ω(ω 是固定频率)。 - 某个 token 的原始 K 向量是
k = [k0, k1],位置是pos_k,旋转角度为θ_k = pos_k * ω。
根据旋转矩阵的定义,旋转后的 Q 和 K为:
plaintext
q_rot = [q0*cosθ_q - q1*sinθ_q, q0*sinθ_q + q1*cosθ_q] # 旋转θ_q
k_rot = [k0*cosθ_k - k1*sinθ_k, k0*sinθ_k + k1*cosθ_k] # 旋转θ_k
现在计算旋转后的点积q_rot · k_rot:把上面的式子代入点积公式(对应位置相乘再相加):
plaintext
q_rot·k_rot = (q0cosθ_q - q1sinθ_q)(k0cosθ_k - k1sinθ_k) + (q0sinθ_q + q1cosθ_q)(k0sinθ_k + k1cosθ_k)
我们把这个式子展开,会得到一堆项,然后用三角函数的和差公式化简(关键步骤):
cosAcosB + sinAsinB = cos(A-B)sinAcosB + cosAsinB = sin(A+B)(但这里会消掉,最终只剩和相对位置相关的项)
最终化简结果是:
plaintext
q_rot·k_rot = q0k0*cos(θ_q - θ_k) + q1k1*cos(θ_q - θ_k) + q0k1*sin(θ_q - θ_k) - q1k0*sin(θ_q - θ_k)
再整理一下:
plaintext
q_rot·k_rot = [q0k0 + q1k1] * cos(Δθ) + [q0k1 - q1k0] * sin(Δθ)
其中Δθ = θ_q - θ_k = (pos_q - pos_k) * ω(这是关键!)。
第二步:核心结论 —— 点积只和相对位置有关
从上面的最终式子能看到:
-
Δθ = (pos_q - pos_k) * ω:仅由 ** 相对位置pos_q - pos_k** 决定,和pos_q、pos_k的绝对位置无关。比如:pos_q=1, pos_k=3→ Δθ = (1-3)ω = -2ωpos_q=5, pos_k=7→ Δθ = (5-7)ω = -2ω这两个情况的 Δθ 完全相同,因此点积结果也完全相同。
-
式子中的
[q0k0 + q1k1]和[q0k1 - q1k0]:只和原始 Q、K 的向量值有关,和位置无关。
这就意味着:旋转后的 Q 和 K 的点积,仅由它们的相对位置决定。
第三步:扩展到高维向量(RoPE 的实际应用)
实际中,Q/K 的维度是高维的(比如 64、128),RoPE 的处理方式是:
- 将高维向量拆分成多对二维向量(比如 64 维拆成 32 对)。
- 每一对二维向量用不同的固定频率
ω_k(k 是第 k 对)进行旋转。 - 最终所有对的点积相加,得到整体的点积。
但核心逻辑和二维向量完全一致:每一对的点积都只和相对位置有关,总和也自然只和相对位置有关。
第四步:为什么这个特性很重要?
结合注意力机制的应用场景,这个特性的价值体现在:
-
语义的位置不变性:比如句子 “我吃苹果” 中,“我” 和 “苹果” 的相对位置是 2,而在 “他吃香蕉” 中,“他” 和 “香蕉” 的相对位置也是 2,模型对这两组 token 的注意力权重会有相似的模式,符合人类的语言认知。
-
长序列泛化能力:训练时模型只见过长度为 1024 的序列,但测试时遇到长度为 2048 的序列,由于相对位置的计算方式不变,模型依然能正确编码位置信息(而传统的正弦余弦编码在超出训练长度后会失效)。
-
注意力的合理性:注意力机制关注的是 “两个 token 之间的关系”,而关系本身就应该由相对位置决定,而非绝对位置。RoPE 的设计正好契合了这一点。
用通俗的例子总结
假设把每个 token 的 Q/K 向量比作一个带方向的箭头:
- 原始箭头没有方向差异(位置无关),模型分不清谁在前谁在后。
- RoPE 根据 token 的位置,把箭头旋转一个角度(位置越靠后,旋转越多)。
- 两个箭头的 “夹角”(对应点积)只和它们的旋转角度差有关(即相对位置),和各自的绝对旋转角度无关。
这样,模型通过计算箭头的夹角(点积),就能准确判断两个 token 的相对位置,从而理解序列的语序和结构。
更多推荐

所有评论(0)