【DeepSeek-R1背后的技术】系列八：位置编码介绍（绝对位置编码、RoPE、ALiBi、YaRN）

【DeepSeek背后的技术】系列博文：
第1篇：混合专家模型（MoE）
第2篇：大模型知识蒸馏（Knowledge Distillation）
第3篇：强化学习（Reinforcement Learning, RL）
第4篇：本地部署DeepSeek，断网也能畅聊！
第5篇：DeepSeek-R1微调指南
第6篇：思维链（CoT）
第7篇：冷启动
第8篇：位置编码介绍（绝对位置编码、RoPE、ALiBi、YaRN）
第9篇：MLA（Multi-Head Latent Attention，多头潜在注意力）

Transformer 模型在处理序列数据时，其自注意力机制使得模型能够全局地捕捉不同元素之间的依赖关系，但这样做的代价是丧失了序列中的元素顺序信息。由于自注意力机制并不考虑元素在序列中的位置，所以在输入序列的任何置换下都是不变的，这就意味着模型无法区分序列中元素的相对位置。在许多自然语言处理任务中，词语之间的顺序是至关重要的，所以需要一种方法来让模型捕获这一信息。这就是位置编码（Positional Encoding）的角色所在。

本文主要介绍常见的绝对位置编码（sinusoidal）、旋转位置编码（Rotary Position Embedding，RoPE）、相对位置编码ALiBi（Attention with Linear Biases）以及YaRN（Yet another RoPE extensioN method）。

1 绝对位置编码

1.1 原理介绍

为了解决时序的问题，Transformer的作者用了一个绝妙的办法：位置编码(Positional Encoding)，在论文 《Attention Is All You Need》 中，作者通过sin函数和cos函数交替来创建 positional encoding，其计算positional encoding的公式如下

其中，pos是指每个token在整个序列中的位置，相当于是0, 1, 2, 3…(看序列长度是多大，比如10，比如100)，d_model代表位置向量的维度(也是词embedding的维度，transformer论文中设置的512维) ，i是embedding向量的位置下标对2求商并取整(可用双斜杠 // 表示整数除法，即求商并取整)，它的取值范围是 [0, …, d_model/2]，例如

即：
2i是指向量维度中的偶数维，即第0维、第2维、第4维…，第510维，用sin函数计算；
2i+1 是向量维度中的奇数维，即第1维、第3维、第5维…，第511维，用cos函数计算。

1.2 实现代码

class PositionalEncoding(nn.Module):  #@save
    """Positional encoding."""
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        # Create a long enough P
        self.P = torch.zeros((1, max_len, num_hiddens))
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        self.P[:, :, 0::2] = torch.sin(X)
        self.P[:, :, 1::2] = torch.cos(X)

    def forward(self, X):
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        return self.dropout(X)

1.3 可视化

绝对位置编码的可视化如下图所示：

原文如下，只用了一小节来介绍：

2 RoPE

原文：Roformer: Enhanced Transformer With Rotray Position Embedding

2.1 原理介绍

旋转位置编码是在位置编码上删除了绝对位置嵌入，而在网络的每一层增加了由苏剑林等人(2021)提出的旋转位置嵌入(RoPE)，其思想是采用绝对位置编码的形式 实现相对位置编码，且RoPE主要借助了复数的思想。

为了能利用上 token 之间的相对位置信息，假定 query 向量 q_m 和 key 向量 k_n 之间的内积操作可以被一个函数 g 表示，该函数 g 的输入是词嵌入向量 x_m ， x_n 和它们之间的相对位置 m-n ：

接下来的目标就是找到一个等价的位置编码方式，从而使得上述关系成立。具体推导过程参考：
https://zhuanlan.zhihu.com/p/642884818
https://blog.csdn.net/v_JULY_v/article/details/134085503
https://kexue.fm/archives/8265

假定现在词嵌入向量的维度是两维 d=2 ，函数g可以表示如下：

将2维推广到任意维度，可以表示如下：

2.2 高效计算

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

可以看到，RoPE 形式上和 Sinusoidal 位置编码有点相似，只不过 Sinusoidal 位置编码是加性的，而 RoPE 可以视为乘性的。在 θ_i 的选择上，RoPE 同样沿用了 Sinusoidal 位置编码的方案，即 θ_i = 10000^-2i/d，它可以带来一定的远程衰减性。

2.3 代码实现

在LLAMA中的实现：

# 生成旋转矩阵
def precompute_freqs_cis(dim: int, seq_len: int, theta: float = 10000.0):
    # 计算词向量元素两两分组之后，每组元素对应的旋转角度\theta_i
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    # 生成 token 序列索引 t = [0, 1,..., seq_len-1]
    t = torch.arange(seq_len, device=freqs.device)
    # freqs.shape = [seq_len, dim // 2] 
    freqs = torch.outer(t, freqs).float()  # 计算m * \theta

    # 计算结果是个复数向量
    # 假设 freqs = [x, y]
    # 则 freqs_cis = [cos(x) + sin(x)i, cos(y) + sin(y)i]
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs) 
    return freqs_cis

# 旋转位置编码计算
def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    # xq.shape = [batch_size, seq_len, dim]
    # xq_.shape = [batch_size, seq_len, dim // 2, 2]
    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 2)
    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 2)
    
    # 转为复数域
    xq_ = torch.view_as_complex(xq_)
    xk_ = torch.view_as_complex(xk_)
    
    # 应用旋转操作，然后将结果转回实数域
    # xq_out.shape = [batch_size, seq_len, dim]
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(2)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(2)
    return xq_out.type_as(xq), xk_out.type_as(xk)

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__()

        self.wq = Linear(...)
        self.wk = Linear(...)
        self.wv = Linear(...)
        
        self.freqs_cis = precompute_freqs_cis(dim, max_seq_len * 2)

    def forward(self, x: torch.Tensor):
        bsz, seqlen, _ = x.shape
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)

        xq = xq.view(batch_size, seq_len, dim)
        xk = xk.view(batch_size, seq_len, dim)
        xv = xv.view(batch_size, seq_len, dim)

        # attention 操作之前，应用旋转位置编码
        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
        
        # scores.shape = (bs, seqlen, seqlen)
        scores = torch.matmul(xq, xk.transpose(1, 2)) / math.sqrt(dim)
        scores = F.softmax(scores.float(), dim=-1)
        output = torch.matmul(scores, xv)  # (batch_size, seq_len, dim)
  # ......

2.4 外推性

• 旋转编码 RoPE 可以有效地保持位置信息的相对关系，即相邻位置的编码之间有一定的相似性，而远离位置的编码之间有一定的差异性。这样可以增强模型对位置信息的感知和利用。这一点是其他绝对位置编码方式（如正弦位置编码、学习的位置编码等）所不具备的，因为它们只能表示绝对位置，而不能表示相对位置。

• 旋转编码 RoPE 可以通过旋转矩阵来实现位置编码的外推，即可以通过旋转矩阵来生成超过预训练长度的位置编码。这样可以提高模型的泛化能力和鲁棒性。这一点是其他固定位置编码方式（如正弦位置编码、固定相对位置编码等）所不具备的，因为它们只能表示预训练长度内的位置，而不能表示超过预训练长度的位置。

• 旋转编码 RoPE 可以与线性注意力机制兼容，即不需要额外的计算或参数来实现相对位置编码。这样可以降低模型的计算复杂度和内存消耗。这一点是其他混合位置编码方式（如Transformer-XL、XLNet等）所不具备的，因为它们需要额外的计算或参数来实现相对位置编码。

3 ALiBi

ALiBi的思想核心是：不给词向量加入位置嵌入向量，而是用一个和query和key之间的距离成比例的一个“惩罚项”来偏置query-key的注意力得分。效果：可以加快11%的训练速度，以及减少11%的内存使用。

方法核心如下图所示：

左边的第一项，是自注意力得分，关于q和k的内积的，这个和传统transformer里面的一样了。第二项是一个相对距离的矩阵，例如

• q1和 k1之间的距离是0，所以对应位置就是0；
• q2和k1，是相对位置偏移为k的索引1 - q的索引2，得到1-2 = -1，就对应到了中间矩阵的取值为-1了；
• 以此类推，相对距离矩阵的中间对角线上都是0，然后左下角的取值都是对应的k的索引-q的索引了。

第三项，是个坡度m，按照论文中的描述，其做法是：

例如，8个heads的时候，m的取值为：1/2, 1/4, 1/8, 1/16, 1/32, 1/64, 1/128, 1/256；如果是16个heads，则m的取值为：1/sqrt(2), 1/2, 1/(2*sqrt(2)), 1/4, …, 1/256，相当于追加了一半的1/sqrt(2)到原来的8个head的每个m的取值。

扩展到一般情况就是，对于n个head的话，m的取值就是2^-8/n，即 2^-8/1，2^-8/2，2^-8/3，……，2^-8/n 这样的m个坡度了。

整体公式就是：

对于第i个query来说，其可见范围，在自回归设置下，只能是从左边开始1， 2， ……， i个位置，而他们之间的相对距离就是：0, -1, -2, …, 1-i 了，即：k的索引 减去 q的索引。k的索引，遍历1, 2, …, i；而q的索引，取值为i。如此，就有了上面的公式了。

4 YaRN

YaRN是基于NTK-aware方法的进一步拓展，YaRN= NTK-aware + NTK-by-parts + Dynamic NTK。通过结合温度缩放和NTK-by-parts插值技术，全面提升长文本外推能力。它核心解决的问题是线性内插导致的self-attention 点积的值增大。由于线性内插会改变旋转向量转动的幅度，原来距离较远的q,k点积由于旋转幅度变小，他们的点积结果会增大，进而导致Softmax操作过于“锐化”，使得注意力分布集中于少数位置，削弱模型对全局上下文的关注能力。 Yarn在 NTK-by-parts 基础上，引入注意力温度因子 t 来调整注意力分布：

优点：

• 超低训练成本（0.1%预训练数据）。
• 几乎兼容目前所有主流的Transformer实现。
• 性能优越，无论是否微调均能在128k上下文中表现出色。

5 四种位置编码的详细对比

5.1 绝对位置编码（Absolute Positional Encoding）

• 原理：通过预定义的正弦/余弦函数或可学习的参数，为每个位置生成唯一的编码向量，直接与词嵌入相加。
• 计算方式：静态编码，直接叠加到输入向量上。
• 优势：
  • 简单直观：易于实现，广泛应用于早期Transformer模型（如BERT、GPT-2）。
  • 绝对位置感知：明确区分每个位置的唯一性。
• 劣势：
  • 长序列泛化差：难以处理超出训练长度的序列。
  • 相对位置间接：模型需自行学习相对位置关系，效率较低。
• 适用场景：短文本任务或对位置敏感但长度固定的场景。

5.2 RoPE（Rotary Position Embedding）

• 原理：通过旋转矩阵将位置信息融入查询（Query）和键（Key）向量，利用复数域旋转改变注意力得分。
• 计算方式：动态调整Q/K向量的相位，生成相对位置感知的注意力。
• 优势：
  • 相对位置编码：直接建模位置差异，提升长序列处理能力。
  • 无需额外参数：旋转矩阵由公式生成，节省参数空间。
  • 兼容性：适配线性注意力机制，计算高效。
• 劣势：
  • 实现复杂：涉及复数运算和矩阵旋转，对硬件优化要求高。
  • 长度外推有限：虽优于绝对编码，但超长序列仍需改进。
• 适用场景：语言模型（如LLaMA、GPT-J）及需要高效相对位置建模的任务。

5.3 ALiBi（Attention with Linear Biases）

• 原理：在注意力得分中添加与位置差成反比的线性偏置，抑制远距离关注。
• 计算方式：固定偏置项（如 $-m\cdot |i-j|$，$m$为预设斜率），无需学习。
• 优势：
  • 高效长序列处理：显式减少远距离依赖，适合超长文本。
  • 零额外参数：偏置固定，训练稳定且计算开销低。
  • 强外推性：在未训练长度上表现优异（如支持16k→100k token推理）。
• 劣势：
  • 灵活性不足：固定偏置可能忽略关键远距离信息。
  • 任务适应性差：对需要复杂位置关系的任务（如解析嵌套结构）效果有限。
• 适用场景：长文档理解、代码生成等需处理超长序列的任务。

5.4 YaRN（Yet another RoPE Extension）

• 原理：基于RoPE改进，通过动态调整旋转基的频率或插值方法，增强长序列外推能力。
• 计算方式：引入缩放因子（如温度参数）调整旋转角度，适配更长上下文。
• 优势：
  • 扩展RoPE能力：显著提升长序列外推性能（如支持64k→128k token）。
  • 动态适应性：通过插值或波长缩放，平衡远近位置关系。
• 劣势：
  • 实现复杂：需调整旋转机制，增加计算和调参成本。
  • 依赖RoPE基础：无法脱离RoPE独立使用。
• 适用场景：需处理超长序列且已采用RoPE的模型（如LLaMA-2扩展）。

5.5 对比总结

5.6 优劣势总结

• 绝对位置编码：简单但缺乏灵活性，适合短序列任务。
• RoPE：平衡相对位置建模与效率，主流选择但需优化实现。
• ALiBi：长序列王者，计算高效但牺牲灵活性。
• YaRN：RoPE的进化版，专攻超长序列外推，代价是复杂性和计算成本。

根据任务需求选择：短序列通用任务选绝对编码或RoPE，超长文本选ALiBi或YaRN，平衡场景选RoPE。