RoFormer_旋转位置嵌入

详细解读RoFormer—带有旋转位置嵌入的增强型transformer

在看复旦的moss的代码的时候，碰到了这个“旋转位置编码”rotary position embedding，正好找这个时间把这个的思想来源，主要方法，以及带来的好处，等，学习一下。

这是追一科技的苏剑林大神的文章，目前引用量有150左右了。

https://arxiv.org/pdf/2104.09864.pdf

RoFormer：https://huggingface.co/docs/transformers/model_doc/roformer 【huggingface上，已集成了】

2021年4月的，两年前的文章了，之前一直错过了。

方法：

rotary position embedding - RoPE，旋转位置嵌入

1. 用一个旋转矩阵rotation matrix来对绝对位置进行编码，于此同时，meanwhile，

2. 在自注意力机制中导入显式的位置依赖。

好处：

1. 自由的序列长度，

2. 随着相对位置的增大，而逐步延缓退化（=衰减）的inter-token dependency，

3. 用相对位置编码来“武装”线性自注意力。

在本文中，我们（=作者们，下同）介绍了一种新方法，即旋转位置嵌入 (rotary, position embedding - RoPE)，以将位置信息利用到 PLMS (预训练模型) 的学习过程中。

具体来说，RoPE 使用旋转矩阵对绝对位置进行编码 ，同时将显式的相对位置依赖性纳入自注意公式中。【核心的两个点，一个是“旋转矩阵”，一个是“显式的相对位置依赖”】

请注意，所提出的 RoPE 具有一系列有价值的属性，优先于现有方法，包括：

1. 序列长度灵活性、

2. 随着相对距离的增加而衰减的tokens间的依赖性，

3. 以及，为线性自注意力，配备（配装的）的相对位置编码的能力。

在多个长文本分类基准数据集上的实验结果表明，与基线替代方案相比，具有“旋转位置嵌入的增强型transformer”，即 RoFormer，可以提供更好的性能，从而证明了所提出的 RoPE 的有效性。

简而言之，我们的贡献有以下三方面：

• 【方法核心】我们研究了现有的相对位置编码方法，发现它们大多是基于将位置编码添加到上下文表示的分解思想构建的。我们引入了一种新方法，即旋转位置嵌入 (RoPE)，以将位置信息利用到 PLMS 的学习过程中。关键思想是通过将“上下文表示”与具有明确理论解释的“旋转矩阵”相乘，来，编码“相对位置”。

• 【硬核背景】我们研究了RoPE 的属性并表明它可以，随着“相对距离”的增加而，衰减，这是自然语言编码所需要的（直观理解：相对距离越远，则彼此之间的关联度越低，依赖度越低，这也符合马尔可夫假设）。我们善意地认为，以前基于相对位置编码的方法与“线性自注意力”不兼容。

• 【实验支撑】我们在各种长文本基准数据集上评估提议的RoFormer。我们的实验表明，与其替代方案相比，它始终能够实现更好的性能。GitHub 上提供了一些使用预训练语言模型的实验：https://github.com/ZhuiyiTechnology/roformer。

注意力机制

一个token: w_i

该token的d-维度的词嵌入向量：x_i

N个token的一个句子：S_N = w_1, w_2, ..., w_N，

三个线性层f_q, f_k, f_v变换；这里的f，可以是“位置编码+线性变换”的“复合函数”（后续会展开）

然后就是，计算q的m号词，和，k的n号词，之间的相似度，以及把这个相似度扔给v作为“权重”：

后续计算attention score和，以attention score来reweight value v的过程

上面是使用softmax函数来计算的。

绝对位置编码 

原始版本的，transformer里面用的，绝对位置编码公式

在位置编码向量中，2t偶数部分，是用sin函数；2t+1的奇数部分，是用cos函数。

例如，一个token，如果用16维度的位置向量表示它的话，是：

[sin(0),

cos(1/10000^(2/16)),

sin(2/10000^(4/16)),

cos(2/10000^(4/16)), ... ]

正弦和余弦，交叉进行。

相对位置编码

这个涉及到的公式，就多了一些，尝试在论文中的公式上，增加一些明确的注释。

2018年的一个，相对位置编码，的版本

上面，q中居然没有位置编码。。。奇怪，

k和v中，分别有一个矩阵，可学习的，是关于相对位置编码的，有|r_max - r_min + 1|个向量，每个向量的维度和x_n一样，都是d维度。即，我们为每个约束范围内的相对位置，分配一个可学习的向量。这也是下标r的由来。

https://aclanthology.org/N18-2074.pdf

然后，就是transformer-xl中的，比较经典的，相对位置编码了：

 

transformer-xl中使用的相对位置编码

这里面，query中的一个位置m，拥有一个相对位置向量p_m；

key中的一个位置n，拥有一个相对位置向量p_n。

更进一步，key ideas:

相对位置编码的改进点，两个：一个是用p_{m-n}，一个是用两个可训练向量u, v（他们已经和具体的位置，无关了，不像p_n, p_m这样的，还需要下标m和n的。。。）。

后面，作者们又提到了三个其他的改进方向：

引入可训练的偏置向量（任意两个位置之间，有一个特定的向量），以及公式9里面的，分别为token之间的相似度和位置之间的相似度，使用不同的映射矩阵

下面是第三个：

每个相对距离，都引入一个可训练的向量。

公式化目标：

目标：构造一个g函数，它有三个输入；和左边的两个f_q以及f_k实现相同的功能！

旋转位置嵌入

2D的情况：

上面的\theta是个预先设置好的常量。如果（W_k * x_n)里面有实数部，虚数部，那么W_k这个矩阵里面应该有复数的设定的。。。TODO

上面这个2d的情况，说实话，看不太懂。

不懂的地方有：

1. W_q, W_k里面都是复数吗？即，计算的细节，还没有搞懂；（一会儿实数，一会儿复数的，最终f和g的输出，都是复数？还是说，复数，就是为了方便叙事，指的是2d，二维空间而已？）

1. 和token之间的关系相关的（x_m和x_n）相似度；

2. 外部的e，是依据m和n的差值的一种方向+角度的“相对位置偏移信息”。

这两者结合起来，也就实现了“加持了相对位置信息的”两个token之间的相似度的一种“度量”。

这里因为有使用Re()函数，就是把中括号里面的复数的实数部分取出来，作为g()函数的返回值即可。

相比之下，为啥f_q和f_k都不取复数的实部呢？？？TODO

论文中有如下一句话：

混合进来相对位置信息的方法，是显而易见的：

翻译一下，就是：“仿射变换后的词嵌入向量”乘以“极坐标向量”，就是糅合了相对位置偏移的“词向量表示”结果了。

一般形式

上面的截屏里面，原始论文中有几个typo。

核心的“角度变换”矩阵，是d/2个2*2的子矩阵拼接起来的，只有对角线上，以及对角线上方和下方，三条斜线上有取值，其他地方都是0。

从对角线上来看：

对角线，上下两条“副”对角线

对角线上的取值，比较简单：

图1

看看图1的细节：

5步糅合位置信息到token的query和key；相当于用极坐标的“角度向量”来做一个位置偏置

RoPE的性质

长距离衰减，即：

inner-product will decay when the relative position increases，【相对距离增大，则内积变小】

也就是说，

a pair of tokens with a long relative distance should have less connection. 【一对儿tokens，如果距离越远，则其联系越少】

公式借鉴的不准确，引用格式有问题

关于这个引用格式，其实是写paper的基本功，强烈建议苏神[@苏剑林]，看看如下的文章：

https://ftp.yz.yamagata-u.ac.jp/pub/CTAN/macros/latex/contrib/natbib/natnotes.pdf

如何引用

natbib包，是neurips template中使用的。

https://media.neurips.cc/Conferences/NeurIPS2021/Styles/neurips_2021.pdf

引用格式，neurips论文模板中的

简单说一下应该使用的公式：（线性注意力的）：

线性注意力机制的核心公式

核心就是先让K和V内部计算一下乘积，然后这个是可以复用到每个query的向量上的。

从而避免了直接计算Q*K^T的这个二次方复杂度的“瓶颈”。

更新Attention的计算

融入了旋转位置编码的新的attention score的计算公式

作者们说，这个分母是特别设计的，就是防止risk of dividing zero。多少有点牵强了，这个+ epsilon，即可防止除法的溢出了。个人建议写论文的时候，还是尽量保证理论上的正确，然后提一句说，实现的时候，可以折衷，效果也不差。。。【如果是特别长的篇幅，chatgpt的8k的那种context，这个是否也这么搞？】

【目前莫测复旦的Moss里面，是直接softmax的，即分母还是引入了位置偏移的。。。】

理论解释

2D下推演RoPE

  

用径向和角度，来刻画f()和g()两个函数

上面的公式(25)是初始条件。

 m=n=0的时候的情况，角度函数只和位置相关，和位置上的具体的token无关

n=m+1的时候，n是来自k的位置（序号），m是来自q的位置

接着往下看吧：

这里又引入了一个γ，思路上。。。多少有点迷失了。。。

往下，就是矩阵形式的表示了：

这是扩展到一般情况了。

多少感觉哪里卡住了。。。引入复数之后，目前还没有看到之前提到的Re[]（复数取实部）的应用。。。

简易计算

【需要注意的是，这里的**-x_{d-1}，其实应该是-x_d**；

而x_d，应该是**x_{d-1}**：顺序是：-x2, x1, -x4, x3, ..., 所以，最后两个元素，应该是先-x_d，后是x_{d-1}！】

三条“对角线”（主对角线+两个副对角线）下的，简易计算

直观理解就是：

基于三条对角线的简易计算方法。

这个在复旦大学Moss中的代码使用，可以参考：

[代码学习]复旦大学MOSS的推理算法代码-part 5-模型前向forward

https://zhuanlan.zhihu.com/p/629110594

“旋转位置编码”部分。

RoPE的长距离衰减

下面这个性质比较重要，目的是让距离更远的两个token的基于循环位置信息的词嵌入，下，计算出来的相似度，衰减，更小。

长距离下的，相似度的衰减的证明

继续往下看：

有些地方还是不理解。。。需要结合代码和例子来看懂了。。。

看看图，

随着距离的变远，相对上界（两个token的词嵌入向量的内积相似度）也在变小。。。

往后的实验部分就先不看了。

局限性

这是来自作者们自己的分析。

尽管我们（=作者们）提供了理论基础和有希望的实验证明，但我们的方法受到以下事实的限制：

• 尽管我们在数学上将相对位置关系格式化为 2D 子空间下的旋转，但缺乏详尽的解释：“为什么它比包含其他位置编码策略的基线模型收敛得更快”（对啊，为什么呢。。。）。

• 虽然我们已经证明我们的模型对 intern-token products(内积)，具有长期衰减的有利特性，第（3.3）节，这类似于现有的位置编码机制，但（这里，感觉不应该是转折，应该是进一步啊。。。）我们的模型在长文本上显示出，优于同行模型的性能，我们 没有提出一个忠实的解释。我们提出的 RoFormer 建立在基于 Transformer 的基础设施之上，这需要硬件资源来进行预训练。  whaosoft aiot http://143ai.com  

【即，相对的背后的更深层次的理论分析，还不够，还可以再深挖。。。】

下次逐行学习代码。下次一定。。。

待续。