大模型系列——旋转位置编码和长度外推

绝对位置编码

旋转位置编码

论文中有个很直观的图片展示了旋转变换的过程：

对于“我”对应的d维向量，拆分成d/2组以后，每组对应一个角度 $\theta$ ，若 $\theta$ 1对应的向量为(x1,x2)，应用旋转位置编码，相当于这个分量旋转了m $\theta$ 1角度。

结合transformer4.42.4版本，qwen2源码分析如下：

1、定义和缓存cos、sin

class Qwen2RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()

        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # Build here to make `torch.jit.trace` work.
        self._set_cos_sin_cache(
            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
        )

    def _set_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len
        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)

        freqs = torch.outer(t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

        return (
            self.cos_cached[:seq_len].to(dtype=x.dtype),
            self.sin_cached[:seq_len].to(dtype=x.dtype),
        )

在上面这段代码中，inv_freq对应的是各分量的旋转角度，长度为d/2

        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)

        freqs = torch.outer(t, self.inv_freq)

这里的t为提前把所有可能的位置id 都先取好，并与对应的角度相乘，对应公式中的m $\theta$ ，计算出来的矩阵freqs维度为（self.max_seq_len，d/2）。这里outer函数计算如下：

torch.outer

import torch
t = torch.tensor([1,2,3])
inv_freq = torch.tensor([0.1,0.2,0.3])
f = torch.outer(t, inv_freq)
tensor([[0.1000, 0.2000, 0.3000],
        [0.2000, 0.4000, 0.6000],
        [0.3000, 0.6000, 0.9000]])

emb = torch.cat((f,f), dim=-1)
emb
tensor([[0.1000, 0.2000, 0.3000, 0.1000, 0.2000, 0.3000],
        [0.2000, 0.4000, 0.6000, 0.2000, 0.4000, 0.6000],
        [0.3000, 0.6000, 0.9000, 0.3000, 0.6000, 0.9000]])

这样取每一行，即对应这个id下的所有m $\theta$ 值。

通过cat，拼接出来的emb后半部分和前半部分是一致的，维度变成d

2、apply_rotary_pos_emb

def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    sin = sin[position_ids].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这个函数，我们模拟下：

q = torch.tensor([0.1000, 0.2000, 0.3000, 0.4000, 0.5000, 0.6000])

rotate_q = rotate_half(q) 
rotate_q 
tensor([-0.4000, -0.5000, -0.6000,  0.1000,  0.2000,  0.3000])

可以看出，q的最后一维后半部分取反后放到了前面，这里和论文中的公式是有区别的。问了下kimi解释如下：

原始论文中可能采用了不同的子空间划分方法，例如将维度0, 1, 2, 3划分为{0, 1}和{2, 3}两组。而在HuggingFace的实现中，可能采用了不同的划分方式，例如将维度0, 1, 2, 3划分为{0, 2}和{1, 3}两组。这两种划分方法都是合理的，只是子空间的划分不同，但最终都能达到相同的旋转编码效果。

因此，即使rotate_half函数的实现与论文中的公式看起来不一样，但由于其最终效果相同，这种实现差异是可以接受的。重要的是理解旋转位置编码的核心思想，即通过旋转操作将位置信息编码到模型中，而不是纠结于具体的实现细节。

3、attention forward

query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

if past_key_value is not None:
    cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}  # Specific to RoPE models
    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)

# repeat k/v heads if n_kv_heads < n_heads
key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups)

attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

参考：

（1）十分钟读懂旋转编码（RoPE）：https://zhuanlan.zhihu.com/p/647109286

（2）旋转位置编码RoPE最通俗易懂解释，https://zhuanlan.zhihu.com/p/701576694

ROPE优化——长度外推

本节来自：Transformer升级之路：16、“复盘”长度外推技术 - 科学空间|Scientific Spaces

顾名思义，免训练长度外推，就是不需要用长序列数据进行额外的训练，只用短序列语料对模型进行训练，就可以得到一个能够处理和预测长序列的模型，即“Train Short, Test Long”。那么如何判断一个模型能否用于长序列呢？最基本的指标就是模型的长序列Loss或者PPL不会爆炸，更加符合实践的评测则是输入足够长的Context，让模型去预测答案，然后跟真实答案做对比，算BLEU、ROUGE等，LongBench就是就属于这类榜单。

但要注意的是，长度外推应当不以牺牲远程依赖为代价——否则考虑长度外推就没有意义了，倒不如直接截断文本——这意味着通过显式地截断远程依赖的方案都需要谨慎选择，比如ALIBI以及《Transformer升级之路：7、长度外推性与局部注意力》所列举的大部分方案，还有带显式Decay的线性RNN，这些方案当序列长度足够大时都表现为局部注意力，即便有可能实现长度外推，也会有远程依赖不足的风险，需要根据自己的场景斟酌使用。

如何判断在长度外推的同时有没有损失远程依赖呢？比较严谨的是像《Transformer升级之路：12、无限外推的ReRoPE？》最后提出的评测方案，准备足够长的文本，但每个模型只算每个样本最后一段的指标，如下图所示：

比如，模型训练长度是4K，想要看外推到16K的效果，那么我们准备一个16K tokens的测试集，4K的模型输入每个样本最后4K tokens算指标，8K模型输入每个样本最后8K tokens但只算最后4K tokens算指标，12K模型输入每个样本最后12K tokens但只算最后4K tokens算指标；依此类推。这样一来，不同长度的模型算的都是同一段tokens的指标，不同的只是输入的Context不一样，如果远程依赖得以有效保留，那么应该能做到Context越长，指标越好。

问题：如果ABC和MNOP没有强关联关系呢？如何评判优劣？

先说个题外话，为什么如今大部分LLM的位置编码都选择了RoPE呢？笔者认为主要有几点原因：

1、RoPE不带有显式的远程衰减，这对于旨在Long Context的模型至关重要；

2、RoPE是一种真正的位置编码，通过不同频率的三角函数有效区分了长程和短程，达到了类似层次位置编码的效果，这也是Long Context中比较关键的一环；

3、RoPE直接作用于Q、K，不改变Attention的形式，与Flash Attention更契合，更容易Scale Up。

相比之下，诸如ALIBI、KERPLE等，虽然有时也称为位置编码，但它们实际上只是一种Attention Bias，没有太多位置信息，且不适用于Encoder，能用于Decoder大体上是因为Decoder本身的下三角Mask就已经有较为充分的位置Bias了，额外的Attention Bias只是锦上添花。此外它们无法在单个头内有效区分长程和短程，而是要通过在不同头设置不同的Decay因子来实现，这也意味着它们用于单头注意力（比如GAU）的效果会欠佳。

说这么多优缺点的对比，看起来像是“王婆卖瓜，自卖自夸”，其实不然，这只是为了跟大家交换一下观点，因为之前也有读者提出过相同的问题。作为RoPE的提出者，笔者对RoPE的理解不见得一定比大家深刻，毕竟当时提出RoPE的初衷纯粹是好玩，当时的想法是有效就很不错了，能媲美Learnable的绝对位置编码就是非常好的消息了。所以，既然是“意料之外”，那么“作者本人也没多透彻的认识”这件事，也是“情理之中”了。

好像又把话题扯偏了。简单来说，其实上两节的内容主要是想表达的观点是：目前看来，RoPE对于Long Context来说是足够的，所以研究RoPE的长度外推是有价值的，以及我们在选择长度外推方案时，不应牺牲远程依赖的能力。

1、窗口截断

在本站最早讨论长度外推的《Transformer升级之路：7、长度外推性与局部注意力》一文中，我们判断长度外推是一个预测阶段的OOD（Out Of Distribution）的问题，尽管用今天的视角看，这篇文章的一些评述已经显得有点过时，但这个根本判断是依然还算正确，放到RoPE中，就是推理阶段出现了没见过的相对距离。为此，一个看上去可行的方案是引入Sliding Window的Attention Mask，如下图左所示：

当然，由于强行截断了窗口外的注意力，所以这个方案并不满足“不牺牲远程依赖的能力”的原则，但我们可以只将它作为一个Baseline看待。很遗憾的是，即便做出了如此牺牲，这个方案却是不Work的——连最基本的PPL不爆炸都做不到！

答案可能让人意外：开头的几个Token很重要，不能扔掉。所以最后可用的Window Mask应该如上图右（LM-Infinite这篇论文管它叫“ΛΛ-Mask”）。

2、位置内插-PI

一位网名为“kaiokendev”的网友在他的博客《https://kaiokendev.github.io/til#extending-context-to-8k》中提出了一个非常朴素的解决办法——“位置内插”——将预测的长文本的位置编码乘上因子LtrainLtest��，缩放到训练长度范围内，如下式所示（式中的位置都是相对位置）。没过多久，Meta在论文《Extending Context Window of Large Language Models via Positional Interpolation》中也发布了同样的方法，命名为“Positional Interpolation（PI）”，并补充了较为充分的实验结果。

然而，位置内插并不算长度外推方案，至少不是免训练的长度外推方案，因为位置内插之后同样会有PPL爆炸的问题。原因也不难理解，尽管位置内插避免了远处的位置越界问题，但这同时压缩了邻近Token的距离，严重扰乱了模型的局部分辨率，而众所周知语言模型本身就是一个非常依赖于局部关系的任务，所以扰乱了局部自然就没法预测准了。

不过，这也并非说位置内插就没有价值了。我们知道，需要长度外推的读者，无外乎是两种情况：一种是没有资源去做长文本微调，希望能够从短文本模型直接得到一个可用的长文本模型，这种需求对长度外推的效果要求会比较高，位置内插就不适合他们了；另一种是有资源去做长文本微调，研究长度外推纯粹是为了得到一个更好的初始化模型，这种情况对模型修改带来的初始损失容忍度比较高，只要能够通过微调快速弥补回损失掉的效果即可，位置内插正好是属于此类方法。Meta的论文显示，经过PI之后，仅需1000步左右的长文本训练，就可以得到一个行之有效的长文本模型，这比不做任何修改直接微调的训练效率高出很多。

3、保近压远-ReRope

直接外推的问题是远处越界，而位置内插的问题是局部失真，看上去两者是互补的，能不能集两者之长呢？这就是《Transformer升级之路：12、无限外推的ReRoPE？》所提出的Leaky ReRoPE，以及它的极限版本ReRoPE。

如果将内插的因子k�取到无穷大，这就得到极简的ReRoPE，它在窗口外的位置编码都变为w�，意味着对于任意长的序列都不会越界，即理论上具备无限外推的潜力！事实上，Leaky ReRoPE和ReRoPE的表现确实都非常好，从Loss来看，它们能做到几乎不损失训练长度内的效果，并且实现了长度外推，且Context越长，Loss越低，说明它们在外推的同时还确实保证了远程依赖。

Leaky ReRoPE和ReRoPE的主要问题在于它们的代码实现稍微有点麻烦。跟Attention Bias类的位置编码不同，RoPE没法通过先构造相对位置矩阵然后才计算相对位置编码的方式来实现（那样效率太低），只能通过绝对位置编码的方式来实现相对位置编码，这意味着它只能实现线性增长的相对位置，而Leaky ReRoPE和ReRoPE的相对位置是分段线性的，这意味着朴素地实现的话，需要算两次Attention矩阵（得到两段不同的线性）然后将它们拼接起来，这样效率无疑明显降低了。

不过，好消息是当前主流的Attention加速手段如Flash Attention都是将Attention分块计算的，比如每128长度为一块，这样当序列足够长时，分段线性的块占比非常少（只有窗口边界附近），如下式所示，只有红绿混色的块才需要重复计算Attention，剩下同色的块都只需要计算一次，所以结合分块计算Attention的话，Leaky ReRoPE和ReRoPE所增加的计算成本时几乎可以忽略的。此前读者 @chu-tianxiang 在评论区也分享了一个基于Triton的实现，大家有兴趣的可以参考一下。

月初Arxiv上提交了一篇论文《LLM Maybe LongLM: Self-Extend LLM Context Window Without Tuning》，其中提出了一种名为“Self-Extend”的免训练长度外推方法，它实际上就是在Leaky ReRoPE的基础上加了Round运算（四舍五入），使得每个相对位置都变回整数，进一步减轻相对位置的OOD问题。论文报告的效果也很好，这进一步肯定了Leaky ReRoPE的有效性。

4、转圈视角-YaRN

备注：i->0时为高频，转的越快。

结论：YaRN效果略逊于ReRope。我感觉是内插的时候过于随意，这种修改没有保证窗口内的局部注意力不变，也许YaRN可以和ReRope结合。

备注：这里的缩放因子，一会会提到

就像NTK-RoPE、YaRN的作者Bowen Peng曾经的观点，高频学习到的是局部的相对距离，低频学习到的是远程的绝对距离，两者都很重要，它们之间更像是一种层次的关系；

5、首个外推方案-NTK-ROPE

相比YaRN本身，YaRN的作者Bowen Peng的故事也许更加称得上“引人入胜”，他早前所提出的NTK-RoPE是RoPE的第一个免训练的长度外推方案，本系列的两篇博客《Transformer升级之路：10、RoPE是一种β进制编码》和《Transformer升级之路：11、将β进制位置进行到底》都直接受启发于它。虽然从目前来看，NTK-RoPE的效果不见得多好（相比YaRN、ReRoPE等），但它首次显示了免训练长度外推的可能性，具有里程碑式的意义，甚至可以说，后续的所有长度外推相关研究，都直接或者间接得益于NTK-RoPE打开了大家的想象力。

不过，尽管NTK-RoPE效果上不如YaRN，但对于前面提到的第二种有资源去做长文本微调的读者，可能会更喜欢NTK-RoPE，因为他们只是为了得到一个更好的初始化模型，反正都是要微调，NTK-RoPE与YaRN的初始效果差异他们并不会太在意，相比之下他们更乐意选择实现更简单的NTK-RoPE了，比如CodeLLAMA就是在LLAMA2的基础上将base改为10的6次方，然后继续训练的。此外，Meta在其论文《Effective Long-Context Scaling of Foundation Models》中，将NTK-RoPE改称为RoPE-ABF（Adjusted Base Frequency），相比神秘的NTK，ABF的名称能更直观体现出它的含义。

6、拒绝交税-Dynamic Scaling和CLEX

备注：从公式可以看出，核心思想就是Ltrain训练窗口长度内的位置不缩放，之外的缩放

在llama的官方代码中，有一个DynamicNTKScalingRotary实现如下（结合以上公式就比较容易看懂了）：

class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding):
    """LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""

    def forward(self, x, position_ids):
        # difference to the original RoPE: inv_freq is recomputed when the sequence length > original length
        seq_len = torch.max(position_ids) + 1
        if seq_len > self.max_position_embeddings:
            base = self.base * (
                (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
            ) ** (self.dim / (self.dim - 2))
            inv_freq = 1.0 / (
                base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(x.device) / self.dim)
            )
            self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO joao: this may break with compilation

        cos, sin = super().forward(x, position_ids)
        return cos, sin

另外llama官方代码中还有一个线性scaling，也记录一下：

class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):
    """LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""

    def forward(self, x, position_ids):
        # difference to the original RoPE: a scaling factor is aplied to the position ids
        position_ids = position_ids.float() / self.scaling_factor
        cos, sin = super().forward(x, position_ids)
        return cos, sin

问题：窗口内的是否真的做到了拒绝交税，窗口外的是否没有丢失远程信息？

7、另起炉灶 HWFA和Key Norm

除了Dynamic Scaling外，“拒绝交税”的另一个思路是“另起炉灶”，通过重新设计预训练时所用的模型架构，使得它具备训练完成后就可以不做任何修改实现长度外推的潜力，在这个系列的文章中，笔者有两篇相关的探讨，分别是在《Transformer升级之路：9、一种全局长度外推的新思路》所提到HWFA（Hybird Window-Full Attention），以及在《Transformer升级之路：15、Key归一化助力长度外推》所验证的Key Norm。