RoPE原理

不讲从0-1的推导原理

背景是Transformer的positional embedding直接添加在word_embedding上

$$\underset{word-embedding}{\underbrace{\left[\begin{array}{c}1\\ 2\end{array}\right]}}+\underset{positional-embedding}{\underbrace{\left[\begin{array}{c}0.3\\ 0.6\end{array}\right]}}=\left[\begin{array}{c}1.3\\ 2.6\end{array}\right]$$
发现对word_embedding的模长和角度都发生了变化，虽然有一定程度的位置区分度，但是也很大程度上引入了一部分噪声。

$$\begin{array}{rl}Q& =W_q(X_m+P_m)\\ K& =W_k(X_n+P_n)\\ Q{K}^T& \approx (X_m+P_m)(X_n+P_n{)}^T\\ & =\underset{词向量间}{\underbrace{X_m{X}_n^T}}+\underset{噪声}{\underbrace{X_m{P}_n^T+P_m{X}_n^T}}+\underset{位置编码间}{\underbrace{P_m{P}_n^T}}\end{array}$$

因此引入RoPE，一种将位置信息以向量旋转的方式，不改变模长的条件下进行引入的方式

首先定义q@[bs,seq,n_head,head_dim]、k/v@[bs,seq,n_kv_head,head_dim]

基于

$$(\underset{real}{\underbrace{a}}+i\underset{imaginary}{\underbrace{b}})(\cos \theta +i\sin \theta )=\underset{real-after-rotate\theta }{\underbrace{(a\cos \theta -b\sin \theta )}}+i\underset{imaginary-after-rotate\theta }{\underbrace{(a\sin \theta +b\cos \theta )}}$$

这个公式，实现如下的向量旋转

$$\left[\begin{array}{c}a\\ b\end{array}\right]\to \left[\begin{array}{c}a\cos \theta -b\sin \theta \\ a\sin \theta +b\cos \theta \end{array}\right]$$

所以也可以通过$q->q_r+j{q}_i->q_r^{\prime }+j{q}_i^{\prime }=>q^{\prime }$这个方式实现对于query向量的旋转

因此有q_r@[bs,seq,n_head,head_dim//2]以及q_i@[bs,seq,n_head,head_dim//2]的拆分

$$\underset{q_r}{\underbrace{{\left[\begin{array}{ccc}{q}_r^{(1,1)}& \cdots & q_r^{(1,d^{\prime })}\\ \cdots & q_r^{(m,k)}& \cdots \\ q_r^{(seq,1)}& \cdots & q_r^{(seq,d^{\prime })}\end{array}\right]}_{seq\times d^{\prime }}}}\ast f(\underset{freq}{\underbrace{{\left[\begin{array}{ccc}1\cdot 10000^{-\frac{2\cdot 1}{d}}& \cdots & 1\cdot 10000^{-\frac{2\cdot d^{\prime }}{d}}\\ \cdots & \cos (m{\theta }_k)& \cdots \\ seq\cdot 10000^{-\frac{2\cdot 1}{d}}& \cdots & seq\cdot 10000^{-\frac{2\cdot d^{\prime }}{d}}\end{array}\right]}_{seq\times d^{\prime }}}})$$
其中$d^{\prime }=\frac{d}{2}$

相当于实现了

$$\begin{array}{rl}{q}_r^{\prime }& =q_r\ast \cos (freq)-q_i\ast \sin (freq)\\ q_i^{\prime }& =q_r\ast \sin (freq)+q_i\ast \cos (freq)\end{array}$$

苏神在博客中写道

$$\underset{R_m}{\underbrace{\left[\begin{array}{ccccccc}\cos m{\theta }_0& -\sin m{\theta }_0& 0& 0& \cdots & 0& 0\\ \sin m{\theta }_0& \cos m{\theta }_0& 0& 0& \cdots & 0& 0\\ 0& 0& \cos m{\theta }_1& -\sin m{\theta }_1& \cdots & 0& 0\\ 0& 0& \sin m{\theta }_1& \cos m{\theta }_1& \cdots & 0& 0\\ \cdots & \cdots & \cdots & \cdots & \cdots & \cdots & \cdots \\ 0& 0& 0& 0& \cdots & \cos m{\theta }_{d/2-1}& -\sin m{\theta }_{d/2-1}\\ 0& 0& 0& 0& \cdots & \sin m{\theta }_{d/2-1}& \cos m{\theta }_{d/2-1}\end{array}\right]}}\left[\begin{array}{c}{q}_0\\ q_1\\ q_2\\ q_3\\ \cdots \\ q_{d-2}\\ q_{d-1}\end{array}\right]$$
并且定义

$$\begin{array}{rl}\left[\begin{array}{c}{q}_0^{\prime }\\ q_1^{\prime }\end{array}\right]& =\left[\begin{array}{cc}\cos m{\theta }_0& -\sin m{\theta }_0\\ \sin m{\theta }_0& \cos m{\theta }_0\end{array}\right]\left[\begin{array}{c}{q}_0\\ q_1\end{array}\right]\\ \left[\begin{array}{c}{q}_2^{\prime }\\ q_3^{\prime }\end{array}\right]& =\left[\begin{array}{cc}\cos m{\theta }_1& -\sin m{\theta }_1\\ \sin m{\theta }_1& \cos m{\theta }_1\end{array}\right]\left[\begin{array}{c}{q}_2\\ q_3\end{array}\right]\\ & \cdots \end{array}$$
去实现q=[q0 q1 ... q_{d-1}] -> q'=[q'0 q'1 ... q'{d-1}]的旋转

[!important]

这里容易误解的点在于，苏神是对$f(q,m)=\left[\begin{array}{cc}\cos m\theta & -\sin m\theta \\ \sin m\theta & \cos m\theta \end{array}\right]\left[\begin{array}{c}{q}_0\\ q_1\end{array}\right]$的一个特例，一定要理解的是m表示的是当前的token在整个sequence中的位置。所以这里相当于

$$q={\left[\begin{array}{ccc}{q}^{(1,1)}& \cdots & q^{(1,d)}\\ \cdots & q^{(i,j)}& \cdots \\ q^{(seq,1)}& \cdots & q^{(seq,d)}\end{array}\right]}_{seq\times d}$$

在第m行，把一行数据拿出来

$${\left[\begin{array}{ccccc}{q}^{(m,1)}& q^{(m,2)}& q^{(m,3)}& \cdots & q^{(m,d)}\end{array}\right]}_{1\times d}$$

并把它称作为$[q_0,q_1,q_2,\cdots ,q_{d-1}]$

因此不要忘记要对除m行以外的所有$q^{(i,j)}$都进行上述的旋转操作！

因此，推导到这里，实际上，这和我上文就是一致的了：

我们从元素级别入手：苏神在博客中的$q_0\in {\mathbb{R}}^{bs\times 1\times n-head\times 1}$，切记这里的$q_0$实则是$q^{(m,0)}$，和它交互的是$\cos m{\theta }_0$与$\sin m{\theta }_0$；而我们的推导当中$q_r^{(m,k)}\in {\mathbb{R}}^{bs\times 1\times n-head\times 1}$，与它交互的是$\cos m{\theta }_k$与$\sin m{\theta }_k$。从元素级别可以看到每一个元素的交互都是一样的。

[!important]

事实上，拓展到高维后，有结论如下：

对于任意维度$d$，任意第$k$组$k=0,1,...,d/2-1$：

$q_{2k}=q_r[k]$第 k 组的第一个元素 = 第 k 个复数的实部

$q_{2k+1}=q_i[k]$第 k 组的第二个元素 = 第 k 个复数的虚部

我们可以把整体进行一个缩小，由于各方法对q@[bs,seq,n_head,head_dim]中bs,n_head不改动，后续高维推导中省略这2个维度，所以这里$\mathbf{q}\in {\mathbb{R}}^{seq\times d}$简化为如下所示

$$\mathbf{q}=\left[\begin{array}{cccc}{q}^{(0,0)}& q^{(0,1)}& \cdots & q^{(0,d-1)}\\ q^{(1,0)}& q^{(1,1)}& \cdots & q^{(1,d-1)}\\ \cdots & \cdots & \cdots & \cdots \\ q^{(seq-1,0)}& q^{(seq-1,1)}& \cdots & q^{(seq-1,d-1)}\end{array}\right]$$

根据上述结论可知：

$$\begin{array}{rl}{\mathbf{q}}_{\mathbf{r}}& =\left[\begin{array}{cccc}{q}^{(0,0)}& q^{(0,2)}& \cdots & q^{(0,d-2)}\\ q^{(1,0)}& q^{(1,2)}& \cdots & q^{(1,d-2)}\\ \cdots & \cdots & \cdots & \cdots \\ q^{(seq-1,0)}& q^{(seq-1,2)}& \cdots & q^{(seq-1,d-2)}\end{array}\right]\\ {\mathbf{q}}_{\mathbf{i}}& =\left[\begin{array}{cccc}{q}^{(0,1)}& q^{(0,3)}& \cdots & q^{(0,d-1)}\\ q^{(1,1)}& q^{(1,3)}& \cdots & q^{(1,d-1)}\\ \cdots & \cdots & \cdots & \cdots \\ q^{(seq-1,1)}& q^{(seq-1,3)}& \cdots & q^{(seq-1,d-1)}\end{array}\right]\end{array}$$

写到这里，也更加清晰，苏神的推导中取$\mathbf{q}$逐$m$行进行$R_m$矩阵旋转：

$$\begin{array}{rl}{q}^{\prime (m,2i)}& =q^{(m,2i)}\cdot \cos m{\theta }_i-q^{(m,2i+1)}\cdot \sin m{\theta }_i\\ q^{\prime (m,2i+1)}& =q^{(m,2i)}\cdot \sin m{\theta }_i+q^{(m,2i+1)}\cdot \cos m{\theta }_i\end{array}$$

而工程代码中总结规律发现上述方法效率太低，既然：

1. $q^{(m,2i)}$和$q^{(m,2i+1)}$总是有固定的$\cos m{\theta }_i$/$\sin m{\theta }_i$和$-\sin m{\theta }_i$/$\cos m{\theta }_i$要相乘；
2. $q^{(m,2i)}$和$q^{(m,2i+1)}$恰好对应${\mathbf{q}}_{\mathbf{r}}$和${\mathbf{q}}_{\mathbf{i}}$中的一行；

不如直接整合成一个矩阵，实现

$$\begin{array}{rl}\underset{{\mathbf{q}}_{\mathbf{r}}^{\prime }}{\underbrace{\left[\begin{array}{cccc}{q}^{\prime (0,0)}& q^{\prime (0,2)}& \cdots & q^{\prime (0,d-2)}\\ q^{\prime (1,0)}& q^{\prime (1,2)}& \cdots & q^{\prime (1,d-2)}\\ \cdots & \cdots & \cdots & \cdots \\ q^{\prime (seq-1,0)}& q^{\prime (seq-1,2)}& \cdots & q^{\prime (seq-1,d-2)}\end{array}\right]}}& =\underset{{\mathbf{q}}_{\mathbf{r}}}{\underbrace{\left[\begin{array}{cccc}{q}^{(0,0)}& q^{(0,2)}& \cdots & q^{(0,d-2)}\\ q^{(1,0)}& q^{(1,2)}& \cdots & q^{(1,d-2)}\\ \cdots & \cdots & \cdots & \cdots \\ q^{(seq-1,0)}& q^{(seq-1,2)}& \cdots & q^{(seq-1,d-2)}\end{array}\right]}}\ast \underset{\cos \mathbf{m\theta }}{\underbrace{\cos (\left[\begin{array}{cccc}0{\theta }_0& 0{\theta }_1& \cdots & 0{\theta }_{d/2-1}\\ 1{\theta }_0& 1{\theta }_1& \cdots & 1{\theta }_{d/2-1}\\ \cdots & \cdots & \cdots & \cdots \\ (seq-1){\theta }_0& (seq-1){\theta }_1& \cdots & (seq-1){\theta }_{d/2-1}\end{array}\right])}}-\underset{{\mathbf{q}}_{\mathbf{i}}}{\underbrace{\left[\begin{array}{cccc}{q}^{(0,1)}& q^{(0,3)}& \cdots & q^{(0,d-1)}\\ q^{(1,1)}& q^{(1,3)}& \cdots & q^{(1,d-1)}\\ \cdots & \cdots & \cdots & \cdots \\ q^{(seq-1,1)}& q^{(seq-1,3)}& \cdots & q^{(seq-1,d-1)}\end{array}\right]}}\ast \underset{\sin \mathbf{m\theta }}{\underbrace{\sin (\left[\begin{array}{cccc}0{\theta }_0& 0{\theta }_1& \cdots & 0{\theta }_{d/2-1}\\ 1{\theta }_0& 1{\theta }_1& \cdots & 1{\theta }_{d/2-1}\\ \cdots & \cdots & \cdots & \cdots \\ (seq-1){\theta }_0& (seq-1){\theta }_1& \cdots & (seq-1){\theta }_{d/2-1}\end{array}\right])}}\\ \underset{{\mathbf{q}}_{\mathbf{i}}^{\prime }}{\underbrace{\left[\begin{array}{cccc}{q}^{\prime (0,1)}& q^{\prime (0,3)}& \cdots & q^{\prime (0,d-1)}\\ q^{\prime (1,1)}& q^{\prime (1,3)}& \cdots & q^{\prime (1,d-1)}\\ \cdots & \cdots & \cdots & \cdots \\ q^{\prime (seq-1,1)}& q^{\prime (seq-1,3)}& \cdots & q^{\prime (seq-1,d-1)}\end{array}\right]}}& =\underset{{\mathbf{q}}_{\mathbf{r}}}{\underbrace{\left[\begin{array}{cccc}{q}^{(0,0)}& q^{(0,2)}& \cdots & q^{(0,d-2)}\\ q^{(1,0)}& q^{(1,2)}& \cdots & q^{(1,d-2)}\\ \cdots & \cdots & \cdots & \cdots \\ q^{(seq-1,0)}& q^{(seq-1,2)}& \cdots & q^{(seq-1,d-2)}\end{array}\right]}}\ast \underset{\sin \mathbf{m\theta }}{\underbrace{\sin (\left[\begin{array}{cccc}0{\theta }_0& 0{\theta }_1& \cdots & 0{\theta }_{d/2-1}\\ 1{\theta }_0& 1{\theta }_1& \cdots & 1{\theta }_{d/2-1}\\ \cdots & \cdots & \cdots & \cdots \\ (seq-1){\theta }_0& (seq-1){\theta }_1& \cdots & (seq-1){\theta }_{d/2-1}\end{array}\right])}}+\underset{{\mathbf{q}}_{\mathbf{i}}}{\underbrace{\left[\begin{array}{cccc}{q}^{(0,1)}& q^{(0,3)}& \cdots & q^{(0,d-1)}\\ q^{(1,1)}& q^{(1,3)}& \cdots & q^{(1,d-1)}\\ \cdots & \cdots & \cdots & \cdots \\ q^{(seq-1,1)}& q^{(seq-1,3)}& \cdots & q^{(seq-1,d-1)}\end{array}\right]}}\ast \underset{\cos \mathbf{m\theta }}{\underbrace{\cos (\left[\begin{array}{cccc}0{\theta }_0& 0{\theta }_1& \cdots & 0{\theta }_{d/2-1}\\ 1{\theta }_0& 1{\theta }_1& \cdots & 1{\theta }_{d/2-1}\\ \cdots & \cdots & \cdots & \cdots \\ (seq-1){\theta }_0& (seq-1){\theta }_1& \cdots & (seq-1){\theta }_{d/2-1}\end{array}\right])}}\end{array}$$

具体复杂度对比这里就不做了，工程代码实现了从$O(n^2)\to O(n)$的简化

此外，由于RoPE的

$$\begin{array}{rl}Q& =W_q(R_m{X}_m)\\ K& =W_k(R_n{X}_n)\\ Q{K}^T& \approx (R_m{X}_m)(R_n{X}_n{)}^T\\ & =R_m{X}_m{X}_n^T{R}_n^T\end{array}$$
没有引入类似positional embedding中的噪声项，因此也更加稳定。

最后由衷说一声，苏神牛逼！

代码实现

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    # 获取输入张量的形状：批量大小、序列长度、键/值对头的数量、每个头的维度大小
    bs, slen, n_kv_heads, head_dim = x.shape
    
    # 如果重复次数为1，则不需要重复，直接返回原始张量
    if n_rep == 1:
        return x
    
    # 对张量进行扩展和重塑操作以重复键值对
    return (
        x[:, :, :, None, :]  # 在第四个维度（头的维度前）添加一个新的维度
        .expand(bs, slen, n_kv_heads, n_rep, head_dim)  # 将新添加的维度扩展到n_rep大小，实现重复的效果
        .reshape(bs, slen, n_kv_heads * n_rep, head_dim)  # 重新塑形，合并键/值对头的数量和重复次数的维度
    )
    
# 注意：此处的dim应为 dim//n_head，因为我们是对每个head进行旋转嵌入
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    # torch.arange(0, dim, 2)[: (dim // 2)].float()生成了一个从0开始，步长为2的序列，长度为dim的一半
    # 然后每个元素除以dim，再取theta的倒数，得到频率
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    # 生成一个从0到end的序列，长度为end
    t = torch.arange(end, device=freqs.device)
    # 计算外积，得到一个二维矩阵，每一行是t的元素乘以freqs的元素
    freqs = torch.outer(t, freqs).float()
    # 计算频率的余弦值，得到实部
    freqs_cos = torch.cos(freqs)
    # 计算频率的正弦值，得到虚部
    freqs_sin = torch.sin(freqs)
    return freqs_cos, freqs_sin
    
def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
    # 获取x的维度数
    ndim = x.ndim
    
    # 断言，确保1在x的维度范围内
    assert 0 <= 1 < ndim
    
    # 断言，确保freqs_cis的形状与x的第二维和最后一维相同
    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
    
    # 构造一个新的形状，除了第二维和最后一维，其他维度都为1，这样做是为了能够将freqs_cis与x进行广播操作
    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
    
    # 将freqs_cis调整为新的形状，并返回
    return freqs_cis.view(shape)
    
def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cos: torch.Tensor,
    freqs_sin: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:

    # 将查询和键张量转换为浮点数，并重塑形状以分离实部和虚部
    xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1)
    xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1)

    # 重新塑形频率张量以进行广播
    freqs_cos = reshape_for_broadcast(freqs_cos, xq_r)
    freqs_sin = reshape_for_broadcast(freqs_sin, xq_r)

    # 应用旋转，分别计算旋转后的实部和虚部
    xq_out_r = xq_r * freqs_cos - xq_i * freqs_sin
    xq_out_i = xq_r * freqs_sin + xq_i * freqs_cos
    xk_out_r = xk_r * freqs_cos - xk_i * freqs_sin
    xk_out_i = xk_r * freqs_sin + xk_i * freqs_cos

    # 将最后两个维度合并，并还原为原始张量的形状
    xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(3)
    xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(3)

    return xq_out.type_as(xq), xk_out.type_as(xk)

class Attention(nn.Module):
    def __init__(self, args: ModelConfig):
        super().__init__()
        # 根据是否指定n_kv_heads，确定用于键（key）和值（value）的头的数量。
        self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
        # 确保总头数可以被键值头数整除。
        assert args.n_heads % self.n_kv_heads == 0

        # 模型并行处理大小，默认为1。
        model_parallel_size = 1
        # 本地计算头数，等于总头数除以模型并行处理大小。
        self.n_local_heads = args.n_heads // model_parallel_size
        # 本地键值头数，等于键值头数除以模型并行处理大小。
        self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
        # 重复次数，用于扩展键和值的尺寸。
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        # 每个头的维度，等于模型维度除以头的总数。
        self.head_dim = args.dim // args.n_heads

        # 定义权重矩阵。
        self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
        # 输出权重矩阵。
        self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)

        # 定义dropout。
        self.attn_dropout = nn.Dropout(args.dropout)
        self.resid_dropout = nn.Dropout(args.dropout)
        # 保存dropout概率。
        self.dropout = args.dropout

        # 检查是否使用Flash Attention（需要PyTorch >= 2.0）。
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
        if not self.flash:
            # 若不支持Flash Attention，则使用手动实现的注意力机制，并设置mask。
            print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
            # 创建一个上三角矩阵，用于遮蔽未来信息。
            mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
            mask = torch.triu(mask, diagonal=1)
            # 注册为模型的缓冲区
            self.register_buffer("mask", mask)

    def forward(self, x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor):
        # 获取批次大小和序列长度，[batch_size, seq_len, dim]
        bsz, seqlen, _ = x.shape

        # 计算查询（Q）、键（K）、值（V）。
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
        # 调整形状以适应头的维度。
        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

        # 应用旋转位置嵌入（RoPE）。
        xq, xk = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin)

        # 对键和值进行扩展以适应重复次数。
        xk = repeat_kv(xk, self.n_rep)
        xv = repeat_kv(xv, self.n_rep)

        # 将头作为批次维度处理。
        xq = xq.transpose(1, 2)
        xk = xk.transpose(1, 2)
        xv = xv.transpose(1, 2)

        # 根据是否支持Flash Attention，选择实现方式。
        if self.flash:
            # 使用Flash Attention。
            output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None, dropout_p=self.dropout if self.training else 0.0, is_causal=True)
        else:
            # 使用手动实现的注意力机制。
            scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
            assert hasattr(self, 'mask')
            scores = scores + self.mask[:, :, :seqlen, :seqlen]
            scores = F.softmax(scores.float(), dim=-1).type_as(xq)
            scores = self.attn_dropout(scores)
            output = torch.matmul(scores, xv)

        # 恢复时间维度并合并头。
        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)

        # 最终投影回残差流。
        output = self.wo(output)
        output = self.resid_dropout(output)
        return output