旋转位置编码 RoPE

xllm

已于 2024-09-02 14:51:04 修改

阅读量273

点赞数 2

文章标签：机器学习人工智能

于 2024-08-31 21:14:49 首次发布

本文链接：https://blog.csdn.net/sail1234ss/article/details/141758590

版权

使用 RoPE 目的

$\textbf{Attention}(K,Q,V)=SoftMax(\frac{QK^T}{\sqrt{d_k}})V$

在计算 Attention 过程中，Q、K 都含有位置信息，计算 $QK^T$ 时，为了 Q 中行向量 $q_m$ 与 K 的列向量 $k_n$ 的内积保留相对位置信息

如何实现

位置信息

$q_m$ 和 $k_n$ 包含的位置信息为

$p_{i,2t}=sin(\frac{i}{10000^{\frac{2t}{d}}}) \\\\\\ p_{i,2t+1}=cos(\frac{i}{10000^{\frac{2t}{d}}}) \\$

代码实现

# position 就对应 token 序列中的位置索引 i
# hidden_dim 就对应词嵌入维度大小 d
# seq_len 表示 token 序列长度
def get_position_angle_vec(position):
    return [position / np.power(10000, 2 * (hid_j // 2) / hidden_dim) for hid_j in range(hidden_dim)]

# position_angle_vecs.shape = [seq_len, hidden_dim]
position_angle_vecs = np.array([get_position_angle_vec(pos_i) for pos_i in range(seq_len)])

# 分别计算奇偶索引位置对应的 sin 和 cos 值
position_angle_vecs[:, 0::2] = np.sin(position_angle_vecs[:, 0::2])  # dim 2t
position_angle_vecs[:, 1::2] = np.cos(position_angle_vecs[:, 1::2])  # dim 2t+1

# positional_embeddings.shape = [1, seq_len, hidden_dim]
positional_embeddings = torch.FloatTensor(position_angle_vecs).unsqueeze(0)

如何把位置信息加入到 $q_m$ 、 $k_n$

$f_q(x_m,m)=(W_qx_m)e^{im\theta} \\ f_k(x_n,n)=(W_kx_n)e^{in\theta} \\$

其中

$e^{im\theta}=\cos (m\theta) + i\sin (m\theta) \\ e^{in\theta}=\cos (n\theta) + i\sin (n\theta) \\ e^{i(m-n)\theta}=\cos ((m-n)\theta) + i\sin ((m-n)\theta) \\$

def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    # xq.shape = [batch_size, seq_len, dim]
    # xq_.shape = [batch_size, seq_len, dim // 2, 2]
    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 2)
    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 2)
    
    # 转为复数域
    xq_ = torch.view_as_complex(xq_)
    xk_ = torch.view_as_complex(xk_)
    
    # 应用旋转操作，然后将结果转回实数域
    # xq_out.shape = [batch_size, seq_len, dim]
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(2)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(2)
    return xq_out.type_as(xq), xk_out.type_as(xk)