LLama2源码分析——Rotary Position Embedding分析

参考：一文看懂 LLaMA 中的旋转式位置编码（Rotary Position Embedding）

原理推导参考自上文，以下结合huggingface代码分析公式计算过程

1 旋转角度计算

计算公式如下，其中d为词嵌入维度，这部分和论文原文一样
$${\theta }_j=10000^{-2(j-1)/d},j\in [1,2,\dots ,d/2]$$

# 计算词向量元素两两分组之后，每组元素对应的旋转角度
# 维度：[dim / 2]
inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))

2 计算整个seq的cos_sin矩阵

def _set_cos_sin_cache(self, seq_len, device, dtype):
    self.max_seq_len_cached = seq_len
    # 生成token长度序列
    t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
    # 计算两个矩阵的外积，结果维度[seq_len, dim // 2]
    freqs = torch.einsum("i,j->ij", t, self.inv_freq)
    # 类似[[0, 2, 4, ..., 0, 2, 4, ...], ...]形式，旋转角度两两一组相同
    emb = torch.cat((freqs, freqs), dim=-1)
    self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
    self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

3 计算旋转式位置编码

f q ( x m , m ) = ( W q x m ) e i m θ f k ( x n , n ) = ( W k x n ) e i n θ \begin{aligned}f_q(x_m,m)&=(W_qx_m)e^{im\theta} \\f_k(x_n,n)&=(W_kx_n)e^{in\theta}\end{aligned} fq​(xm​,m)fk​(xn​,n)​=(Wq​xm​)eimθ=(Wk​x