位置编码

总阅读量 加载中... 次

位置编码（Positional Encoding）是 Transformer 架构中的关键组件之一。由于 Transformer 采用自注意力机制进行信息交互，其本身不具备处理序列顺序的能力——无论输入序列如何打乱，注意力计算的输出结果都是相同的。位置编码的引入正是为了解决这一问题，使模型能够区分序列中不同位置的 token。

本文围绕 Transformer 中的位置编码技术展开综述，介绍绝对位置编码与相对位置编码两类主流方法。

1. 背景与问题定义

在 Transformer 的自注意力计算中：

Attention (Q, K, V) = Softmax (\frac{Q K^{T}}{\sqrt{d_{k}}}) V

公式本身对输入序列的顺序是对称的——交换输入序列中两个 token 的位置，注意力矩阵的值不会发生任何变化。这意味着，如果不额外引入位置信息，模型将无法区分"我爱你"和"你爱我"这类语序敏感的任务。

位置编码的核心目标是为序列中的每个位置 $i$ 生成一个独特的向量 $P E_{i}$ ，将其与 token 的嵌入向量相加，从而使模型能够感知位置信息。

绝对位置编码是最早被提出的位置编码方案，为序列中的每个绝对位置生成一个固定的编码向量。

2. 可学习位置编码

Learnable 位置编码（可学习位置编码）最早由 Bahdanau 等人在论文《Neural Machine Translation by Jointly Learning to Align and Translate》（2014）中提出，随后被广泛应用于早期神经机器翻译模型。

其核心思想是将位置编码视为可学习的参数矩阵 $P E \in R^{L \times d_{model}}$ ，通过反向传播自动学习最优的位置表示。

{Input}_{i} = {TokenEmbedding}_{i} + {PositionEmbedding}_{i}

2.1 代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18class LearnablePositionalEncoding(torch.nn.Module):
    """
    Learnable Positional Encoding
    论文来源：Bahdanau et al., 2014
    """
    def __init__(self, d_model: int, max_len: int = 5000):
        super(LearnablePositionalEncoding, self).__init__()
        self.position_embeddings = torch.nn.Embedding(max_len, d_model)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        :param x: Tensor, shape = [batch_size, seq_len, d_model]
        :return: Tensor, shape = [batch_size, seq_len, d_model]
        """
        batch_size, seq_len, _ = x.shape
        # 生成位置索引 [0, 1, 2, ..., seq_len-1]
        positions = torch.arange(seq_len, device=x.device).expand(batch_size, seq_len)
        return x + self.position_embeddings(positions)

3. 正弦位置编码

Sinusoidal 位置编码（正弦位置编码）是 Transformer 架构奠基之作 Attention Is All You Need 中提出的一种绝对位置编码技术。它不需要通过数据训练去学习，而是直接利用不同频率的正弦和余弦函数组合，为序列中的每个位置生成一个固定的、唯一的特征向量。

3.1 核心数学公式

假设我们需要为一个长度为 $L$ 的文本序列生成位置编码，每个位置的位置编码向量维度为 $d$ （通常与 Embedding 维度相同，且为偶数）。

对于序列中的第 $p o s$ 个位置（ $p o s \in [0, L - 1]$ ），以及该位置编码向量中的第 $i$ 个维度（ $i \in [0, d - 1]$ ），其计算公式定义如下：

P E_{(p o s, 2 i)} = \sin (\frac{p o s}{10000^{\frac{2 i}{d}}})

P E_{(p o s, 2 i + 1)} = \cos (\frac{p o s}{10000^{\frac{2 i}{d}}})

公式拆解：

$p o s$ ：当前 Token 在句子中的绝对位置（如第 0 个词、第 1 个词…）。
$2 i$ / $2 i + 1$ ：代表位置编码向量中的索引。偶数维度使用正弦（ $\sin$ ），奇数维度使用余弦（ $\cos$ ）。
$10000^{\frac{2 i}{d}}$ ：这是一个分母，通常定义为波长控制参数。随着维度索引 $i$ 的增大，分母呈指数级变大，这意味着波长越来越长，频率越来越低。

3.2 完整的计算步骤（矩阵视角）

在工程实现中，通常不会用双重循环去逐个元素计算，而是采用矩阵化操作。以下是完整的全矩阵计算流程：

第一步：构建位置向量 (Position Vector)

创建一个形状为 $(L, 1)$ 的列向量，代表序列中每个点的绝对位置：

p o s = [\begin{matrix} 0 \\ 1 \\ 2 \\ ⋮ \\ L - 1 \end{matrix}]

第二步：计算逆频率向量 (Inverse Frequency Vector)

根据维度的前半部分（偶数项），计算缩放系数（也叫角速度 $ω_{i}$ ）。共有 $d / 2$ 个项：

inv_freq = [\frac{1}{10000^{0}}, \frac{1}{10000^{\frac{2}{d}}}, \frac{1}{10000^{\frac{4}{d}}}, \dots, \frac{1}{10000^{\frac{d - 2}{d}}}]

其形状为 $(1, d / 2)$ 的行向量。

第三步：计算角度矩阵 (Angles Matrix)

将位置列向量与逆频率行向量进行外积（Matrix Outer Product）相乘，得到一个形状为 $(L, d / 2)$ 的角度矩阵。矩阵中的每一个元素都是 $p o s \cdot ω_{i}$ ：

angles = p o s \times inv_freq = [\begin{matrix} 0 \cdot ω_{0} & 0 \cdot ω_{1} & \dots & 0 \cdot ω_{d / 2 - 1} \\ 1 \cdot ω_{0} & 1 \cdot ω_{1} & \dots & 1 \cdot ω_{d / 2 - 1} \\ ⋮ & ⋮ & ⋱ & ⋮ \\ (L - 1) \cdot ω_{0} & (L - 1) \cdot ω_{1} & \dots & (L - 1) \cdot ω_{d / 2 - 1} \end{matrix}]

第四步：应用 $\sin$ 和 $\cos$ 并拼接

对角度矩阵的每一个元素求正弦，得到 $\sin (angles)$ ，形状为 $(L, d / 2)$ 。
对角度矩阵的每一个元素求余弦，得到 $\cos (angles)$ ，形状为 $(L, d / 2)$ 。

最后，将这两个矩阵横向拼接。在标准实现中，通常是以交错（Interleave）或者前后拼接的方式组合。

如果采用交错拼接，最终得到的位置编码矩阵 $P E$ 形状为 $(L, d)$ ：

P E = [\begin{matrix} \sin (0 \cdot ω_{0}) & \cos (0 \cdot ω_{0}) & \sin (0 \cdot ω_{1}) & \cos (0 \cdot ω_{1}) & \dots \\ \sin (1 \cdot ω_{0}) & \cos (1 \cdot ω_{0}) & \sin (1 \cdot ω_{1}) & \cos (1 \cdot ω_{1}) & \dots \\ ⋮ & ⋮ & ⋮ & ⋮ & ⋱ \end{matrix}]

3.3 直观理解：它长什么样？

如果我们将最终的 $P E$ 矩阵可视化为一张热力图：

纵轴（位置 $p o s$ ）：从上往下代表句子的第一个词到最后一个词。
横轴（维度 $d$ ）：从左往右。
左侧（低维区域）： $\sin$ 和 $\cos$ 的频率非常高，图形呈现密集的正弦波浪，能够精确捕捉近距离、微观的位置差异。
右侧（高维区域）：频率极低，图形随着 $p o s$ 的增加变化非常缓慢，甚至趋于直线，用于捕捉长距离、宏观的全局位置框架。

这种设计使得每一个位置的 $d$ 维向量都是独一无二的"数字指纹"。

3.4 为什么要这样设计？（设计巧妙之处）

原论文指出，选择正弦和余弦函数的核心原因在于：它允许模型轻松学习到相对位置关系。

根据三角函数的和差化积公式：

\sin (α + β) = \sin α \cos β + \cos α \sin β

\cos (α + β) = \cos α \cos β - \sin α \sin β

这意味着，对于任何固定的相对偏移量 $k$ ，位置 $p o s + k$ 的位置编码向量 $P E_{p o s + k}$ ，都可以表示为位置 $p o s$ 的位置编码向量 $P E_{p o s}$ 的线性变换。

也就是说，Attention 机制在计算两个相距为 $k$ 的词的注意力时，它们之间的点积可以通过线性矩阵变换，转化成只与相对距离 $k$ 相关的变换，从而让模型既具备绝对位置感，又具备动态的相对距离感。

最后，生成好这个 $(L, d)$ 的 $P E$ 矩阵后，在 Transformer 中直接将其与原始词向量矩阵（Token Embeddings）进行按位置相加（Element-wise Add），便完成了位置信息的注入：

Input to Transformer = Token Embeddings + P E

3.5 代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28import torch
import math

class SinusoidalPositionalEncoding(torch.nn.Module):
    """
    Sinusoidal Positional Encoding
    论文来源：Vaswani et al., 2017
    """
    def __init__(self, d_model: int, max_len: int = 5000):
        super(SinusoidalPositionalEncoding, self).__init__()
        # 创建位置编码矩阵 [max_len, d_model]
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # [max_len, 1]
        # 计算除数 10000^(2i/d_model)
        div_term = torch.exp(
            torch.arange(0, d_model, 2, dtype=torch.float) * (-math.log(10000.0) / d_model)
        )  # [d_model/2]
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数维度
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数维度
        pe = pe.unsqueeze(0)  # [1, max_len, d_model] 便于批量相加
        self.register_buffer('pe', pe)  # 注册为 buffer，不参与梯度更新

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        :param x: Tensor, shape = [batch_size, seq_len, d_model]
        :return: Tensor, shape = [batch_size, seq_len, d_model]
        """
        return x + self.pe[:, :x.size(1), :]

3.6 绝对位置编码对比

特性	Learnable	Sinusoidal
论文来源	Bahdanau et al., 2014	Vaswani et al., 2017
参数量	$L \times d_{model}$	无需额外参数
外推能力	有限（受 max_len 限制）	理论支持（周期函数）
灵活性	可学习、更灵活	固定、不可调优
效果	通常略优于 Sinusoidal	表现稳定

绝对位置编码在处理长序列时存在明显的局限性：一方面，模型难以直接学习位置之间的相对关系；另一方面，当测试序列长度超过训练集最大长度时，编码效果急剧下降。相对位置编码通过直接编码 token 之间的相对距离来解决这些问题。

4. RoPE 位置编码

旋转位置编码（Rotary Position Embedding，RoPE）由 Su 等人在论文《RoFormer: Enhanced Transformer with Rotary Position Embedding》（2022）中提出，是目前大语言模型（如 LLaMA、Baichuan、Mistral 等）中最主流的位置编码技术之一。它的核心思想是：通过在复数空间中旋转 Query 和 Key 向量，将绝对位置信息自然地转化为相对位置信息。

ROPE

4.1 核心数学思想

RoPE 的核心目标是：通过一个变换函数 $f (x, pos)$ ，将位置信息编码到向量中，使得变换后的向量进行点积运算时，结果只依赖于相对位置。

具体来说，假设输入两个特征向量 $x_{q}$ 和 $x_{k}$ ，分别处于绝对位置 $m$ 和 $n$ 。我们希望找到一个变换函数 $f (x, pos)$ ，满足：

⟨ f (x_{q}, m), f (x_{k}, n) ⟩ = g (x_{q}, x_{k}, m - n)

RoPE 巧妙地利用了二维空间中的旋转矩阵来实现这一特性。在二维空间中，旋转一个向量不改变其长度，只改变其角度。

在 RoPE 中，这个变换函数的具体形式为：

f (x, pos) = R_{Θ, pos}^{d} \cdot x

其中 $R_{Θ, pos}^{d}$ 是一个分块对角旋转矩阵，将位置 $pos$ 编码为旋转角度。

4.2 完整的计算步骤

在实际的 Transformer 架构中，假设多头注意力机制中某个头的维度为 $d$ （ $d$ 通常为偶数，如 64 或 128）。RoPE 的具体计算流程如下：

第一步：拆分二维子空间

因为旋转是在二维平面上定义的，RoPE 会将一个 $d$ 维的 Query（或 Key）向量拆分为 $d / 2$ 个独立的二维子空间。

对于位置为 $m$ 的 $d$ 维向量 $q_{m} = [q_{0}, q_{1}, q_{2}, q_{3}, \dots, q_{d - 2}, q_{d - 1}]^{T}$ ，拆分为：

[\begin{matrix} q_{0} \\ q_{1} \end{matrix}], [\begin{matrix} q_{2} \\ q_{3} \end{matrix}], \dots, [\begin{matrix} q_{d - 2} \\ q_{d - 1} \end{matrix}]

第二步：计算每个子空间的旋转角度 $θ_{i}$

RoPE 为不同的二维子空间分配不同的旋转基数（类似于 Transformer 原始正弦位置编码）。

对于第 $i$ 个二维子空间（ $i \in [0, 1, \dots, \frac{d}{2} - 1]$ ），其旋转角速度 $θ_{i}$ 计算公式为：

θ_{i} = {base}^{- \frac{2 i}{d}} (通常 base = 10000)

在位置 $m$ 处，该子空间实际旋转的角度为：

ω_{i} = m θ_{i}

第三步：应用旋转变换（矩阵相乘）

对于第 $i$ 个二维子空间 $[\begin{matrix} q_{2 i} \\ q_{2 i + 1} \end{matrix}]$ ，乘以其对应的二维旋转矩阵 $R_{m}^{(i)}$ ：

[\begin{matrix} {\tilde{q}}_{2 i} \\ {\tilde{q}}_{2 i + 1} \end{matrix}] = [\begin{matrix} \cos (m θ_{i}) & - \sin (m θ_{i}) \\ \sin (m θ_{i}) & \cos (m θ_{i}) \end{matrix}] [\begin{matrix} q_{2 i} \\ q_{2 i + 1} \end{matrix}]

展开后的计算结果为：

${\bar{q}}_{2 i} = q_{2 i} \cos (m θ_{i}) - q_{2 i + 1} \sin (m θ_{i})$
${\bar{q}}_{2 i + 1} = q_{2 i} \sin (m θ_{i}) + q_{2 i + 1} \cos (m θ_{i})$

将所有子空间旋转后的结果重新拼接，就得到了带有位置 $m$ 信息的全维度向量 ${\tilde{q}}_{m}$ 。

4.3 整体矩阵表示

如果将上述过程写成一个整体的大矩阵乘以 $q_{m}$ ，RoPE 的全局变换矩阵 $R_{Θ, m}^{d}$ 是一个分块对角矩阵：

{\tilde{q}}_{m} = R_{Θ, m}^{d} q_{m} = [\begin{matrix} \cos m θ_{0} & - \sin m θ_{0} & 0 & 0 & \dots & 0 & 0 \\ \sin m θ_{0} & \cos m θ_{0} & 0 & 0 & \dots & 0 & 0 \\ 0 & 0 & \cos m θ_{1} & - \sin m θ_{1} & \dots & 0 & 0 \\ 0 & 0 & \sin m θ_{1} & \cos m θ_{1} & \dots & 0 & 0 \\ ⋮ & ⋮ & ⋮ & ⋮ & ⋱ & ⋮ & ⋮ \\ 0 & 0 & 0 & 0 & \dots & \cos m θ_{\frac{d}{2} - 1} & - \sin m θ_{\frac{d}{2} - 1} \\ 0 & 0 & 0 & 0 & \dots & \sin m θ_{\frac{d}{2} - 1} & \cos m θ_{\frac{d}{2} - 1} \end{matrix}] [\begin{matrix} q_{0} \\ q_{1} \\ q_{2} \\ q_{3} \\ ⋮ \\ q_{d - 2} \\ q_{d - 1} \end{matrix}]

4.4 高效的工程实现

在代码实现中，直接进行大型稀疏矩阵乘法效率极低。为了加速，通常会利用逐元素相乘（Element-wise Product）的 trick。

我们将原向量 $q_{m}$ 复制并做一个"交错取反"的操作，定义一个辅助向量 $q_{m}^{half}$ ：

q_{m}^{half} = [- q_{1}, q_{0}, - q_{3}, q_{2}, \dots, - q_{d - 1}, q_{d - 2}]^{T}

同时，我们将所有的 $\cos (m θ_{i})$ 和 $\sin (m θ_{i})$ 复制扩展成与 $q$ 维度相同的 $d$ 维向量：

$C_{m} = [\cos (m θ_{0}), \cos (m θ_{0}), \cos (m θ_{1}), \cos (m θ_{1}), \dots]^{T}$
$S_{m} = [\sin (m θ_{0}), \sin (m θ_{0}), \sin (m θ_{1}), \sin (m θ_{1}), \dots]^{T}$

最终的 RoPE 计算可以精简为两组向量的哈达玛积（Hadamard Product，即对应位置直接相乘）之和：

{\tilde{q}}_{m} = q_{m} ⊙ C_{m} + q_{m}^{half} ⊙ S_{m}

提示： 在很多主流模型（如 LLaMA）的代码实现中，为了内存对齐和计算方便，二维切分的方式不是 $[0, 1], [2, 3]$ 邻近切分，而是前后对半切分（前 $\frac{d}{2}$ 维与后 $\frac{d}{2}$ 维两两配对）。虽然切分方式不同，但数学本质完全一致。

4.5 为什么能表示相对位置？

当计算位置 $m$ 的 Query 和位置 $n$ 的 Key 的点积时：

⟨ {\tilde{q}}_{m}, {\tilde{k}}_{n} ⟩ = (R_{Θ, m}^{d} q_{m})^{T} (R_{Θ, n}^{d} k_{n}) = q_{m}^{T} (R_{Θ, m}^{d})^{T} R_{Θ, n}^{d} k_{n}

由于旋转矩阵是正交矩阵，满足 $(R_{m})^{T} = R_{- m}$ ，且矩阵乘法满足角度相加减：

(R_{Θ, m}^{d})^{T} R_{Θ, n}^{d} = R_{Θ, - m}^{d} R_{Θ, n}^{d} = R_{Θ, n - m}^{d}

因此：

⟨ {\tilde{q}}_{m}, {\tilde{k}}_{n} ⟩ = q_{m}^{T} R_{Θ, n - m}^{d} k_{n}

结论： 经过 RoPE 编码后，Query 和 Key 的点积结果只取决于它们的相对距离 $n - m$ （或 $m - n$ ），这完美实现了相对位置编码的效果，同时保留了绝对位置的线性变换能力。

4.6 代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51import torch
import math

class RoPEPositionalEncoding(torch.nn.Module):
    """
    Rotary Position Embedding
    论文来源：Su et al., 2022
    """
    def __init__(self, dim: int, base: int = 10000):
        super(RoPEPositionalEncoding, self).__init__()
        self.dim = dim
        self.base = base
        # 预计算旋转角度 inv_freq[i] = theta_i
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, seq_len: int, device: torch.device = None):
        """
        生成旋转角度矩阵
        :param seq_len: 序列长度
        :param device: 计算设备
        :return: 旋转角度矩阵 [seq_len, dim/2]
        """
        if device is None:
            device = self.inv_freq.device
        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        # 外积得到 [seq_len, dim/2]
        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)  # [seq_len, dim]
        return emb

    @staticmethod
    def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
        """
        应用旋转编码到输入向量
        :param x: 输入向量 [batch, num_heads, seq_len, dim]
        :param cos: 余弦值 [seq_len, dim]
        :param sin: 正弦值 [seq_len, dim]
        :return: 旋转后的向量
        """
        # 将 cos/sin 扩展为与 x 相同的维度
        cos = cos.unsqueeze(0).unsqueeze(0)  # [1, 1, seq_len, dim]
        sin = sin.unsqueeze(0).unsqueeze(0)
        # 旋转公式：x_rot = x * cos + rotate(x) * sin
        # 其中 rotate(x) = [-x[..., 1::2], x[..., ::2]]
        x1 = x[..., :x.size(-1)//2]
        x2 = x[..., x.size(-1)//2:]
        return torch.cat([
            x1 * cos[..., :x1.size(-1)] - x2 * sin[..., :x2.size(-1)],
            x1 * sin[..., :x1.size(-1)] + x2 * cos[..., :x2.size(-1)]
        ], dim=-1)

5. YaRN 位置编码

YaRN（Yet another RoPE extensioN method）是目前大语言模型（LLM）用于**长文本上下文扩展（Context Window Extension）**最先进、效果最好的技术之一。

它的核心痛点在于：标准的 RoPE 无法直接推广到比预训练更长的序列（会外推失败）。YaRN 基于"NTK-by-parts"插值法，通过对不同特征维度（频率）进行分段非线性缩放，并引入温度修正，使得模型只需极少的数据微调，就能将上下文窗口扩展 16 倍以上（例如从 4k 扩展到 64k/128k）。

5.1 核心理论：频率分段

RoPE 的本质是不同维度有着不同的旋转"波长"。YaRN 的核心观察是：不同波长的维度，在扩展上下文时扮演的角色完全不同。

若预训练的最大长度为 $L$ ：

高频维度（波长 $λ < L$ ）：这些维度在预训练的窗口内已经旋转了不止一圈，它们主要编码局部绝对位置和紧密相对关系。如果盲目对它们进行压缩（插值），会破坏原有的局部感知力。因此，YaRN 对高频区域不压缩（保留外推）。
低频维度（波长 $λ > L$ ）：这些维度在整个预训练窗口内连一圈都没转完，其点积不具备跨越全周期的相对位置区分度，主要编码全局宏观框架。如果直接外推，模型会遇到未见过的角度组合导致崩塌。因此，YaRN 对低频区域进行完整的线性压缩（插值）。
中频维度：处于两者之间，YaRN 采用平滑过渡（Smooth Blend）。

5.2 完整的计算步骤

假设我们要将上下文从原长度 $L$ 扩展到新长度 $L^{'} = s \cdot L$ （其中 $s$ 为扩展倍率，如 $s = 4$ 或 $s = 16$ ）。

第一步：计算原始的波长 $λ_{i}$

对于特征维度为 $d$ 的每个二维子空间索引 $i \in [0, 1, \dots, \frac{d}{2} - 1]$ ，原始的角速度和波长分别为：

θ_{i} = {base}^{- \frac{2 i}{d}}

λ_{i} = \frac{2 π}{θ_{i}} = 2 π \cdot {base}^{\frac{2 i}{d}}

第二步：定义分段边界（超参数 $α$ 和 $β$ ）

YaRN 引入了两个超参数 $α$ 和 $β$ 来界定频率：

当 $λ_{i} < α L$ 时，判定为高频（纯外推）。
当 $λ_{i} > β L$ 时，判定为低频（纯插值）。

为了方便计算，通常将其转化为维度索引边界（即根据上面的波长公式反推出具体的维度 $i$ ）：

高频边界对应的维度索引 $r_{fast}$ ：

r_{fast} = \frac{d}{2} \cdot \frac{\ln (α L / 2 π)}{\ln (base)}

低频边界对应的维度索引 $r_{slow}$ ：

r_{slow} = \frac{d}{2} \cdot \frac{\ln (β L / 2 π)}{\ln (base)}

第三步：计算渐变比例函数 $γ_{i}$

对于当前的维度索引 $i$ ，引入一个斜坡函数（Ramp Function） $γ_{i}$ ，用于在插值与外推之间做平滑过渡：

γ_{i} = {\begin{cases} 0 & if i < r_{fast} (高频，完全外推) \\ 1 & if i > r_{slow} (低频，完全插值) \\ \frac{i - r_{fast}}{r_{slow} - r_{fast}} & otherwise (中频，线性渐变) \end{cases}

第四步：计算 YaRN 的新频率（角速度） $θ_{i}^{'}$

利用 $γ_{i}$ ，将原始频率 $θ_{i}$ 与线性缩放后的频率 $\frac{θ_{i}}{s}$ 进行加权混合。YaRN 改进的实际修正频率公式为：

θ_{i}^{'} = (1 - γ_{i}) \cdot θ_{i} + γ_{i} \cdot \frac{θ_{i}}{s}

解读： 当 $γ_{i} = 0$ 时， $θ_{i}^{'} = θ_{i}$ （保留原样不压缩）；当 $γ_{i} = 1$ 时， $θ_{i}^{'} = θ_{i} / s$ （频率变慢，完美容纳 $s$ 倍长文本）。

第五步：注意力温度修正（Attention Temperature Scaling）

这是 YaRN 相比之前方法的另一大关键创新。当位置编码被压缩后，不同位置向量之间的平均点积值会变小（因为角度靠得更近了）。这会导致 Attention 矩阵的熵增加，分布变得极其平缓，从而丧失了注意力集中的能力。

为了抵消这一副作用，YaRN 引入了一个注意力温度放大系数 $t$ ：

t = 0.1 + \frac{1.16 \cdot \ln (s)}{\ln (10)}

在计算应用完位置编码的 Query ( $q$ ) 和 Key ( $k$ ) 的点积时，通过乘以系数 $\sqrt{\frac{1}{t}}$ （或者直接将对应的 $q, k$ 向量除以 $\sqrt{t}$ ）来放大点积差异，强行恢复注意力矩阵的动态范围：

Attention Score = Softmax (\frac{1}{t} \cdot \frac{q^{T} k}{\sqrt{d_{head}}})

5.3 最终矩阵应用

最终，对于任何一个绝对位置为 $m$ 的 Token，其对应的第 $i$ 个二维子空间，所采用的旋转矩阵不再是标准的 $m θ_{i}$ ，而是替换为了：

[\begin{matrix} {\tilde{q}}_{2 i} \\ {\tilde{q}}_{2 i + 1} \end{matrix}] = \frac{1}{\sqrt{t}} [\begin{matrix} \cos (m θ_{i}^{'}) & - \sin (m θ_{i}^{'}) \\ \sin (m θ_{i}^{'}) & \cos (m θ_{i}^{'}) \end{matrix}] [\begin{matrix} q_{2 i} \\ q_{2 i + 1} \end{matrix}]

5.4 总结与意义

YaRN 位置编码计算流程的伟大之处在于：

局部高精度：保证了相邻 Token 之间的短距离高频相对位置信息完全不失真。
全局自适应：长距离的信息通过动态插值被优雅地压缩在新窗口内。
无痛扩展：通过温度因子 $t$ 解决了点积分布涣散的工程难题。

这也使得 LLaMA-2/3、Mistral 等主流开源模型在需要扩展长文本时，普遍首选基于 YaRN 变体的解决方案。

5.5 代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44import torch
import math

class YaRNPositionalEncoding(torch.nn.Module):
    """
    Yet another RoPE extensioN
    论文来源：Peng et al., 2023
    """
    def __init__(self, dim: int, base: int = 10000, original_max_position_embeddings: int = 2048,
                 scaling_factor: float = 1.0):
        super(YaRNPositionalEncoding, self).__init__()
        self.dim = dim
        self.base = base
        self.scaling_factor = scaling_factor
        # 原始训练的最大长度
        self.original_max_position_embeddings = original_max_position_embeddings
        # 重新计算缩放后的 base
        self.base = base * (scaling_factor ** (dim / (dim - 2)))
        # 预计算旋转角度
        inv_freq = 1.0 / (self.base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, seq_len: int, device: torch.device = None):
        """生成旋转角度矩阵（支持缩放）"""
        if device is None:
            device = self.inv_freq.device
        # 对位置索引进行缩放
        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        t = t / self.scaling_factor  # 线性缩放
        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        return emb

    @staticmethod
    def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
        """应用旋转编码（与 RoPE 相同）"""
        cos = cos.unsqueeze(0).unsqueeze(0)
        sin = sin.unsqueeze(0).unsqueeze(0)
        x1 = x[..., :x.size(-1)//2]
        x2 = x[..., x.size(-1)//2:]
        return torch.cat([
            x1 * cos[..., :x1.size(-1)] - x2 * sin[..., :x2.size(-1)],
            x1 * sin[..., :x1.size(-1)] + x2 * cos[..., :x2.size(-1)]
        ], dim=-1)

6. 总结

本文综述了位置编码技术的发展历程，从最早的可学习位置编码（Bahdanau et al., 2014），到 Transformer 的标准方案 Sinusoidal（Vaswani et al., 2017），再到近年来针对长序列优化的相对位置编码（RoPE、YaRN），研究者们始终在解决一个核心问题：如何让模型有效感知序列中的位置信息。

方法	论文来源	类型	外推能力
Learnable	Bahdanau et al., 2014	绝对	有限
Sinusoidal	Vaswani et al., 2017	绝对	理论支持
RoPE	Su et al., 2022	相对	强
YaRN	Peng et al., 2023	相对	最强

当前主流大语言模型普遍采用 RoPE 或 YaRN 作为位置编码方案。YaRN 通过 NTK 感知缩放技术，在长文本处理方面展现出显著优势，已成为大模型长上下文扩展的推荐方案。