PaLM中ROPE位置编码实现源码解析

1、源码

import torch
from einops import rearrange
from torch import einsum, nn

class RotaryEmbedding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim)) # [d/2]
        # inv_freq形式化为 [theta_0, theta_1, ..., theta_(d/2-1)]
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, max_seq_len, *, device):
        # 计算m
        seq = torch.arange(max_seq_len, device=device, dtype=self.inv_freq.dtype) # [length]

        freqs = einsum("i , j -> i j", seq, self.inv_freq)
        # freqs形式化为 [m*theta_0, m*theta_1, ..., m*theta_d/2],其中 m=0,1,...,length-1

        # return结果形式化为 [m*theta_0, m*theta_1, ..., m*theta_(d/2-1),m*theta_0, m*theta_1, ..., m*theta_(d/2-1)],其中 m=0,1,...,length-1
        return torch.cat((freqs, freqs), dim=-1) # [length, d]


def rotate_half(x):
    # x为q或k， 形式化为[q0, q1, .., qd-1]
    # x: [bs, head, length, d]--> [bs, head, length, 2, d/2]
    # 下式将x进行划分，前半部分形式化为[q0, q1, .., q(d/2-1)]
    x = rearrange(x, "... (j d) -> ... j d", j=2)
    # x1形式化为[q0, q1, .., q(d/2-1)]
    # x2形式化为[q(d/2）, q(d/2+1）, .., q(d-1)]
    x1, x2 = x.unbind(dim=-2)
    return torch.cat((-x2, x1), dim=-1) # [-q(d/2）, -q(d/2+1）, .., -q(d-1), q0, q1, .., q(d/2-1)]

def apply_rotary_pos_emb(pos, t):
    # t: [bs, head, length, d], [q0, q1, .., qd-1]
    # pos: [length, d], [m*theta_0, m*theta_1, ..., m*theta_(d/2-1),m*theta_0, m*theta_1, ..., m*theta_(d/2-1)]
    rotate_half(t)
    # 以第一个为例，q0*cos(m*theta_0) - q(d/2)*sin(m*theta_0)
    #  第二个，q1*cos(m*theta_1) - q(d/2+1)*sin(m*theta_1)
    return (t * pos.cos()) + (rotate_half(t) * pos.sin())


if __name__=='__main__':
    # (bs, head, length, d)
    q = torch.randn((2, 12, 10, 32))  # q=[q0, q1, .., qd-1]
    k = torch.randn((2, 12, 10, 32))
    v = torch.randn((2, 12, 10, 32))
    print('q:', q[0][0][0])
    print('k:', k[0][0][0])
    rotary_emb = RotaryEmbedding(dim=32)
    pos_emb = rotary_emb(max_seq_len=10, device=torch.device('cpu'))  # [length, d]
    q_new, k_new = map(lambda t: apply_rotary_pos_emb(pos_emb, t), (q, k))
    print()

2、公式
$$\left(\begin{array}{c}{q}_0\\ q_1\\ ...\\ q_{d/2-1}\\ q_{d/2}\\ ...\\ q_{d-2}\\ q_{d-1}\end{array}\right)\ast \left(\begin{array}{c}cos(m{\theta }_0)\\ cos(m{\theta }_1)\\ ...\\ cos(m{\theta }_{d/2-1})\\ cos(m{\theta }_0)\\ ...\\ cos(m{\theta }_{d/2-2})\\ cos(m{\theta }_{d/2-1})\end{array}\right)+\left(\begin{array}{c}-q_{d/2}\\ -q_{d/2+1}\\ ...\\ -q_{d-1}\\ q_0\\ ...\\ q_{d/2-2}\\ q_{d/2-1}\end{array}\right)\left(\begin{array}{c}sin(m{\theta }_0)\\ sin(m{\theta }_1)\\ ...\\ sin(m{\theta }_{d/2-1})\\ sin(m{\theta }_0)\\ ...\\ sin(m{\theta }_{d/2-2})\\ sin(m{\theta }_{d/2-1})\end{array}\right)$$
3、图形
观察上图，可以发现$q_0$和$q_{d/2}$相互作用，生成新的$q_0^{new}$和$q_{d/2}^{new}$，拆解后可以得到下式

$$q_0^{new}=q_0\ast cos(m{\theta }_0)-q_{d/2}\ast sin(m{\theta }_0)$$
$$q_{d/2}^{new}=q_0\ast sin(m{\theta }_0)+q_{d/2}\ast cos(m{\theta }_0)$$
也即向量$(q_0^{new},q_{d/2}^{new})$由向量$(q_0,q_{d/2})$逆时针旋转$m{\theta }_0$得到
可于下面链接中LLaMA中ROPE实现做对比
LLaMA中ROPE位置编码实现源码解析