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

前言

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


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

位置编码知识准备

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

相对位置编码对两个单词之间的相对位置进行建模,并且将相对位置信息加入到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在做什么,该公式是将一个原始的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变化的计算示意图如下

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

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


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

在之前介绍的sin-cos位置编码中Transformer系列:快速通俗理解Transformer的位置编码,我们知道sin-cos位置编码因为三角函数的性质,使得它可以表达相对位置信息,具体而言是:给定距离,任意位置的位置编码都可以表达为一个已知位置的位置编码的关于距离的线性组合,而RoPE的位置编码也是同样的思路,采用绝对位置编码实现相对距离的表达,区别如下

  • 实现相对位置能力的途径不同: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的计算方式,下面结合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和θ相乘的结果

θ只生成了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]的二维矩阵,模拟如下

紧接着作者分别用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的代码逻辑它实现的计算公式实际为

该公式和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能够自动感知到相对位置,而不需要对Attention的结构进行改造。由于Attention是对两个位置的token的向量进行运算,因此只需要在Attention之前,对Query和Key向量进行绝对位置编码改造即可,跟Value没有关系,我们设一个m位置的token1,它的Query为q,n位置的token2,它的Key为k,改造函数为f,则经过注入位置信息改造之后的新向量为

这样改造的目的是使得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的函数,且满足关于m的等差数列关系,将求解的R和θ代入改造函数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位置逆时针旋转角度
00度
157度
2114度
3171度

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

全文完毕。

学习策略调整建议

鉴于当前市场的积极态势,对于初学者而言,学习LLM不应仅仅停留在理论层面,更应注重实践能力和创新思维的培养。以下是一些针对当前行情的学习策略建议。

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

学习大型人工智能模型,如GPT-3、BERT或任何其他先进的神经网络模型,需要系统的方法和持续的努力。既然要系统的学习大模型,那么学习路线是必不可少的,下面的这份路线能帮助你快速梳理知识,形成自己的体系。

L1级别:AI大模型时代的华丽登场

L2级别:AI大模型API应用开发工程

L3级别:大模型应用架构进阶实践

L4级别:大模型微调与私有化部署

一般掌握到第四个级别,市场上大多数岗位都是可以胜任,但要还不是天花板,天花板级别要求更加严格,对于算法和实战是非常苛刻的。建议普通人掌握到L4级别即可。

以上的AI大模型学习路线,不知道为什么发出来就有点糊,高清版可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费

在这里插入图片描述

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

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

img

三、大模型经典PDF籍

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

img

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

img

作为普通人,入局大模型时代需要持续学习和实践,不断提高自己的技能和认知水平,同时也需要有责任感和伦理意识,为人工智能的健康发展贡献力量。

  • 20
    点赞
  • 27
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值