Published on

LLM 长上下文扩展之路:Transformer 位置编码的持续演进

Authors
Table of Contents

positional_encoding_clean.png

为什么是 sin 和 cos?

Transformer 模型在处理序列数据时取得了革命性的成功,其核心机制之一是自注意力(Self-Attention),它允许模型在处理序列中的每个元素时,都能考虑到序列中所有其他元素的信息。然而,自注意力机制本身是位置无关的,这意味着它无法区分序列中词语的顺序。为了解决这个问题,Transformer 引入了位置编码 (Positional Encoding)

在原始的 Transformer 论文中,采用的是一种特殊的绝对位置编码 (Absolute Positional Encoding)。它的核心思想是:

  1. 利用不同频率的 sin 和 cos 函数来生成一个与词嵌入 (Word Embedding) 维度相同的向量 PE。
  2. 这个 PE 向量会被直接加到输入词嵌入 xx 上,形成带有位置信息的新嵌入 x=x+PEx' = x + PE

位置编码 PE(pos,k)PE(pos, k) 的计算方式如下:

  1. 偶数维度 (k=2ik = 2i) 的位置编码: PE(pos,2i)=sin(pos100002idmodel)PE(pos, 2i) = sin\left(\frac{pos}{10000^{\frac{2i}{d_{model}}}}\right)

  2. 奇数维度 (k=2i+1k = 2i+1) 的位置编码: PE(pos,2i+1)=cos(pos100002idmodel)PE(pos, 2i+1) = cos\left(\frac{pos}{10000^{\frac{2i}{d_{model}}}}\right)

其中:

  • pospos 是序列中的位置(从 0 开始)。
  • ii 是向量维度的索引,它决定了当前维度对的频率,取值范围为 0,1,,dmodel/210, 1, \dots, d_{model}/2 - 1
  • dmodeld_{model} 是词嵌入的维度(即模型的隐藏层大小)。

这个公式中的 100002idmodel10000^{\frac{2i}{d_{model}}} 项,实际上是为每个维度对 (2i,2i+1)(2i, 2i+1) 定义了一个独特的频率 ωi=1100002idmodel\omega_i = \frac{1}{10000^{\frac{2i}{d_{model}}}}。通过这种方式,不同维度的位置编码能够捕捉到不同尺度的位置信息。

positional_encoding_labeled.png 相关可视化代码见文末附录。

那么,为什么 Transformer 要选择这种基于sin和cos函数的周期性方法,而不是简单地使用一个递增的数字(例如 1, 2, 3, …)来代表位置呢?这种函数编码提供了哪些关键的数学特性?

序列长度可扩展性

首先,一个直观但重要的原因是可扩展性 (Extensibility)。我们希望模型能够处理任意长度的序列,即使这个长度超过了在训练过程中见过的最大长度。如果仅仅使用递增的数字,当序列长度超出预设的最大值时,我们就无法为新的位置提供有效的编码。

而 sin 和 cos 函数是周期性的,并且可以为任何位置 pospos 生成一个确定的 PE 向量。这意味着无论序列有多长,我们总能为每个位置生成一个独特且一致的位置编码,从而完美地解决了处理任意长度序列的问题。

通过将这些不同频率(不同尺度)的波形信号组合在一个向量里,就为每个位置生成了一个独一无二、且信息丰富的“坐标”,远比简单地用 0, 1, 2, 3... 这样的数字来表示位置要强大得多。

隐式地编码相对位置信息

除了可扩展性,使用sin和cos函数还有一个更关键的数学原因,这直接关系到 Transformer 的自注意力机制

尽管我们将 PE 的值加到了词嵌入上,但 sin 和 cos 的巧妙设计让模型能够隐式地学习到相对位置信息。任意位置 pos+k 的位置编码的每个维度上的值,都可以通过一个固定的线性变换从 pos 位置编码的对应维度值中计算出来。这对于注意力机制计算两个词之间的相对距离至关重要。

回想一下三角函数中的角度差cos公式cos(αβ)=cosα×cosβ+sinα×sinβcos(\alpha - \beta) = cos\alpha \times cos\beta + sin\alpha \times sin\beta

在原始 Transformer 的位置编码中,PE 的每个维度 ii 上的值,都可以看作是 posωipos \cdot \omega_i 的 sin 或 cos 值(其中 ωi\omega_i 是一个预设的频率,也就是 100002idmodel10000^{\frac{2i}{d_{model}}})。当计算注意力得分时,Query 向量 QQ 和 Key 向量 KK 都包含了这种形式的绝对位置编码。

为了理解为什么 QKQ \cdot K 的点积会自然地产生类似于 cos(αβ)cos(\alpha - \beta) 的项,我们首先简化问题,暂时忽略线性层 WQ,WKW_Q, W_K 以及词嵌入 xx 的影响,只关注位置编码自身的点积PE1PE2PE_1 \cdot PE_2

原始的 PE 设计是将词嵌入维度 dmodeld_{model} 分成 dmodel/2d_{model}/2 对(偶数和奇数维度)。每一对维度 (2i2i2i+12i+1) 都使用相同的频率 ωi\omega_i,分别应用 sin 和 cos 函数。

让我们关注其中任意一对维度的点击:

维度位置 pos1pos_1 (PE 1)位置 pos2pos_2 (PE 2)
2i2isin(α)sin(\alpha)sin(β)sin(\beta)
2i+12i+1cos(α)cos(\alpha)cos(β)cos(\beta)

其中 α=pos1ωi\alpha = pos_1 \cdot \omega_iβ=pos2ωi\beta = pos_2 \cdot \omega_i

计算这两个维度上的点积: Contribution=[sin(α)sin(β)]+[cos(α)cos(β)]Contribution = [sin(\alpha) \cdot sin(\beta)] + [cos(\alpha) \cdot cos(\beta)]

根据我们回顾的角度差cos公式:这使得 PE1PE2PE_1 \cdot PE_2 的这一小部分贡献正好等于 cos(αβ)cos(\alpha - \beta)

由于 αβ=(pos1ωi)(pos2ωi)=ωi(pos1pos2)\alpha - \beta = (pos_1 \cdot \omega_i) - (pos_2 \cdot \omega_i) = \omega_i \cdot (pos_1 - pos_2),这个贡献只与两个位置的相对距离 pos1pos2pos_1 - pos_2 有关(同一向量维度上的频率 ωi\omega_i 是一样的)。

总结来说,传统的 sin-cos 位置编码通过将 sin 和 cos 函数成对放置,使得 PE 向量的自身点积自然地将绝对位置信息(pos1pos_1pos2pos_2)转换成了与相对距离(pos1pos2pos_1 - pos_2)相关的项。换句话说,PE(pos+k)PE(pos+k) 的每个维度上的值都可以表示为 PE(pos)PE(pos) 对应维度上的 sin 和 cos 值的线性组合。

这使得 Transformer 模型在计算注意力得分时,能够隐式地捕捉并利用词语之间的相对位置关系。

传统 PE 的局限性

尽管传统的位置编码设计巧妙,但这种“加在词向量上,然后让模型隐式地从点积中学习相对位置”的方式,也存在一些潜在的局限性:隐式的学习还是容易丢失信息,这正是传统 PE 的核心问题,也是 Rotary Position Embeddings (RoPE) 想要解决的痛点,这些局限性促使了后续研究中更先进的位置编码方法(如 RoPE)的出现,以期更明确、更有效地融入相对位置信息。

尽管传统 PE 的三角函数设计非常巧妙,但它有以下几个主要限制,都源于它的隐式性

  1. 信息耦合与混淆 (Coupling): 传统的做法是简单地将位置向量 PE 到词向量 xx 上 (x=x+PEx′=x+PE)。这使得词汇内容信息位置信息紧密地耦合在一起,模型必须通过后续的线性层 WQWQ ​,WKWK ​ 来学习如何将两者解耦 (disentangle)。这个解耦过程本身增加了模型的学习负担,并可能导致位置信息的丢失或混淆。
  2. 缺乏对相对位置的直接建模: 模型拿到的是 xpos1​​+PE1x_{pos1}​​+PE_1​xpos2​​+PE2x_{pos2}​​+PE_2 ​。它必须通过点积 QKQ⋅K 间接地推导出 pos1​−pos2​ 这种相对关系。这不如直接将相对位置编码到向量中效率高。
  3. 长序列泛化性差 (Extrapolation): 当序列长度远超训练集时,模型很难保证它在训练中学习到的那种复杂的、间接的相对位置关系在新长度下仍然有效。

RoPE:用旋转实现显式建模

rope_rotated_clean.png

RoPE (Rotary Position Embeddings) 的核心思想在于显式地将相对位置信息编码到 Query 和 Key 向量中,并通过旋转操作确保自注意力机制的点积结果只依赖于词语之间的相对距离,而非绝对位置。

这与传统的位置编码形成了鲜明对比。传统方法是将位置编码直接到词嵌入上,然后期望模型通过学习来隐式地从点积中提取相对位置信息。而 RoPE 则通过旋转的方式,更直接、更优雅地解决了这个问题。

核心思想

  1. 告别加法,拥抱旋转: 传统 PE 将位置信息通过加法融入词嵌入 x=x+PEx' = x + PE。RoPE 摒弃了这种方式,而是对 Query 和 Key 向量本身进行基于位置的旋转
  2. 显式建模相对位置: RoPE 的设计目标是让 QQKK 向量在经过位置 mmnn 的旋转后,其点积 QmTKnQ_m^T K_n 能够直接表示为只与相对距离 mnm-n 相关的函数。
  3. 利用 2D 旋转的特性: RoPE 的灵感来源于 2D 平面上的旋转操作,它将词向量的维度两两分组,在每个 2D 子空间中应用旋转。这种旋转操作天然地与三角函数相关联。

image.png

2D 旋转矩阵

RoPE 将 dmodeld_{model} 维度的 QQKK 向量拆分成 dmodel/2d_{model}/2 个 2D 向量小块。对于每个小块,应用一个基于位置 mm 和特定频率 ωi\omega_i 的旋转操作:

RmQ=[cos(mω)sin(mω)sin(mω)cos(mω)][q0q1]R_m Q = \begin{bmatrix} \cos(m\omega) & -\sin(m\omega) \\ \sin(m\omega) & \cos(m\omega) \end{bmatrix} \begin{bmatrix} q_0 \\ q_1 \end{bmatrix}

最终 QQKK 向量的完整旋转就是对所有 2D 小块的独立旋转的集合。

RmR_m 是一个2D 旋转矩阵 (2D Rotation Matrix),它的作用是将一个 2D 向量在二维平面上绕原点旋转一个特定的角度。在 RoPE 中,RmR_m 表示对一个 2D 向量进行与位置 mm 相关的旋转。这里的旋转角度是 mωm\omega(其中 ω\omega 是一个预设的频率,它会因维度对而异,通常写作 ωi\omega_i)。

NOTE

这个矩阵的形式是数学上标准且唯一的 2D 旋转矩阵。

一个 2D 向量 (x,y)(x, y) 绕原点逆时针旋转角度 θ\theta 后的新坐标 (x,y)(x', y') 可以通过以下矩阵乘法得到:

[xy]=[cosθsinθsinθcosθ][xy]\begin{bmatrix} x' \\ y' \end{bmatrix} = \begin{bmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{bmatrix} \begin{bmatrix} x \\ y \end{bmatrix}

将这个公式应用到 RoPE 中:

  • 旋转角度 θ\theta 在 RoPE 中,旋转角度是由当前位置 mm 和一个频率 ω\omega 决定的,即 θ=mω\theta = m\omega。这个频率 ω\omega 确保了不同维度对的旋转角度不同,从而能够编码不同尺度的位置信息。
  • 被旋转的向量: 在 RoPE 中,被旋转的是 Query 向量 QQ(或 Key 向量 KK)的某个 2D 小块。

这个矩阵的特性是,它是一个正交矩阵1,并且其行列式为 12。正交矩阵在向量点积计算中具有非常好的性质,它能保持向量的长度不变,并且在计算点积时能自然地引入角度差(相对位置)信息,这正是 RoPE 所需要的。

在 RoPE 的上下文中,Query 向量 QQ(以及 Key 向量 KK)的维度是 dmodeld_{model}。RoPE 的处理方式是将这个 dmodeld_{model} 维的向量两两分组

q0q_0q1q_1 代表了 Query 向量 QQ 中某一对连续的维度分量。

例如,如果 QQ 向量是 [Q0,Q1,Q2,Q3,,Qdmodel1][Q_0, Q_1, Q_2, Q_3, \dots, Q_{d_{model}-1}],那么第一组就是 (Q0,Q1)(Q_0, Q_1),此时 q0=Q0,q1=Q1q_0 = Q_0, q_1 = Q_1

RoPE 对 QQ 向量的每个这样的 2D 小块独立地应用一个 2D 旋转矩阵。每个 2D 小块会使用一个不同的频率 ωi\omega_i 来计算旋转角度 mωim\omega_i,从而实现对整个 QQ 向量的“分块旋转”。

这种分块旋转的设计,使得 RoPE 能够将高维向量的位置编码问题分解为多个独立的 2D 旋转问题,并且每个 2D 旋转都能够显式地编码相对位置信息。

rope_rotated_labeled.png 相关可视化代码见文末附录。

点积与相对距离

RoPE 的数学优越性在于它能通过矩阵代数,证明旋转后的点积只取决于位置差 nmn-m

我们从注意力得分的点积开始:

Score(m,n)=QmKn=(RmQ)T(RnK)Score(m, n) = Q'_m \cdot K'_n = (R_m Q)^T (R_n K)

利用旋转矩阵 RR 的两个关键性质:

  1. 转置是逆旋转:RmT=RmR_m^T = R_{-m}
  2. 连续旋转相加:RaRb=Ra+bR_a R_b = R_{a+b}

代入公式进行化简:

Score(m,n)=QmKn=(RmQ)T(RnK)=QTRmTRnK=QT(RmRn)K=QT(Rnm)K\begin{aligned} Score(m, n) &= Q'_m \cdot K'_n \\ &= (R_m Q)^T (R_n K) \\ &= Q^T R_m^T R_n K \\ &= Q^T (R_{-m} R_n) K \\ &= Q^T (R_{n-m}) K \end{aligned}

结论:旋转后的 QQKK 的点积,完美地等于 QQ 和一个只旋转了相对距离 nmn-mKK 的点积。这在数学上保证了 RoPE 实现了线性地依赖于相对位置的注意力机制。

传统 PE 必须在训练时看到绝对位置 m 的编码 PEmPE_m ​,然后学习如何在点积中推导出相对关系。一旦遇到训练中没见过的绝对位置(例如序列长度 L 以外的 L+1),模型将遇到完全陌生的 PE 模式,导致性能急剧下降。

RoPE 的注意力得分取决于相对距离 k=n−m,而不取决于绝对位置 m 或 n。因为 RkR_k​ 由 sin 和 cos 函数定义,对于任何距离 k 都有明确的数学定义。模型只需泛化它在训练中看到的相对距离关系,而不是泛化全新的绝对位置。这使得 RoPE 能够更自然、更可靠地将学到的相对位置关系泛化(外推)到比训练时更长的序列上。

RoPE 已经成为现代大型语言模型(LLM)的标准配置之一,LLaMA、Qwen、DeepSeek 以及 Mistral、Gemma 等主流架构都使用 RoPE。

RoPE 长上下文扩展之路:从 PI 到 YaRN

尽管 RoPE 在理论上具有无限的外推能力,但在实际应用中,当我们尝试将模型的上下文窗口从训练时的 2K 扩展到 16K 或 32K 时,仍然会遇到性能下降,其无法泛化至超出训练时所见的上下文窗口范围

RoPE 的数学原理保证了它只依赖相对距离,但问题出在模型训练的数据范围上:

  1. 频率与训练范围: 尽管 sin / cos 函数是连续的,但模型只在相对距离 kk 小于训练上下文长度 LtrainL_{train} 的范围内见过 RkR_k 的旋转模式。
  2. 高频混沌: 在 RoPE 的高频维度(对应短距离关系)上,当计算一个很长的相对距离 klongk_{long} 时,旋转角度 θ=klongωhigh\theta = k_{long} \cdot \omega_{high} 会变得非常大。
  3. 性能下降: 这导致了 sin 和 cos 产生模型在训练中从未见过的、高度振荡的旋转模式。模型无法可靠地将这些“陌生”的模式与准确的注意力得分联系起来,从而导致长序列性能急剧下降。

解决这个问题(RoPE 长上下文扩展性能下降)的核心思路是:对 RoPE 使用的频率 (ωi\omega_i) 进行缩放 (Scale Down)。

目的是将 RoPE 的频率 ωi\omega_i 降低,使得当我们处理一个新的长距离时,由此产生的旋转角度 θnew\theta_{new} 能够映射回模型在训练时看到的相似角度 θtrain\theta_{train}

线性插值 (PI)

论文提出 Position Interpolation(PI),将输入位置索引按比例下调,使最长位置回到模型预训练的原始上限,从而在计算 RoPE 时始终处于已训练范围内。相比直接外推,这种“插值”避免注意力分数失控,理论上其上界比外推小约 600 倍,显著更稳定。

原理: 假设训练长度 LtrainL_{train},目标长度 LnewL_{new},缩放因子 s=Lnew/Ltrains = L_{new} / L_{train}。线性缩放会统一地将所有频率 ωi\omega_i 缩放 1/s1/sωi=ωi/s\omega_i' = \omega_i / s image.png

上面是位置插值(Position Interpolation)方法的示意图。

假设有一个 Llama 模型,它在训练时使用的上下文窗口长度是2048。

  • 左上角的图 展示了LLM模型的正常使用情况:输入的位置索引(蓝色点)都在模型预训练的范围之内。
  • 右上角的图 展示了长度外推(length extrapolation)的情况,这时模型需要处理一些它在训练时从未见过的位置(红色点),例如将上下文长度扩展到4096。
  • 左下角的图 展示了位置插值(Position Interpolation)的方法:我们直接将位置索引(蓝色和绿色点)从 [0, 4096] 的范围缩小到 [0, 2048],从而强制它们回到模型预训练的范围内。

所有的相对距离都被均匀压缩。该方法简单、架构无改动、工程易落地,能有效地防止注意力分数爆炸。PI 以极低成本将 RoPE 模型上下文扩展至 32k,同时保留短上下文能力。其稳定性源于插值而非外推的理论优势,实践中微调≤1000步即可获得显著长序列收益。

但它也等比例地降低了高频维度(负责短距离、细粒度语法信息的维度)的频率。这使得模型在处理长文本时,短距离关系的分辨率被降低,从而可能影响其在局部语法和语义上的准确性。

非线性缩放 (YaRN)

像 RoPE 这样的嵌入可看作一组频率特征,短距离语法依赖高频信息,而长距离结构依赖低频信息。在 NTK 视角下,模型更容易学习低频、难以捕捉高频。如果把所有维度等比例拉伸,可能削弱高频分辨率,从而损害近邻相对位置的识别,描述最小距离的旋转不能太小,否则网络无法检测到

也就是说,模型(尤其是 Transformer 及其 RoPE 变体)在处理长序列时,更容易捕捉和泛化长距离(低频)的粗糙结构,但更容易丢失和难以泛化短距离(高频)的细粒度信息。

这正是 “NTK-aware” 方法通过调整频率基数来保护高频信息的动机。

NOTE

简单理解什么是 Neural Tangent Kernel (NTK):

在训练的开始阶段,网络的学习行为(即参数更新的轨迹)可以被一个固定不变的核函数(Kernel Function)所完全描述,这个函数就称为 Neural Tangent Kernel,NTK 衡量了网络关于其参数的梯度之间的相似度相关性

NTK把一个足够宽的神经网络在梯度下降下的训练过程近似为线性模型在一个“核”上的学习,这个核随参数几乎不变,被称为神经切线核。它定义了输入之间的相似度:相似度越高,训练信号在函数空间里传播得越强。结果是,网络的函数收敛行为可以用核回归精确描述,解释了为何宽网络在小步长训练时呈现“惰性学习”(参数变动小、函数近线性)以及如何优先拟合低频组件。

因此, YaRN (Yet another RoPE extensioN method):在“线性位置缩放”的基础上,按频率与维度“分段”处理并保留高频信息,结合以下三项改进:

  • NTK-aware:改变 RoPE 频率基数,让低频维度缩放得更多、高频维度缩得更少;
  • NTK-by-parts(分段):依据各维波长与原上下文比值,只对低频维度插值、高频维度不插值,中间维度线性过渡;
  • Dynamic Scaling + 温度缩放:推理时随当前序列长度动态更新缩放因子,并对注意力 logits 施加长度相关的温度缩放以抑制过度自信。

NTK-aware

不改架构与注意力代码,仅将 RoPE 的频率基数 1000010000 替换为更大的 bb',让低频维度被更多缩放而高频维度少动:

  • 原始频率定义:θi=100002i/dmodel\theta_i = 10000^{-2i/d_{model}}
  • NTK-aware 基数变换:b=10000sdmodel/(dmodel2)b' = 10000 \cdot s^{d_{model}/(d_{model}-2)}ss 是缩放因子。
  • 替换后频率:θi=b2i/dmodel\theta_i' = b'^{-2i/d_{model}}

直观效果:低频(大波长)维度被更强插值以扩展远距离感知;高频(小波长)维度更少被改变,以维持近邻排序与细粒度位置识别。

NTK-aware 方法的局限性在于它会使部分位置编码维度超出预训练范围,导致微调效果不如 PI,并且实际的上下文扩展比例需要通过经验调整,不能完全依赖理论缩放因子。

NTK-by-parts

NTK‑by‑parts 是按频率“分段”插值 RoPE:低频强插值、高频不插值、中频平滑过渡。

  • 定义每个维度的波长与“相对尺度”: λd=2πθd\lambda_d = \frac{2\pi}{\theta_d}, r(d)=Lλdr(d) = \frac{L}{\lambda_d}, 其中 LL 是原始上下文长度、θd\theta_d 是 RoPE 频率。
  • 直觉:
    • 高频维度(波长短,r(d)r(d) 大)主要用于分辨近邻位置;对它们插值会压缩局部距离、使近邻变“模糊”,因此应尽量不插值。
    • 低频维度(波长长,r(d)r(d) 小)承载长程位置;对它们插值可安全扩大感知范围。

作为纯插值方案,训练时不产生“越界”频率,实践中微调效果优于仅基数变换的 NTK‑aware,通常更利于保持短上下文性能。

Dynamic Scaling

在长上下文扩展方法里都有一个缩放因子 ss 用来把位置或频率“压缩”进预训练范围。Dynamic Scaling(动态缩放)的核心是:不要在整个推理过程中固定一个 s,而是每次前向计算都依据当前序列长度 l′ 重新设置 s,从而让模型在不同长度下都用更合适的插值力度3

  • 固定法 (传统做法): s=L/Ls = L'/L,推理全程恒定,超过训练的 LL' 时性能会突然恶化。
  • 动态法 (推荐): s=max(1,l/L)s = \max(1, l'/L),每次前向用当前序列长度 ll' 相对原始长度 LL 的比值,当长度变长时逐步增加 ss,形成“渐进插值”。

这里 LL 是模型预训练的原始上下文上限 (如 4k),ll' 是当前步的序列长度 (生成时每步 +1)。

为什么要用动态缩放 

  • 平滑退化:长度未达到训练的扩展上限时不过度缩放,接近上限时插值力度才逐步加大,避免在临界点“断崖式”性能下降。
  • 零微调增益:在完全未做长窗微调的原模型(L=LL'=L)上,动态缩放已能显著缓解困惑度爆炸,延长可用推理长度。
  • 通用叠加:它是推理时策略,可与 PI、NTK‑by‑parts、YaRN 叠加;与 NTK‑aware 结合常称 “Dynamic NTK”。

WARNING

在 Transformer 推理阶段,Query 与 Key 向量会先经过线性映射,再按当前位置施加RoPE 旋转,随后参与注意力内积。若使用上下文扩展方法(如PI、NTK‑by‑parts、YaRN、Dynamic Scaling),也是在这一刻对位置或频率按当前序列长度应用相应缩放。做 KV 缓存时应缓存未施加 RoPE 的 K/V,并在每次前向以一致的尺度对新旧向量施加RoPE,避免因尺度变化导致错配。

因为动态 ss 会改变每个 token 的 RoPE 旋转结果,因此必须缓存未施加 RoPE 的 KV,在每次前向时用当下的 ss 对本步和缓存中的 KV 向量施加一致的 RoPE;若缓存的是已 RoPE 的 KV,就无法与新的 ss 对齐,会导致注意力错配。

温度缩放

YaRN 在注意力 softmax 前引入温度 tt,对 Q/K 进行等效缩放以稳定长上下文的熵与困惑度。在超长上下文下,所有位置被“压缩”后,注意力分数分布容易变得偏平或偏陡,导致困惑度波动。通过温度因子 tt 对注意力 logits 做统一缩放,使分布保持在更合适的熵水平,从而稳定性能: softmax(qmTkntD).softmax\left(\frac{q_m^T k_n}{t\sqrt{|D|}}\right).

总的来说,YaRN 是一种对 RoPE 插值的改进方案,几乎零改动即可把上下文扩展到超长并保持原有能力。YaRN 以分段频率插值加注意力温度缩放为核心,在不修改模型架构的前提下替代 PI,稳定扩展至超长上下文同时在基准上保持近似原性能。

总结

Transformer 通过位置编码(PE)解决自注意力机制的位置无关性。原始 PE 利用不同频率的 sin/cos 函数,提供序列长度可扩展性并隐式编码相对位置信息,但存在信息耦合和长序列泛化性差的局限。

旋转位置嵌入(RoPE)通过对 Query/Key 向量进行旋转,显式地将相对位置信息编码到注意力机制中,理论上具有无限外推能力。然而,在实际超长上下文应用中,RoPE 仍面临性能下降,原因是模型未在训练中见过极端长距离的高频旋转模式。

为解决此问题,主要有两种扩展方法:

  1. 线性插值(PI):通过统一缩放位置索引,将长序列“压缩”回训练范围,简单有效但可能牺牲短距离高频信息的精度。
  2. 非线性缩放(YaRN):在 PI 基础上改进,结合了 NTK-aware(调整频率基数保护高频)、NTK-by-parts(按频率分段插值)和动态缩放(根据当前长度调整缩放因子)及温度缩放(稳定注意力分数),旨在更精细地平衡长短距离信息,实现更稳定、高效的超长上下文扩展。

自 Transformer 的提出已经约 8 年了,但在这期间,Transformer 的改进从未停止过,正如其位置编码的不断迭代一样。尽管人们不断呼吁下一个超越 Transformer 模型架构,也有很多学者尝试提出替代性的架构,例如基于 Mamba 等思路的模型,试图在长序列建模、计算效率或记忆机制上取得突破。

但是 Transformer 依然凭借其通用性、可扩展性以及强大的表示能力,稳坐深度学习的头把交椅。如今,“如何取代 Transformer” 与 “如何让 Transformer 更高效、更具归纳偏置” 已成为研究者们思考的两条并行路线。

附录

原始 transformer 位置编码可视化代码:

import numpy as np
import matplotlib.pyplot as plt

def create_positional_encoding(max_seq_len, d_model):
    """
    创建Transformer位置编码矩阵

    Args:
        max_seq_len: 最大序列长度(位置数)
        d_model: 模型维度(向量维度)

    Returns:
        位置编码矩阵,形状为 (max_seq_len, d_model)
    """
    # 初始化位置编码矩阵
    pe = np.zeros((max_seq_len, d_model))

    # 位置索引
    position = np.arange(0, max_seq_len).reshape(-1, 1)

    # 计算分母项
    div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))

    # 偶数位置使用sin,奇数位置使用cos
    pe[:, 0::2] = np.sin(position * div_term)
    pe[:, 1::2] = np.cos(position * div_term)

    return pe

# 创建位置编码矩阵
max_seq_len = 1000  # 位置数量
d_model = 64  # 向量维度

pe = create_positional_encoding(max_seq_len, d_model)

# 创建图形,带坐标轴和标签
fig, ax = plt.subplots(figsize=(14, 10))

# 绘制热力图
im = ax.imshow(pe, cmap='viridis', aspect='auto', interpolation='nearest')

# 设置标题
ax.set_title('Standard Transformer Positional Encoding', fontsize=16, pad=20)

# 设置坐标轴标签
ax.set_xlabel('Vector Index (Dimension)', fontsize=14)
ax.set_ylabel('Position', fontsize=14)

# 添加colorbar
cbar = plt.colorbar(im, ax=ax)
cbar.set_label('Encoding Value', fontsize=12)

# 设置坐标轴刻度
ax.set_xticks(np.arange(0, d_model, 8))
ax.set_yticks(np.arange(0, max_seq_len, 100))

# 保存图片
plt.savefig('positional_encoding_labeled.png',
            bbox_inches='tight',
            dpi=300)

print("带标签的位置编码可视化图片已保存为: positional_encoding_labeled.png")
print(f"矩阵形状: {pe.shape}")
print(f"位置数量: {max_seq_len}, 向量维度: {d_model}")

plt.close()

RoPE 作用在原始的 QK 向量上:

import numpy as np
import matplotlib.pyplot as plt

def create_rope_rotated_vectors(max_seq_len, d_model, base=10000.0):
    """
    创建RoPE旋转后的向量可视化
    这里我们可视化一个初始向量经过RoPE旋转后在不同位置的表现

    Args:
        max_seq_len: 最大序列长度(位置数)
        d_model: 模型维度(向量维度)
        base: 基数,默认10000

    Returns:
        旋转后的向量矩阵
    """
    # 创建一个初始向量(模拟一个token的query/key向量)
    # 使用一个简单的模式:前半部分为1,后半部分为-1
    initial_vector = np.ones(d_model)
    initial_vector[d_model//2:] = -1

    # 初始化结果矩阵
    rotated_vectors = np.zeros((max_seq_len, d_model))

    # 计算频率
    dim_indices = np.arange(0, d_model, 2)
    freqs = 1.0 / (base ** (dim_indices / d_model))

    # 对每个位置应用RoPE旋转
    for pos in range(max_seq_len):
        # 复制初始向量
        vec = initial_vector.copy()

        # 对每一对维度应用旋转
        for i in range(0, d_model, 2):
            freq = freqs[i // 2]
            angle = pos * freq

            # 提取这一对维度
            x = vec[i]
            y = vec[i + 1]

            # 应用旋转矩阵
            vec[i] = x * np.cos(angle) - y * np.sin(angle)
            vec[i + 1] = x * np.sin(angle) + y * np.cos(angle)

        rotated_vectors[pos] = vec

    return rotated_vectors

# 创建RoPE旋转后的向量矩阵
max_seq_len = 1000
d_model = 64

rope_rotated = create_rope_rotated_vectors(max_seq_len, d_model)

# 创建图形,带坐标轴和标签
fig, ax = plt.subplots(figsize=(14, 10))

# 绘制热力图
im = ax.imshow(rope_rotated, cmap='viridis', aspect='auto', interpolation='nearest')

# 设置标题
ax.set_title('RoPE Applied: Rotated Vector at Different Positions', fontsize=16, pad=20)

# 设置坐标轴标签
ax.set_xlabel('Vector Index (Dimension)', fontsize=14)
ax.set_ylabel('Position', fontsize=14)

# 添加colorbar
cbar = plt.colorbar(im, ax=ax)
cbar.set_label('Rotated Vector Value', fontsize=12)

# 设置坐标轴刻度
ax.set_xticks(np.arange(0, d_model, 8))
ax.set_yticks(np.arange(0, max_seq_len, 100))

# 保存图片
plt.savefig('rope_rotated_labeled.png',
            bbox_inches='tight',
            dpi=300)

print("带标签的RoPE旋转向量可视化图片已保存为: rope_rotated_labeled.png")
print(f"矩阵形状: {rope_rotated.shape}")
print("这个可视化展示了一个初始向量在不同位置经过RoPE旋转后的效果")

plt.close()

欢迎关注微信公众号👏

微信公众号

欢迎微信扫码加入我的付费知识星球👏

知识星球

Footnotes

  1. 正交矩阵是一个方阵,它的列向量(或行向量)都是单位向量,并且它们之间两两相互垂直(正交)

  2. 行列式为 1 意味着这个旋转操作不会改变向量的长度,也不会放大或缩小它所作用的二维空间。

  3. 想起训练阶段的自适应学习率。