LLMs组件系列：Attention排列不变性与位置编码（超详细）

💡 Transformer注意力可并行计算的特性既是其相较于RNN的优势，也使其丢失了时间步语义信息，所以位置编码 (positional encodings) 是Transformer这种并行结构必须要有的，否则机器就会已读乱回。但位置编码原理是什么，工程上怎么实现呢？网上有许多讲解位置编码的方案，但很多都为了文章篇幅而舍弃掉一些前置信息，如果没有基础看起来还是比较吃力的。为此，本文用篇幅换取详细度，详细介绍了当前四种主流的位置编码方案。具体行文上，先在第一部分介绍了Attention机制的排列不变性，也即为什么需要位置编码。之后将当前位置编码方案分为绝对位置编码和相对位置编码分别详述，并附工程代码。

1 Attention排列不变性与语义序列的冲突

1.1 排列不变性

上图是一个非First Iteration的Attention点积机制。其中query_states是input_embed在经过q_proj投影矩阵后获得的Query隐状态向量（只有一个当前token，position_ids为3），key_states和value_states则是当前token在k_proj和v_proj投影后与kv_cache中之前token隐状态cat的结果（position_ids为0-3）。attn_weights则是QK注意力点积获得的注意力评分，在First Iteration中是一个方阵（对角线下方有效），在这里则直接取方阵最后一行，图中：

• 上部分为QK点积，结果为attn_weights评分向量（两个position_ids组合为编号，例如query_states编号为3，与key_states中第一列编号为0的向量点积后获得attn_weights向量中第一个编号为30的值）；
• 下部分为注意力权重与值向量的点积，同理，attn_outputs向量中第一个元素A为attn_weights向量与value_states矩阵第一列点积的结果，表示上写为attn_outputs每个编号和value_states第一列中每个编号（0a…3a）拼接相加的结果。

可见Attention机制的输出总体上可以看作是对query_states进行线性变换的结果（从长度为6降维到4），可视作对value_states进行带权垂直压缩的结果。具体来说，attn_outputs每个维度（ABCD）都是注意力评分矩阵作为权重（30, 31, 32, 33）与value矩阵对应行（0,1,2,3）进行加权求和的结果，这个对应关系是固定的，就算将token序列打乱，在下部分注意力权重与值向量的点积的过程中，还是各自token的权重与各自token的value向量相乘后再所有token的value值向量相加的结果，值不变。如下图所示，将token序列0，1，2，3逆转为3，2，1，0再执行一样的过程，attn_outputs结果的值不会改变。

上图中的结果说明，在Attention结构中，就算token的序列打乱，query_states的变换结果attn_outputs也不会变，这种情况在First Iteration中一致，只不过当query_states有多行时，attn_outputs也在对应的行上与其一致。这就是Attention的排列不变性，本质是token进行全组合相乘然后求和，而全组合本质是一个无序的集合。

1.2 语义的序列要求与位置编码概念

上述结果说明了Attention的本质是对token进行组合求积并求和，与序列无关。这很好啊，既然是组合的集合，那么我就可以同时对所有组合进行计算了，这也是Attention可并行计算的根本原因。之前的LSTM需要按照时间步将每个token串行放入进行计算，前一个token的计算结果既可以当作当前时间步的输出，也得传递给下个时间步作为输入。相比之下简直就是串行到并行的飞跃！

但人家LSTM这样串行当然是有其道理的，那就是语义的时序信息。在Attention中，既然token顺序和结果无关，那岂不是说我是狗的主人、狗是我的主人、主人是我的狗、我xxxxx这几句话等价了？简直是倒反天罡！快归快，不能乱来对不对。那怎么办呢？有没有哪种方法既保留Transformer快速并行矩阵计算的特点，又能捕获词序信息？

那就是位置编码Positional Encoding，Transformer原论文中对这一点有比较清楚的叙述：

Since our model contains no recurrence and no convolution, in order for the model to make use of the order of the sequence, we must inject some information about the relative or absolute position of the tokens in the sequence. To this end, we add “positional encodings” to the input embeddings at the bottoms of the encoder and decoder stacks.

即，位置编码，或者说位置性的编码，就是对输入序列中的token表示进行位置信息（绝对或相对）的注入。位置编码经历了范式上的编码，逐渐发展成如今适用于Transformer的编码，具体可参见Transformer学习笔记一：Positional Encoding（位置编码） - 猛猿的文章 - 知乎。简单来说，位置编码经过以下几个阶段

• 用整型值标记位置：较为通用且直观，但无界且量纲上与浮点隐藏状态不匹配，不适用于Attention。
• 用[0,1]范围浮点数表示：机械切分，无法外推
• 二进制编码：用n位二进制的隐藏状态就可以表示$2^n$个词，但维度上与Embedding不匹配会有空值，且离散
• 用周期函数如三角函数sin：范围固定[-1,1]浮点数，连续，可外推，三角函数和差化积、积化和差的特性可以编码相对位置信息。如下，但这种表示是非线性的，有待改进。

$$sin(x+\Delta x)=\sin (x)\cos (\Delta x)+\cos (x)\sin (\Delta x)=[\begin{array}{cc}\sin x& \cos x\end{array}][\begin{array}{c}\cos \Delta x\\ \sin \Delta x\end{array}]$$

• 三角函数坐标系（一对sin和cos）：具备以上所有优点。为了表示线性变换，可以引入三角函数坐标系，将输入的x理解为角度$\theta$。那么直角坐标系中的一个位置就可以写作一对正余弦三角函数$[\begin{array}{c}\sin x\\ \cos x\end{array}]$。此时其经过一个线性变换可以表示任意相对位置的旋转（即$x+\Delta x$视作$x$的旋转），$[\begin{array}{c}\sin x+\Delta x\\ \cos x+\Delta x\end{array}]$，如下：
  $$[\begin{array}{c}\sin (x+\Delta x)\\ \cos (x+\Delta x)\end{array}]=[\begin{array}{cc}\cos \Delta x& \sin \Delta x\\ -\sin \Delta x& \cos \Delta x\end{array}][\begin{array}{c}\sin x\\ \cos x\end{array}]$$

Transformer原文紧接着还有两句：

The positional encodings have the same dimension dmodel as the embeddings, so that the two can be summed. There are many choices of positional encodings, learned and fixed [8].

也就是说，位置编码一般和Embedding的shape一样，注入的方式就是加到Embedding上。此外，在当时已经有了一些位置编码方案，分为习得的learned的和固定的Fixed。当然了这种分类方式在当时条件下是成立的，但从我们现在的视角来看，位置编码首先应该按照是否是绝对的位置编码进行分类：

• 绝对位置编码：一个位置一个定死的编码，即只要position ids一样，那么token的embedding的所有分量上加的编码信息值就一样；实现上是直接加到embedding上。
• 相对位置编码：根据token的相对位置进行编码，实现上是分别成到query和key向量上，并通过QK点积实现相对性。

2 绝对位置编码

在Embedding层加入位置编码信息。

2.1 Learned Positional Embedding(学习式编码)：习得Learned

在出现时间上早于Transformer的Sinusoidal编码，也在其原始论文中被讨论过，之后也比Sinusoidal应用更广。其基本原理是，将位置变量作为一个$shape=[最大可编码长度,embedding维度]$的可学习的参数矩阵，与模型一起训练。早期的BERT, Roberta, GPT, GPT2使用的就是这种，最早使用则可溯源至2017年Facebook的一篇论文Convolutional sequence to sequence learning。Transformer原始论文中也提到了这个方法并用于对比实验，实验结果是Transformer的位置编码方案（Sinusoidal）与这种可学习的方案的结果差不多，但这种方案无法外推extrapolate至超过最大可编码长度的序列（目前也有研究表明可以通过其他方式让学习式编码具备外推的能力）。

2.2 Sinusoidal(正弦曲线编码)：固定Fixed

Transformer论文中给出的基于三角函数余计算、固定的位置编码。即不是通过学习习得一个$shape=[最大可编码长度,embedding维度]$的参数矩阵，而是通过一个函数计算任意长度的序列的位置编码与embedding后的结果相加，也就是具备可外推性质。但实际上这个方法并没有很优越，只有Transformer, Transformer XL在使用。

编码函数$PE$如下。$PE$函数接收两个参数，$pos$是token在序列中的位置编号，$2i(+1)$是token嵌入向量每个元素的编号，最大为嵌入维度$d_{model}-1$。这个函数对输入序列中每个token、每个token的嵌入空间中每个嵌入值都赋予一个特定、通用、绝对的位置编码。

$$\begin{gathered} PE(pos,2i)=sin(pos/10000^{\frac{2i}{d_{model}}}) \\ PE(pos,2i+1)=cos(pos/10000^{\frac{2i}{d_{model}}}) \end{gathered}$$

式子有点抽象，看一下原文的解释：

where pos is the position and i is the dimension. That is, each dimension of the positional encoding corresponds to a sinusoid. The wavelengths form a geometric progression from 2π to 10000 · 2π. We chose this function because we hypothesized it would allow the model to easily learn to attend by relative positions, since for any fixed offset k, $P{E}_{pos+k}$ can be represented as a linear function of $P{E}_{pos}$.

重点有二：

• 对于Dimension，在横向的Embedding分量维数值等差增长时（例如对一个128维的Embedding维度，从第0维到127维），关于pos的三角函数波长呈现一个从2 $\pi$到 10000 · 2 $\pi$的几何级数；
• 对于Position，其能允许模型简单地基于相对位置施加关注，即$P{E}_{pos+k}$可以由$P{E}_{pos}$线性表示。

（1）首先，让我们来看波长呈现几何级数是什么意思

如上图，对一个token长度为5、嵌入维度为128（$d_{model}$）的输入序列进行位置编码，形式上要生成一个shape=[5, 128]的位置编码矩阵用以与Embedding矩阵相加。128个维数分量的正弦函数的输入$pos/10000^{\frac{2i}{d_{model}}}$其实可以转化为标准形式观察波长函数：

$$\frac{2\pi \cdot pos}{2\pi \cdot 10000^{\frac{2i}{d_{model}}}}=\frac{2\pi x}{\lambda }$$

其中$\lambda =2\pi \cdot 10000^{\frac{2i}{d_{model}}}$为波长，$i$取值从0到$\frac{d_{model}}{2}$，128个维度相邻两个为一组（共享一个波长），那么总共有64组（上图中彩色部分及其下方对应的波长）。那么这个位置编码正弦函数的波长是关于dimension的一个等比数列，即呈几何级数（The wavelengths form a geometric progression from 2π to 10000 · 2π）。当$d_{model}=128$时，波长数列公比为$10000^{\frac{1}{64}}\approx 1.15$，如下图。下图画出了embedding维度为128的Sinusoidal位置编码在前7个dimension的值情况，红色为偶数维度，蓝色为奇数维度。可以观察到奇偶数的波长一样但存在一个位移，奇数和偶数维度的波长则递增，这就使得不同维度的编码频率是有规律且唯一的。

此外，上图展现的一个重要规律是，在Dimension层面，维数分量从0-127不断变大时，每个分量上对于所有Position Ids处理的三角函数的波长越大，这会导致数值变化非常小，如下。颜色表示分量上的数值（[-1,1]范围），越小越蓝，越大越红，横轴为128个维度，竖轴为50个token。可以看到从左往右（0到127维分量）颜色越单一，即数值上变化越小。

（2）其次，基于相对位置施加关注是什么意思？

让我们先在形式上稍稍改变$PE$函数，将其参数列表中的维度分量变量i去掉并合为一个式子，即$PE(pos)$。其可以自动对pos位置的token的所有embedding分量进行两两一组的分块矩阵的编码，此时将原式中$pos/10000^{\frac{2i}{d_{model}}}$简化为$pos\cdot {\lambda }_i$，其中$i$和原式中保持一致两两共享，如下。

$$PE(pos)=\left(\begin{array}{cccc}\left[\begin{array}{cc}sin(pos\cdot {\lambda }_0),& cos(pos\cdot {\lambda }_0)\end{array}\right],& \left[\begin{array}{cc}sin(pos\cdot {\lambda }_1),& cos(pos\cdot {\lambda }_1)\end{array}\right],& ......,& \left[\begin{array}{cc}sin(pos\cdot {\lambda }_{d_{model}/2-1}),& cos(pos\cdot {\lambda }_{d_{model}/2-1})\end{array}\right]\end{array}\right)$$

此时

$$PE(pos+\Delta k)=\left(\begin{array}{cccc}\left[\begin{array}{cc}sin((pos+\Delta k)\cdot {\lambda }_0),& cos((pos+\Delta k)\cdot {\lambda }_0)\end{array}\right],& \left[\begin{array}{cc}sin((pos+\Delta k)\cdot {\lambda }_1),& cos((pos+\Delta k)\cdot {\lambda }_1)\end{array}\right],& ......,& \left[\begin{array}{cc}sin((pos+\Delta k)\cdot {\lambda }_{d_{model}/2-1}),& cos((pos+\Delta k)\cdot {\lambda }_{d_{model}/2-1})\end{array}\right]\end{array}\right)$$

正如前一部分提到的基于三角函数坐标系组合使用sin和cos的位置编码可以实现用线性变换表示相对位置，如下式：

$$[\begin{array}{c}\sin (x+\Delta x)\\ \cos (x+\Delta x)\end{array}]=[\begin{array}{cc}\cos \Delta x& \sin \Delta x\\ -\sin \Delta x& \cos \Delta x\end{array}][\begin{array}{c}\sin x\\ \cos x\end{array}]$$

现在非常清晰了，可以将下式的线性变换规则套用到上面两个式子。套用结果如下（分块矩阵的转置是内外都转置）：

$$\begin{gathered} PE(pos+\Delta k{)}^T=T_{\Delta k}\ast PE(pos{)}^T= \\ \left(\begin{array}{ccc}\left[\begin{array}{cc}\cos (\Delta k\cdot {\lambda }_0)& \sin (\Delta k\cdot {\lambda }_0)\\ -\sin (\Delta k\cdot {\lambda }_0)& \cos (\Delta k\cdot {\lambda }_0)\end{array}\right]& ....& 0\\ ...& ...& ...\\ 0& ...& \left[\begin{array}{cc}\cos (\Delta k\cdot {\lambda }_{d_{model}/2-1})& \sin (\Delta k\cdot {\lambda }_{d_{model}/2-1})\\ -\sin (\Delta k\cdot {\lambda }_{d_{model}/2-1})& \cos (\Delta k\cdot {\lambda }_{d_{model}/2-1})\end{array}\right]\end{array}\right)\left(\begin{array}{c}\left[\begin{array}{c}sin(pos\cdot {\lambda }_0)\\ cos(pos\cdot {\lambda }_0)\end{array}\right]\\ ......\\ \left[\begin{array}{c}sin(pos\cdot {\lambda }_{d_{model}/2-1})\\ cos(pos\cdot {\lambda }_{d_{model}/2-1})\end{array}\right]\end{array}\right) \end{gathered}$$

除此之外，Sinusoidal还有一个相对位置编码的特点，即任何两个Pos位置的编码结果向量的点积结果仅仅取决于两个pos的差值$\Delta k$，用到的还是三角函数和差化积、积化和差的定理，如下：

$$\begin{gathered} PE(pos+\Delta k)\ast PE(pos{)}^T=\sum _{i=0}^{\frac{d_{model}}{2}-1}sin((pos+\Delta k)\cdot {\lambda }_i)sin(pos\cdot {\lambda }_i)+cos((pos+\Delta k)\cdot {\lambda }_i)cos(pos\cdot {\lambda }_i)= \\ \sum _{i=0}^{\frac{d_{model}}{2}-1}\cos ((pos+\Delta k-pos)\cdot {\lambda }_i)=\sum _{i=0}^{\frac{d_{model}}{2}-1}\cos (\Delta k\cdot {\lambda }_i) \end{gathered}$$
请注意，这个特性在Sinusoidal中没有被充分利用，但在苏神的RoPE中被进一步使用了（见下）。至此，Sinusoidal的细节已拆解完毕，最后总结以下，Sinusoidal有以下功能：

• 基本要求：有界连续函数唯一性编码：使用三角函数进行编码，编码值在[-1,1]之间且连续。使用token序列编号pos和嵌入维度分量序号i输入编码函数$PE$进行唯一性编码。
• 嵌入维度层面（i）：相邻两个i一组，嵌入维度的分量数值i（0-127）数值越大，三角函数的波长越长，且随i的变大呈几何级数增长，导致越大的i会使得pos的数值能造成的编码数值变化越小。
• 相对位置信息编码（pos）：进阶使用三角函数，借鉴三角函数坐标轴形式，在嵌入维度层面让相邻两个i作为一组波长相同的sin与cos。此时所有pos的token的相同emnbedding分量i（一列）使用一个三角函数，在token位置层面$PE(pos+\Delta k)$可以用$PE(pos)$经过线性变换得到（左乘变换矩阵），两者的点积也仅由$\Delta k$这个Pos的相对位置差决定。

3 相对位置编码

在Attention前加入位置编码信息。Llama/GLM/Qwen/Baichuan7B用Rotary，BLOOM/ BloombergGPT/Falcon/Baichuan13B使用ALiBi

3.1 RoPE旋转位置编码

苏剑林在RoFormer中提出，位置编码是通过旋转向量来实现的。这种旋转依赖于对应位置的正弦和余弦值。
与Sinusoidal非常像，但不是直接加到embedding上，而是逐位乘到QK向量上，并通过点积实现相对位置编码，具体的原理可以惨遭Sinusoidal中演示的最后一个特性，即两个位置编码的点积结果由相对位置决定。为了让这个结果更清楚一点，下面我们深入展开。

首先，让我回顾一下Sinusoidal中使用的那个$T_{\Delta k}$线性变换矩阵。其作用是对输入的hidden_states进行两两一组的编码，在Sinusoidal中为了便于理解，我们将其两两一组的特性显式地表示为分块矩阵的形式，下面我们将展示其原本的样子：

• 设QK的hidden_size，即head_dim设为$d_{head}$；设token序列中的位置序号为$pos$。
• 复用Sinusoidal中$PE$编码函数中三角函数内pos的系数，由于这个系数的变量是嵌入维度的分量i，所以我们将其单独命名为$\Theta$并将其表示为一个长度为$\frac{d_{head}}{2}$向量形式（两两一组），即
  $$\Theta =\left\{\begin{array}{c}{\theta }_i=10000^{-\frac{2(i-1)}{d_{head}}},i=\in [1,2,...,\frac{d_{head}}{2}]\end{array}\right\}$$
• 将Sinusoidal中使用的关于位置pos和维度分量i的矩阵$T_{\Delta k}$线性变换矩阵的转置展开（将显式的分块矩阵[]去掉）为一个$shape=[d_{head},d_{head}]$的方阵，并重新命名为RoPE旋转位置编码矩阵$R_{\Theta ,pos}$，表示这个矩阵的要素由关于维度分量i的函数$\Theta ({\theta }_i)$和位置pos构成，如下。
  $$R_{\Theta ,pos}^{d_{head}}=\left(\begin{array}{cccccccc}\cos pos{\theta }_1& -\sin pos{\theta }_1& 0& 0& ...& 0& 0& \\ \sin pos{\theta }_1& \cos pos{\theta }_1& 0& 0& ...& 0& 0& \\ 0& 0& \cos pos{\theta }_2& -\sin pos{\theta }_2& ...& 0& 0& \\ 0& 0& \sin pos{\theta }_2& \cos pos{\theta }_2& ...& 0& 0& \\ ......& ......& ......& ......& ...& ......& ......& \\ 0& 0& 0& 0& ...& \cos pos{\theta }_{d_{model}/2}& -\sin pos{\theta }_{d_{model}/2}& \\ 0& 0& 0& 0& ...& \sin pos{\theta }_{d_{model}/2}& \cos pos{\theta }_{d_{model}/2}& \end{array}\right)$$

如上，此时还是可以套用Sinusoidal中的相对位置线性变换特性 $PE(pos+\Delta k)=PE(pos)R_{\Theta ,\Delta k}^{d_{head}}$，但我们的目的却是将其用于Sinusoidal中最后一条关于点积取决于相对位置的特性。因为RoPE与Sinusoidal最大的区别在于，前者将位置编码作用于Attention机制中将input_embed分别经由q_proj和k_proj投影矩阵变化后获得的query_states和key_states，而后者直接作用于input_embed。这一点不同带来的结果就是，Sinusoidal中演示的最后一个特性大放异彩。为了直观体现，设经过进入Attention机制的经过embedding的输入$input\_embed$维矩阵为$x$，其形为$n$ * $d_{model}$，其中n为token序列长度，$x_m$为其中第m个token的嵌入向量；设QK的投影权重矩阵q_proj和k_proj为$W_q,W_k$，形状均为$d_{model}$*$d_{head}$，也就是将嵌入维度映射到$head\_dim$维度。则RoPE的针对q或k的特定序列位置m的编码函数如下：
$$RoP{E}_{q,k}(x_m,m)=x_m{W}_{q,k}{R}_{\Theta ,m}^{d_{head}}$$
之后在Attention中就是QK点积，为了更清楚看到query_states和key_states向量中不同位置pos的向量点积的结果，我们令query_states中的pos(m值)为a，key_states中的pos(m值)为b，则a和b两个位置的向量点积结果为
$$RoP{E}_q(x_a,a)RoP{E}_k(x_b,b{)}^T=x_a{W}_q{R}_{\Theta ,a}^{d_{head}}(x_b{W}_k{R}_{\Theta ,b}^{d_{head}}{)}^T=x_a{W}_q{R}_{\Theta ,a}^{d_{head}}{R}_{\Theta ,b}^{d_{head}T}{W}_k^T{x}_b^T$$

由于$R_{\Theta ,a}^{d_{head}}$为一个对角分块矩阵，所以上式中$R_{\Theta ,a}^{d_{head}}{R}_{\Theta ,b}^{d_{head}T}$实际上是对角线上的所有shape=[2,2]的分块矩阵进行点积。例如，其中第一组分块矩阵的乘积结果如下：
$$\begin{gathered} \left[\begin{array}{cc}\cos a{\theta }_1& -\sin a{\theta }_1\\ \sin a{\theta }_1& \cos a{\theta }_1\end{array}\right]\left[\begin{array}{cc}\cos b{\theta }_1& \sin b{\theta }_1\\ -\sin b{\theta }_1& \cos b{\theta }_1\end{array}\right]=\left[\begin{array}{cc}\cos a{\theta }_1\cos b{\theta }_1+\sin a{\theta }_1\sin b{\theta }_1& \cos a{\theta }_1\sin b{\theta }_1-\sin a{\theta }_1\cos b{\theta }_1\\ \sin a{\theta }_1\cos b{\theta }_1-\cos a{\theta }_1\sin b{\theta }_1& \sin a{\theta }_1\sin b{\theta }_1+\cos a{\theta }_1\cos b{\theta }_1\end{array}\right]=\left[\begin{array}{cc}\cos (a-b){\theta }_1& -\sin (a-b){\theta }_1\\ \sin (a-b){\theta }_1& \cos (a-b){\theta }_1\end{array}\right]= \\ {R}_{\Theta ,a-b}^{d_{head}} \end{gathered}$$

乘积后结果与之前形态一致，其他位置的分块矩阵类似。所以，QK点积的结果如下。

$$RoP{E}_q(x_a,a)RoP{E}_k(x_b,b{)}^T=x_a{W}_q{R}_{\Theta ,a-b}^{d_{head}}{W}_k^T{x}_b^T$$

值得一提的是，旋转编码矩阵$R_{\Theta ,m}^{d_{head}}$是一个正交矩阵，因此可以确保在位置编码过程中的稳定性，简单理解就是，query_states和key_states中相同pos位置的token的隐向量点积后为0，只有不同位置才表示位置信息。

具体工程实现上，大模型一般和Sinusoidal一样先缓存一个token序列足够长的的缓存这些值后，模型在每次处理输入序列时，可以快速应用这些位置编码。下面跟随工程代码了解具体实现。

class RotaryEmbedding( nn.Module ):
    def __init__(self, head_dim, max_position_embeddings=151643, base = 10000 ):
        super().__init__()
        self.base = base
        self.head_dim = head_dim
        self.max_position_embeddings = max_position_embeddings
        '''
        inv_freq逆频率向量，就是前面RoPE部分的第一个式子，一个长度为d_head/2的\Theta向量。base默认为10000
        '''
        self.inv_freq = 1.0 / ( base ** ( torch.arange( 0, self.head_dim, 2, dtype=torch.int64).float() / self.head_dim ) )
        self.cos_cached, self.sin_cached = self._set_cos_sin_cache( self.max_position_embeddings )
    
    def _set_cos_sin_cache(self, max_position_embeddings ):
        """
        （1）足够长（151643）的序列的position ids向量t，和inv_freq向量外积，double，变成一个shape=[max_position_embedding * head_dim]的矩阵，形式如下
        [
        [0* \theta1, 0* \theta2,...,  0* \theta_(d_head/2), 0* \theta1, 0* \theta2,...,  0* \theta_(d_head/2)],
        [1* \theta1, 1* \theta2,...,  1* \theta_(d_head/2), 1* \theta1, 1* \theta2,...,  1* \theta_(d_head/2)],
        ......,
        [151642* \theta1, 151642* \theta2,...,  151642* \theta_(d_head/2), 151642* \theta1, 151642* \theta2,...,  151642* \theta_(d_head/2)],
        ]
        （2）将double_freqs作为输入逐位计算正弦和余弦值并分别返回两个矩阵进行缓存备用。
        （3）请注意，这里实现的功能等价于前面的旋转位置编码矩阵RoPE_{q,k}(x_m,m)，但为了工程上快速计算，而取巧了。具体见下。
        """
        t = torch.arange(0, max_position_embeddings)
        freqs = torch.outer( t, self.inv_freq)
        double_freqs = torch.cat( (freqs, freqs), dim=-1)
        cos_cached = double_freqs.cos().to(dtype)
        sin_cached = double_freqs.sin().to(dtype)
        return cos_cached, sin_cached

    def forward(self, x, seq_len = None):
        """
        返回和当前对话序列长度等长的正弦余弦序列，若序列长度大于max_position_embeddings，则根据序列长度重新生成正余弦序列
        """
        if seq_len == None:
            seq_len = 1
        if seq_len > self.max_position_embeddings:
            self._set_cos_sin_cache( seq_len )
        return (
            self.cos_cached[:seq_len].to(dtype=x.dtype),
            self.sin_cached[:seq_len].to(dtype=x.dtype)
        )
    @classmethod
    def rotary_inputs( cls, x ):
      '''
      分别将query_states或key_states按d_head维度一分为二，后半部分取符号前后调换位置再cat合并。本质是语义向量复数化。将前半部分划为实部，后半部分划为虚部。
      '''
        x1, x2 = x[..., : x.shape[-1]//2], x[..., x.shape[-1]//2: ]
        return torch.cat( ( -x2, x1 ), dim=-1)

    @classmethod
    def apply_rotary_pos_emb(cls, query, key, cos, sin, position_ids=None, unsqueeze_dim=1 ):
        '''
        下式对query和key进行旋转。对上述提到的分别对query_states或key_states向量进行位置编码的操作进行取巧，具体操作是：
        （1）cos和sin分别是预先计算的足够长的token序列的self.cos_cached和self.sin_cached中对应于本次编码长度的截取部分，shape均为[seq_len, d_head]，其中d_head维度为两组[0, d_head/2]部分的拼接。也可以表示为sin_double=[sin, sin], cos_double=[cos, cos]
        （2）query和key分别是被计算的query_states或key_states，shape均为[seq_len, d_head]
        （3）位置信息注入的方式是逐位积，将query向量切一半，表示为query=[q_1, q_2]，则q_embed=[q_1*cos-q_2*sin, q_2*cos+q_1*sin]，同样的k_embed=[k_1*cos-k_2*sin, k_2*cos+k_1*sin]
        '''
        cos = cos[ position_ids ].unsqueeze( unsqueeze_dim ) 
        sin = sin[ position_ids ].unsqueeze( unsqueeze_dim )
        q_embed = query * cos + cls.rotary_inputs(query) * sin 
        k_embed = key * cos + cls.rotary_inputs(key) * sin
        return q_embed, k_embed

因此，当上面代码中apply_rotary_pos_emb返回的q_embed与k_embed的转置进行点积时，实际上得到矩阵结果为，
$$\left[\begin{array}{c}{q}_1,q_2\end{array}\right]\left[\begin{array}{cc}cos& sin\\ -sin& cos\end{array}\right]\left[\begin{array}{cc}cos& -sin\\ sin& cos\end{array}\right]\left[\begin{array}{c}{k}_1,k_2\end{array}\right]$$

当涉及到矩阵中具体行列的值时，同$RoP{E}_q(x_a,a)RoP{E}_k(x_b,b{)}^T$结果，位置相同时为0，不相同时表示为相对位置编码结果。这种相对位置关系可以写成Sinusoidal中提到的线性变换关系，由于是维度分量两两一组为sin和cos，输入值可以理解为角度，所以RoPE又称为旋转位置编码。这种代码在旋转矩阵的可扩建上以及乘积的位置相对性上都是可以外推的。

3.2 ALiBi线性偏置注意力，2021

之前的两个位置编码方式都是对embedding或者query_states加上或者逐位乘一个基于三角函数的位置编码向量，而Attention with Linear Biases (ALiBi)(论文Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation)的核心思想是用一个和query_states, key_states的特定token之间的距离成比例的一个惩罚项来偏置query_states和key_states点积后的注意力权重得分。原始论文宣称可以加快11%的训练速度，以及减少11%的内存使用。

简单来说，该方法认为RoPE通过三角函数来编码相对位置的思路过于冗杂，既然最终结果是要表示不同两个token的相对位置信息，何不直接根据相隔的距离直接编码？两个token在序列中相隔一个位置，就对qk点积后的注意力得分矩阵中的位置-1，相差2就-2，简单明了。如下图：

此外，为了对多个head的情况也加以区分，还引入了一个坡度概念m，当有n个head时，m的计算公式为：

$$m=\left\{\begin{array}{c}{2}^{\frac{-8}{i}},i\in [1,2,...,n]\end{array}\right\}$$
也就是说，在不同的head的注意力评分矩阵中，根据head的序号i，对位置差这一乘法统一乘以一个坡度 $2^{\frac{-8}{i}}$。
仔细一想，确实，这个方法简单粗暴有效，目前已经和RoPE并称两大位置编码方案。但如果自行设计一个大模型，到底用什么position encoding其实还没有一个公认的最优方案，这两者（RoPE和ALiBi）目前看都可尝试，具体情况具体分析。