ROPE编码

参考https://www.zhihu.com/tardis/bd/art/647109286def apply_rotary_pos_emb(x: torch.Tensor, pos: torch.Tensor) - torch.Tensor: apply rotary position embedding x: [b, num_heads, num_tokens, head_dim] pos: [b, num_tokens, 2] returns: [b, num_heads, num_tokens, head_dim] b, num_heads, num_tokens, head_dim x.shape assert head_dim % 2 0, head_dim必须为偶数。 rotary_dim head_dim // 2 pos pos.unsqueeze(1) # [b, 1, num_tokens, 2] freq_base 10000.0 rotary_idx torch.arange(rotary_dim, devicex.device, dtypex.dtype) inv_freq 1.0 / (freq_base ** (rotary_idx / rotary_dim)) # [rotary_dim] angle_x pos[..., 0:1] * inv_freq # [b, 1, num_tokens, rotary_dim] angle_y pos[..., 1:2] * inv_freq # [b, 1, num_tokens, rotary_dim] angle angle_x angle_y # [b, 1, num_tokens, rotary_dim] freqs_cis torch.polar(torch.ones_like(angle), angle) # [b, 1, num_tokens, rotary_dim] x_ x.reshape(*x.shape[:-1], -1, 2) # [b, num_heads, num_tokens, rotary_dim, 2] x_ torch.view_as_complex(x_) # [b, num_heads, num_tokens, rotary_dim] x_rot torch.view_as_real(x_ * freqs_cis).reshape(*x.shape[:-1], -1) # [b, num_heads, num_tokens, head_dim] return x_rot def apply_rotary_pos_emb_2d_split(x: torch.Tensor, pos: torch.Tensor) - torch.Tensor: x: [b, h, t, d] pos: [b, t, 2], pos[..., 0]x, pos[..., 1]y return: [b, h, t, d] b, h, t, d x.shape assert d % 4 0, 2D RoPE通常要求head_dim能被4整除 d_half d // 2 # x-part / y-part d_quarter d // 4 # complex pairs per axis x_pos pos[..., 0] # [b, t] y_pos pos[..., 1] # [b, t] inv_freq 1.0 / ( 10000.0 ** (torch.arange(d_quarter, devicex.device, dtypex.dtype) / d_quarter) ) # [d/4] angle_x x_pos[:, None, :, None] * inv_freq[None, None, None, :] # [b,1,t,d/4] angle_y y_pos[:, None, :, None] * inv_freq[None, None, None, :] # [b,1,t,d/4] freqs_x torch.polar(torch.ones_like(angle_x), angle_x) # [b,1,t,d/4] freqs_y torch.polar(torch.ones_like(angle_y), angle_y) # [b,1,t,d/4] x1 x[..., :d_half] # [b,h,t,d/2] x2 x[..., d_half:] # [b,h,t,d/2] x1_c torch.view_as_complex(x1.reshape(b, h, t, d_quarter, 2)) # [b,h,t,d/4] x2_c torch.view_as_complex(x2.reshape(b, h, t, d_quarter, 2)) # [b,h,t,d/4] x1_rot torch.view_as_real(x1_c * freqs_x).reshape(b, h, t, d_half) x2_rot torch.view_as_real(x2_c * freqs_y).reshape(b, h, t, d_half) return torch.cat([x1_rot, x2_rot], dim-1)a.主要是因为PE在原理上有一个问题所以提出了RoPEb.我们自然的想到在LLM领域有一个叫RoPE的东西这是一种相对位置编码也是在建模相对关系只不过建模的是一个句子里不同位置的单词之间的相对位置关系c.RoPE的特点就是我在query/key上加入绝对特征绝对位置在attention的计算过程里这种关系会巧妙地转变为相对位置特征这和我们的目标不谋而合这样就省去了显示计算相对关系矩阵的开销d.虽然在Attention is all you need里提出了上述正余弦对PE但是在Bert, GPT等文章里仍然选择使用learnable PE也即用一组初始化的可学习向量用作PE而且效果看起来都要更好一些e.因为PE的有一个原理上的问题在q和k上增加了旋转在v上面没有Rope原理解析核心原理就是将attention中的qk矩阵相乘等价于复数相乘。而q中带上旋转角度也是一个复数相乘。从二维旋转矩阵推导出复数指数形式是理解 RoPE 为何能高效实现“相对位置建模”的关键。我们可以分三步走从矩阵乘法到复数乘法再到欧拉公式。第一步二维旋转矩阵 (Matrix Form)在二维空间中将一个向量x[x1,x2]T\mathbf{x} [x_1, x_2]^Tx[x1​,x2​]T顺时针旋转θ\thetaθ角度标准的旋转矩阵RθR_\thetaRθ​定义为Rθ(cos⁡θ−sin⁡θsin⁡θcos⁡θ)R_\theta \begin{pmatrix} \cos \theta -\sin \theta \\ \sin \theta \cos \theta \end{pmatrix}Rθ​(cosθsinθ​−sinθcosθ​)当我们把位置信息mmm注入 Query 向量qqq时f(q,m)Rmθ⋅q(cos⁡mθ−sin⁡mθsin⁡mθcos⁡mθ)(q1q2)(q1cos⁡mθ−q2sin⁡mθq1sin⁡mθq2cos⁡mθ)f(q, m) R_{m\theta} \cdot q \begin{pmatrix} \cos m\theta -\sin m\theta \\ \sin m\theta \cos m\theta \end{pmatrix} \begin{pmatrix} q_1 \\ q_2 \end{pmatrix} \begin{pmatrix} q_1 \cos m\theta - q_2 \sin m\theta \\ q_1 \sin m\theta q_2 \cos m\theta \end{pmatrix}f(q,m)Rmθ​⋅q(cosmθsinmθ​−sinmθcosmθ​)(q1​q2​​)(q1​cosmθ−q2​sinmθq1​sinmθq2​cosmθ​)第二步映射到复数空间 (Complex Mapping)二维向量[q1,q2]T[q_1, q_2]^T[q1​,q2​]T可以完美映射到复平面上的一个复数qq1iq2q q_1 iq_2qq1​iq2​。同样地旋转矩阵RmθR_{m\theta}Rmθ​在复数空间中对应的是复数旋转算子eimθe^{im\theta}eimθ。根据复数乘法法则(q1iq2)⋅(cos⁡mθisin⁡mθ)(q_1 iq_2) \cdot (\cos m\theta i \sin m\theta)(q1​iq2​)⋅(cosmθisinmθ)展开得到(q1cos⁡mθ−q2sin⁡mθ)i(q1sin⁡mθq2cos⁡mθ) (q_1 \cos m\theta - q_2 \sin m\theta) i(q_1 \sin m\theta q_2 \cos m\theta)(q1​cosmθ−q2​sinmθ)i(q1​sinmθq2​cosmθ)你会发现这个复数的实部和虚部刚好对应第一步中旋转矩阵变换后的两个分量。所以f(q,m)⟷q⋅eimθf(q, m) \longleftrightarrow q \cdot e^{im\theta}f(q,m)⟷q⋅eimθ第三步计算内积与相对位置 (Dot Product)现在我们要计算位置mmm的qqq和位置nnn的kkk的内积。在向量空间中内积是⟨q,k⟩\langle \mathbf{q}, \mathbf{k} \rangle⟨q,k⟩而在复数空间中两个复数向量的内积等于q⋅k∗q \cdot k^*q⋅k∗的实部Re\text{Re}Re。带入旋转后的复数⟨f(q,m),f(k,n)⟩Re[(qeimθ)⋅(keinθ)∗]\langle f(q, m), f(k, n) \rangle \text{Re} \left[ (q e^{im\theta}) \cdot (k e^{in\theta})^* \right]⟨f(q,m),f(k,n)⟩Re[(qeimθ)⋅(keinθ)∗]利用共轭性质KaTeX parse error: Double superscript at position 13: (ab)^ a^ b^̲以及(einθ)e−inθ(e^{in\theta})^ e^{-in\theta}(einθ)e−inθ**Re[q⋅eimθ⋅k∗⋅e−inθ]\text{Re} \left[ q \cdot e^{im\theta} \cdot k^* \cdot e^{-in\theta} \right]Re[q⋅eimθ⋅k∗⋅e−inθ]指数合并Re[q⋅k∗⋅ei(mθ−nθ)]Re[q⋅k∗⋅ei(m−n)θ]\text{Re} \left[ q \cdot k^* \cdot e^{i(m\theta - n\theta)} \right] \text{Re} \left[ q \cdot k^* \cdot e^{i(m-n)\theta} \right]Re[q⋅k∗⋅ei(mθ−nθ)]Re[q⋅k∗⋅ei(m−n)θ]DROPEDynamic Rotary Positional Embedding动态旋转位置编码. 为什么需要 “Dynamic” (D)普通的 RoPE如 Llama 中使用的通常处理的是离散的整数索引第 1 个词、第 2 个词……。但在自动驾驶中输入往往是连续的物理坐标比如障碍物在 BEV 下的x15.5m,y3.2mx15.5m, y3.2mx15.5m,y3.2m。DRoPE 的核心改进连续性建模 它不直接使用整数 index而是将物理距离作为输入。动态频率缩放 为了让模型能处理更广阔的感知范围比如从 50 米扩展到 200 米它会动态调整旋转的基底频率Base Frequency这类似于大语言模型中的 NTK-aware Scaling。2. DRoPE 在 2D/3D 空间中的运作在你之前的代码逻辑中角度是x⋅θy⋅θx \cdot \theta y \cdot \thetax⋅θy⋅θ。DRoPE 可能会引入自适应权重Anglefscale(dist)⋅(x⋅θxy⋅θy)\text{Angle} f_{scale}(dist) \cdot (x \cdot \theta_x y \cdot \theta_y)Anglefscale​(dist)⋅(x⋅θx​y⋅θy​)通过这种方式模型可以感知距离的非线性 远处的障碍物对当前规划的影响应该逐渐减弱DRoPE 可以通过旋转频率的控制让远距离的向量在内积时自然产生更大的衰减。坐标系对齐 动态地根据自车的速度Ego-speed或朝向调整编码相位保持空间的一致性。# 伪代码引入动态缩放 dist torch.norm(pos[..., :2], dim-1, keepdimTrue) # 根据距离动态调整频率基数例如距离越远旋转越慢保留更多全局信息 dynamic_inv_freq inv_freq / (1 torch.log(1 dist)) angle_x pos[..., 0:1] * dynamic_inv_freq angle_y pos[..., 1:2] * dynamic_inv_freq