浅析RoPE旋转位置编码的远程衰减特性

为什么${\theta }_i$的取值会造成远程衰减性

旋转位置编码的出发点为：通过绝对位置编码的方式实现相对位置编码。

对词向量$\mathbf{q}$添加绝对位置信息$m$，希望找到一种函数$f$，使得：
$$<f(\mathbf{q},m),f(\mathbf{k},n)>=g(\mathbf{q},\mathbf{k},m-n)$$
假设词向量是二维的，借用复数来进行求解（具体求解过程参考：https://spaces.ac.cn/archives/8265），最终得到一种可行解：
$$\begin{array}{rl}f(\mathbf{q},m)& =\mathbf{q}{e}^{im\theta }\\ & =\left(\begin{array}{cc}cosm\theta & -sinm\theta \\ sinm\theta & cosm\theta \end{array}\right)\left(\begin{array}{c}{q}_0\\ q_1\end{array}\right)\end{array}$$
扩展到多维：

$f(\mathbf{q},m)={\mathbf{R}}_m\mathbf{q}$
$${\mathbf{R}}_m=\left(\begin{array}{ccccccc}cosm{\theta }_0& -sinm{\theta }_0& 0& 0& \cdots & 0& 0\\ sinm{\theta }_0& cosm{\theta }_0& 0& 0& \cdots & 0& 0\\ 0& 0& cosm{\theta }_1& -sinm{\theta }_1& \cdots & 0& 0\\ 0& 0& sinm{\theta }_1& cosm{\theta }_1& \cdots & 0& 0\\ ⋮& ⋮& ⋮& ⋮& \ddots & ⋮& ⋮\\ 0& 0& 0& 0& \cdots & cosm{\theta }_{d/2-1}& -sinm{\theta }_{d/2-1}\\ 0& 0& 0& 0& \cdots & sinm{\theta }_{d/2-1}& cosm{\theta }_{d/2-1}\end{array}\right)$$
相当于左乘一个旋转矩阵，或者说高维向量，每两维一组，分别旋转一个角度，且不改变模长。

显然，$({\mathbf{R}}_m\mathbf{q}{)}^T({\mathbf{R}}_n\mathbf{k})={\mathbf{q}}^T{\mathbf{R}}_m^T{\mathbf{R}}_n\mathbf{k}={\mathbf{q}}^T{\mathbf{R}}_{n-m}\mathbf{k}$，这样Attention就包含相对位置信息了。


下面分析为什么${\theta }_i$的取值会造成远程衰减性

远程衰减性指的是，对于两个词向量，如果两者相对距离较近，那么它们的注意力分数应该偏高，反之应该偏低。

假设$\mathbf{q}$和$\mathbf{k}$均为ones向量，则$({\mathbf{R}}_m\mathbf{q}{)}^T({\mathbf{R}}_n\mathbf{k})={\mathbf{q}}^T{\mathbf{R}}_{n-m}\mathbf{k}=2{\sum }_{i=0}^{d/2-1}cos(n-m){\theta }_i$，设相对距离$n-m$为$x$，则相对距离为$x$的向量之间注意力得分：
$$g(x)=2\sum _{i=0}^{d/2-1}cosx{\theta }_i$$
如果任意${\theta }_i=0$，则$g(x)=d$，无论相对距离多大，注意力得分都相等

如果任意${\theta }_i=1$，则$g(x)=dcosx$，随着相对距离增大，注意力得分呈周期性变化，但不会震荡衰减：

而作者在${\theta }_i$的选择上，沿用了Sinusoidal位置编码的方案，即${\theta }_i=10000^{-2i/d}$，它会带来一定的远程衰减性。

每个${\theta }_i$，$cosx{\theta }_i$的周期大小$T_i$等于$\frac{2\pi }{{\theta }_i}=\frac{2\pi }{10000^{-2i/d}}=2\pi \ast 10^{8i/d}$，所以$i$越大，$T_i$越大，最小周期为$T_0=2\pi$，最大周期为$T_{d/2-1}=2\pi \ast 10^{(4-\frac{8}{d})}$。

如果对于所有的$x$，$x<\frac{1}{4}{T}_{d/2-1}=\frac{\pi }{2}\ast 10^{(4-\frac{8}{d})}$，也就是说，$cosx{\theta }_{d/2-1}$处于单调递减区间（下方的蓝色区间）

由于前面的$cosx{\theta }_i$呈周期变化，而周期变化的函数 + 单调递减的函数 = 震荡递减的函数。因此，注意力得分$g(x)$随着相对距离$x$的增大而震荡减小。

比如在LLaMA中，$d=4096$，$\frac{1}{4}{T}_{d/2-1}$近似于$10^4$，由于实际应用中，最大序列长度一般不会大于$10^4$，所以相对距离$x<\frac{1}{4}{T}_{d/2-1}$一般是成立的，当然，也可以增大${\theta }_i=10000^{-2i/d}$中的10000，这样$T_{d/2-1}$会变得更大。

下面举几个例子：

import numpy as np 
import matplotlib.pyplot as plt 
from matplotlib.axes import Axes

def create_sin_cos_cache(max_num_tokens, head_size):
    theta = 10000 ** (-np.arange(0, head_size, 2) / head_size)  #(128,)
    # theta = np.ones(head_size//2)
    theta = theta.reshape(-1, 1).repeat(2, axis=1).flatten()  #(256,)

    pos = np.arange(0, max_num_tokens) #(512,)
    table = pos.reshape(-1, 1) @ theta.reshape(1, -1)  # [max_num_tokens, head_size] 512*256

    sin_cache = np.sin(table) #(512, 256)
    sin_cache[:, ::2] = -sin_cache[:, ::2]

    cos_cache = np.cos(table)
    return sin_cache, cos_cache

def rotate_half(vec):
    return vec.reshape(-1, 2)[:, ::-1].flatten()

def rotary(vec, pos, sin_table, cos_table):
    #vec.shape=256,是原始的q向量（q=Wx），
    #rotate_half(vec) 是处理过后的q向量
    #cos_table.shape=512*256
    return vec * cos_table[pos] + rotate_half(vec) * sin_table[pos]

def plot(plt_obj: Axes, pic_index, query_index=0, head_size=256, max_num_tokens=8192, step=1):
    q_vec = np.ones(head_size) #(256,)
    k_vec = np.ones(head_size) #(256,)
    sin_table, cos_table = create_sin_cos_cache(max_num_tokens, head_size) #(512, 256), (512, 256)

    rotated_q_vec = rotary(q_vec, query_index, sin_table, cos_table) #(256,)
    #如果query_index=0，则rotated_q_vec全为1
    #rotated_q_vec 是旋转后的q，即波浪q
    k_indices = np.arange(0, max_num_tokens, step) #(512,)
    rotated_k_vecs = rotary(k_vec, k_indices, sin_table, cos_table) #(512, 256)
    attn_scores = (rotated_k_vecs @ rotated_q_vec) / np.sqrt(head_size) #(512,)

    plt_obj.plot(k_indices, attn_scores)
    plt_obj.set_title(f"Figure {pic_index}: query_index={query_index}, d={head_size}")
    plt_obj.set_xlabel("key index")
    plt_obj.set_ylabel("attention score")

plt.rcParams.update({
    "font.sans-serif": ["Times New Roman", ],
    "font.size": 10
})

_, axes = plt.subplots(nrows=2, ncols=2, figsize=(10, 10))
plot(axes[0, 0], 1, query_index=0, head_size=256, max_num_tokens=10000)
# plot(axes[0, 1], 2, query_index=32, max_num_tokens=128)
# plot(axes[1, 0], 3, query_index=0, max_num_tokens=6553)
# plot(axes[1, 1], 4, query_index=0, head_size=8, max_num_tokens=65535)
plt.show()

当$d=4$时，最大周期$T_{d/2-1}$是628，下面的示例$x$会超过$\frac{1}{4}{T}_{d/2-1}$，因此$g(x)$呈周期性，并不是震荡减小

当$d=256$时，下面的示例$x$不超过$\frac{1}{4}{T}_{d/2-1}=14617$，因此震荡减小。