避开复数推导，我们还可以怎么理解RoPE（重制版）

zenRRan

于 2024-10-12 21:52:05 发布

阅读量17

点赞数

原文链接：https://mp.weixin.qq.com/s?__biz=MzI3ODgwODA2MA==&mid=2247532518&idx=3&sn=11e208ad36a0cfb9bd41b7fe1214eff6&chksm=eaa7f18ce2684868a0172166d4d491bc26ae70f42ff95e49dd7ae4c27e0ecf8943e6f9347131&scene=126&sessionid=0

版权

来自：大猿搬砖简记

写在前面：前段日子公众号发过从圆盘旋转视角理解RoPE的文章，发出后反响比较好。有读者在后台问，是否可以延续圆盘旋转的视角，重新理解NTK-RoPE呢？因为直接从数学推导的角度看实在太抽象了，不能想象到相比于RoPE，它到底是怎样起到（特别是在长文本推理和continue pretrain）上的改进作用的。所以这篇文章的第三部分，我延续了圆盘旋转的逻辑，更加可视化地解读NTK-RoPE，并辅以直观的数学推导。除此以外的部分和上一篇文章保持一致～（都怪公众号不能原文修改）

以下是正文：

大家好，今天的这篇文章，我想避开复数的推导，从一些全新的、更好玩、更可视化的角度，来探究RoPE的原理和各种性质。

这里所说的“可视化”，不仅仅是大家熟悉的“空间向量的旋转”，而是：

具体能让你在调控RoPE的超参时，可以在脑海里快速绘制出一副图，预估你的调参对模型效果的大致影响
或者是当你想探寻衰减性和外推性时，你的脑海里不再仅有代表结果的那一副曲线图，你能动态地绘制出这些重要的性质是怎么一步步产生的。

诸如此类。而当你看完这篇文章，你就能站在几何的角度去理解复数推导的过程了（复数的运算本身就具有几何意义，本文也会给出一定解读）。

【全文目录如下】

一、原始Transformer函数式位置编码
1.1 从旋转的角度理解原理
1.2 这个位置编码为什么得不到人们的青睐

二、RoPE

2.1 在做一件什么事
2.2 旋转角度：二维空间
2.3 旋转角度：高维空间
（1）可以使用一个很小的吗
（2）钟表视角下的高维旋转
2.4 理解衰减性：从傅立叶变换的角度（快乐版）
2.5 理解外推性：基数的选择
（1）可视化位置编码的训练过程
（2）基数的选择

三、圆盘视角下的NTK-RoPE(有公式但快乐版)

一、原始Transformer函数式位置编码

1.1 从旋转的角度解读

transformer位置编码原理我们在这篇文章中详细讲过，这里我们对它进行一个快速的回顾解读。了表达方便，我们先以二维特征空间为例，根据transformer PE的构造方法，我们有：

位置的编码为
位置的编码为

其中是我们设定好的常数。

那么根据：

我们可以把和的关系拆解成：

从这个拆解关系我们可以直观看出，由于是一个表示顺时针旋转的正交矩阵（正交意味着旋转不改变向量的模长，只改变方向），则其实就是以为基础，顺时针旋转这个角度而来。

为了直观体会到这点，我们以为基础，画图看一下不同的是如何旋转而来的：

在直观理解了如何从推导至的基础上，我们来探究的性质：

由于是预先设定好的一个常数，所以当我们假设某个t固定不变，然后慢慢增大时，逐渐变小，这也意味着相距较远的两个位置编码的内积越小，即内积可以用于反馈两个位置向量在绝对位置上的远近

但是，细心的你一定发现了，如果我保持红色不变，顺时针慢慢转动绿色PE时，可能会出现下图的情况，即图中所示的两个绿向量和红向量的内积是一样的，但是左侧绿向量明明距离红向量更远，此时，我们似乎无法从内积大小判断两个位置向量的远近：

那该怎么办呢？我们有一个粗暴但有效的解决方法：让每次位置变动时的转动角度小一些，不就可以了吗？由于我们转动角度为，这意味着只要我们尽量把设置得小一些（这就意味着调大了的周期），让绿线旋转的幅度小一些，使得不管有多少个位置向量，绿线都在第一和第四象限内移动，不就可以了吗？这就是关于的一个最简单的直观解释，在后文中，我们还会结合更细致的内容，继续探寻在更高维的位置向量特种空间中的作用。

1.2 缺陷

从1.1节的介绍中，我们提取到两个关于transformer原始位置编码的重要信息：

在设置合适的值的前提下，每个位置都能取到唯一的位置编码（绝对性）
一个位置编码可以由另一个位置编码旋转而来（相对性），且在设置合适的值的前提下，两个位置编码的内积大小可以反应位置的远近，内积越小，距离越远（衰减性）。

我们固定住某个t，变动，来可视化一下的变动趋势:

如图：

横轴表示
纵轴表示固定某个的情况下，改变后得到的和的内积
d表示不同的hidden_size（例如在1.1节中，我们就假设d = 2）

从图中我们可以发现：

在固定某个的情况下，两个位置编码的内积具有对称性（很简单，对照1.1的圆圈，想象cos函数的对称性。更严谨的推导参见开头引用的文章的3.3节（1）部分）
在固定某个的情况下，两个位置编码的内积具有远程衰减性（在上面实验图对应的原始论文中被称为“距离意识（distance-aware）”），即两个位置编码相距越远，距离越小。

看起来，transformer的原始位置编码应该已经足够好了，可是为什么在很长的一段时间里，人们还是普遍用可学习式的位置编码，甚至还有很多实验证明了transformer的这种位置编码对最终的效果没有起到实质性帮助呢？

这是因为，我们上述的一切分析，都是在原始位置编码上的结果，我们一般把位置编码和输入层的token相加，然后让他们继续去做接下来的计算。可是，transformer架构是复杂的，更详细地说，当token向量进入attention层时，起作用的还是这个部分吗？

为了更详细探究这个问题，我们假设：

分别为两个不同位置的原始token向量，其尺寸为(hidden_size, 1)
分别为两个不同位置的原始PE向量，其尺寸为 (hidden_size, 1)
分别为尺寸为(hidden_size, hidden_size)的Q、K矩阵

那么数据过attention部分可以表示成：

我们只关注其中和两个位置变量都相关的部分，也就是：

从中我们不难发现，经过attention层后，位置编码真正起作用的不再是，而是引入了线性变化后的。那么再引入这种线性变化后，位置编码还能保持上述所说的绝对性、相对性和远距离衰减性这种优良性值吗？我们同样用实验的方式来细看这一点。

由于本质上可以合成一种线性变化，所以我们可以随机初始化一个线性矩阵来代替它，在我们的实验中，我们做了三组试验：

，其中是我们随机初始化的一个线性矩阵
，其中是我们随机初始化的一个线性矩阵

同样，我们固定住某个我们固定住某个t，变动，来可视化一下这三组试验的结果：

上图中的橘线和绿线即表示两位置编码内积间引入线性变化后的结果，可以发现相比于标准的蓝线（），原始位置编码的优良性质（远程衰减性等）都受到了极大程度的破坏。

总结来看，虽然原始transfomer位置编码本身考虑了绝对性、相对性和远程衰减性，但是由于位置编码经过attention层后，最终起作用的形式是而不是，而进一步我们通过实验直观证明了插入一个线性变化会极大破坏位置编码设计之初的各种优良性质，所以早期transformer的这种函数式的位置编码并没有得到大家的青睐。

二、RoPE

在第一部分的分析中，我们已经知道attention层的计算(）会破坏掉输入层位置编码的优良性质，那么我们自然而然会想到：如果我直接在attention层中融入位置信息，也就是我直接把位置编码作用于，这样我不就能维持位置编码优良性质不变吗？（在接下来的讲解中，为了表达方便，我们直接用下标m和n表示两个位置）

2.1 在做一件什么事

这里都是尺寸为(hidden_size, 1)的向量。

我们知道，当我们计算时，我们是在做attention score的计算，其结果表示m位置的token和n位置token间的相关性分数。现在让我们切换一下视角，这个相关性分数其实就是两个特征向量之间的内积，在一定程度上衡量了两个向量之间的相似性。

现在我们希望把位置编码的信息直接引入中，这也就意味着，我们希望根据|n-m|的结果，给这个内积计算一定的惩罚：

当|n-m|较小时，我们希望拉进的距离
当|n-m|较大时，我们希望拉远的距离

你可能觉得有些抽象，不要紧，我们马上仿照1.1中的方式，给出可视化的解释

旋转角度：二维空间

假设原始的特征向量如下：

模仿1.1的方式，如果我们想让具备位置信息，我们可以分别把他们旋转m和n度。在1.1中我们采用顺时针旋转，在这里我们使用逆时针旋转（没有特殊原因，只是为了和原始rope旋转方式贴合），那么我们可以得到如下结果，其中分别表示旋转后的结果：

可以发现，旋转过后，随着|n-m|差值的变大，在保持向量模长不变的情况下，我们拉远了两者之间的距离，也即降低了它们的内积，最后达到降低（惩罚）attention score的效果。

同时，和1.1中一样，我们同样需要引入用于参数，通过拉长三角函数周期的方式确保不同位置的向量不发生碰撞。我们以上图中的为例，当保持其模长不变的情况下，它的逆时针运动轨迹是一个圆，这也意味着一个较小的m和一个较大的m代表的向量可能会重合，所以我们需要使用这种形式，对q每次旋转的幅度加以控制。

总结起来，在二维特征空间下，我们定义了一个逆时针正交旋转矩阵(正如1.1所说，正交旋转矩阵保证了模长不变，即维护q，k原始的特征，只单纯做旋转使其拥有绝对位置信息):

其中表示绝对位置，例如0，1，2...

接着我们将其作用到attention部分的计算上，则有：

端详上面这个等式，我们可以发现：

：两者分别赋予绝对位置信息
：绝对位置信息相乘的结果拥有了相对位置的概念

2.3 旋转角度：高维空间

(1) 可以使用一个很小的吗

到目前为止，我们都在讨论二维特征空间下的位置编码，那么到了高维特征空间，位置编码应该如何设计，比如此时我们有：

其中就是我们所说的hidden_size。

我们先再次把目光聚焦回二维特征空间上，这次我们只看（因为也是同理）

正如2.2节所说，我们通过旋转的方式赋予原始绝对位置的信息，由于是一个逆时针的圆周旋转，这也意味着一个较小的m可能和一个较大的m代表的向量重合，为了避免这一点，我们引入一个较小的调控参数，使得旋转角度变为，放慢旋转的步调。

你可能会想：这个方法真是不错。如果我面对的是一排数量很多的token（长文本），那么我只要把设置得尽可能小，我就能保证不同位置的绝对位置向量一定是唯一的了！

但是，过小的会产生一个问题：由于旋转角度十分微弱，不同位置的向量几乎重合，这就使得“位置”这个信息带来的帮助有限了。

(2) 钟表视角下的高维旋转

当你看见上图中，向量逆时针做圆周运动时，不知道是否让你想起了日常生活中一个很熟悉的物体：钟表。对于一块钟表：

秒针：走得最快
分针：走得中等
时针：走得最慢

而“时刻”这个东西，就是由这三个频率各不相等的圆周运动组成的。

当两块表摆在你面前时，即使走得最快的秒针重合了，但是如果走得较慢的分针和时针不一样，他们照样能代表不同的时刻。

现在，让我们把高维的位置编码想象成一个“时刻”，那么如下图所示：

在这幅图中：

横向表示不同位置的位置编码（m =0和m=1）
纵向可以看成是3个不同频率的逆时针旋转运动（秒针、分针、时针），你可以发现，从m=0变到m=1，秒针旋转的幅度最大，时针最小。

更具体一点来说，假设这三个旋转运动对应的旋转矩阵分别是：

则有。

我们继续做一些延展，对于向量：

两两为一组（），每一组都可以看成是一个旋转速度不一样的指针，我们可以用个指针来表示的位置向量，如下图所示：

因此在高维度空间中，用于表示绝对位置的旋转矩阵最终可以表示成：

由于以上的是一个稀疏矩阵，大部分位置为0，为了节省算力，在代码开发时我们一般采用下列方式：

为了满足我们所说的的特性，同时又不能让太大，我们设:

其中，10000这个数决定了的大小，我们称其为基数（base）。在后面的章节中，我们回来讨论基数应该如何选择。

2.4 理解衰减性：从傅立叶变换角度（快乐版）

（⚠️本节没有令人抓狂的数学公式，所以是快乐版，大家可以放心食用）

傅立叶变换的基本思想，是一个函数可以用无穷多个周期性的线性组合来逼近，我们通过一张经典的示例图来可视化这句话：

在图中，红色方波就是我们想做傅立叶变换的原始函数，它可以被分解成无穷个频率不一致的蓝色三角函数(sin)的线性和。再形象一点解释的话，这些频率不一致的sin函数加总后，在各个方向上的趋势相互影响，最终生成了红色方波有的地方平，有的地方陡峭的模样。（当然，并不是只有方波可以被分解，这里只是给出了一个例子）

现在，让我们回到那个逆时针转动的圆盘上。我们画出两个空的二维坐标系。一个坐标系我们放着这个圆（圆心在原点），记录下它每次转动的角度，和对应的y轴坐标。然后我们在另一个坐标系上绘制出角度和y的关系，你会发现自己得到了一个周期性的三角函数，如下图所示（快快回忆起高中的知识）：

我们只需关注第一行的图即能帮助我们理解上面说的内容。第二～三行的图是在从圆周旋转的角度帮助我们理解方波是如何形成的。图片来自wiki百科。

理解到这一步，我们结合“旋转角度”和“傅立叶变换”角度，再来看一次我们的RoPE，不难理解，一块钟表的旋转，可以被转换成一个周期性的三角函数。

现在再让我们回头看看这个内积计算（假设我们已经把RoPE的信息赋给q和k了），当向量中的各项相加时：

是不是就相当于各个不同的周期的三角函数在做线性组合？
那么线性组合的结果，是不是就可以看成是某种具有周期性的波形？（就像时分秒针组成了一个时刻一样）
当我们固定住某个m，改变n，去计算这个内积。如果|n-m|越大，内积越小，这不就是我们所说的“衰减性”吗？那现在我们再来看下面张图（注意箭头），你是不是能体会到衰减性是怎么来的？

那么再进一步，你可能在实验中会发现，对于的base，它设置得越大，衰减曲线越平缓。现在你对照着上图，大base引起小，也即相当于加大了每个蓝线分量的周期（把它们变得更宽更平缓了），那么红线的趋势自然也就拉平了。

（注意，这里并不是说内积就长成方波的样子，只是结合这个图例，给大家一个感性的理解）

现在，你是不是已经能渐渐认识到了RoPE的魅力，并且能从图像化的角度去思考它的超参意义了呢？

2.5 理解外推性：基数的选择

（1）可视化理解位置编码的训练过程

在2.3节中，我们提过原始RoPE中，的设置为：

你的每一个值就控制着一块圆盘的转动速度，一共有d/2个圆盘，这里我们再把这张图放一次：

结合这张图，我们再来看下旋转矩阵的定义，其中每个[]表示一个圆盘：

我们单独拎出一个圆盘来看：

在二维空间中，这个圆盘表示逆时针方向的旋转，结合我们学过的极坐标的知识，它可以被画成一个单位圆，也就是指针将在这个单位圆上每次做角度为的逆时针旋转（这就是复数的几何意义，理解了这一点，也能从可视化的角度去看复数推导了，不理解的话直接忽略）

所以，当我们训练位置编码时，我们其实就是在训练这d/2个转速不一的单位圆盘：

（2）基数的选择

理论上，当我们把的基数设得很大时，每个圆盘的转速都很慢，这样就可以保证不管有多少个token，它们的绝对位置编码都不会重复，也就是如果仅从编码表示的范围上看，较大的基数已经能够允许我们灵活表示不同长度的token了。

但是，从模型效果上来看呢？例如：假设我想训练一个能处理长文本的模型，我直接用“长文本数据 + 大基数”去做这件事，效果会怎么样呢？

我们可以从（1）中的示意图的角度来可视化地思考这个问题：当我们使用“长文本数据+大基数”时，每个圆盘的旋转速度都变慢了，当我们完成训练后，我们会发现大部分圆盘都没有转完一周（被训练完），也就是我们的位置编码信息训练不到位（没有收敛），这时如果我们想用这个模型做推理，就可能得不到好效果。

同样地，如果我们直接用一个“短文本数据 + 小基数”的方式去训练模型，那么虽然有更多的圆盘旋转过一周了，但是并不是圆盘上的每一个点都能训练到（见下图）。所以如果我们直接用这个模型去做长文本的推理，也得不到好效果

这时，我们可以换一个角度思考：如果我能尽量让每个圆盘都转一周（当然这个过程并不是把圆周上的每一个点都走过），把大致的信息学到。然后再去微调圆盘上点和点之间的缝隙，尽量把这个圆周的细节都填满（在填的过程中你可能会经过之前已经学过的点，有了先验知识，那效果就更好了），那不就能快速学到尽可能完整的位置信息，加速模型收敛了吗？

把这个思想用于实践，就得到了目前训练长文本的一个常用方法：

现用“小基数 + 短数据”做训练（让每个圆盘都尽量转满一圈）
再用“大基数 + 长文本”做微调（弥补圆盘上的空隙）

所以，当你想探究位置编码的外推性，或者想研究基数选择对外推性的影响时，不妨在脑海里可视化这d//2个圆盘，想想如何让圆盘上尽可能多的点位被训练到。

三、NTK-RoPE（有公式但快乐版）

让我们来延续圆盘训练的视角，来可视化理解一下NTK-RoPE的设计原理和运作流程。

假设现在我们已经过一次pretrain，那么对于不同转速的圆盘，其被训练过的圆周长度也是不一样的，对于高频（i更靠近0）的圆盘来说，它被训练过的圆周长度越长。而对于低频(i更靠近d//2-1)的圆盘来说，它被训练过的圆周长度越短，如下图所示：

在这个基础上，现在我们想使用更长的文本做continue-pretrain或者推理，这个时候，从直觉上来说，我们肯定希望圆盘能实现下面的要求：

（1）尽量不要偏移已经训练过的圆周范围。例如，对于图中第一个圆盘，我们就在pretrain走过的绿点的圆周范围内，按照2.5（2）中所说的方式做细节填充（蓝色点），这个操作也就等于尽可能利用已经训练好的位置相关信息。
（2）尽量学到比pretrain更多的位置信息。虽然我们希望尽可能实现（1），但同时圆盘上那些没训练过的圆周位置，我们就不管了吗？如果我们引入了更长的文本，我们当然希望在保守训练的同时，能学到一些新知识。

读到这里你可能已经察觉到，（1）和（2）之间需要有一个balance，那怎么做呢？

我们再回来端详这个圆盘，假设当我们从位置50变动到位置10000时，你希望这d//2个圆盘的指针怎么变动，才能更加【突显】出两个位置之间的差异？我们回想一下，越靠后的圆盘转速越慢，意味着它对【绝对位置】的变动不敏感，所以，如果我们想反映出这个【绝对位置】上的差异，我们最好让靠前的圆盘走得多一些：也就是，靠前的圆盘，我们希望在数据长度变更时，它能尽量学到pretrain以外的位置知识，走更多的圆弧。那么同理，对于靠后的圆盘，它的指针转速小，精度高，我们希望它遵循pretrain中学到的规律，限制在原始的圆弧内，不要做新的尝试。那么如果说大角度转速对应着【绝对位置】，那么小角度转速当然对应着【相对位置】间的精细学习，因此，靠后的圆盘其实是在控制【相对位置】的学习能力。

好，我们来做个总结，对于长文本的continue-pretrain或者inference：