详解旋转位置编码

首先通过一个具体的数字例子，详细说明如何使用旋转位置编码（Rotary Position Embedding, RoPE）。

---

1. 输入数据

假设：

• 输入张量 x：形状为 (batch_size=2, seq_len=3, num_heads=1, head_dim=4)，值为：
  $$x=\left[\begin{array}{c}[[1,2,3,4]],\\ [[5,6,7,8]],\\ [[9,10,11,12]]\end{array}\right]$$
• 位置索引 t：序列长度为 3，位置索引为 [0, 1, 2]。
• 旋转频率 theta：设为 10000.0。

---

2. 预计算旋转编码 pos_cis

(1) 计算频率向量

• 向量维度 dim = 4，因此 dim // 2 = 2。
• 频率向量公式：
  $${\theta }_j=\frac{1}{{\text{base}}^{j/d}}$$
  其中 $j=0,1$，$d=4$，base = 10000.0。

计算：
$$\begin{gathered} {\theta }_0=\frac{1}{10000^{0/4}}=1.0 \\ {\theta }_1=\frac{1}{10000^{1/4}}\approx 0.5623 \end{gathered}$$
因此，频率向量为：
$$\text{freqs}=[1.0,0.5623]$$

(2) 计算旋转角度

• 位置索引 t = [0, 1, 2]。
• 旋转角度公式：
  $$m{\theta }_j=t\cdot {\theta }_j$$
  计算：
  $$t\cdot \text{freqs}=\left[\begin{array}{cc}0\cdot 1.0& 0\cdot 0.5623\\ 1\cdot 1.0& 1\cdot 0.5623\\ 2\cdot 1.0& 2\cdot 0.5623\end{array}\right]=\left[\begin{array}{cc}0.0& 0.0\\ 1.0& 0.5623\\ 2.0& 1.1246\end{array}\right]$$

(3) 生成复数形式的旋转编码

使用 torch.polar 生成旋转编码：
$$\text{pos\_cis}=e^{i\cdot m{\theta }_j}$$
计算：
$$\text{pos\_cis}=\left[\begin{array}{cc}{e}^{i\cdot 0.0}& e^{i\cdot 0.0}\\ e^{i\cdot 1.0}& e^{i\cdot 0.5623}\\ e^{i\cdot 2.0}& e^{i\cdot 1.1246}\end{array}\right]=\left[\begin{array}{cc}1.0+i\cdot 0.0& 1.0+i\cdot 0.0\\ 0.5403+i\cdot 0.8415& 0.8472+i\cdot 0.5314\\ -0.4161+i\cdot 0.9093& 0.4284+i\cdot 0.9036\end{array}\right]$$

---

3. 应用旋转位置编码

(1) 将输入张量 x 转换为复数形式

• x 的形状为 (2, 3, 1, 4)，将其转换为复数形式：
  $$x_q=\left[\begin{array}{c}[1+i\cdot 2,3+i\cdot 4],\\ [5+i\cdot 6,7+i\cdot 8],\\ [9+i\cdot 10,11+i\cdot 12]\end{array}\right]$$

(2) 调整 pos_cis 的形状

• pos_cis 的形状为 (3, 2)，调整为 (1, 3, 1, 2)，以便与 x 广播。

(3) 应用旋转编码

将 pos_cis 与 x_q 逐元素相乘：
$$x_q^{\prime }=x_q\cdot \text{pos\_cis}$$
计算：
$$x_q^{\prime }=\left[\begin{array}{cc}(1+i\cdot 2)\cdot (1.0+i\cdot 0.0)& (3+i\cdot 4)\cdot (1.0+i\cdot 0.0)\\ (5+i\cdot 6)\cdot (0.5403+i\cdot 0.8415)& (7+i\cdot 8)\cdot (0.8472+i\cdot 0.5314)\\ (9+i\cdot 10)\cdot (-0.4161+i\cdot 0.9093)& (11+i\cdot 12)\cdot (0.4284+i\cdot 0.9036)\end{array}\right]$$
逐元素计算结果：
$$x_q^{\prime }=\left[\begin{array}{cc}1.0+i\cdot 2.0& 3.0+i\cdot 4.0\\ -2.6248+i\cdot 7.3479& -0.3776+i\cdot 10.1706\\ -13.5123+i\cdot 2.0013& -6.5604+i\cdot 15.1404\end{array}\right]$$

(4) 将复数形式转换回实数形式

将 x_q' 转换回实数形式，形状为 (2, 3, 1, 4)：
$$x_q^{\prime }=\left[\begin{array}{c}[[1.0,2.0,3.0,4.0]],\\ [[-2.6248,7.3479,-0.3776,10.1706]],\\ [[-13.5123,2.0013,-6.5604,15.1404]]\end{array}\right]$$

---

4. 最终结果

应用旋转位置编码后，x 的值变为：
$$x^{\prime }=\left[\begin{array}{c}[[1.0,2.0,3.0,4.0]],\\ [[-2.6248,7.3479,-0.3776,10.1706]],\\ [[-13.5123,2.0013,-6.5604,15.1404]]\end{array}\right]$$

---

5. 总结

通过这个具体的数字例子，我们展示了旋转位置编码的完整计算过程：

1. 预计算旋转编码 pos_cis。
2. 将输入张量 x 转换为复数形式。
3. 应用旋转编码，将位置信息融入 x。
4. 将结果转换回实数形式。

旋转位置编码的核心思想是通过复数旋转将位置信息融入查询和键向量中，从而增强模型对序列位置的感知能力。

好的！以下是代码对应的数学公式和原理解释，帮助你更好地理解旋转位置编码（Rotary Position Embedding, RoPE）的实现。

---

1. 旋转位置编码的数学原理

旋转位置编码的核心思想是通过复数旋转将位置信息融入查询（Q）和键（K）向量中。具体来说，对于位置 m 和 n，查询和键向量会被旋转一个与位置相关的角度，从而在计算注意力分数时引入位置信息。

(1) 旋转公式

对于向量 x 和位置 m，旋转位置编码的公式为：
$$\text{RoPE}(x,m)=x\cdot e^{im\theta }$$
其中：

• $x$ 是输入向量。
• $m$ 是位置索引。
• $\theta$ 是旋转角度，由频率向量决定。

(2) 频率向量

频率向量 ${\theta }_j$ 的计算公式为：
$${\theta }_j=\frac{1}{{\text{base}}^{j/d}}$$
其中：

• $j$ 是维度索引。
• $d$ 是向量的维度。
• $\text{base}$ 是一个常数（代码中的 theta）。

---

2. 代码的数学公式

以下是代码中每一步对应的数学公式。

(1) precompute_pos_cis 函数

1. 计算频率向量 ${\theta }_j$：
   $${\theta }_j=\frac{1}{{\text{base}}^{j/d}}$$
   代码实现：

   freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
   

2. 计算旋转角度 $m{\theta }_j$：
   $$m{\theta }_j=t\cdot {\theta }_j$$
   代码实现：

   freqs = torch.outer(t, freqs).float()
   

3. 生成复数形式的旋转编码：
   $$\text{pos\_cis}=e^{im{\theta }_j}$$
   代码实现：

   pos_cis = torch.polar(torch.ones_like(freqs), freqs)
   

(2) apply_rotary_emb 函数

1. 将查询和键向量转换为复数形式：
   $$x_q=x_q+i{x}_q^{\prime },x_k=x_k+i{x}_k^{\prime }$$
   代码实现：

   xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
   xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
   

2. 应用旋转位置编码：
   $$x_q^{\prime }=x_q\cdot e^{im{\theta }_j},x_k^{\prime }=x_k\cdot e^{im{\theta }_j}$$
   代码实现：

   xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)
   xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)
   

---

3. 公式与代码的对应关系

数学公式	代码实现
${\theta }_j=\frac{1}{{\text{base}}^{j/d}}$	freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
$m{\theta }_j=t\cdot {\theta }_j$	freqs = torch.outer(t, freqs).float()
$\text{pos\_cis}=e^{im{\theta }_j}$	pos_cis = torch.polar(torch.ones_like(freqs), freqs)
$x_q=x_q+i{x}_q^{\prime }$	xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
$x_q^{\prime }=x_q\cdot e^{im{\theta }_j}$	xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)

---

4. 示例

假设：

• 输入向量 $x_q=[1,2,3,4]$。
• 位置索引 $m=1$。
• 频率向量 ${\theta }_j=[0.1,0.2]$。

(1) 计算旋转角度

$$m{\theta }_j=[1\cdot 0.1,1\cdot 0.2]=[0.1,0.2]$$

(2) 生成复数形式的旋转编码

$$\text{pos\_cis}=e^{i[0.1,0.2]}=[\cos (0.1)+i\sin (0.1),\cos (0.2)+i\sin (0.2)]$$

(3) 应用旋转位置编码

$$x_q^{\prime }=[1+i2,3+i4]\cdot [\cos (0.1)+i\sin (0.1),\cos (0.2)+i\sin (0.2)]$$

---

5. 总结

通过数学公式和代码的对应关系，可以看出旋转位置编码的核心是通过复数旋转将位置信息融入查询和键向量中。具体步骤包括：

1. 计算频率向量 ${\theta }_j$。
2. 计算旋转角度 $m{\theta }_j$。
3. 生成复数形式的旋转编码 $e^{im{\theta }_j}$。
4. 将旋转编码应用到查询和键向量上。

这种方法能够在不增加额外参数的情况下，显著提升模型对序列位置的感知能力。