RoPE (Rotary Position Embedding) 深度解析：旋转位置编码的数学之美
6516 字
33 分钟
RoPE (Rotary Position Embedding) 深度解析：旋转位置编码的数学之美
1. 引言：位置编码的必要性
1.1 Transformer 的位置不变性
┌─────────────────────────────────────────────────────────────┐│ Transformer 的位置不变性 │├─────────────────────────────────────────────────────────────┤│ ││ 问题：Self-Attention 是位置无关的操作 ││ ││ 输入序列: [A, B, C, D] ││ ││ Attention(Q, K, V) 的计算: ││ ││ Q = XW_q ──┐ ││ K = XW_k ──┼──> Attention(Q, K, V) ││ V = XW_v ──┘ ││ ││ 其中 X = [x₀, x₁, x₂, x₃] ││ ││ 问题：Attention(Q, K, V) 对输入顺序是不变的！ ││ ││ Attention([A,B], [C,D], ...) ││ = Attention([C,D], [A,B], ...) ││ ││ 但 "狗咬人" ≠ "人咬狗" ││ ││ 解决：给 token 添加位置信息 ││ │└─────────────────────────────────────────────────────────────┘1.2 位置编码发展历程
class PositionEncodingHistory: """ 位置编码发展历程 """
def timeline(self): """ 发展时间线 """ return { "2017": "Original Transformer (Vaswani et al.)", "2018": "Learnable Position Embedding (BERT, GPT-2)", "2018": "Relative Position Embedding (Shaw et al.)", "2020": "T5 Relative Position Bias (Raffel et al.)", "2021": "ALiBi (Press et al.) - 无需训练的位置编码", "2022": "RoPE (Su et al.) - 旋转位置编码", "2023": "NTK-Aware Scaling (YaRN) - 扩展上下文", "2024": "CoMayRoPE, RoPE-AFA - 进一步扩展", }
def taxonomy(self): """ 位置编码分类 """ return """ 位置编码分类:
┌─────────────────────────────────────────────────────┐ │ 位置编码 │ ├─────────────────────────────────────────────────────┤ │ │ │ 绝对位置编码 (Absolute Position Encoding) │ │ ├── Sinusoidal (Transformer 原始) │ │ └── Learnable (BERT, GPT) │ │ │ │ 相对位置编码 (Relative Position Encoding) │ │ ├── Shaw et al. (2018) │ │ ├── T5 Bias (2020) │ │ ├── ALiBi (2021) │ │ └── RoPE (2022) │ │ │ └─────────────────────────────────────────────────────┘
关键区分: → 是否编码绝对位置 → 是否在 Attention 中引入相对偏差 → 外推能力如何 """1.3 位置编码的核心挑战
class PositionEncodingChallenges: """ 位置编码的核心挑战 """
def extrapolation(self): """ 外推问题 (Extrapolation)
模型在训练时见到的最长序列长度为 L_train， 但推理时可能需要处理更长的序列 L_infer > L_train。 """ return """ 外推问题示例:
训练: 最大 2048 tokens
推理需求: - 输入 3000 tokens - 位置 2500 > 2048 (训练未见过的位置)
问题: ┌─────────────────────────────────────────────────────┐ │ 绝对位置: PE[2500] 不存在，需要外推 │ │ 相对位置: 位置差 2500-2048 = 492，超出训练范围 │ └─────────────────────────────────────────────────────┘
好的位置编码应该: 1. 对未见过的位置也能给出合理编码 2. 保持位置之间的相对关系 3. 不需要重新训练模型 """
def length_generalization(self): """ 长度泛化 """ return { "sinusoidal": "理论上可以外推，但实际效果差", "learnable": "无法外推，泛化能力差", "relative": "相对位置可能越界", "rope": "相对位置编码，外推性好", "alibi": "基于距离的线性衰减，外推性最好", }1.4 本系列文章关联
| 文章 | 关联 |
|---|---|
| Attention Mechanism | Self-Attention 基础 |
| LLaMA Architecture | RoPE 在 LLaMA 中的应用 |
| Long Context Extension | 上下文扩展技术 |
2. 经典位置编码回顾
2.1 Sinusoidal 位置编码
class SinusoidalPositionEncoding: """ Transformer 原始的 Sinusoidal 位置编码 """
def formula(self): """ 公式
PE(pos, 2i) = sin(pos / 10000^(2i/d_model)) PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
其中: - pos: 位置 (0, 1, 2, ...) - i: 维度索引 (0, 1, 2, ..., d_model/2) - d_model: 模型维度 """ return ''' import numpy as np
def get_sinusoidal_position_encoding( seq_len, d_model ): """生成 Sinusoidal 位置编码"""
positions = np.arange(seq_len)[:, np.newaxis] div_term = np.exp( np.arange(0, d_model, 2) * (-np.log(10000.0) / d_model) )
pe = np.zeros((seq_len, d_model)) pe[:, 0::2] = np.sin(positions * div_term) pe[:, 1::2] = np.cos(positions * div_term)
return pe
# 示例 # position 0: [0, 1, 0, 1, 0, 1, ...] # position 1: [sin(1), cos(1), sin(1/10000), ...] # ... '''
def properties(self): """ 特性 """ return { "advantage_1": "每个位置的编码是唯一的", "advantage_2": "可以表示任意长度的位置 (理论)", "advantage_3": "不同位置编码有线性关系", "advantage_4": "周期特性可以从短周期合成长周期",
"disadvantage_1": "参数量大 (d_model × max_len)", "disadvantage_2": "需要为每个位置学习 embedding", "disadvantage_3": "训练时未见过的位置泛化差", }
def geometric_intuition(self): """ 几何直觉 """ return """ 几何解释:
每个维度 i 对应一个旋转:
dim 0,1: 频率 ω₀ = 1/10000^0 = 1 dim 2,3: 频率 ω₁ = 1/10000^(2/d) dim 4,5: 频率 ω₂ = 1/10000^(4/d) ...
PE(pos) = [sin(ω₀·pos), cos(ω₀·pos), sin(ω₁·pos), cos(ω₁·pos), ...]
这相当于在多个圆上以不同频率旋转:
ω₀ (快): ●→○→●→○ (位置每增加 2π 重复) ω₁ (慢): ●→→→→→○ (位置每增加 2π/ω₁ 重复) ...
不同频率的组合可以唯一标识每个位置。 """2.2 Learnable 位置编码
class LearnablePositionEncoding: """ 可学习的位置编码 (BERT, GPT-2) """
def approach(self): """ 方法
将位置编码作为可学习的参数: PE ∈ ℝ^(max_len × d_model)
训练时通过反向传播学习。 """ return ''' import torch import torch.nn as nn
class LearnablePositionalEncoding(nn.Module): """可学习的位置编码"""
def __init__(self, max_len, d_model): super().__init__() self.pe = nn.Embedding(max_len, d_model)
# 可选的: 初始化 nn.init.normal_(self.pe.weight, mean=0, std=0.02)
def forward(self, x): """ x: (batch, seq_len, d_model) """ batch_size, seq_len, _ = x.shape
# 创建位置索引 [0, 1, 2, ..., seq_len-1] positions = torch.arange( seq_len, device=x.device ).unsqueeze(0).expand(batch_size, -1)
return x + self.pe(positions) '''
def issues(self): """ 问题 """ return { "issue_1": "需要预先确定最大长度", "issue_2": "超出范围的位置无法处理", "issue_3": "泛化能力差（未见过的位置）", "issue_4": "参数量大", }2.3 相对位置编码
class RelativePositionEncoding: """ 相对位置编码 """
def shaw_encoding(self): """ Shaw et al. (2018) 相对位置编码 """ return """ Shaw 的方法:
在 Attention Score 中添加相对位置偏差:
Attention(Q, K, V)_{i,j} = softmax( Q_i · K_j + b_{j-i} ) · V_j
其中 b_{j-i} 是位置差 j-i 的可学习嵌入。
限制: 相对位置范围限制在 [-k, k] 内。 """
def alibi(self): """ ALiBi (Attention with Linear Biases) """ return """ ALiBi 方法:
优点: 完全不需要位置编码，天然支持外推！
Attention Score 修改:
Attention(Q, K, V)_{i,j} = softmax( Q_i · K_j - |j-i|·m ) · V_j
其中 m 是与头有关的斜率因子:
m_h = 1 / 2^(8h/H) (h = 0, 1, 2, ..., H-1)
特点: → 越远的 token 注意力越低 → 线性衰减 (linear bias) → 外推性好 """
def comparison(self): """ 对比 """ return { "method": ["Sinusoidal", "Learnable", "Shaw", "T5", "ALiBi", "RoPE"], "参数可学习": ["否", "是", "是", "是", "否", "否"], "绝对位置": ["是", "是", "否", "否", "否", "否"], "相对位置": ["隐式", "隐式", "显式", "显式", "隐式", "显式"], "外推能力": ["一般", "差", "一般", "一般", "好", "好"], }3. RoPE 核心思想
3.1 旋转的直觉
class RoPEMotivation: """ RoPE 的直观理解 """
def the_core_idea(self): """ 核心思想
RoPE 的核心洞察是：使用旋转矩阵来编码位置信息。
关键观察： - Attention(Q, K, V) 中的 Q 和 K 决定注意力分数 - 如果 Q 和 K 都进行相同角度的旋转，那么它们的点积只取决于相对位置！ """ return """ 旋转点积的不变性:
给定向量 q 和 k，我们希望:
⟨Rot(θ_q)·q, Rot(θ_k)·k⟩ = ⟨q, k⟩ (仅依赖于相对位置)
这正是 RoPE 要实现的目标！
几何直觉:
q = (1, 0) k = (0.8, 0.6) ↗ (旋转 θ_q) ↗ (旋转 θ_k)
Rot(θ_q)·q = (cos θ_q, -sin θ_q) Rot(θ_k)·k = (0.8 cos θ_k - 0.6 sin θ_k, 0.8 sin θ_k + 0.6 cos θ_k)
点积 = (cos θ_q)(0.8 cos θ_k - 0.6 sin θ_k) + (-sin θ_q)(0.8 sin θ_k + 0.6 cos θ_k) = 0.8 cos(θ_q - θ_k) + 0.6 sin(θ_q - θ_k) = k · Rot(-θ_q)q ← 只依赖于相对角度！ """
def why_rotation_works(self): """ 为什么旋转有效 """ return """ 关键性质: 旋转保持内积的相对结构
对于 2D 旋转:
⟨Rot(θ)·a, Rot(θ)·b⟩ = ⟨a, b⟩
证明: ⟨Rot(θ)·a, Rot(θ)·b⟩ = ⟨a, Rot(-θ)·b⟩ (旋转是正交变换) = ⟨a, b⟩ (当 θ=0 时)
但我们想要的是: ⟨Rot(θ_a)·a, Rot(θ_b)·b⟩ 只依赖于 θ_b - θ_a
设 θ_a = 0, θ_b = θ: ⟨a, Rot(θ)·b⟩ = ⟨Rot(-θ)·a, b⟩
这正是 RoPE 的设计基础！ """3.2 复数表示
class ComplexNumberRepresentation: """ 复数表示 """
def complex_mapping(self): """ 复数映射
将 2D 向量映射到复数: [x, y] → x + yi
其中 i = √(-1) """ return ''' import numpy as np
def vector_to_complex(x, y): """将 2D 向量转为复数""" return complex(x, y)
def complex_to_vector(z): """将复数转为 2D 向量""" return [z.real, z.imag]
def rotate_2d(x, y, theta): """2D 旋转""" cos_t, sin_t = np.cos(theta), np.sin(theta) x_rot = x * cos_t - y * sin_t y_rot = x * sin_t + y * cos_t return x_rot, y_rot
# 复数形式 def rotate_complex(z, theta): """复数旋转 = 乘以 e^(iθ)""" return z * np.exp(1j * theta) '''
def group_property(self): """ 群的性质
旋转在复数乘法下构成一个群 (圆群)。 """ return """ 旋转群的性质:
1. 单位元: θ = 0 (乘以 1) Rot(0)·v = v
2. 闭合性: 旋转后还是旋转 Rot(θ₁)·Rot(θ₂) = Rot(θ₁ + θ₂)
3. 结合性: (Rot(θ₁)·Rot(θ₂))·Rot(θ₃) = Rot(θ₁)·(Rot(θ₂)·Rot(θ₃))
4. 逆元: Rot(θ)⁻¹ = Rot(-θ)
这些性质保证了: → 位置编码是良定义的 → 可以合成任意的相对位置 → 不需要截断或限制范围 """3.3 从 2D 到多维
class MultiDimensionalExtension: """ 扩展到多维 """
def pairwise_rotation(self): """ 两两旋转
对于 d_model 维的向量，每 2 维构成一个"旋转对"。 """ return """ 将 d_model 维向量分成 d_model/2 个 2D 对:
q = [q₀, q₁, q₂, q₃, q₄, q₅, ..., q_{d-2}, q_{d-1}] └──────┘└──────┘└──────┘ └──────┘ pair 0 pair 1 pair 2 pair d/2-1
第 i 个 pair 的旋转角度: θ_i = pos × base^i
其中 base 通常取 10000。
这样: - 低频维度 (i 较小) → 较快的旋转 - 高频维度 (i 较大) → 较慢的旋转
这与 Sinusoidal 编码的设计思想一致！ """
def rotation_matrix(self): """ 旋转矩阵 """ return ''' import torch
def get_rope_matrix(seq_len, dim, base=10000): """ 生成 RoPE 旋转矩阵
返回一个形状为 (seq_len, dim/2) 的角度数组 """ # 角度 = position × base^{2i/dim} positions = torch.arange(seq_len) indices = torch.arange(0, dim, 2) thetas = torch.pow(base, -indices / dim) angles = positions[:, None] * thetas[None, :]
return angles # (seq_len, dim/2)
def apply_rope(q, cos, sin): """ 应用 RoPE 旋转
q: (batch, heads, seq_len, dim) cos, sin: (seq_len, dim/2) """ # 将 q 分成 pairs q_half = q.view(*q.shape[:-1], -1, 2) # (..., dim/2, 2)
# 旋转: q' = q * cos - rotate(q) * sin # 其中 rotate(q) = [-q[..., 1], q[..., 0]] q_rot = torch.stack([-q_half[..., 1], q_half[..., 0]], dim=-1)
q_prime = q_half * cos.unsqueeze(-2) + q_rot * sin.unsqueeze(-2)
return q_prime.flatten(-2) '''4. RoPE 数学推导
4.1 基本公式
class RoPEMathematicalDerivation: """ RoPE 数学推导 """
def formal_definition(self): """ 形式化定义
对于位置 m 处的 query 向量 q_m 和位置 n 处的 key 向量 k_n:
RoPE(q_m) = W_q · x_m · R_m RoPE(k_n) = W_k · x_n · R_n
其中 R_m 是旋转矩阵。 """ return """ 符号说明:
d: 模型维度 (必须为偶数) m, n: 位置索引 x_m: 位置 m 的输入向量 (d 维) q_m: 位置 m 的 query 向量 k_n: 位置 n 的 key 向量 W_q, W_k: query/key 投影矩阵 R_m: 位置 m 的旋转矩阵 (d×d)
旋转矩阵 R_m 是块对角的:
R_m = diag(R_m,1, R_m,2, ..., R_m,d/2)
其中每个 R_m,i 是 2×2 旋转矩阵:
R_m,i = [[cos(m·θ_i), -sin(m·θ_i)], [sin(m·θ_i), cos(m·θ_i)]]
频率 θ_i = base^{-2i/d} = base^{-2i/d} """
def attention_with_rope(self): """ 带 RoPE 的 Attention """ return """ Attention 计算:
A_{m,n} = ⟨q_m, k_n⟩ / √d = ⟨W_q x_m R_m, W_k x_n R_n⟩
展开为 pairs: = Σ_{i=0}^{d/2-1} ⟨R_m,i · q_{m,i}, R_n,i · k_{n,i}⟩
根据旋转内积性质: = Σ_{i=0}^{d/2-1} ⟨q_{m,i}, R_{n-m,i} · k_{n,i}⟩
关键：结果只依赖于相对位置 (n-m)！
这意味着: → RoPE 自然实现了相对位置编码 → 不需要在 Attention 分数上添加位置偏差 → 位置信息通过旋转隐式编码 """4.2 旋转的几何意义
class GeometricInterpretation: """ 几何意义 """
def pairwise_visualization(self): """ 两两配对可视化 """ return """ 二维旋转的几何意义:
位置 m: ┌─────────────────────────────────────────────────────┐ │ │ │ y │ │ ↑ │ │ │ │ │ │ ● q_m │ │ │ ↗ │ │ │ m·θ │ │ │ ↙ │ │ │ ○ origin │ │ └──────────────────→ x │ │ │ │ q_m = (cos(m·θ), sin(m·θ)) │ │ 旋转角度 = m·θ │ │ │ └─────────────────────────────────────────────────────┘
不同位置的向量旋转不同角度: → 位置 0: 旋转 0° → 位置 1: 旋转 θ° → 位置 2: 旋转 2θ° → ... → 位置 m: 旋转 m·θ° """
def relative_position_emergence(self): """ 相对位置的涌现 """ return """ 点积 → 相对位置:
q_m · k_n = ⟨R_m·q, R_n·k⟩
当 q 和 k 是同一个向量时: q_m · q_n = ⟨R_m·q, R_n·q⟩ = ⟨q, R_{n-m}·q⟩ = q · Rot(n-m)·q
这只依赖于 (n-m)，即相对位置！
验证: 位置差 = 1: ⟨q, Rot(θ)·q⟩ = cos(θ) 位置差 = 2: ⟨q, Rot(2θ)·q⟩ = cos(2θ) 位置差 = 3: ⟨q, Rot(3θ)·q⟩ = cos(3θ)
余弦值只与位置差有关！ """4.3 频率设置
class FrequencySetting: """ 频率设置 """
def base_selection(self): """ base 参数选择
θ_i = base^{-2i/d}
常见选择: - LLaMA: base = 10000 - GPT-NeoX: base = 10000 - GPT-4: base = 220 (非官方) """ return { "base_10000": { "description": "Transformer 原始选择", "min_frequency": "10000^{-1} = 0.0001", "max_frequency": "10000^{0} = 1", }, "base_500": { "description": "更大的 base，压缩高频", "min_frequency": "500^{-1} = 0.002", }, "base_effect": { "large_base": "低频维度的旋转更慢，适合短距离", "small_base": "低频维度的旋转更快，适合长距离", }, }
def wavelength_calculation(self): """ 波长计算 """ return """ 波长 = 2π / θ_i
对于 base = 10000, d_model = 4096:
dim pair 0: θ₀ = 10000^0 = 1 wavelength = 2π ≈ 6.28 tokens
dim pair 1: θ₁ = 10000^{-1} = 0.0001 wavelength = 2π × 10000 ≈ 62832 tokens
dim pair 2: θ₂ = 10000^{-2} = 10^{-8} wavelength ≈ 6.28 × 10^8 tokens
...
结论: → 低维度的波长很短 (短距离信息) → 高维度的波长很长 (长距离信息) → 这允许模型同时捕获短程和长程依赖 """5. 实现细节
5.1 核心实现
class RoPEImplementation: """ RoPE 实现 """
def precomputation(self): """ 预计算旋转角度 """ return ''' import torch import math
def precompute_freqs_cis( seq_len: int, dim: int, base: float = 10000.0, device: torch.device = None, ) -> torch.Tensor: """ 预计算旋转角度的 cos 和 sin 值
Args: seq_len: 序列长度 dim: 注意力头维度 (必须是偶数) base: 频率基数 device: 计算设备
Returns: freqs_cis: 形状 (seq_len, dim//2)，复数形式的旋转因子 """ # 计算每个维度的频率 # θ_i = base^{-2i/d} = base^{-i/(d/2)} freqs = 1.0 / ( base ** (torch.arange(0, dim, 2, device=device).float() / dim) )
# 计算每个位置的相位 # positions = [0, 1, 2, ..., seq_len-1] t = torch.arange(seq_len, device=device).float()
# 计算角度: θ * position # freqs_cis[i, j] = e^{i * position[i] * freq[j]} freqs = torch.outer(t, freqs)
# 转换为复数形式: e^{iθ} = cos(θ) + i·sin(θ) freqs_cis = torch.polar( torch.ones_like(freqs), # magnitude = 1 freqs )
return freqs_cis # (seq_len, dim//2) '''
def apply_to_queries_keys(self): """ 应用 RoPE 到 Query 和 Key """ return ''' def apply_rotary_pos_emb( q: torch.Tensor, k: torch.Tensor, freqs_cis: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """ 应用旋转位置编码到 q 和 k
Args: q: query 张量 (batch, heads, seq_len, dim) k: key 张量 (batch, heads, seq_len, dim) freqs_cis: 预计算的旋转因子 (seq_len, dim//2)
Returns: q, k: 应用 RoPE 后的张量 """ batch, heads, seq_len, dim = q.shape
# 重塑为复数形式 # q: (batch, heads, seq_len, dim//2, 2) # → (batch, heads, seq_len, dim//2) q_float = q.view(batch, heads, seq_len, dim // 2, 2) k_float = k.view(batch, heads, seq_len, dim // 2, 2)
# 转换为复数 # q: (batch, heads, seq_len, dim//2) 复数 q_complex = torch.view_as_complex(q_float) k_complex = torch.view_as_complex(k_float)
# 广播旋转因子并相乘 # freqs_cis: (seq_len, dim//2) # q_complex: (batch, heads, seq_len, dim//2) q_rotated = q_complex * freqs_cis.unsqueeze(0).unsqueeze(0) k_rotated = k_complex * freqs_cis.unsqueeze(0).unsqueeze(0)
# 转回实数 q_out = torch.view_as_real(q_rotated).flatten(-2) k_out = torch.view_as_real(k_rotated).flatten(-2)
return q_out, k_out '''
def efficiency_considerations(self): """ 效率考虑 """ return """ 优化要点:
1. 预计算 - 旋转角度可以预计算并缓存 - 不需要每次前向传播重新计算
2. 内存友好 - freqs_cis 的大小: seq_len × dim/2 - 对于 seq_len=2048, dim=128: ~256KB - 完全可以放在显存中
3. 计算量 - 额外的复数乘法: O(batch × heads × seq_len × dim) - 与 Attention 计算相比可以忽略
4. CUDA Kernel 融合 - 可以将 RoPE 和 QKV 投影融合 - 减少内存访问 """5.2 与 Attention 集成
class AttentionIntegration: """ 与 Attention 集成 """
def full_attention_forward(self): """ 完整的 Attention 前向 """ return ''' def attention_forward( x: torch.Tensor, # (batch, seq_len, d_model) freqs_cis: torch.Tensor, # (seq_len, d_model//2) W_qkv: torch.Tensor, # (3, d_model, d_model) W_o: torch.Tensor, # (d_model, d_model) ) -> torch.Tensor: """带 RoPE 的 Attention 前向传播"""
batch, seq_len, d_model = x.shape
# QKV 投影 qkv = torch.einsum("bsd,hdk->bhsdk", x, W_qkv) q, k, v = qkv.unbind(dim=2)
# 应用 RoPE (在 d_model 维度分成 pairs) q = reshape_for_rope(q, freqs_cis.shape[-1]) k = reshape_for_rope(k, freqs_cis.shape[-1]) q, k = apply_rope(q, k, freqs_cis)
# Attention 计算 # Scaled dot-product attention scale = 1.0 / math.sqrt(q.shape[-1]) attn_weights = torch.einsum("bhqd,bhkd->bhqk", q, k) * scale attn_weights = F.softmax(attn_weights, dim=-1)
# Value 加权 attn_output = torch.einsum("bhqk,bhkd->bhqd", attn_weights, v)
# Output 投影 output = torch.einsum("bhqd,dh->bqd", attn_output, W_o)
return output '''
def causal_mask_consideration(self): """ 因果掩码考虑 """ return """ 与因果 Attention 的结合:
RoPE 不修改 Attention 的掩码逻辑。 因果掩码在 softmax 之前应用:
1. 计算 attention scores scores = q · k^T (已应用 RoPE)
2. 应用因果掩码 scores = scores + mask (mask[i,j] = 0 if j<=i else -inf)
3. Softmax attn_weights = softmax(scores, dim=-1)
RoPE 和因果掩码是正交的，可以独立使用。 """5.3 显存与计算分析
class MemoryComputationAnalysis: """ 显存与计算分析 """
def memory_overhead(self): """ 显存开销 """ return { "freqs_cis": "seq_len × (dim/2) × 8 bytes (复数 float32)", "example_7b": { "seq_len": 4096, "head_dim": 128, "freqs_cis_size": "4096 × 64 × 8 = 2 MB per layer", "total": "32 layers × 2 MB = 64 MB", }, "comparison": "与模型参数相比可以忽略 (~14GB for 7B)", }
def flops_overhead(self): """ FLOPs 开销 """ return { "rope_operation": "每个 token 每个头的复数乘法", "flops_per_layer": "2 × batch × heads × seq_len × (dim/2)", "attention_flops": "4 × batch × heads × seq_len² × (dim/2)",
"ratio": { "rope_flops": "O(seq_len)", "attention_flops": "O(seq_len²)", "overhead": "~1% for long sequences", }, }6. RoPE 在主流模型中的应用
6.1 LLaMA 系列
class LLaMARoPE: """ LLaMA 中的 RoPE """
def llama_architecture(self): """ LLaMA 架构 """ return """ LLaMA 架构使用 RoPE:
┌─────────────────────────────────────────────────────┐ │ LLaMA Layers (L = 32/40/80) │ ├─────────────────────────────────────────────────────┤ │ │ │ 每层包含: │ │ 1. Self-Attention + RoPE │ │ 2. RMSNorm │ │ 3. Feed-Forward Network │ │ │ └─────────────────────────────────────────────────────┘
关键参数: - LLaMA 7B: 32 layers, d_model=4096, n_heads=32, head_dim=128 - LLaMA 13B: 40 layers, d_model=5120, n_heads=40, head_dim=128 - LLaMA 70B: 80 layers, d_model=8192, n_heads=64, head_dim=128 """
def rope_configuration(self): """ RoPE 配置 """ return { "base": 10000, "max_position_embeddings": { "llama_1": 2048, "llama_2_7b_13b": 4096, "llama_2_70b": 4096, "llama_3_8b": 8192, "llama_3_70b": 128 * 1024, # 128K with RoPE scaling },
"rope_scaling": { "type": "linear", # or "yarn" for llama 3 "factor": 2.0, # 扩展倍数 }, }6.2 GPT-NeoX
class GPTNeoXRoPE: """ GPT-NeoX 中的 RoPE """
def config(self): """ GPT-NeoX 配置 """ return """ GPT-NeoX 使用与 LLaMA 相同的 RoPE 配置:
base = 10000 max_position = 2048 (原始)
GPT-NeoX-20B: - 44 layers - d_model = 6144 - n_heads = 64 - head_dim = 96 """
def attention_implementation(self): """ 实现 """ return ''' class RotaryEmbedding(nn.Module): """GPT-NeoX 风格的 RoPE"""
def __init__(self, dim, max_position, base=10000): super().__init__() self.dim = dim self.base = base self.max_position = max_position
# 预计算 self._compute_freqs()
def _compute_freqs(self): """计算频率""" inv_freq = 1.0 / ( self.base ** ( torch.arange(0, self.dim, 2).float() / self.dim ) )
t = torch.arange(self.max_position * 2) freqs = torch.outer(t, inv_freq)
# 复数形式 self.freqs_cis = torch.polar( torch.ones_like(freqs), freqs )
def forward(self, q, k): """应用 RoPE""" seq_len = q.shape[2] return apply_rotary_pos_emb( q, k, self.freqs_cis[:seq_len] ) '''6.3 Qwen 系列
class QwenRoPE: """ Qwen 中的 RoPE """
def qwen2_config(self): """ Qwen2 配置 """ return { "rope_type": "yarn", "original_max_position": 32768, "extracted_max_position": 512 * 1024, # 512K
"yarn_parameters": { "beta_fast": 32, # 快衰减开始的位置 (相对于 original_max) "beta_slow": 1, # 慢衰减的权重 }, }
def yarn_explanation(self): """ YaRN 解释 """ return """ YaRN (Yet another RoPE extensioN):
目标: 扩展 RoPE 支持更长的上下文
方法: 调整频率，使得在原始上下文之外也有合理的旋转。
具体: - 对高频维度不修改 - 对低频维度进行缩放
效果: - 在原始上下文内保持不变 - 在扩展上下文中仍有意义 """7. RoPE 扩展技术
7.1 NTK-Aware Scaling
class NTKAwareScaling: """ NTK-Aware Scaling (Neural Tangent Kernel) """
def motivation(self): """ 动机
简单的线性 RoPE 缩放（增加 max_position）在扩展时效果不佳。 原因：低频维度（负责长距离依赖）的波长变得太短。 """ return """ 线性缩放的问题:
原始: θ_i = base^{-2i/d} 缩放: θ_i = (base/s)^(-2i/d) (s > 1)
问题: 缩放后波长变为 λ/s，所有频率都变快。
对于长距离依赖: - 需要缓慢旋转（长波长） - 但缩放后旋转变快 - 导致模型困惑
NTK-Aware 方法: - 不均匀缩放 - 高频维度缩放多 - 低频维度缩放少 - 保持长距离能力 """
def formula(self): """ 公式 """ return """ NTK-Aware RoPE:
位置 m 的旋转角度:
θ_i' = (base · s)^{(-2i/d)} · λ^{(2i/d)} 当 m > original_max = base^{(-2i/d)} 当 m ≤ original_max
其中: - s: 缩放因子 - λ: NTK 调整参数
简化为: θ_i' = base^{(-2i/d)} · (s · λ^{1/d})^{2i/d}
设 α = s · λ^{1/d} θ_i' = (α · base)^{(-2i/d)}
但只对 i > threshold 的维度应用。 """
def implementation(self): """ 实现 """ return ''' def compute_ntk_freqs( seq_len: int, dim: int, base: float = 10000.0, scale: float = 2.0, original_max: int = 8192, ) -> torch.Tensor: """NTK-Aware 频率计算"""
# 计算基础频率 freqs = 1.0 / ( base ** (torch.arange(0, dim, 2).float() / dim) )
# 位置 positions = torch.arange(seq_len)
# NTK 调整 # 对于超过 original_max 的位置，应用缩放 scale_factor = torch.ones_like(freqs)
# 计算阈值维度 # λ ≈ (original_max / s)^{2i/d} 的反函数 threshold_dim = dim // 2 * (1 - torch.log(scale) / torch.log(original_max / seq_len)).item()
for i in range(dim // 2): if i >= threshold_dim: # 应用 NTK 调整 scale_factor[i] = scale ** (2 * i / dim)
freqs = freqs * scale_factor
# 计算相位 angles = positions[:, None] * freqs[None, :]
return torch.polar(torch.ones_like(angles), angles) '''7.2 YaRN
class YaRNImplementation: """ YaRN 实现 """
def yarn_method(self): """ YaRN 方法
Yet another RoPE extensioN """ return """ YaRN 的核心思想:
1. 位置插值 (Position Interpolation) - 将大位置映射到训练范围 - [0, L] → [0, L/s] - 但这会损害短距离信息
2. 注意力缩放 (Attention Scaling) - 调整 softmax 前的 logits - 补偿插值带来的注意力衰减 - 使用 yarm 因子调整
3. 频段调整 (Band Adjustment) - 对不同频段应用不同策略 - 高频（短距离）保持不变 - 低频（长距离）进行插值 """
def parameters(self): """ YaRN 参数 """ return { "beta_fast": "低频开始缩放的位置 (相对于 original_max)", "beta_slow": "缩放因子，控制衰减速度",
"example_qwen2": { "beta_fast": 32, "beta_slow": 1, "factor": 1e-5, },
"tuning": { "small_beta_fast": "更早开始缩放", "large_beta_slow": "更平滑的衰减", }, }7.3 其他扩展
class OtherExtensions: """ 其他扩展技术 """
def comet(self): """ CoMeT (Compound Memory Transformers) """ return """ CoMeT 方法:
思想: 结合不同 base 的 RoPE
RoPE_1 (base = 10000): 短距离依赖 RoPE_2 (base = 1000000): 长距离依赖
q_final = concat([RoPE_1(q), RoPE_2(q)], dim=-1)
优点: - 同时捕获多个尺度的依赖 - 自然处理长序列 """
def longrope(self): """ LongRoPE """ return """ LongRoPE:
1. 非均匀位置插值 - 不同维度的缩放因子不同 - 基于该维度的波长
2. 渐进式扩展 - 先在短序列上微调 - 再在长序列上微调
3. 位置解耦 - 将位置分成多个范围 - 每个范围独立编码 """8. 与其他位置编码对比
8.1 详细对比
class PositionEncodingComparison: """ 位置编码对比 """
def feature_comparison(self): """ 特性对比 """ return { "sinusoidal": { "外推能力": "一般", "实现复杂度": "低", "可学习": "否", "相对位置": "隐式", "计算开销": "低", }, "learnable": { "外推能力": "差", "实现复杂度": "低", "可学习": "是", "相对位置": "隐式", "计算开销": "低", }, "alibi": { "外推能力": "好", "实现复杂度": "中", "可学习": "否", "相对位置": "隐式", "计算开销": "低", }, "rope": { "外推能力": "好", "实现复杂度": "中", "可学习": "否", "相对位置": "显式", "计算开销": "中", }, }
def attention_pattern(self): """ Attention 模式对比 """ return """ 不同位置编码的 Attention 模式:
绝对位置编码: - Attention 分布与绝对位置相关 - 可能对位置过度依赖
相对位置编码 (Shaw): - Attention 分布与相对位置相关 - 需要额外的位置嵌入
ALiBi: - 固定的注意力衰减 - 线性偏置 - 自然倾向于近距离
RoPE: - 通过旋转编码位置 - 不添加额外的偏置项 - 保持标准 attention 结构 """8.2 适用场景
class UseCaseRecommendation: """ 适用场景推荐 """
def recommendation_table(self): """ 推荐 """ return { "短序列 + 高精度": "Sinusoidal 或 Learnable", "长序列 + 自然外推": "RoPE + NTK Scaling", "简单实现 + 好外推": "ALiBi", "预训练 + 微调": "RoPE (大多数开源模型)", "特定领域优化": "根据任务特性选择", }9. 实践指南
9.1 配置建议
class ConfigurationAdvice: """ 配置建议 """
def dimension_selection(self): """ 维度选择 """ return """ RoPE 维度:
head_dim 必须是偶数（推荐 64, 96, 128, 256）
常见配置: - LLaMA: head_dim = 128 - GPT-NeoX: head_dim = 96 - Falcon: head_dim = 64
head_dim 影响: - 波长数量: d/2 个不同的波长 - 计算效率: 2 的幂次更高效 """
def base_selection_guide(self): """ base 选择指南 """ return { "base_10000": "标准选择，大多数模型使用", "base_500": "更快的低频旋转，适合更长的上下文", "base_100000": "非常慢的旋转，适合超长距离依赖",
"rule_of_thumb": "base 越大，最长波长越短", "empirical": "base=10000 在大多数任务上表现良好", }
def memory_considerations(self): """ 显存考虑 """ return """ 显存使用:
freqs_cis 的大小: = seq_len × (head_dim/2) × 8 bytes (BF16)
示例: - seq_len=2048, head_dim=128: 2048 × 64 × 8 = 1 MB
- seq_len=100K, head_dim=128: 100000 × 64 × 8 = 48 MB
建议: - 预计算并缓存 - 超过 1M tokens 时考虑按需计算 """9.2 常见问题
class CommonIssues: """ 常见问题 """
def nan_loss(self): """ Loss NaN """ return { "symptom": "训练过程中 loss 突然变为 NaN", "causes": [ "RoPE 角度计算溢出", "数值精度问题", "学习率过大", ], "solutions": [ "检查 freqs_cis 计算", "使用 float32", "降低学习率", "添加 gradient clipping", ], }
def poor_long_context(self): """ 长上下文效果差 """ return { "symptom": "模型在短序列上好，长序列上差", "causes": [ "base 太小", "未使用 RoPE scaling", "训练数据缺乏长序列", ], "solutions": [ "增大 base", "应用 YaRN 或 NTK scaling", "使用长序列数据微调", ], }10. 核心公式汇总
10.1 RoPE 定义
其中 。
10.2 应用 RoPE
其中 。
10.3 旋转后的内积
根据旋转性质，这等价于：
仅依赖于相对位置 。
11. 总结
11.1 RoPE 核心要点
┌─────────────────────────────────────────────────────────────┐│ RoPE 核心要点 │├─────────────────────────────────────────────────────────────┤│ ││ 1. 数学基础 ││ → 旋转矩阵保持内积结构 ││ → 两两配对的 2D 旋转 ││ → 通过旋转编码位置信息 ││ ││ 2. 核心优势 ││ → 自然实现相对位置编码 ││ → 不需要在 Attention 添加偏置 ││ → 外推能力好 ││ → 实现简洁高效 ││ ││ 3. 频率设计 ││ → θ_i = base^{-2i/d} ││ → 低维度：高频（短距离） ││ → 高维度：低频（长距离） ││ ││ 4. 扩展技术 ││ → NTK-Aware Scaling: 非均匀缩放 ││ → YaRN: 注意力缩放 + 频段调整 ││ → 支持 128K+ 上下文 ││ ││ 5. 应用广泛 ││ → LLaMA, GPT-NeoX, Qwen, Falcon ││ → 大多数开源 LLM 采用 RoPE ││ │└─────────────────────────────────────────────────────────────┘11.2 选择建议
位置编码选择建议:
RoPE: ✓ 大多数场景的首选 ✓ 预训练模型 (LLaMA, Qwen) ✓ 需要长上下文
ALiBi: ✓ 简单实现 ✓ 对称注意力 ✓ 外推性要求高
Sinusoidal: ✓ 原版 Transformer ✓ 不需要微调 ✓ 快速原型
Learnable: ✓ 小数据集 ✓ 固定长度 ✓ 简单任务11.3 未来方向
RoPE 的未来发展:
1. 更长的上下文 → 1M+ tokens 的支持 → 更高效的扩展方法
2. 自适应频率 → 根据输入动态调整 → 任务相关的频率选择
3. 与其他机制的结合 → Flash Attention + RoPE → Sparse Attention + RoPE
4. 理论分析 → NTK 理论支持 → 表达能力分析推荐阅读
- RoPE Paper (Su et al., 2022) —— 原始论文
- RoPE 原理解释博客 —— 作者博客
- YaRN Paper —— 上下文扩展技术
- LLaMA 论文 —— RoPE 的实际应用
参考资料
- Su, J., et al. (2022). “RoFormer: Enhanced Transformer with Rotary Position Embedding.” arXiv<2104>.09864.
- Press, O., et al. (2021). “Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation.” ICLR 2022.
- Vaswani, A., et al. (2017). “Attention Is All You Need.” NeurIPS.
- Touvron, H., et al. (2023). “LLaMA: Open and Efficient Foundation Language Models.” Meta Research.
- Black, S., et al. (2022). “GPT-NeoX-20B: An Open-Source Autoregressive Language Model.” arXiv.
- Peng, B., et al. (2023). “YaRN: Efficient Context Window Extension of Large Language Models.” arXiv.
- Chen, S., et al. (2023). “NTK-Aware Scaled RoPE Allows LLaMA Models to Have Extended (8K+) Context Size.” Blog.
- Team, Q. (2024). “Qwen2 Technical Report.” arXiv.
- Dao, T., et al. (2022). “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.” NeurIPS.
文章分享
如果这篇文章对你有帮助，欢迎分享给更多人！
RoPE (Rotary Position Embedding) 深度解析：旋转位置编码的数学之美
https://aiattnstudio.link/posts/rope/欢迎来到Federico的个人博客
57文章
7分类
404标签