如何拓展大模型上下文长度

最近需要将 Qwen2.5 Math 7B 的模型上下文从 4k 扩展到 16k，记录一下是怎么实现的。这里采用的是 NTK aware 的方法，原因是足够简单且可行。

RoPE

Qwen 系列模型使用了 RoPE 位置编码，所以首先回顾一下 RoPE 的方法：对于给定的 Query ($q$) 和 Key ($k$) 向量，它们在送入注意力计算之前，会根据其在序列中的绝对位置 $i$，将每两个维度视为复平面上的一个点，进行旋转。由于旋转的相对性，天然就有相对位置编码的特性。

对一个维度为 $d$ 的向量 $x=[x_0,x_1,\dots,x_{d-1}]$（在序列中的位置为 $m$），将第 $2i$ 和 $2i+1$ 位作为一个二维向量进行旋转：

其中 $\theta_{i}$ 是一个和位置 $i$ 相关的角频率，计算方式为：

$b$ 为基频，在 Qwen2.5 Math 7B 中为 $10000$。在实际代码中，旋转矩阵 $\mathbf{R}_{\mathbf{\Theta},m}^{d}$ 非常稀疏，需要使用如下方式进行计算以提高计算效率：

在 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)
  # ......

对位置在 $m$ 处的 query $q_{m}$ 和位置在 $n$ 处的 key $k_{n}$，计算内积就变成了

的确只和相对位置有关。

对 $\theta_{i}$，当 $i$ 较小时（$x$ 的低维部分），$\theta_{i}\approx 1$，是高频信息（即在 $m\theta_{i}$ 旋转速度更快），用于捕捉短距离注意力关系；而当 $i$ 较大时（$x$ 的高维部分），是低频信息，用于捕捉长距离注意力关系。总的来说就是低维高频、高维低频。至此，对于 RoPE 的理解就足够了，可以开始推导如何扩展上下文长度了。

扩展上下文长度

虽然好像 RoPE 编码的定义使其好像很容易进行上下文扩展（扩展上下文长度就相当于增大 $m$ 的取值范围），但是注意，由于基频是预先定义好的，所以所有维度的频率都是固定的，面对更长的长度，$m$ 的取值范围增大，$m\theta_{i}$ 的变化范围也会增大，多出来的范围就是模型并没有学习过的，对高频短距离信息尤为如此；从另一个视角来看，上下文长度的增加会带来更多的低频长距离信息，这是模型没有学习过的。因此必须做出一些必要的调整以适应增加的低频长距离信息。与此同时，需要注意的是，高频的短距离信息几乎不应该有变化，因此调整时也要注意到这些信息的保存。

由于低频、高频信息和 $\theta_{i}$ 直接相关，而 $\theta_{i}=b^{-2i/d}$ 中的 $d$ 是随模型固定的，$i$ 应该尽量避免变化，所以自然的想法是调整原有的基频 $b$。从上面说到的，上下文长度扩展的程度对低频信息的引入有很大的关联，因此引入一个变量 $s$ 为新、旧最大上下文长度的比值，并定义新的基频

这里 $k$ 是一个待定的常数。带入 $\theta_{i}$ 可得

对高频信息，当 $i\to 0$ 时，$s^{-2ik/d}\to 1$，故 $\theta'_{i}\approx\theta_{i}$，满足之前讨论到的“高频的短距离信息几乎不应该有变化”。

而对低频信息，即当 $i\to \frac{d}{2}$ 时（注意 $i$ 的取值范围是 $[0, d / 2-1]$），由于扩展 $s$ 倍上下文，所以希望能够在低频部分进行 $s^{-1}$ 的线性插值（角频率缩小 $s$ 倍），即希望 $s^{-2ik/d}\approx s^{-1}$，得到

那么也就得到了最终的基频修改公式：

在实际应用中，发现由于 $d$ 通常比较大，即使是取 $b'=b\cdot s$ 也能有不错的效果。并且原作者提到的，对 LLaMA 7B 的模型，你甚至不需要做任何的微调就能适应新的上下文长度！²

扩展 Qwen2.5 Math 7B 上下文长度

回到正题，终于可以扩展上下文长度了！Qwen2.5 Math 7B 的默认参数配置 ³ 为：

{
  "architectures": [
    "Qwen2ForCausalLM"
  ],
  "attention_dropout": 0.0,
  "bos_token_id": 151643,
  "eos_token_id": 151643,
  "hidden_act": "silu",
  "hidden_size": 3584,
  "initializer_range": 0.02,
  "intermediate_size": 18944,
  "max_position_embeddings": 4096,
  "max_window_layers": 28,
  "model_type": "qwen2",
  "num_attention_heads": 28,
  "num_hidden_layers": 28,
  "num_key_value_heads": 4,
  "rms_norm_eps": 1e-06,
  "rope_theta": 10000,
  "sliding_window": 4096,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.44.0",
  "use_cache": true,
  "use_mrope": false,
  "use_sliding_window": false,
  "vocab_size": 152064
}

这里的 rope_theta 就是基频 $b$，由于使用了多头注意力机制，所以实际的向量维度为 hidden_size / num_attention_heads，即 $128$，新旧最大上下文长度之比 $s=16\text{k} / 4\text{k}=4$，那么可以算得

如果取 $s$ 的指数为 $1$，那么

将 max_position_embeddings 和 rope_theta 的值分别修改为 $16384$ 和 $40000$ 即可。

结束……了？

当然没有！实际上，上面所推导的方法正是 NTK Aware 方法 ²，基于这个方法，还有一些其他的变体，例如 Dynamic NTK Interpolation⁴、NTK-by-parts Interpolation⁵、YaRN⁶ 等，可以阅读 这篇博客 来迅速了解这些方法。

另外，上面的推导仅仅只是对 NTK aware 算法的一个近似求解，具体的推导过程可以看下面的内容（对数学要求较高，可以跳过不看）。

神经内切核

首先我们需要了解 NTK（Neural Tangent Kernel，神经内切核）的理论 ⁷。对于一个由参数 $\mathbf{\theta}\in \mathbb{R}^{P}$ 参数化的神经网络 $f(\mathbf{x};\mathbf{\theta})$，其神经内切核 $\mathbf{\Theta}(\mathbf{x},\mathbf{x}')$ 是一个衡量输入 $\mathbf{x}$ 和 $\mathbf{x}'$ 之间关系的核函数，定义为

根据 Jacot 等 ⁷ 的研究，可以知道在无限宽度和无穷小学习率的极限下，神经网络的训练动态可以用这个核函数来精确描述。模型在函数空间中的演化等价于一个核回归问题，其使用的核就是 NTK。因此，NTK 的性质（尤其是其谱特性）决定了模型的学习能力和泛化行为。

NTK Aware 方法

RoPE 核函数

将 RoPE 视为一个特征映射 $\phi:\mathbb{N}\to \mathbb{C}^{d/2}$，将一个位置索引 $m$ 映射到一个 $\frac{d}{2}$ 维的复数向量：

其中 $\theta_{i}=b^{-2i/d}$，$(d / 2)^{-1/2}$ 用于归一化。

然后将注意力机制中 RoPE 的贡献抽象为一个核函数 $K(m,n)$，用于衡量位置 $m$ 和 $n$ 之间的相似性，有

只和 $m-n$ 有关，为了简单，记核函数 $K(m,n)=K(\Delta m)=K(\Delta m; b)$，$f(x)=\cos(\Delta m\cdot x)$，则 $K(\Delta m)=\sum_{i=0}^{d/2-1}f(\theta_{i})$。另外也可以发现，任意两个位置之间的相似性和所有频率 $\{ \theta_{i} \}_{i=0,1,\dots,d/2-1}$ 相关，也就是说，

NOTE

模型内部对位置差异的感知，是由所有频率分量共同贡献的结果

在进行上下文长度扩展前后，希望核函数的性质尽可能保持不变，也就是说

NOTE

在新的、扩展后的上下文中，任意两个位置 $m$ 和 $n$ 的核函数值 $K(\Delta m;b')$，与在原始上下文中、按比例缩放后的位置 $m / s, n / s$ 的核函数值 $K\left(\frac{m}{s}, \frac{n}{s};b \right)$ 尽可能保持一致：

即

这里 $\theta_{i}'=b'^{-2i/d}$。那么问题就转化为如何计算或近似出上式的左右两边（实际上只需要左边就行了），当向量维度 $d$ 足够大时，求和可以转化为积分近似，即：

代入上式得到

谱密度函数

然而，求解这个积分等式非常困难，需要通过其他途径进行分析。将核函数 $K(m,n)$ 生成的核矩阵 $[K]_{mn}$ 视为一个大型矩阵，从随机矩阵理论的视角，这个矩阵的谱特性（如谱密度、谱半径）与核函数的分析性质（如其傅里叶变换）紧密相关。保持谱特性的稳定是实现稳健上下文扩展的关键。我们来计算核函数 $K(m,n)$ 的谱密度 $\hat{K}(\omega;b)$：

代入狄拉克函数

并由于 $\delta(x)=\delta(-x)$，得到

只考虑正频率 $\omega>0$，由于 $b>1$ 且 $0<x<1$，有 $b^{-x}>0$，则 $\delta(\omega+b^{-x})=0$，上式简化为：

然后分析这个谱的支撑集。为了使谱密度 $\hat{K}(\omega;b) > 0$，则必然存在 $x\in[0,1]$，使得 $\omega-b^{-x}=0$，解得 $x=-\ln \omega / \ln b$。为了让 $x\in[0,1]$，可以得到

那么就得到：

IMPORTANT

理想化的 RoPE 核（在无限长序列上）其谱密度函数的支撑集为 $\left[\frac{1}{b}, 1 \right]$

接着我们使用狄拉克函数的性质：

其中 $x_{i}$ 为 $f(x)=0$ 的根。代入后可以得到谱密度：

它在低频端点 $\omega=\frac{1}{b}$ 处取最大值，在高频端点 $\omega=1$ 处取最小值。至此，谱密度函数的性质已经分析完毕。

非理想情况

然而以上是理想情况——因为我们不知不觉中，在计算谱密度的时候，假定了傅里叶变换的存在性，即 $\Delta m$ 可以取得 $\pm \infty$，也即序列长度可以达到 $+\infty$。这显然是不可能的。我们不得不回到现实情况中来，也即序列长度最多只能到达 $L$。此时核函数 $K(m,n)$ 的谱矩阵是一个 $L\times L$ 的矩阵，更准确的说，是一个托普利茨矩阵（常对角矩阵）⁸，记为 $T_{L}(f)$：

CAUTION

除此之外，还有一个让这种理想情况不可能达到的原因，想想是什么？（在后面会有揭晓）

这个矩阵的行为完全由生成函数 $f$，即上一步计算得到的谱密度函数 $\hat{K}(\omega;b)$ 决定 ⁹。

接下来，我们需要使用一个研究大型托普利茨矩阵行列式渐近行为的强大的引理：强斯格极限定理 (Strong Szegő Limit Theorem) ¹⁰，如下

NOTE

若生成函数 $f(\theta)$ 满足某些条件，例如 $f>0$ 且足够光滑，则：

$$\lim_{ L \to \infty } \frac{\det(T_{L}(f))}{G(f)^{L}} = E(f)$$

其中：

• $G(f)$ 是 $f$ 的几何平均数 (Geometric Mean)，代表了行列式的体行为 (bulk behavior)。对于充分大的 $L$，有 $\det(T_{L}(f))\approx G(f)^{L}$。

$$G(f)=\exp\left( \frac{1}{2\pi}\int _{0}^{2\pi}\ln f(e^{i\theta}) \, \mathrm{d}\theta \right)$$

• $E(f)$ 是一个和 $f$ 的光滑性相关的边界项，依赖于 $\ln f$ 的傅里叶系数

$$E(f)=\exp\left( \sum_{k=1}^{\infty} k\left\lvert\widehat{\ln f}_{k} \right\rvert^{2} \right)$$

其中 $\widehat{\ln f}_{k}$ 是 $\ln f$ 的第 $k$ 个傅里叶系数。

将上述引理应用于解释扩展上下文长度的稳定性质：回顾我们求出的谱密度函数：

明显，在两侧端点 $\frac{1}{b}$ 和 $1$ 处，$f$ 均出现了跳变，结合强斯格极限定理，可以得到如下结论。

IMPORTANT

由于谱密度函数 $f$ 天然具有跳变不连续的性质，在简单的拓展上下文长度 $L$ 时，模型的离散频率分辨率 $\sim \frac{1}{L}$ 和谱密度函数 $f$ 两侧端点位置之间的关系发生了改变，系统进入了一个对这种不连续性极其敏感的区域。从而使矩阵 $T_L(f)$ 出现了离群特征值，同时 $L\to \infty$ 时 $E(f)$ 没有良好的收敛性质，从而导致了简单扩展上下文长度 $L$ 的崩溃。

离群特征值的出现原因

可以看到，在上面的结论中，我们使用了一些描述性的语言，以便于当前能顺利理解。接下来，我们顺理成章的需要研究：

这种不连续性是如何导致托普利茨矩阵 $T_{L}(f)$ 的离群特征值的出现，以及这一现象何时发生？ （但是这太困难了，已经超出了我的知识范围，只好请教 gemini 大人了😭，贴在下面等有缘人来确认正确性）

首先让我们描述离群特征值是什么。根据 Fisher–Hartwig 猜想 ¹¹（已被证明）以及 Harold Widom¹² 的后续工作，可知对于一个具有跳变不连续的生成函数 $f(\omega)$，其对应的托普利茨矩阵 $T_{L}(f)$ 的谱在 $L\to \infty$ 时呈现如下结构：

• 连续谱 (Continuous Spectrum)：绝大多数（几乎所有）的特征值会密集地分布在由 $f(\omega)$ 的值域构成的区间内。这个区间我们称之为体谱 (bulk spectrum)。
• 离散谱 (Discrete Spectrum)：在生成函数 $f(\omega)$ 的不连续点处，可能会有一个或多个特征值脱离体谱，成为孤立的离群点。

不稳定性就等价于这些离群特征值的出现和它们的病态行为。

接下来，是最后的推导（虽然并不详细）。

Harold Widom¹² 对特定生成函数，即可以表示为一个光滑函数和一个区间指示函数的乘积的函数，证明了积分方程定理（Integral Equation Theorem for Outliers）：当 $L\to \infty$ 时，大多数特征值位于所谓的体谱区间内（由函数值域决定），而离群特征值的渐近行为，可以被一个定义在有限区间上的积分算子 $\mathcal{K}_L$ 的特征值精确描述。该积分算子的核一般与 sinc 核形式（如 $\sin(c(x-y)) / (x-y)$）相关。这意味着可以将一个复杂的、$L\times L$ 离散矩阵的谱问题，转化成一个更易于分析的连续积分算子 $\mathcal{K}_{L}$ 的问题。准确来说，$T_{L}(f)$ 的行列式的渐近行为被 $\mathcal{K}_{L}$ 的 Fredholm 行列式 ¹³ $\det(I-\mathcal{K}_{L})$ 描述。

通过对这个等效积分算子 $\mathcal{K}_{L}$ 的研究，数学家们发现它的谱（特别是其最大特征值）存在相变 (Phase Transition)¹⁴ 现象。相变的发生与否，由一个关键的无量纲参数控制，我们称之为序参量，记为 $\beta$。

在进行接下来的推导之前，首先需要对原始 RoPE 核进行适当的“修正”：回归原始 RoPE 核：

对 $i=0$ 这一项，即 $\cos(\Delta m)$ 这一项，因为它并非平方可积，所以不满足连续谱分析的前提，应该排除这一项，即稳定部分由 $i\in \{ 1,2,\dots,d/2-1 \}$ 这些项构成。这样做改变了生成函数的精确形状和支撑集，但可以验证，其定性行为（如边界不连续性）依然存在。那么有效维度就只有 $d_{\text{eff}}=d-2$ 个了（排除了第 $0,1$ 两个维度）。

进行了这样的修正之后，$x=2k/d$ 的范围就缩小为了 $[2 / d, (d-2) / d]$，对应的频率 $\omega=b^{-x}$ 的范围变为 $[b^{-(d-2)/d}, b^{-2/d}]$，也就是新的支撑集。

通过分析与算子 $\mathcal{K}_{L}$ 相关的潘勒韦方程 ¹⁵ 的解的渐近行为，并对支撑集下边界 $b^{-(d-2)/d}$ 进行分析，得到与这个边界相关的系统特征长度 $\lambda_{\text{char}}$，和基频 $b$、有效维数 $d_{\text{eff}}$ 的关系为：

这个特征长度可以被理解为系统内部结构能够保持相干的最大距离。定义上面提到的序参量 $\beta$ 为系统外在尺寸 $L$ 和这个内在特征长度 $\lambda_{\text{char}}$ 的比值：

Harold Widom 和他的合作者一起证明了：离群特征值的出现，精确地发生在序参量 $\beta$ 跨越一个特定的临界值 $\beta_{c}$​ 的时候，即，当 $\beta <\beta_{c}$​ 时，系统处于稳定状态，没有离群特征值；当 $\beta>\beta_{c}$​ 时，系统失稳，离群特征值从体谱中分裂出来。因此，系统保持在临界稳定状态的条件是：

那么为了能够稳定的扩展上下文，就需要有

$C$ 是一个常数。重新整理得到：

将新旧基频 $b'$ 和 $b$ 代入 $L'=s\cdot L$，得到

那么就有

总结

至此，艺术已成。

至此，我们从理论上证明了：

NOTE

当扩展上下文长度 $L$ 时，基频 $b$ 需要满足 $b'=b\cdot s^{d/(d-2)}$ 的变化方式，才能保证模型不会崩溃。

回顾我们在非理想情况一节中使用的描述性语言：

由于谱密度函数 $f$ 天然具有跳变不连续的性质，在简单的拓展上下文长度 $L$ 时，模型的离散频率分辨率 $\sim \frac{1}{L}$ 和谱密度函数 $f$ 两侧端点位置之间的关系发生了改变，系统进入了一个对这种不连续性极其敏感的区域。从而使矩阵 $T_L(f)$ 出现了离群特征值，同时 $L\to \infty$ 时 $E(f)$ 没有良好的收敛性质，从而导致了简单扩展上下文长度 $L$ 的崩溃。

现在可以理解：

• 系统进入了一个对这种不连续性极其敏感的区域：就是指随着 $L$ 的增大，$\beta > \beta_c$，从而导致离群特征值的产生，也正是相变现象。
• 模型的离散频率分辨率 $\sim \frac{1}{L}$ 和谱密度函数 $f$ 两侧端点位置之间的关系发生了改变：
  • 对长度 $L$ 有限的离散序列，频率空间就是由离散傅里叶变换 (DFT) 的频率点构成，即 $\Omega_L = \{ \omega_k = \frac{2\pi k}{L} | k=0,1,2,\dots, L-1 \}$，格点分辨率 $\Delta \omega = \frac{2\pi}{L}\sim \frac{1}{L}$
  • 随着 $L$ 的增大，如果 $f$ 满足良好的性质，那么 $T_L(f_b)$ 的特征值会越来越精确地填充到体谱中；然而 $f$ 存在跳跃不连续点，所以这种情况不可能发生，从而导致离群特征值出现的必然性。说“关系”发生改变，也正是说离散特征值不能精确填充到体谱中去，逐渐变成了离群特征值。

总结一下我们的证明结论：

NOTE

对于一个由固定的谱密度函数 $f_b$（其支撑集 $I_b$ 的边界存在跳变不连续）生成的托普利茨算子/矩阵序列 $\{ T_L(f_b) \}_{L\in \mathbb{Z}^+}$，存在一个由上下文长度 $L$ 和系统参数 $(b ,d)$ 共同决定的无量纲序参量 $\beta \propto L \cdot b^{-(d-2)/d}$。当 $L$ 增大并使得 $\beta$ 超过一个临界阈值 $\beta_c$ 时，系统发生相变：$T_L(f_b)$ 的谱结构发生定性改变，一个或多个离群特征值会从体谱 $I_b$ 中分裂出来。这个相变点 $L_{\text{crit}} \propto b^{(d-2)/d}$ 定义了在不修改基数 $b$ 的情况下，上下文长度能够保持稳定的上限。因此，当扩展上下文长度 $L$ 时，基频 $b$ 需要满足 $b'=b\cdot s^{d/(d-2)}$ 的变化方式，才能保证模型不会崩溃。