有哪些免费自学设计软件的网站_企业网站优化兴田德润怎么样_网络营销师怎么考_谷歌广告联盟怎么做

探秘Transformer系列之（17）— RoPE（下）

因为本文字数超限，顾分为上下两篇发出，上篇参见。

探秘Transformer系列之（17）— RoPE（上）

0x03 性质

本节来学习RoPE的一些主要特性以及业界思考。

3.1 相关性

旋转编码 RoPE 有如下特点：

• 计算$q{k}^T$点积时，保留了词语的相对位置信息（不会因词语的绝对位置发生改变），这样可以有效地保持位置信息的相对关系。
• 相邻位置的编码之间有一定的相似性，即便在旋转后，相邻的位置仍然会有相似的嵌入。而距离较远的编码之间有一定的差异性。这样可以增强模型对位置信息的感知和利用。
• 语义相似的Token平均来说获得更多的注意力。即，当 k和q相近时，不管它们的相对距离n-m多大，其注意力 $q^T{R}_{n-m}k$平均都应该更大，至少要比随机的两个token更大。

3.2 周期性

因为旋转一圈的弧度是$2\pi$ ，所以RoPE中的向量旋转就像时钟一样，每组分量的旋转都具有周期性。因为每组分量的旋转弧度都随着位置索引的增加而线性增加，所以越靠后的分组，它的正弦函数的周期越大、频率越低，它的旋转速度越慢。整体频率可以对应到低频，以及高频上。

所以我们接下来就有一个问题：随着位置的增大，旋转角度是否会重复？具体解答如下。

• 在任意第k个子空间，只要${\theta }_k$公式中不含有$\pi$，那么就不会出现周期性重复。
• 如果每个子空间都不会出现周期性重复，整体更不会重复。

3.3 $\beta$进制

苏神认为RoPE是β进制编码，原文如下。

位置n的旋转位置编码（RoPE），本质上就是数字n的β进制编码！

对于一个10进制的数字n，如果希望得到其的β进制表示的（从右往左数）第m位数字，方法是

$$\lfloor \frac{n}{{\beta }^{m-1}}\rfloor mod\beta$$

也就是先除以${\beta }^{k-1}$次方，然后求模（余数）。而RoPE可以改写为
$$[cos(\frac{n}{{\beta }^0}),sin(\frac{n}{{\beta }^0}),cos(\frac{n}{{\beta }^1}),sin(\frac{n}{{\beta }^1}),\cdots ,cos(\frac{n}{{\beta }^{d/2-1}}),sin(\frac{n}{{\beta }^{d/2-1}}]$$

其中，$\beta =10000^{2/d}$。模运算的最重要特性是周期性，cos,sin刚好也是周期函数。所以，除掉取整函数这个无关紧要的差异外，RoPE（或者说Sinusoidal位置编码）其实就是数字n的β进制编码！

3.4 对称性

对照三角函数编码性质，对于RoPE编码，位置m的token A对于位置n的Token B的注意力影响，和位置为2n-m上的token C对于Token B的注意力影响一样。尤其当位置m的Token与位置2n-m的Token相同时，有如下表达式

$$g(x_m,x_n,m-n)=g(x_n,x_{2n-m},n-(2n-m))$$

这证明RoPE编码也是符合对称性，没有学习到方向的差异。

3.5 频域

$\theta$的大小决定了对应维度的单调性，也赋予了这些维度上的参数不同的学习倾向。$\theta$ 就对应到了傅里叶变换中的频率这一概念。$\theta$ 较大时，注意力计算结果仅在相对距离 t−s 较小时保持一致的单调性，之后陷入波动，本质上就是高频信号； $\theta$ 较小时，注意力计算结果能在相对距离 t−s 较大时仍然能保持一致的单调性，波动较为平缓，本质上就是低频信号。

论文“SCALING LAWS OF ROPE-BASE"指出，如果用$q_t{k}_s$表示 s 位置的token对 t 位置token的语义相似度，$q_t{k}_s$是一个二维时域信号，有 t,s 两个时域维度。语义相似度 $q_t{k}_s$ 就是由不同频域维度上的语义相似度分量$q_t^{(n)}{k}_s^{(n)}$组合而成的，每个维度对应一个频段 ${\theta }_n$ ，高频分量对应局部语义影响，低频分量对应长上下文语义影响。从频域到时域，最朴素的转换方式就是傅里叶逆变换，通过$e^{i(s-t){\theta }_n}$将不同频段的分量组合。由于是为了获取 s 位置对 t 位置的位置信息，所以变换对象是 $q_t^{(1)}{k}_s^{(1)}...q_t^{(d)}{k}_s^{(d)}$，变换的目标维度是原始二维时域的对角线方向，即 s−t 方向。

论文“Round and Round We Go! What makes Rotary Positional Encodings useful?”也揭示了RoPE不同频率成分在模型学习中的作用：高频用于位置注意力，低频用于语义注意力。

我们可以计算出每个维度的ROPE对应的波长（Wavelength）是${\lambda }_d=\frac{2\pi }{{\theta }_d}=2\pi b^{\frac{2d}{|D|}}$，，其中 |D| 是维度的总数，b是base。波长描述了嵌入在该维度上完成一次完整旋转（2π）所需的标记数量。波长与RoPE嵌入的频率有关，且在不同维度上可能有所不同。

当给定一个长度L，有的维度会出现周期比L更长，可以假设，当出现这个情况的时候，所有的位置都能获得一个唯一的编码，也就是绝对位置都保留了下来。反之，周期比较短的维度只能保留相对位置信息。

3.6 高频低频

RoPE中，向量旋转就像时钟一样，因为旋转一圈的弧度是$2\pi$，所以每组分量的旋转都具有周期性。RoPE以角度${\theta }_i$对每个二维向量（维度对$(q_i,q_{i+1})$）分别进行旋转，旋转角的取值与三角式位置编码相同，即采样频率 $\theta$ 乘上token下标（$m{\theta }_i=m\times bas{e}^{-2i/d}$），旋转完将所有切分拼接，就得到了含有位置信息的特征向量。这里 ${\theta }_i=10000^{-2i/d}$ ，沿用了 Transformer 最早的 Sinusoidal 位置编码的方案。它可以带来一定的远程衰减性。每个位置旋转的角度不一样。

$$\left(\begin{array}{ccccccc}cos(m{\theta }_1)& -sin(m{\theta }_1)& 0& 0& ...& 0& 0\\ sin(m{\theta }_1)& cos(m{\theta }_1)& 0& 0& ...& 0& 0\\ 0& 0& cos(m{\theta }_2)& -sin(m{\theta }_2)& ...& 0& 0\\ 0& 0& sin(m{\theta }_2)& cos(m{\theta }_2)& ...& 0& 0\\ 0& 0& 0& 0& ...& 0& 0\\ .& .& .& .& .& .& .\\ .& .& .& .& .& .& .\\ .& .& .& .& .& .& .\\ 0& 0& 0& 0& ...& cos(m{\theta }_{d/2})& -sin(m{\theta }_{d/2})\\ 0& 0& 0& 0& ...& sin(m{\theta }_{d/2})& cos(m{\theta }_{d/2})\end{array}\right)\left(\begin{array}{c}{q}_m^{(1)}\\ q_m^{(2)}\\ q_m^{(3)}\\ q_m^{(4)}\\ .\\ .\\ .\\ q_m^{(d-1)}\\ q_m^{(d)}\end{array}\right)$$
在周期函数中，如$sin(\omega x)$ ，$\omega$ 越大，频率越大。在RoPE中，$\omega$随维度变量 k 增加,$b^{-2k/d}$减小，从而频率降低。
$$sin(\frac{p}{b^{\frac{2k}{d}}}),cos(\frac{p}{b^{\frac{2k}{d}}})$$
我们可得：位置编码的低维对应高频，高维对应低频。对于每组分量，它的旋转弧度随着位置索引的增加而线性增加。越靠后的分组，它的旋转速度越慢，正弦函数的周期越大、频率越低。

• 高频：是RoPE的位置向量，i 比较小（前面的维度）， 𝜃𝑖 较大的时候，周期短，频率高。

• 低频: 是RoPE的位置向量，i 比较小（后面的维度），𝜃𝑖 较小的时候，周期长，频率低。

NTK-RoPE、YaRN的作者Bowen Peng认为：高频学习到的是局部的相对距离，低频学习到的是远程的绝对距离。高频低频两者都很重要，它们之间更像是一种层次的关系；用进制类别来看，低频对应的就是高位，如果只保留低位而去掉高位，那么结果就相当于求模（余数），无法准确表达出位置信息来。

3.7 远程衰减

远程衰减基于一个很朴素的假设：相对距离越远，则彼此之间的关联度越低，依赖度越低。如果位置编码具有远程衰减特性，则可以让位置相近的Token平均来说获得更多的注意力。

表现

RoPE也呈现出远程衰减的性质，具体表现为：对于两个词向量，若它们之间的距离越近，则它们的内积分数越高，反之则越低。即，位置 m 的 Query 向量 $q_m$ 与位置 n 的 Key 向量 $k_n$ 相对距离越远（ |n−m| 越大）， $(R_m{q}_m{)}^T(R_n{k}_n))$ 越小。从下图可以看到，随着相对距离的变大，内积结果有衰减趋势。

从图上也可以看出，在衰减曲线后期会产生很大波动，产生了U形状的注意力模式。对比图如下。

论文”HoPE: A Novel Positional Encoding Without Long-Term Decay for Enhanced Context Awareness and Extrapolation"对此进行了细致的分析，发现在RoPE中，U形状的注意力模式是由特定学习到的组件（learned components）造成的，这些组件也是限制RoPE表达能力和外推能力的关键因素。具体参见下图。

• (a) 表示将RoPE分解为组件（Comps）进行分析（见图上红圈方程式）。上部子图显示了每个组件对整体注意力逻辑的贡献。我们用红色突出显示了一些具有突出模式（patterns）的组件，即“激活”组件，用蓝色突出显示了低频组件。下部子图展示了整体注意力逻辑，以及“激活”组件的组合效应。
• (b) 给出了训练期间RoPE不同组件的方差（VAF）。
• © 揭示了外推中的OOD现象是由“激活”组件引起的。两个上部的子图显示了第一层的注意模式，下部的子图则显示了后续层的异常模式。

基于这些发现，该论文提出了一种新的位置编码方法——High-frequency rotary Position Encoding（HoPE）。HoPE通过去除RoPE中的位置依赖组件，保留高频信号，从而理论上也打破了长期衰减的原则。

能否设计不振荡的位置编码？很难，位置编码函数如果不振荡，那么往往缺乏足够的容量去编码足够多的位置信息，也就是某种意义上来说，位置编码函数的复杂性本身也是编码位置的要求。

论证

我们接下来对远程衰减进行论证。

首先，我们用论文中的推导来看，具体参见下图。

其次，有研究人员认为下面公式为RoPE的主要功能项。Transformer位置编码（意义） 河畔草lxr
$$\begin{gathered} C_{RoPE}(t-s)=\frac{1}{d/2}\sum _{n=1}^{d/2}cos(s-t){\theta }_n \\ {\theta }_n=10000^{-2n/d} \end{gathered}$$
$C_{RoPE}(t-s)$ 大致随相对距离 t−s 呈现单调减的关系。但是整体偏置的单调减并不意味着每个维度偏置的单调减： ${\theta }_n$ 的大小决定了维度 2n−1,2n 的单调性，也赋予了这些维度上的参数不同的学习倾向：

• n 较小时， ${\theta }_n$ 较大，趋向于1，仅在相对距离 t−s 较小时保持一致的单调性，之后陷入波动，诱导对应维度刻画较近的位置信息；
• n 较大时， ${\theta }_n$ 较小，趋向于0，能相对距离 t−s 较大时仍然保持一致的单调性，诱导对应维度刻画较远的位置信息。

反过来，不同相对位置的语义信息也会反映在不同的特征维度上：

• 在相对距离 t−s 较小时，所有维度的偏置都接近于1，对应自注意力分布更加关注相邻位置的信息；
• 在相对距离 t−s 较大时，多数维度有正有负相互抵消，只有部分维度的偏置较大，如果两个token对应维度的语义特征高度重合则会予以部分强调，否对应自注意力分布趋近于0。这也正是相对偏置的一大优势，即对相对距离较远的语义关联，没有给予绝对的惩罚，而是给予相对的过滤：虽然通过整体偏置抑制较远距离的信息，但是仍然允许某些特征维度上的语义汇集到自注意力计算中。

基数

对于${\theta }_n=10000^{-2n/d}$，10000这个数决定了 𝜃 的大小，我们称其为基数（base）。base的不同取值会影响注意力远程衰减的程度。因为“随距离衰减”是外推的关键，所以base的性质与大模型的长度外推息息相关，如NTK-Aware Scaled RoPE、NTK-by-parts、Dynamic NTK等长度外推方法，本质上都是通过改变base，从而影响每个位置对应的旋转角度，进而影响模型的位置编码信息，最终达到长度外推的目的。

由于 RoPE 中的 attention 值除了 q,k 本身外，仅和$ R_{n-m} $因子相关，下面考察$ R_{n-m} $ 因子的特点
$$\begin{array}{rlrl}⟨{\mathbf{p}}_{\mathrm{m}},{\mathbf{p}}_{\mathrm{n}}⟩& =\mathrm{Re}\left[{\mathrm{e}}^{\mathrm{i}(\mathrm{m}-\mathrm{n}){\theta }_0}+{\mathrm{e}}^{\mathrm{i}(\mathrm{m}-\mathrm{n}){\theta }_1}+\cdots +{\mathrm{e}}^{\mathrm{i}(\mathrm{m}-\mathrm{n})\theta \ast \mathrm{d}/2-1}\right]& =\frac{\mathrm{d}}{2}\cdot \mathrm{Re}\left[\sum \ast {\mathrm{i}=0}^{\mathrm{d}/2-1}{\mathrm{e}}^{\mathrm{i}(\mathrm{m}-\mathrm{n})10000^{-\mathrm{i}/\mathrm{d}/2)}}\frac{1}{\mathrm{d}/2}\right]& \approx \frac{\mathrm{d}}{2}\cdot \mathrm{Re}\left[{\int }_0^1{\mathrm{e}}^{\mathrm{i}(\mathrm{m}-\mathrm{n})\cdot 10000^{-\mathrm{t}}\mathrm{dt}}\right]\end{array}$$

那么问题就变成了积分 $\int_0^1 \mathrm{e}^{\mathrm{i}(\mathrm{m}-\mathrm{n}) \cdot 10000^{-\mathrm{t}} }\mathrm{dt} $的渐进估计问题，通过一下函数计算积分值与位置距离的关系就可以分析出不同 base 值的影响。

• base=1，完全失去远程衰减特性。
• base 越小，衰减得越快且幅度也更大。太小的base会破坏注意力远程衰减的性质，例如base=10或100时，注意力分数不再随着相对位置的增大呈现出震荡下降的趋势。
• base 越大，衰减得越慢且幅度也越小。这也是为什么训练更长的窗口，要把base改大的原因。所以现在业界的主流做法都是窗口变长后，base也要跟着变大做适配。苹果就在其模型中用了很大的基数。输入序列越长，base就需要越大，让未充分训练过的窗口强行衰减变慢，本身也是降低崩的概率的一种方式。

平滑性

另外，embedding维度和衰减曲线的平滑程度成正相关，维度越高，衰减曲线越平滑。外推性的基本前提是函数的“光滑性”。外推性就是局部推断整体，它依赖的就是给定函数的高阶光滑性（高阶导数存在且有界）。但是三角函数编码或RoPE不具备这种性质。它们是一系列正余弦函数的组合，这算是关于位置编码k的高频振荡函数，而不是直线或者渐近趋于直线之类的函数，所以基于它的模型往往外推行为难以预估。

3.8 外推

尽管RoPE可以理论上可以编码任意长度的绝对位置信息，并且通过旋转矩阵（三角计算）来生成超过预训练长度的位置编码，并且RoPE也具有远程衰减特性（“随距离衰减”是外推的关键）。RoPE仍然存在外推问题（length extrapolation problem），即对于基于RoPE的大语言模型，测试长度超过训练长度之后，模型的效果会有显著的崩坏，具体表现为语言建模困惑度急剧攀升。远程衰减属性导致在更长文的外推中，RoPE编码的作用影响也在衰减，效果在逐步变差。

我们将在后面专门写一篇来做具体分析。

0x04 实现

4.1 基础Torch知识

torch.outer

torch.outer(a, b) 计算两个 1D 向量 a 和 b 的外积，生成一个二维矩阵，其中每个元素的计算方式为：result[i,j]=𝑎[i]×𝑏[j]。即，result矩阵的第 i 行、第 j 列的元素等于向量 a 的第 i 个元素与向量 b 的第 j 个元素的乘积。

外积（outer product）是指两个向量 a 和 b 通过外积操作生成的矩阵：𝐴=𝑎⊗𝑏。其中 𝑎⊗𝑏 生成一个矩阵，行数等于向量 𝑎 的元素数，列数等于向量 𝑏的元素数。

torch.matmul

当输入张量的维度大于 2 时，torch.matmul将执行批量矩阵乘法。

torch.polar

torch.polar()函数会构造一个复数张量，用法是torch.polar(abs, angle, *, out=None) → Tensor。其元素是极坐标对应的笛卡尔坐标，绝对值为 abs，角度为 angle。 out=abs⋅cos(angle)+abs⋅sin(angle)⋅j。

torch.repeat_interleave

torch.repeat_interleave()函数会返回一个具有与输入相同维度的重复张量。

torch.view_as_complex

把一个tensor转为复数形式，要求这个tensor的最后一个维度形状为2。

torch.view_as_real

把复数tensor变回实数，可以看做是是刚才操作的逆变换。

4.2 在Transformer中的位置

不同于原始 Transformer 的绝对位置编码，RoPE位于多头注意力机制的内部，直接作用于每个头完成变换的query和key，而且每个头使用相同的RoPE（RoPE的输入参数只有位置和维度，跟头无关），这也意味着在 transformer中的每一层都要加入RoPE。

4.3 llama3

lama中对RoPE的实现采用复数的公式来计算$f_q(x_m,m)=(W_q{x}_m)e^{im\theta }$。该方式速度较快，但不方便后续修改。

具体而言，是把每个向量（Key或者Query）两维度一组切分，分成元素对$(q^1,q^2),(q^3,q^4),...$，每对都解释为二维向量。然后RoPE以角度${\theta }_i$对每个二维向量（维度对$(q_i,q_{i+1})$）分别进行旋转，旋转角的取值与三角式位置编码相同，即采样频率 $\theta$ 乘上token下标（$m{\theta }_i=m\times bas{e}^{-2i/d}$），旋转完将所有切分拼接，就得到了含有位置信息的特征向量。

$$\left(\begin{array}{ccccccc}cos(m{\theta }_1)& -sin(m{\theta }_1)& 0& 0& ...& 0& 0\\ sin(m{\theta }_1)& cos(m{\theta }_1)& 0& 0& ...& 0& 0\\ 0& 0& cos(m{\theta }_2)& -sin(m{\theta }_2)& ...& 0& 0\\ 0& 0& sin(m{\theta }_2)& cos(m{\theta }_2)& ...& 0& 0\\ 0& 0& 0& 0& ...& 0& 0\\ .& .& .& .& .& .& .\\ .& .& .& .& .& .& .\\ .& .& .& .& .& .& .\\ 0& 0& 0& 0& ...& cos(m{\theta }_{d/2})& -sin(m{\theta }_{d/2})\\ 0& 0& 0& 0& ...& sin(m{\theta }_{d/2})& cos(m{\theta }_{d/2})\end{array}\right)\left(\begin{array}{c}{q}_m^{(1)}\\ q_m^{(2)}\\ q_m^{(3)}\\ q_m^{(4)}\\ .\\ .\\ .\\ q_m^{(d-1)}\\ q_m^{(d)}\end{array}\right)$$
这里 ${\theta }_i=10000^{-2i/d}$ ，沿用了 Transformer 最早的 Sinusoidal 位置编码的方案。它可以带来一定的远程衰减性。每个位置旋转的角度不一样。

总体

其总体代码和公式对应如下图所示。

在实现 RoPE 算法之前，需要注意：为了方便代码实现，在进行旋转之前，需要将旋转矩阵转换为极坐标形式，嵌入向量（q、k）需要转换为复数形式。完成旋转后，旋转后的嵌入需要转换回实数形式，以便进行注意力计算。

准备旋转矩阵

precompute_freqs_cis()函数会生成旋转矩阵，即 给定维度预计算频率θ。θ 完全由 Q、K、V 的向量长度 d 决定。位置 m 对应我们的 query 长度，实际代码中由 max_position_embeddings 参数决定，可以理解为模型支持的最长 query 的长度，因此 max 有了，m 的范围也就有了。结合上面的信息，针对一个固定了最长 query 长度 m 和向量维度 d 的 LLM，我们可以提前将其对应的旋转变换矩阵构造完成。

freqs = torch.outer(t, freqs)的矩阵如下。
$$freqs=\left[\begin{array}{cccccc}1{\theta }_1& 1{\theta }_2& 1{\theta }_3& ...& 1{\theta }_{d/2-1}& 1{\theta }_{d/2}\\ 2{\theta }_1& 2{\theta }_2& 2{\theta }_3& ...& 2{\theta }_{d/2-1}& 2{\theta }_{d/2}\\ ⋮& ⋮& ⋮& \ddots & ⋮& ⋮\\ ⋮& ⋮& ⋮& ⋮& ⋮& ⋮\\ m{\theta }_1& m{\theta }_2& m{\theta }_3& ...& m{\theta }_{d/2-1}& m{\theta }_{d/2}\end{array}\right]$$
结合这个 Rd 的变换矩阵，分别执行 cos 和 sin，便可以得到我们计算所需的全位置全维度的变换矩阵。

torch.polar之后的 freqs 如下。
$$freqs=\left[\begin{array}{ccccc}cos({\theta }_1)+i\cdot sin({\theta }_1)& cos({\theta }_2)+i\cdot sin({\theta }_2)& ...& cos({\theta }_{d/2})+i\cdot sin({\theta }_{d/2})& \\ cos(2{\theta }_1)+i\cdot sin(2{\theta }_1)& cos(2{\theta }_2)+i\cdot sin(2{\theta }_2)& ...& cos(2{\theta }_{d/2})+i\cdot sin(2{\theta }_{d/2})& \\ ⋮& ⋮& \ddots & ⋮& \\ cos(m{\theta }_1)+i\cdot sin(m{\theta }_1)& cos(m{\theta }_2)+i\cdot sin(m{\theta }_2)& ...& cos(m{\theta }_{d/2})+i\cdot sin(m{\theta }_{d/2})& \end{array}\right]$$

具体代码如下。

# 生成旋转矩阵
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):# 根据维度 d 生成旋转角度θ向量。计算词向量元素两两分组之后，每组元素对应的旋转角度 θ_i，由于是将向量两两旋转应用 RoPE，所以共有 dim/2 个 θ。θ 完全由 Q、K、V 的向量长度 dim 决定# freqs 长度是 dim/2，一半的维度。2表示是偶数这里 θ 完全由 Q、K、V 的向量长度 d 决定，即 dim维度，取0，2，4...等维度freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) # 生成 token 序列索引 t = [0, 1,..., seq_len-1]，即拿到所有位置对应的ID，就是论文中常说的m或者nt = torch.arange(end, device=freqs.device, dtype=torch.float32)# 计算m * θ。将旋转角度和 `token` 位置索引相乘，即求向量的外积，结果是一个矩阵，该矩阵包含了每个位置和每个维度对应的旋转角度，即每个元素代表位置t在第i维上的旋转角度（频率）freqs = torch.outer(t, freqs) # freqs的形状是 [seq_len, dim // 2]，具体参见上面公式。# 将上一步的结果写成复数的形式𝑒^{𝑖𝑚𝜃},模是1，幅角是freqs。freqs_cis的大小为(seqlen, dim//2)# 假设 freqs = [x, y]，则 freqs_cis = [cos(x) + sin(x)i, cos(y) + sin(y)i]    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64return freqs_cis

precompute_freqs_cis()函数用如下方式进行调用。

class Transformer(nn.Module):def __init__(self, params: ModelArgs):super().__init__()self.params = paramsself.vocab_size = params.vocab_sizeself.n_layers = params.n_layersself.tok_embeddings = VocabParallelEmbedding(params.vocab_size, params.dim, init_method=lambda x: x)self.layers = torch.nn.ModuleList()for layer_id in range(params.n_layers):self.layers.append(TransformerBlock(layer_id, params))self.norm = RMSNorm(params.dim, eps=params.norm_eps)self.output = ColumnParallelLinear(params.dim, params.vocab_size, bias=False, init_method=lambda x: x)# 预先计算出来选择矩阵，乘以2是为了动态扩展self.freqs_cis = precompute_freqs_cis(params.dim // params.n_heads,params.max_seq_len * 2,params.rope_theta,)

实现

apply_rotary_emb()方法用于将 cos、sin 的旋转矩阵应用到原始的 query 和 key 向量上，这样在 Attention 内积时，就会为 query 和 key 引入位置信息。

# 为了匹配q和k，需要对角度进行扩展  
# freqs_cis维度是[seq len, dim/2]
def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):ndim = x.ndimassert 0 <= 1 < ndim# 需要确保形状和x的形状匹配，即是(x.shape[1]=seq len, x.shape[-1]=dim/2)assert freqs_cis.shape == (x.shape[1], x.shape[-1])# x的第二维和最后一维保留，其他维度置为1shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]return freqs_cis.view(*shape) # [1,S,1,head_dim//2] def apply_rotary_emb(xq: torch.Tensor,xk: torch.Tensor,freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:"""作用: 将q,k向量分别与旋转向量相乘,得到旋转后的q,k向量q/k_rotated输入: x_q(torch.Tensor): 实际上是权重 W_q * 词嵌入向量值, 来自上一个线性层的输出, 形状为 [batch_size, seq_len, n_heads, head_dim]或者[batch_size, seq_len, dim]x_k(torch.Tensor): 实际上是权重 W_k * 词嵌入向量值, 来自上一个线性层的输出, 形状为 [batch_size, seq_len, n_heads, head_dim]或者[batch_size, seq_len, dim]freqs_cis (torch.Tensor): 频率复数张量, 形状为 [max_seq_len, head_dim]输出: 施加了旋转编码后的q和k"""    # 实数域张量转为复数域张量。将一个大小为n的向量xq_两两组合形成复数来计算，需要增加维度，把最后一维变成2，即把最后一维的两个实数作为一个复数的实部和虚部来构建一个复数。 # 计算过程q:[batch_size,atten_heads,seq_len,atten_dim]->q_complex:[b,a_h,s,a_d//2,2]->[b,a_h,s,a_d//2]->[b,a_h,s,a_d//2,2]# [:-1]意思是从第一维到倒数第二维；*是为了展开列表；-1, 2表示把最后一维展开成两维：x/2和2，即最后一维是2； # xq_.shape = [batch_size,atten_heads,seq_len,atten_dim//2,2]，如果不考虑多头，则是[batch_size, seq_len, dim // 2, 2]xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) # 复数形式张量xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) # 复数形式张量# freqs_cis 的形状必须与 xq 和 xk 相匹配，因此我们需要将 freqs_cis 的形状从 [max_seq_len, head_dim] 调整为 [1, max_seq_len, 1, head_dim]。即，旋转矩阵（freqs_cis）的维度在序列长度（seq_len，维度 1）和头部维度（head_dim，维度 3）上需要与嵌入的维度一致。  freqs_cis = reshape_for_broadcast(freqs_cis, xq_)# 通过复数乘法实现向量旋转操作，然后将结果转回实数域。这是幅度不变，角度变换的操作，即把结果恢复成原来的样子，将第三维之后压平，也就是(atten_dim//2,2)->(atten_dim)。位置编码只和向量的序列位置还有向量本身有关，和batch以及注意力头无关，所以只用关注第二维和第四维# xq_out.shape = [batch_size, seq_len, dim]xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)return xq_out.type_as(xq), xk_out.type_as(xk) # 又是实数了

调用

Transformer会调用Transformer层进行RoPE操作。

class Transformer(nn.Module):@torch.inference_mode()def forward(self, tokens: torch.Tensor, start_pos: int):_bsz, seqlen = tokens.shapeh = self.tok_embeddings(tokens)self.freqs_cis = self.freqs_cis.to(h.device)freqs_cis = self.freqs_cis[start_pos : start_pos + seqlen]mask = Noneif seqlen > 1:mask = torch.full((seqlen, seqlen), float("-inf"), device=tokens.device)mask = torch.triu(mask, diagonal=1)# When performing key-value caching, we compute the attention scores# only for the new sequence. Thus, the matrix of scores is of size# (seqlen, cache_len + seqlen), and the only masked entries are (i, j) for# j > cache_len + i, since row i corresponds to token cache_len + i.mask = torch.hstack([torch.zeros((seqlen, start_pos), device=tokens.device), mask]).type_as(h)for layer in self.layers:h = layer(h, start_pos, freqs_cis, mask)h = self.norm(h)output = self.output(h).float()return output

TransformerBlock会直接调用到Attention的forward函数。

class TransformerBlock(nn.Module):def __init__(self, layer_id: int, args: ModelArgs):super().__init__()self.n_heads = args.n_headsself.dim = args.dimself.head_dim = args.dim // args.n_headsself.attention = Attention(args)self.feed_forward = FeedForward(dim=args.dim,hidden_dim=4 * args.dim,multiple_of=args.multiple_of,ffn_dim_multiplier=args.ffn_dim_multiplier,)self.layer_id = layer_idself.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)def forward(self,x: torch.Tensor,start_pos: int,freqs_cis: torch.Tensor,mask: Optional[torch.Tensor],):h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)out = h + self.feed_forward(self.ffn_norm(h))return out

Attention会做如下操作。

def forward(self,x: torch.Tensor,start_pos: int,freqs_cis: torch.Tensor,mask: Optional[torch.Tensor],
):bsz, seqlen, _ = x.shapexq, xk, xv = self.wq(x), self.wk(x), self.wv(x)xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)# attention 操作之前，应用旋转位置编码xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)self.cache_k = self.cache_k.to(xq)self.cache_v = self.cache_v.to(xq)self.cache_k[:bsz, start_pos : start_pos + seqlen] = xkself.cache_v[:bsz, start_pos : start_pos + seqlen] = xvkeys = self.cache_k[:bsz, : start_pos + seqlen]values = self.cache_v[:bsz, : start_pos + seqlen]# repeat k/v heads if n_kv_heads < n_headskeys = repeat_kv(keys, self.n_rep)  # (bs, cache_len + seqlen, n_local_heads, head_dim)values = repeat_kv(values, self.n_rep)  # (bs, cache_len + seqlen, n_local_heads, head_dim)# Q/K/V 对应维度为 [bsz, seq_len, num_heads, head_dim]，transpose 将 seq_len 和 num_heads 的维度调换了，得到的 states 维度为 [bsz, num_heads, seq_len, head_dim]。这个变换是为了将 seq_len x head_dim = 4096 x 8 挪到一起，方便后面的 ⊗ 对位相乘。xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)keys = keys.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)values = values.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)if mask is not None:scores = scores + mask  # (bs, n_local_heads, seqlen, cache_len + seqlen)scores = F.softmax(scores.float(), dim=-1).type_as(xq)output = torch.matmul(scores, values)  # (bs, n_local_heads, seqlen, head_dim)output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)return self.wo(output)

4.4 rotate_half

rotate_half是RoPE中经常使用的方法，我们专门来分析下。rotate_half() 的作用是将输入张量x的一半隐藏维度进行旋转，即进行语义向量复数化，实现向量乘以虚数i，等价于向量逆时针旋转90度。
$$f_q(x_m,m)=(W_1{x}_m)(cos(m\theta )+isin(m\theta ))$$

上述式子继续推导，合并cos和sin就可以发现，$q_t$旋转后的结果就是$q_t$乘上cos，再加上$q_t$翻转维度并取反一维后乘上sin的结果，因此程序里实现叫rotate_half。
$$\left(\begin{array}{c}{q}_m^{(1)}\\ q_m^{(2)}\end{array}\right)=cos(m\theta )\left(\begin{array}{c}{q}_m^{(1)}\\ q_m^{(2)}\end{array}\right)+sin(m\theta )\left(\begin{array}{c}-q_m^{(2)}\\ q_m^{(1)}\end{array}\right)$$
rotate_half其实有两种实现方式。我们首先看看其中一种。具体来说，它将输入张量的后半部分（划为虚部）取负，然后与前半部分（划为实部）拼接，从而实现旋转操作。其流程如下：

1. ‌分割张量‌：假设输入张量x的形状为[batch_size, num_attention_heads, seq_len, head_size]，函数首先将张量x分割成两部分：x1和x2。x1包含前半部分，x2包含后半部分。
2. ‌旋转操作‌：将x2取负，然后将x2与x1拼接在一起。这样，原始张量的后半部分被旋转到了前半部分的位置，实现了旋转效果。
3. ‌拼接‌：最后，将取负后的x2与x1在最后一个维度上拼接，形成最终的旋转位置嵌入张量。

具体代码对应如下。

# 后半部分和前半部分进行了交换，并且将后半部分的符号取反。
# 这个函数很好理解，就是将原始向量从中间劈开分为 A、B 两份，然后拼接为 [-B, A] 的状态：比如 [q0,q1,q2,q3,q4,q5,q6,q7] -> [-q4,-q5,-q6,-q7,q0,q1,q2,q3]
def rotate_half(x):"""Rotates half the hidden dims of the input."""# 前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] # 后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 :] # 后64embedding位置取负号，和前64embedding位置拼接return torch.cat((-x2, x1), dim=-1)def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):cos = cos.unsqueeze(unsqueeze_dim)sin = sin.unsqueeze(unsqueeze_dim)q_embed = (q * cos) + (rotate_half(q) * sin)k_embed = (k * cos) + (rotate_half(k) * sin)return q_embed, k_embed

将 rotate_half() 代入到 apply_rotary_pos_emb()，以 q=[x1,x2] 为例：

q_embed = [x1, x2] * cos + [-x2, x1] * sin = [x1 * cos - x2 * sin, x2 * cos + x1 * sin]

具体参见下图。这里的负号，对应和角公式中的负号。计算旋角 $m\theta$ 的过程此处省略。

然而，上面的代码是HuggingFace的Transformer库的实现，和RoPE论文公式有些许差异，具体为元素位置排列上的差异，在论文中q0的结果是q0和q1这一对元素经过三角函数变换而成的，但是在实际公式中q0是由q0和$q_{d/2+1}$这一对形成的。

• HuggingFace：$[-q_4,-q_5,-q_6,-q_7,q_0,q_1,q_2,q_3]$

• 论文：$[-q_1,q_0,-q_2,q_3,....q_{n-1},q_{n-2}]$

具体近似如下。

其实，这涉及到两个对特征维度进行切分的不同的实现。

按照RoPE论文的思路，就是GPT-J style。实现过程中对特征向量的奇偶维度进行rotate_half操作，相邻两维度一组（ ⊙ 表示对应位相乘，对 $k_s$ 的操作相同）。

由于对奇偶维度旋转需要将维度两两交错，实现较为复杂，后来的研究人员提出，直接将特征维度一切二，这种实现方式称为GPT-NeoX style，实现过程中对特征向量的前后各半进行rotate_half操作。GPT-J style 和 GPT-NeoX style 是等价的，可以互相转化的：GPT-J style中的奇数维度对应GPT-NeoX style的前一半维度，GPT-J style中的偶数维度对应GPT-NeoX style的后一半维度。将GPT-J style的奇数维度抽出来整体拼接在偶数维度之前，就会得到GPT-NeoX style的结果。

两种实现方式只是对应的R矩阵不同，最终都可以实现绝对位置实现相对位置编码的目的。对最终的结果没有影响。因为RoPE对原始向量的改造本质上是以一对元素为单位经过旋转矩阵运算，将所有对的结果进行拼接的过程，而到底是选择连续的元素作为一对，还是其他的挑选方式都是可以的，只要是embedding维度为偶数，且挑选的策略为不重复的一对，最终Attention的内积结果都能感知到相对位置信息，因为Attention满足内积线性叠加性，至于谁和谁一组进行叠加并不重要。

GPT-J sytle

是和原始论文和博客一样，使用的相邻两个为一组。
$$\left[\begin{array}{c}{q}_m^{(1)}\\ q_m^{(2)}\\ q_m^{(3)}\\ q_m^{(4)}\\ ⋮\\ q_m^{(d-1)}\\ q_m^{(d)}\end{array}\right]=\left[\begin{array}{c}cos(m{\theta }_1)\\ cos(m{\theta }_1)\\ cos(m{\theta }_2)\\ cos(m{\theta }_2)\\ ⋮\\ cos(m{\theta }_{d/2})\\ cos(m{\theta }_{d/2})\end{array}\right]\otimes \left[\begin{array}{c}{q}_m^{(1)}\\ q_m^{(2)}\\ q_m^{(3)}\\ q_m^{(4)}\\ ⋮\\ q_m^{(d-1)}\\ q_m^{(d)}\end{array}\right]+\left[\begin{array}{c}sin(m{\theta }_1)\\ sin(m{\theta }_1)\\ sin(m{\theta }_2)\\ sin(m{\theta }_2)\\ ⋮\\ sin(m{\theta }_{d/2})\\ sin(m{\theta }_{d/2})\end{array}\right]\otimes \left[\begin{array}{c}-q_m^{(2)}\\ q_m^{(1)}\\ -q_m^{(4)}\\ q_m^{(3)}\\ ⋮\\ -q_m^{(d)}\\ q_m^{(d-1)}\end{array}\right]$$

GPT-NeoX style

不是相邻两个元素为一组，而是 𝑞0 和 $q_{d/2-1}$ 为一组。

在FlashAttention的源码中就实现了GPT-J sytle 和 GPT-NeoX style的RoPE。

https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/layers/rotary.py

def rotate_half(x, interleaved=False):if not interleaved:x1, x2 = x.chunk(2, dim=-1)return torch.cat((-x2, x1), dim=-1)else:x1, x2 = x[..., ::2], x[..., 1::2]return rearrange(torch.stack((-x2, x1), dim=-1), '... d two -> ... (d two)', two=2)def apply_rotary_emb_torch(x, cos, sin, interleaved=False):"""x: (batch_size, seqlen, nheads, headdim)cos, sin: (seqlen, rotary_dim / 2)"""ro_dim = cos.shape[-1] * 2assert ro_dim <= x.shape[-1]cos = repeat(cos, 's d -> s 1 (2 d)')sin = repeat(sin, 's d -> s 1 (2 d)')return torch.cat([x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin, x[..., ro_dim:]], dim=-1)

0xEE 个人信息

★★★★★★关于生活和技术的思考★★★★★★

微信公众账号：罗西的思考

如果您想及时得到个人撰写文章的消息推送，或者想看看个人推荐的技术资料，敬请关注。

0xFF 参考

Base of RoPE Bounds Context Length Xin Men etc.

LLM时代Transformer中的Positional Encoding MrYXJ

LLM（廿三）：LLM 中的长文本问题 紫气东来

LLaMA中的旋转位置编码（RopE）实现解读 qwdjiq

Long LLM第二篇——why RoPE? 王焱

ROUND AND ROUND WE GO! WHAT MAKES ROTARY POSITIONAL ENCODINGS USEFUL?

RedHerring RedHerring

RoPE外推的缩放法则 —— 尝试外推RoPE至1M上下文 河畔草lxr

RoPE旋转位置编码深度解析：理论推导、代码实现、长度外推 JMXGODLZ

Transformer位置编码（基础） 河畔草lxr

Transformer位置编码（改进） 河畔草lxr

Transformer升级之路：10、RoPE是一种β进制编码

Transformer升级之路：15、Key归一化助力长度外推 苏剑林

Transformer升级之路：16、“复盘”长度外推技术

Transformer升级之路：2、博采众长的旋转式位置编码

Transformer改进之相对位置编码(RPE) Taylor Wu

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

llama源代码逐行分析 bookname

qwen源码解读3-解读QWenAttention模型的调用 programmer

transformers 库提供的 llama rope 实现

中文语言模型研究：(1) 乘性位置编码 PENG Bo

位置编码算法背景知识 Zhang

千问Qwen2 beta/1.5模型代码逐行分析(一) bookname

浅谈LLM的长度外推 uuuuu

深入剖析大模型原理 — Qwen Blog

羡鱼智能：【OpenLLM 009】大模型基础组件之位置编码-万字长文全面解读LLM中的位置编码与长度外推性（上）

羡鱼智能：【OpenLLM 010】大模型基础组件之位置编码-万字长文全面解读LLM中的位置编码与长度外推性（ 中）

让研究人员绞尽脑汁的Transformer位置编码 - 科学空间|Scientific Spaces

LLM：旋转位置编码（RoPE）的通俗理解 莲子

Effective Long-Context Scaling of Foundation Models

HoPE: A Novel Positional Encoding Without Long-Term Decay for Enhanced Context Awareness and Extrapolation