大模型系列：快速通俗理解Transformer旋转位置编码RoPE

猿类崛起@

已于 2024-07-09 11:54:00 修改

阅读量2.4k

点赞数 20

文章标签： transformer 人工智能大数据开发语言学习语言模型自然语言处理

于 2024-06-15 11:05:53 首次发布

本文链接：https://blog.csdn.net/2401_84495872/article/details/139698878

版权

前言

旋转位置编码RoPE（Rotary Position Embedding）是一种Transformer模型中的位置编码策略，它广泛应用于LLama，ChatGLM等大模型，本篇先介绍RoPE的实现步骤和源码，再深入讲解RoPE涉及到的数学原理，力求做到从易到难，学习曲线平滑。

内容摘要

位置编码知识准备
旋转位置编码的本质和计算流程
旋转位置编码如何表达相对位置信息
旋转位置编码的源码分析
旋转位置编码的推导

位置编码知识准备

由于Transformer的Self Attention具有排列不变性，因此需要通过引入位置编码来让模型感知到输入序列中每个单词的位置信息，位置编码分为绝对位置编码和相对位置编码。
绝对位置编码根据单个单词的绝对位置来定义位置编码，每个位置都会分配一个位置编码，将位置编码的表征和单词本身的表征进行融合，再输入给Self Attention，相当于在输入层就把位置信息给弥补上去。绝对位置编码从实现方式上又分为固定式和可学习式，固定式形如原生的Transformer所采用的三角sin-cos位置编码，所谓固定指的是根据一个无参的固定公式就可以推演出位置编码，而可学习式没有固定的位置编码公式，通过初始化位置向量让模型根据上下文数据自适应地学习出来，Bert和GPT采用的可学习式。

Bert的可学习式绝对位置编码和原始输入相加

相对位置编码对两个单词之间的相对位置进行建模，并且将相对位置信息加入到Self Attention模型结构中，形如Transformer-XL，DeBERTa等采用的就是相对位置编码。Self Attention的本质是两个单词信息的内积操作，相对位置编码的思想是对内积的计算方式进行改进，在内积中注入两个单词的相对位置因素。

词和词之间的相对位置

旋转位置编码的本质和计算流程

旋转位置编码RoPE是一种固定式的绝对位置编码策略，但是它的绝对位置编码配合Transformer的Attention内积注意力机制能达到相对位置编码的效果。RoPE的本质是对两个token形成的Query和Key向量做一个变换，使得变换后的Query和Key带有位置信息，进一步使得Attention的内积操作不需要做任何更改就能自动感知到相对位置信息。换句话说，RoPR的出发点和策略用的相对位置编码思想，但是实现方式的确用的是绝对位置编码。
固定式表明RoPE没有额外需要模型自适应学习的参数，因此RoPE是一种高效的编码方式。绝对位置编码表明RoPE给文本的每个位置单词都分配了一个位置表征，和三角sin-cos位置编码一样，RoPE通过token在句子中的位置，token embedding中每个元素的位置，这两个要素一起确定位置编码的表达，先给出RoPE的公式如下

RoPE位置编码公式

RoPE有一定数学推导环节，但是最终的公式并不复杂，因此本篇先从RoPE公式入手介绍RoPE在做什么，该公式是将一个原始的token向量改造为一个注入位置信息之后的新向量的过程。
其中第一项代表某个位置为m的token的原始Query向量，0~d-1代表向量每个位置的元素，d代表向量的维度，第二项为一个同样长度是d的带有cos三角函数的向量，它和Query向量逐位相乘，第三项由原始Query变换而来，第四项和第二项类似区别是将cos替换为sin。
该公式的目的是将原始Query向量改造成一个带有位置信息的新向量，位置信息由参数m和θ进行表征，其中m为token在句子中的位置，θ的下标和向量中各元素的位置直接相关，公式如下

θ和向量中各元素的位置的关系

因此只要给到某个token的输入Query向量，知道token在上下文窗口下处于第几位，就可以将它的Query向量通过RoPE的公式改造为一个新的向量形式，新形成的向量和原向量维度完全一致。以“我爱你”这句话中的第二个词“爱”为例，设词向量的维度d=4，词向量表征为[0.2, 0.1, -0.3, 0.7]，则经过RoPE变化的计算示意图如下

RoPE计算示意图

公式中的第三项由原始向量变换而来，对于原始输入向量，将前后两个元素位置构成一对，交换两者的位置，并且对于偶数位取了相反数，因此每个元素位的注入位置信息的过程，可以看成是该元素和它相邻的元素，分别经过sin，cos三角函数加权求和的结果，比如q0的RoPE结果是q0和q1这一对元素经过三角函数变换的结果。在下文的源码分析中，我们会介绍此处的相邻条件并不是必须的，而是任意不重复的一对都满足这个变换性质。

在Transformer原生的三角sin-cos位置编码中，采用相加的形式将位置编码融入到词向量中，而在RoPE中采用的是类似哈达马积的乘积形式，读者可以将以上RoPE公式做的事情类比于Transformer中原始向量表征和sin-cos位置编码相加的过程。

旋转位置编码如何表达相对位置信息

实现相对位置能力的途径不同：sin-cos位置编码由于三角函数的性质，导致它本身就具备表达相对距离的能力，而RoPE位置编码本身不能表达相对距离，需要结合Attention的内积才能激发相对距离的表达能力
和原输入的融合计算方式不同：sin-cos位置编码直接和原始输入相加，RoPE位置编码采用类似哈达马积相乘的形式

在知识准备模块我们介绍的相对位置编码，其主要的思想是原始输入不变，将相对位置信息注入Attention模块，采用对Attention的网络结构进行修改方式，将位置表征因素也额外的加入Attention计算，使得Attention模块能够把输入层丢失的位置信息弥补回来。
RoPE参考相对位置编码的思想，它也是在Attention模块让模型感知到相对位置，但是它是不改变Attention的结构，反而像绝对位置编码一样在输入层做文章，对输入向量做改造，改造后Attention模块能够重新感知到相对位置，同样能把位置信息弥补回来，因此RoPE可是说是使用绝对位置编码的方式实现了相对位置编码，是两者的融合。
至于为什么RoPE可以通过Attention来激发相对位置信息，原因是带有RoPE位置编码两个token，它们形成的Quey向量和Key向量进入Self Attention层之后，Attention内积的结果可以恒等转化一个函数，该函数只和Quey向量，Key向量，以及两个token位置之差有关，细节推导将在下文的进行介绍，读者先对这个结论有个初步印象。

RoPE配合内积感知相对位置信息

旋转位置编码的源码分析

在前文已经通过公式和一个具体的例子说明了RoPE的计算方式，下面结合HuggingFace的LLaMA大模型实现类LlamaForCausalLM中RoPE的源码再巩固一下。先给到源码实现的步骤，分为三步

1. 初始化cos向量和sin向量：根据给定的上下文窗口大小作为m，多头下每个头的向量的维度大小作为d，生成cos向量和sin向量，也就是RoPE公式中的第二项和第四项。在LLaMA2中上下文窗口为m=4096，每个头下的向量维度为d=128。
1. 截取对应长度的cos向量和sin向量：根据输入Query的实际长度，截取步骤一中生成的cos向量和sin向量，例如上下文窗口为4096，但是实际输入句子长度仅为10，则截取出前10个位置的cos向量和sin向量。
3.使用cos向量和sin向量改造Query和Key：根据步骤二产出的cos向量和sin向量，套用RoPE的公式，对原始Query和Key分别计算出注入位置信息之后的Query和Key。

我们顺着这三个步骤查看LlamaForCausalLM中RoPE的实现，RoPE在Attention操作类LlamaAttention中实现

class LlamaAttention(nn.Module):
    def __init__(self, config: LlamaConfig):
        ...
        # 步骤一：初始化
        self.rotary_emb = LlamaRotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)

    def forward(...):
        ...
        # 步骤二：截取长度
        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
        # 步骤三：改造Query，Key
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
        ...

最关键的三行代码分别对应步骤一二三，在LlamaAttention的初始化模块通过LlamaRotaryEmbedding子模块实现对RoPE的初始化，具体为对公式中的第二项cos向量和第四项sin向量进行初始化。

class LlamaRotaryEmbedding(torch.nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()
        # TODO dim=128, max_position_embeddings=4096, 远程衰减
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
        self.register_buffer("inv_freq", inv_freq)

        # Build here to make `torch.jit.trace` work.
        self.max_seq_len_cached = max_position_embeddings
        # 4096
        t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        # [4096, 64] => [4096, 128]
        emb = torch.cat((freqs, freqs), dim=-1)
        # TODO [1, 1, 4096, 128]
        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)

由于第二项和第四项仅仅是三角函数不同，三角函数的右侧参数是相同的，都是mθ，因此只需要将所有的mθ生成好，再对结果分别取cos和sin即可。在实现上作者通过m向量和θ向量的笛卡尔积相乘构造出来了mθ组合矩阵，核心代码为以下5行，freqs即为mθ的组合结果

        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
        self.register_buffer("inv_freq", inv_freq)
        self.max_seq_len_cached = max_position_embeddings
        t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)

以m=4096，θ=128为例，可以通过m和θ的罗列将这个过程展现出来，每个格子中的结果为m和θ相乘的结果

初始化mθ组合

θ只生成了64种情况，作者将两个freqs在θ拼接，形成了最终的128种情况，代码备注中作者说这个地方和论文的公式不一样，但是最终的效果是相同的，不一样体现在θ下标的排布顺序和论文公式不一样

            # Different from paper, but it uses a different permutation in order to obtain the same calculation
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)

因此最终的mθ组合为一个[4096，128]的二维矩阵，模拟如下

最终mθ组合

紧接着作者分别用cos和sin生成了两个结果向量，并且将它们从二维矩阵变成了四维，原因是在多头注意力中，Query和Key都是四维的形式存在，分别是[batch_size, num_heads, seq_len, head_dim]

        self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
        self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)

初始化完毕之后，在LlamaRotaryEmbedding的forward阶段根据seq_len完成截取操作，对第三维就是上下文窗口m这个维度进行截取

    def forward(self, x, seq_len=None):
        ...
        return (
            # TODO [1, 1, seq_len, emb_size=128]
            self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
            self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
        )

其中seq_len为输入文本的实际长度，在调用的时候它等于Key向量的实际长度，如果每次输入的是一部分token，有前文past_key_value状态，则文本长度会和之前进行拼接相加，最终得到的cos，sin就是截取之后公式中的第二项和第四项

        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        # [batch_size, num_headsm, kv_seq_len, head_dim] => kv_seq_len
        kv_seq_len = key_states.shape[-2]
        if past_key_value is not None:
            kv_seq_len += past_key_value[0].shape[-2]
        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)

进入步骤三，将原始的Query，Key向量，cos，sin输入到apply_rotary_pos_emb中，输出的query_states, key_states就是注入位置信息之后的Query，Key向量结果

query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

在apply_rotary_pos_emb中出现了RoPE公式，第一项为Query，第二项为cos向量，第三项通过rotate_half方法对Query进行变换，第四项为sin向量，通过逐位相乘再相加的形式得到结果，分别对Query和Key用同样的方式进行改造

def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
    ...
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

进一步看rotate_half是否和论文公式中给定的变换一致，答案是否定的，而在前文中对于cos和sin向量的实现和论文也不一致，这两处代码的不一致恰好使得最终的效果和论文一致

def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    # TODO 前64个embedding位置 x=[batch_size, num_heads, seq_len, emb_size] => [batch_size, num_heads, seq_len, emb_size/2]
    x1 = x[..., : x.shape[-1] // 2]
    # TODO 后64个embedding位置 x=[batch_size, num_heads, seq_len, emb_size] => [batch_size, num_heads, seq_len, emb_size/2]
    x2 = x[..., x.shape[-1] // 2 :]
    # TODO 后64embedding位置取负号，和前64embedding位置拼接
    return torch.cat((-x2, x1), dim=-1)

HuggingFace的代码逻辑它实现的计算公式实际为

image.png

该公式和RoPE论文公式在第二，三，四项上都有些许差异，具体为元素位置排列上的差异，在原RoPE公式中q0的结果是q0和q1这一对元素经过三角函数变换而成的，但是在实际公式中q0是由q0和q64这一对形成的，只需要把q1想像成q64则两个公式完全等价，那q1和q64互换对最终的结果影响吗？答案是没有影响，RoPE对原始向量的改造本质上是以一对元素为单位经过旋转矩阵运算，将所有对的结果进行拼接的过程，而到底是选择连续的元素作为一对，还是其他的挑选方式都是可以的，只要是embedding维度为偶数，且挑选的策略为不重复的一对，最终Attention的内积结果都能感知到相对位置信息，因为Attention满足内积线性叠加性，至于谁和谁一组进行叠加并不重要。

代码改版的公式和论文原版的公式

在改造完Query和Key之后，将他们灌入注意力网络，计算注意力权重再携带Value信息，代码如下

attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
attn_output = torch.matmul(attn_weights, value_states)

注意此处的注意力并没有做任何的结构调整，和传统的Transformer注意力的结构一模一样，RoPE的相对位置改造对天然适配下游注意力网络，另外Value信息没有参加RoPE改造，RoPE只对内积过程中的Query和Key做改造。

旋转位置编码的推导

直接使用RoPE的结论在网络结构中使用起来不复杂，RoPE怎么来的需要经过一系列公式推导，其中涉及复数的概念，包括复数的坐标表示和三角表示，复数相乘运算，共轭复数，欧拉公式和旋转矩阵。本篇的讲解会直接引用RoPE的作者博客[Transformer升级之路：2、博采众长的旋转式位置编码

注入位置信息的一般形式

这样改造的目的是使得Attention内积能够自动感知到相对位置信息，即内积可以恒等转化为一个函数，这个函数只和原始的Query，Key，以及两个token之间的距离m-n相关，令g为这个恒等变换函数，则有以下公式

内积恒等变换中的相对位置因素

下面就是要找到一个改造函数f，使得以上这个恒等变换g成立。
作者首先从最简单的二维角度考虑，假设q和k的embedding维度都是2维，将变换后的q，k用复数进行表示，其中第一维为复数的实部，第二维为复数的虚部，以一个[-2.1, 3.2]的二维向量为例，复数形式表示如下

二维向量的复数表示

则两者的内积等于q和k的共轭复数相乘的实部，公式如下

内积转化为复数和共轭复数相乘取实部

公式中的Re代表复数的实部，f*代表共轭复数，这里涉及复数的乘法和共轭复数

复数和共轭复数
复数z的坐标表示为z=a+bi，其中a是复数的实部，b是复述的虚部，z的共轭复数是a-bi，即实部不变，虚部取相反数

复数的乘法
两个复数相乘直接展开相乘即可，z1=a+bi，z2=c+di，则z1×z2=(ac-bd)+(bc+ad)i

根据以上两个性质，等式右侧等于(ac+bd)+(bc-ad)i，其实部为ac+bd，真好为两向量对应位置元素相乘再相加，因此该内积公式成立，等式联立可得

内积恒等变换中的相对位置因素

把实部Re拿掉，f(q,m)和f(k,n)共轭相乘的结果是一个复数，设其结果为g，该复数也必定和q，k，m-n相关，令下式为公式一

复数相乘等式——公式一

我们将三个复数用复数的三角形式表示，表示为向量的模长和幅角形式，令下式为公式二

复数的三角表示——公式二

其中R代表向量的模长，e的iθ次幂为欧拉公式，欧拉公式展开如下

欧拉公式

和向量的模R相乘欧拉公式对应复数的三角表示，其中θ为幅角

复数的三角表示

下面的推导需要用到复数相乘的性质

复数三角形式相乘
复数的三角形式，两个复数相乘，模长相乘，幅角相加。这个可以用三角表示的相乘展开证明，这里举一个例子：复数z=1+√3i，其中模长为，幅叫我为60度，如果z和z相乘，根据性质，相乘的结果映射到坐标系应该模长为4，幅角为120度，因此z×z=-2+2√3i，在坐标系下的可视化如下

复数相乘的性质

复数z再乘以z，在坐标系上相当于将z的模长乘以2，并且逆时针旋转了z的幅角60度。

根据复数相乘的性质，因此等式一左边两个复数相乘的模相乘，角相加，等式右边也是一个复数，因此两边的模和角度应该相等，则有

模和幅角相等

注意第二行为两个θ角度相减，原因是f(k,n)取了共轭复数，因此幅角取负。接下来我们取一个特例m=n=0的时候，令初始化阶段0位置的向量就是向量本身不做任何变化，则对于第一个式子有

m=n=0的模长特例

同样将m=n=0带入第二个式子，则有

m=n=0的幅角特例

可得θ是一个关于位置参数m的函数，且满足关于m的等差数列关系，将求解的R和θ代入改造函数f的三角表示可得f的一般形式

注入位置函数f的一般形式

e的imθ次幂根据欧拉公式展开实际该变换对应着向量的旋转，所以称之为“旋转式位置编码”，改写成矩阵相乘的形式如下

旋转矩阵操作公式

将mθ看作一个参数，将旋转矩阵以函数形式实现，令二维向量坐标为[1, 2]，将其旋转60度的numpy实现如下

>>> import numpy as np

>>> def rotary_matrix(xita):
        matrix = np.array([[np.cos(xita), -np.sin(xita)], [np.sin(xita), np.cos(xita)]])
        return matrix

>>> m = rotary_matrix(np.pi / 3)
>>> one = np.array([[1], [2]])
>>> two = np.matmul(m, one)
>>> print(two)
array([[-1.23205081],
       [ 1.8660254 ]])

以上代码定义个参数xita，若xita等于60度，则代表将原始的二维向量逆时针旋转60度，可以通过两个向量内积除以向量的乘积的模来验证旋转之后两个向量的夹角，首先验证旋转前后向量的模长不变

>>> np.linalg.norm(one)
2.23606797749979
>>> np.linalg.norm(one)
2.23606797749979

旋转之后两个向量的内积除以模乘积等于0.5，因此旋转的夹角为60度

>>> np.dot(one.T, two) / (np.linalg.norm(one) * np.linalg.norm(two))
array([[0.5]])

整个旋转过程可视化如下

旋转矩阵示意图

当向量为二维时，θ下标为0，因此θ的实际结果为1，此时单词位置m控制了旋转的幅度，m越大旋转幅度越大

token	位置	逆时针旋转角度
我	0	0度
爱	1	57度
中	2	114度
国	3	171度
…	…	…

从旋转矩阵的角度，本质上，RoPE是对各个位置的token向量根据自身位置m计算角度做逆时针旋转，在Attention的内积操作中，内积能够感知到旋转之后两个向量之间的夹角，这个夹角就是相对位置信息。
此时二维向量的RoPE得证，由于内积满足线性叠加性，因此任意偶数维的向量都可以表示为二维情形的拼接，因此RoPE的最终公式如下，回到开头介绍RoPE的实现公式

任意偶数维的RoPE公式

如何系统的去学习大模型LLM ？

作为一名热心肠的互联网老兵，我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。

但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的 AI大模型资料 包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

😝有需要的小伙伴，可以V扫描下方二维码免费领取🆓

在这里插入图片描述

一、全套AGI大模型学习路线

AI大模型时代的学习之旅：从基础到前沿，掌握人工智能的核心技能！

二、640套AI大模型报告合集

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。

三、AI大模型经典PDF籍

随着人工智能技术的飞速发展，AI大模型已经成为了当今科技领域的一大热点。这些大型预训练模型，如GPT-3、BERT、XLNet等，以其强大的语言理解和生成能力，正在改变我们对人工智能的认识。那以下这些PDF籍就是非常不错的学习资源。

在这里插入图片描述

四、AI大模型商业化落地方案

阶段1：AI大模型时代的基础理解

目标：了解AI大模型的基本概念、发展历程和核心原理。
内容：
- L1.1 人工智能简述与大模型起源
- L1.2 大模型与通用人工智能
- L1.3 GPT模型的发展历程
- L1.4 模型工程
- L1.4.1 知识大模型
- L1.4.2 生产大模型
- L1.4.3 模型工程方法论
- L1.4.4 模型工程实践
- L1.5 GPT应用案例

阶段2：AI大模型API应用开发工程

目标：掌握AI大模型API的使用和开发，以及相关的编程技能。
内容：
- L2.1 API接口
- L2.1.1 OpenAI API接口
- L2.1.2 Python接口接入
- L2.1.3 BOT工具类框架
- L2.1.4 代码示例
- L2.2 Prompt框架
- L2.2.1 什么是Prompt
- L2.2.2 Prompt框架应用现状
- L2.2.3 基于GPTAS的Prompt框架
- L2.2.4 Prompt框架与Thought
- L2.2.5 Prompt框架与提示词
- L2.3 流水线工程
- L2.3.1 流水线工程的概念
- L2.3.2 流水线工程的优点
- L2.3.3 流水线工程的应用
- L2.4 总结与展望

阶段3：AI大模型应用架构实践

目标：深入理解AI大模型的应用架构，并能够进行私有化部署。
内容：
- L3.1 Agent模型框架
- L3.1.1 Agent模型框架的设计理念
- L3.1.2 Agent模型框架的核心组件
- L3.1.3 Agent模型框架的实现细节
- L3.2 MetaGPT
- L3.2.1 MetaGPT的基本概念
- L3.2.2 MetaGPT的工作原理
- L3.2.3 MetaGPT的应用场景
- L3.3 ChatGLM
- L3.3.1 ChatGLM的特点
- L3.3.2 ChatGLM的开发环境
- L3.3.3 ChatGLM的使用示例
- L3.4 LLAMA
- L3.4.1 LLAMA的特点
- L3.4.2 LLAMA的开发环境
- L3.4.3 LLAMA的使用示例
- L3.5 其他大模型介绍