AI学习记录 - 旋转位置编码

创作不易，有用点赞，写作有利于锻炼一门新的技能，有很大一部分是我自己总结的新视角

1、前置条件：要理解旋转位置编码前，要熟悉自注意力机制，否则很难看得懂，在我的系列文章中有对自注意力机制的画图解释。

先说重要的结论(下面 q向量 和 k向量 是自注意力矩阵诞生的，不懂先去看注意力机制)：

结论1：旋转位置编码本身是绝对位置编码，但是和自注意力机制中的一个qk向量结合之后，就变成了相对位置编码。因为自注意力机制中qk会计算点积，正是恰好这个内积，顺带把旋转位置编码变成了相对位置编码，所以一般说旋转位置编码既包含了绝对位置编码含义，也包含了相对位置编码含义。

结论2：假设没有位置编码这个东西，自注意力机制中，qk向量进行内积的时候，经过反向传播，会逐渐得出词汇与词汇的关联度矩阵，假设10个词汇计算内积，当两个词汇关联度越高，这两个词汇的内积(q * k)越大，重点来了：当对q 和 k叠加上旋转位置编码之后，那不仅仅是两个词汇关联度越高，内积越大，并且当两个词汇位置距离越近，内积也越大。

结论3：原来词向量跟词向量的内积大小只跟词汇的语义相关，内积越大，两个词汇的语义关联度越高。叠加上旋转位置编码后，距离相近的词向量内积也大。当一个句子中，两个词汇距离很远但是语义强相关，那他们的内积就是大；当两个词汇语义没啥关联但是距离很近，内积也是大；当两个词汇距离又近，语义有强相关，内积就是大大的。

2、经过上面的结论，其实我们知道了旋转位置编码在哪个位置起到的作用，就是得出 q 和 k 向量之后。

在说旋转位置编码怎么旋转之前，数学界已经就有了怎么对一个向量进行旋转，举个例子

如果你本身对位置编码不熟悉，在了解旋转位置编码之前，建议先去看我的另一篇博客，有个传统绝对位置编码的解释，旋转位置编码在没和qk叠加之前，其实和绝对位置编码差不多，你会发现他们的公式在某些地方非常的接近。如果这个所谓的旋转位置编码和传统绝对位置编码通过一样的方式叠加到词向量上面，旋转位置编码还是一个绝对位置编码，关键在于叠加方式不一样。当然传统位置编码使用旋转位置编码的叠加方式，也没有产生相对位置含义，所以旋转位置编码的计算公式和他的叠加方式是相互相成的。

传统绝对位置编码公式：

旋转位置编码公式：

对应的图片变化：

3、上面知道如果向量需要旋转，其实需要一个二维向量，但是 q 和 k 都是一维向量，怎么办呢，通过如下叠加，把 q 和 k 向量都按照如下图所示变成二维向量：

然后把q的每一列当成（x，y）取出来，下图所示，一共有8个（x，y），所有的q向量都进行这样子的计算，计算完成之后，我们就说q叠加上了旋转位置编码。

然后又转换回来，这个q叠加上了旋转位置编码

4、我简单提供一个证明，证明在向量在旋转位置编码之后，词汇距离越近，内积就越大，假设两个token的q向量都一样。

假设两个token的初始表示为相同的向量：𝑣=[1,0,1,0]

旋转矩阵为：

下面我们来套用上面说到的公式计算：

当这个向量位置为 1

当这个向量位置为 3

5、最后代码实现，在这里我也是拿某些大佬的，我在这里写了很多print形状，从观察矩阵形状变化去理解比较好

我这里提一下，就是你会发现代码其实有点难以看懂，这是因为涉及到批次计算，多头，导致矩阵代码中做了很多的矩阵变换，但是本质的流程还是我上面所说的，只是在实现过程中，考虑到优化导致的代码难以按照我上面所述的流程看懂，但是本质和上面一样。

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

# %%

def sinusoidal_position_embedding(batch_size, nums_head, max_len, output_dim, device):
    # batch_size = 8
    # nums_head = 12
    # max_len = 10
    # output_dim = 32
    position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(-1)
    ids = torch.arange(0, output_dim // 2, dtype=torch.float)  # 即公式里的i, i的范围是 [0,d/2]
    theta = torch.pow(10000, -2 * ids / output_dim)

    print(position) # [[0.],[1.],[2.],[3.],[4.],[5.],[6.],[7.],[8.],[9.]]
    print(output_dim) # 32
    print(theta) # tensor([1.0000e+00, 5.6234e-01, 3.1623e-01, 1.7783e-01, 1.0000e-01, 5.6234e-02,
                    # 3.1623e-02, 1.7783e-02, 1.0000e-02, 5.6234e-03, 3.1623e-03, 1.7783e-03,
                    # 1.0000e-03, 5.6234e-04, 3.1623e-04, 1.7783e-04])

    print(theta.size()) # torch.Size([16])
    print(position.size()) # torch.Size([10, 1])
    embeddings = position * theta  # 即公式里的：pos / (10000^(2i/d))
    print(embeddings.size()) # torch.Size([10, 16])
    # (max_len, output_dim//2, 2)
    embeddings = torch.stack([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)
    # For example：
    # torch.sin(embeddings) = tensor([[ 0.0000,  0.8415,  0.9093,  0.1411, -0.7568, -0.9589]])
    # torch.cos(embeddings) = tensor([[ 1.0000,  0.5403, -0.4161, -0.9900, -0.6536,  0.2837]])
    # torch.stack = tensor([[[ 0.0000,  1.0000],
    #                        [ 0.8415,  0.5403],
    #                        [ 0.9093, -0.4161],
    #                        [ 0.1411, -0.9900],
    #                        [-0.7568, -0.6536],
    #                        [-0.9589,  0.2837]]])
    print(embeddings.size()) # torch.Size([10, 16, 2])
    embeddings = embeddings.repeat((batch_size, nums_head, *([1] * len(embeddings.shape))))  # 在bs维度重复，其他维度都是1不重复
    print(embeddings.size()) # torch.Size([8, 12, 10, 16, 2])
    # reshape后就是：偶数sin, 奇数cos了
    embeddings = torch.reshape(embeddings, (batch_size, nums_head, max_len, output_dim))



    print(embeddings.size()) # torch.Size([8, 12, 10, 32])
    embeddings = embeddings.to(device)
    return embeddings


# %%

def RoPE(q, k):
    # q,k: (bs, head, max_len, output_dim)
    batch_size = q.shape[0] # batch_size = 8
    nums_head = q.shape[1] # nums_head = 12
    max_len = q.shape[2] # max_len = 10
    output_dim = q.shape[3] # output_dim = 32
    

    pos_emb = sinusoidal_position_embedding(batch_size, nums_head, max_len, output_dim, q.device)
    print(pos_emb.size()) # torch.Size([8, 12, 10, 32])




    # 看rope公式可知，相邻cos，sin之间是相同的，所以复制一遍。如(1,2,3)变成(1,1,2,2,3,3)
    cos_pos = pos_emb[...,  1::2].repeat_interleave(2, dim=-1)  # 将奇数列信息抽取出来也就是cos 拿出来并复制
    sin_pos = pos_emb[..., ::2].repeat_interleave(2, dim=-1)  # 将偶数列信息抽取出来也就是sin 拿出来并复制



    print(cos_pos.size()) # torch.Size([8, 12, 10, 32])
    print(sin_pos.size()) # torch.Size([8, 12, 10, 32])



    q2 = torch.stack([-q[..., 1::2], q[..., ::2]], dim=-1)
    print(q2.size()) # torch.Size([8, 12, 10, 16, 2])
    q2 = q2.reshape(q.shape)  # reshape后就是正负交替了
    print(q2.size()) # torch.Size([8, 12, 10, 32])
    # 更新qw, *对应位置相乘
    q = q * cos_pos + q2 * sin_pos
    print(q.size()) # torch.Size([8, 12, 10, 32])




    k2 = torch.stack([-k[..., 1::2], k[..., ::2]], dim=-1)
    k2 = k2.reshape(k.shape)
    # 更新kw, *对应位置相乘
    k = k * cos_pos + k2 * sin_pos

    return q, k


# %%

def attention(q, k, v, mask=None, dropout=None, use_RoPE=True):
    # q.shape: (bs, head, seq_len, dk)
    # k.shape: (bs, head, seq_len, dk)
    # v.shape: (bs, head, seq_len, dk)

    if use_RoPE:
        q, k = RoPE(q, k)

    d_k = k.size()[-1]

    att_logits = torch.matmul(q, k.transpose(-2, -1))  # (bs, head, seq_len, seq_len)
    att_logits /= math.sqrt(d_k)

    if mask is not None:
        att_logits = att_logits.masked_fill(mask == 0, -1e9)  # mask掉为0的部分，设为无穷大

    att_scores = F.softmax(att_logits, dim=-1)  # (bs, head, seq_len, seq_len)

    if dropout is not None:
        att_scores = dropout(att_scores)

    # (bs, head, seq_len, seq_len) * (bs, head, seq_len, dk) = (bs, head, seq_len, dk)
    return torch.matmul(att_scores, v), att_scores


if __name__ == '__main__':
    # (bs, head, seq_len, dk)
    q = torch.randn((8, 12, 10, 32))
    k = torch.randn((8, 12, 10, 32))
    v = torch.randn((8, 12, 10, 32))

    res, att_scores = attention(q, k, v, mask=None, dropout=None, use_RoPE=True)


    # (bs, head, seq_len, dk),  (bs, head, seq_len, seq_len)
    print(res.shape, att_scores.shape)