第 6 章 长度外推与序列建模中的 NTK

开篇段落

本章致力于应用神经切线核（NTK）理论来分析现代深度学习，特别是 Transformer 模型在处理超长序列时的失败机制——长度外推（Length Extrapolation）问题。我们将首先界定长度外推在统计学习中的 OOD (Out-of-Distribution) 性质，并利用 NTK 的 RKHS (再生核 Hilbert 空间) 视角，精确描述为什么模型在超出训练长度时会失去泛化能力。核心讨论将聚焦于序列模型中的关键组件：位置编码（PE）和自注意力机制如何影响 NTK 的结构和稳定性。最终，我们将分析 RoPE (旋转位置编码) 和 ALiBi 等先进结构如何通过设计稳定的 NTK 特征空间来实现有效的长度外推。

6.1 长度外推与 NTK 的 OOD 视角

6.1.1 泛化悖论与 OOD 挑战

标准的深度学习泛化理论通常假设训练数据 $D_{train}$ 和测试数据 $D_{test}$ 来自相同的分布 $P$。然而，在长度外推问题中，$D_{train}$ 仅包含长度 $L \le L_{train}$ 的序列，而 $D_{test}$ 包含 $L > L_{train}$ 的序列。这是一个典型的协变量偏移（Covariate Shift）问题，即输入 $x$ 的边缘分布 $P(x)$ 发生了变化。

深度网络在 NTK regime 下，学习过程等价于在初始化时的核 $\Theta_0$ 所定义的 RKHS $H_{\Theta_0}$ 中寻找一个最小范数插值解。

$$ f^*(x) = \sum_{i=1}^N \alpha_i \Theta_0(x, x_i) $$

这个解 $f^*$ 的泛化性能，取决于测试样本 $x_{test}$ 在 $H_{\Theta_0}$ 中的表达能力。

6.1.2 RKHS 支持与特征空间崩溃

当长度 $L$ 增加，序列 $x_{long}$ 包含训练集 $D_{train}$ 中未见的绝对位置信息（尤其是绝对位置编码 $p_L$）。

从 NTK 视角看，长度外推失败的本质是：

长度导致的特征空间崩溃 (Feature Space Collapse): 训练数据 $D_{train}$ 张成的特征空间 $S_{train} = \text{span}\{\phi(x_i) | x_i \in D_{train}\}$ 无法包含或有效近似 $x_{long}$ 的特征向量 $\phi(x_{long}) = \nabla_\theta f_\theta(x_{long})$。

在核空间中，这意味着：

  1. 测试点远离支持： $x_{long}$ 在 RKHS $H_{\Theta_0}$ 中的范数 $|!|f(x_{long})|!|_{H_{\Theta_0}}$ 变得非常大，表示其复杂度高且难以用训练数据线性组合表示。
  2. 核值衰减： 对于任何训练样本 $x_{short}$，核函数值 $\Theta_0(x_{long}, x_{short})$ 接近于零。

$$ \text{失败模式： } \Theta(x_{long}, x_{short}) \approx 0 $$

如果核值接近零，则训练得到的系数 $\alpha_i$ 对 $f^*(x_{long})$ 的贡献可以忽略，模型预测退化为随机初始化的输出 $f_{\theta_0}(x_{long})$，即性能崩溃。

6.1.3 绝对位置编码的破坏性

在标准 Transformer 中，输入 $x_i$ 包含嵌入 $e_i$ 和绝对位置编码 $p_i$. $$ \text{Input}_i = e_i + p_i $$ 当 $L_{test} > L_{train}$ 时，如果 $p_L$ 是训练期间从未见过的向量，它将导致梯度 $\nabla_\theta f_\theta(x_{long})$ 发生剧烈、不预测的变化。这种不连续性直接打破了 NTK 线性化近似在 OOD 区域的有效性，并导致 $\Theta(x_{long}, x_{short})$ 极不稳定。

6.2 序列模型的 NTK 近似结构

6.2.1 自注意力机制与梯度流

在一个 $L$ 长的序列上，神经网络函数 $f_\theta(x)$ 是一个维度为 $L \times d_{out}$ 的输出矩阵。我们通常将其展平为向量进行梯度计算。

$$ \phi(x) = \nabla_\theta f_\theta(x) \in \mathbb{R}^P $$

其中 $P$ 是总参数量。

在 Transformer 中，参数 $\theta$ 分布在 $Q, K, V$ 投影矩阵、FFN 权重以及归一化层的参数中。NTK 的计算涉及这些参数梯度之间的内积。

自注意力层的 NTK 结构直观： 自注意力机制使得序列中任意两个位置 $i$ 和 $j$ 相互依赖。这在 NTK 矩阵中体现为：$\Theta(x, x')$ 矩阵是 $L \times L$ 块矩阵的求和，其中每个块 $\Theta_{i, j}(x, x')$ 描述了序列 $x$ 的位置 $i$ 和序列 $x'$ 的位置 $j$ 之间的核相关。

$$ \Theta(x, x') = \sum_{i=1}^L \sum_{j=1}^{L'} \nabla_\theta f_{i}(x)^\top \nabla_\theta f_{j}(x') $$

其中 $f_i(x)$ 是序列 $x$ 在位置 $i$ 的输出。

6.2.2 深度递归与 NNGP vs. NTK

对于深层网络，NTK 遵循递归关系。每一层的输出核 $K^{(l)}$ 和切线核 $\Theta^{(l)}$ 都依赖于前一层 $l-1$ 的核。

$$ \Theta^{(l)}(x, x') = \dots + \text{Interaction Term}(\Theta^{(l-1)}, K^{(l-1)}) $$

在序列模型中，这种递归关系必须考虑到位置编码对每一层输入的注入。如果位置编码在底层就导致了 OOD 区域的特征不稳定，那么这种不稳定性将通过深度递归逐层放大。

NNGP (Neural Network Gaussian Process) 核： NNGP 核 $K(x, x')$ 仅描述初始化分布 $p(\theta)$ 下 $f_\theta(x)$ 和 $f_\theta(x')$ 的协方差。 NTK $\Theta(x, x')$： 描述训练过程中函数空间变化的方向（梯度方向）。

在序列建模中，如果网络很深，NNGP 往往会退化为简单的核（如常数核线性核），导致其在区分长序列结构上效果不佳。NTK 虽然更复杂，但它捕捉了更丰富的、与梯度相关的特征结构，因此对序列长度变化更加敏感。

6.3 序列长度变化对 NTK 谱的影响

6.3.1 NTK Gram 矩阵的长度敏感性

对于一个固定最大长度 $L_{\max}$ 的训练集 $D$，我们可以计算训练 Gram 矩阵 $K = \Theta(D, D)$。该矩阵的特征值 $\lambda_k$ 及其对应的特征函数 $v_k(x)$ 决定了 RKHS $H_\Theta$ 的结构和训练收敛速度。

长度外推中，$L_{test} > L_{train}$ 引入了一个新输入 $x_{long}$。我们关注 $K$ 与 $x_{long}$ 之间的交叉核矩阵 $K_{cross} = \Theta(x_{long}, D)$.

绝对 PE 导致的谱退化： 如果使用绝对 PE，当 $L$ 增加时，$x_{long}$ 的特征向量 $\phi(x_{long})$ 会在特征空间中指向一个与 $S_{train}$ 几乎正交的方向。这导致：

$$ K_{cross} \to \mathbf{0} $$

后果： $x_{long}$ 无法通过 $K$ 的特征函数基 $\{v_k\}$ 进行有效分解。即使训练集 $D$ 成功学习了所有主要的特征函数 $v_k$，它们也不能推广到 $x_{long}$。

6.3.2 长度缩放与稀疏化

在实际应用中，处理长序列可能需要对位置编码进行某种形式的缩放（如长度外推的插值方法）。

长度缩放对 NTK 的影响： 如果我们将训练中使用的 PE 频率进行拉伸或压缩，本质上是在改变输入数据的表示，从而修改了 NTK $\Theta(x, x')$ 的函数形式。

  • 拉伸（如 RoPE Scaling）： 试图使 $p_L$ 的向量在 $L > L_{train}$ 时表现得像 $p_{L'}$ ($L' < L_{train}$) 一样。这是一种通过显式修改输入特征来稳定 NTK 的策略，目的是保持 $K_{cross}$ 矩阵的值不衰减。

ASCII 图示：NTK 矩阵在长度扩展下的稳定性对比

 训练长度 L_train
|-----------|
| A | B     |  <- K_train = 稳定，高秩

| A | B     |  <- K_train = 稳定，高秩
|-----------|
| C | D     |  <- K_cross (绝对PE: C, D -> 0)
|-----------|

 理想外推核 (相对PE)
|-----------|-----------|
| K_short   | K_cross   |  <- K_cross 保持非零且结构稳定

| K_short   | K_cross   |  <- K_cross 保持非零且结构稳定
|-----------|-----------|
| K'_cross  | K_long    |
|-----------|-----------|

6.4 稳定化 NTK 结构：相对位置编码

6.4.1 NTK 的平移不变性需求

为了实现长度外推，序列模型的 NTK 必须具备某种形式的平移不变性（Translation Invariance） 或更精确的相对位置依赖性。

定义： 一个序列核 $\Theta(x, x')$ 是理想长度外推核，如果对于任意 $x, x'$ 和任意平移量 $\tau$，满足：

$$ \Theta(x_\tau, x'_\tau) \approx \Theta(x, x') $$ 其中 $x_\tau$ 表示序列 $x$ 被平移了 $\tau$ 个位置（即所有位置索引 $i \to i+\tau$）。

为什么需要平移不变性？ 如果模型能够识别“位置 1 和位置 5 之间的关系”与“位置 101 和位置 105 之间的关系”是相同的，那么它就不依赖于绝对长度 $L$。

在 NTK 视角下，这意味着梯度特征 $\phi(x)$ 应该要编码序列的局部或相对结构，而不是全局的绝对位置。

6.4.2 ALiBi 的线性偏置与 NTK 稳定化

ALiBi（Attention with Linear Biases）通过将距离偏置 $m \cdot (j-i)$ 直接注入到注意力 Logits 中，显式地去除了对绝对位置的依赖。

$$ \text{Logits}_{i, j} = \frac{Q_i K_j^\top}{\sqrt{d_k}} + \underbrace{m \cdot (j - i)}_{\text{相对偏置}} $$

NTK 分析： 在宽度极限下，梯度流通过注意力分数 $A_{i, j}$。由于 $A_{i, j}$ 的输入（Logits）仅依赖于相对距离 $j-i$，且这种依赖是线性的，因此：

  1. 梯度连续性： 当 $i$ 和 $j$ 超出 $L_{train}$ 时，只要 $j-i$ 保持在合理的范围，Attention 的输入 Logits 仍是训练时见过的函数值域。
  2. 特征平滑： 梯度特征 $\phi(x_{long})$ 能够平滑地从 $\phi(x_{short})$ 扩展，避免了因未见过的绝对 $p_L$ 导致的特征突变。

ALiBi 是通过结构设计强制 NTK 保持稳定的典型案例，它有效地确保了 $K_{cross}$ 矩阵具有与 $K_{train}$ 相似的结构。

6.5 案例研究：旋转位置编码（RoPE）的几何稳定性

6.5.1 RoPE 的数学基础与内积分解

RoPE 采用旋转矩阵 $R_i$ 来编码位置 $i$： $$ Q_i = W_Q R_i e_i $$ 它巧妙地利用了三角函数特性，确保查询 $Q_i$ 和键 $K_j$ 之间的点积能够被分解：

$$ Q_i^\top K_j = (W_Q R_i e_i)^\top (W_K R_j e_j) \propto \cos(\theta(i-j)) $$ （假设 $W_Q=W_K$ 且忽略嵌入 $e$ 的影响）。

核心机制： RoPE 将绝对位置 $i$ 和 $j$ 的信息编码，但其交互结果只依赖于相对距离 $i-j$。

6.5.2 RoPE 与 NTK 几何稳定性

在 NTK 理论中，我们关注的是梯度 $\nabla_\theta f$ 的内积。如果 $Q_i^\top K_j$ 对所有参数的梯度贡献（即 $\nabla_\theta (Q_i^\top K_j)$）也是相对距离 $i-j$ 的函数，那么整个 NTK 矩阵 $\Theta(x, x')$ 就会具备长度稳定性。

RoPE 达成了这一点，因为它在数学上确保了：

  1. 注意力Logits的稳定性： 无论 $i, j$ 多大，Logits 只取决于 $i-j$。
  2. 下游梯度的稳定性： 由于 Softmax 的输入稳定且具有周期性/平滑性（$\cos$ 函数），参数权重 $W_Q, W_K$ 的梯度 $\nabla_{W} Q_i^\top K_j$ 也对 $i-j$ 保持稳定的函数依赖。

这意味着，即使 $i$ 和 $j$ 超过了训练长度，只要它们的相对距离不超出训练期间的相对距离范围，NTK 的局部结构就不会改变。

6.5.3 实践指导：RoPE 外推的限制

虽然 RoPE 极大地提高了外推能力，但它并非完美：

  1. 相对距离限制： RoPE 的外推能力仍然受限于训练中见过的最大相对距离 $|i-j|_{max}$。如果 $|i-j|$ 超出了这个范围，模型的性能仍可能下降，因为 $\cos(\theta(i-j))$ 及其梯度表现出训练期间未见的周期或相位。
  2. 插值缩放： 实际应用中，为了进一步增强 RoPE 的外推能力，通常需要结合长度缩放（如 Llama 2 中的 RoPE Scaling）。这在 NTK 视角下，是通过调整 $\Theta(x, x')$ 的频率偏置，以确保长序列的特征在短序列的特征空间中得到平滑的低频表达。

本章小结

| 概念 | 描述 | NTK 影响 |

概念 描述 NTK 影响
长度外推 序列长度 $L_{test} > L_{train}$ 时的 OOD 泛化。 导致 NTK 矩阵 $K_{cross}$ 衰减，特征空间崩溃。
NTK 视角下的失败 $x_{long}$ 远离 $D_{train}$ 定义的 RKHS $H_\Theta$ 的支持域。 训练得到的函数 $f^*$ 无法对 $x_{long}$ 有效插值或预测。
绝对 PE 依赖于绝对索引的固定向量注入。 在 $L > L_{train}$ 引入不连续的特征 $\phi(x_{long})$，破坏 $\Theta$ 的稳定性。
平移不变性 理想序列核应仅依赖于相对距离。 确保 $\Theta(x_{long}, x_{short})$ 不随绝对长度变化而剧烈衰减。
RoPE/ALiBi 通过设计注意力 Logits，使其依赖于相对距离 $j-i$。 强制 NTK 保持几何/函数形式的稳定性，使得 $\phi(x_{long})$ 与 $\phi(x_{short})$ 保持高相关性。

Rule-of-Thumb 6： 要实现良好的长度外推能力，模型结构必须确保其神经切线核（NTK）关于输入长度是平滑的或相对不变的。在序列模型中，这意味着模型对相对位置的编码必须是鲁棒且可推广的，以防止特征在 OOD 区域发生剧烈旋转或正交化。

练习题

基础题

6.1 绝对 PE 失效的 NTK 描述 解释为什么使用绝对位置编码（PE）的 Transformer 在 $L_{test} > L_{train}$ 时，其经验 NTK 矩阵 $K_{test} = \Theta(D_{long}, D_{short})$ 的元素倾向于趋近于零。请使用“特征向量的正交性”这一概念进行描述。

提示

新的绝对 PE 向量 $p_{L_{train}+k}$ 在高维嵌入空间中通常是随机初始化的。在宽度极限下，如果 $p_{L_{train}+k}$ 与训练中见过的所有 $p_i$ 都不相关，那么 $x_{long}$ 的梯度特征向量 $\phi(x_{long})$ 就会与训练特征空间 $S_{train}$ 近乎正交，导致内积（核值）趋零。

6.2 NTK 与 RKHS 支持域 在 NTK 理论中，长度外推失败意味着 $x_{long}$ 位于训练数据 $D_{train}$ 张成的 RKHS $H_\Theta$ 的支持域之外。请解释，如果 $x_{long}$ 真的完全位于支持域之外，那么通过核回归得到的预测 $f^*(x_{long})$ 将会是什么值？

提示

在 NTK Regime 中，$f^*(x) \approx f_{\theta_0}(x) + \sum \alpha_i \Theta(x, x_i)$. 如果 $\Theta(x_{long}, x_i) \approx 0$ 对于所有 $i$，预测值将退化到网络的初始化输出 $f_{\theta_0}(x_{long})$，即一个未经过训练、近乎随机的基线值。

6.3 卷积网络的平移不变性与 NTK 简单的一维卷积神经网络（CNN，无 PE）在理论上为什么被认为比标准 Transformer 具有更好的长度外推潜力？请从 NTK $\Theta(x, x')$ 的平移不变性角度简要分析。

提示

CNN 的卷积核 $W$ 在整个序列上共享。梯度 $\nabla_W$ 依赖于局部特征窗口，与绝对位置无关。因此，CNN 的 NTK 天然近似于一个平移不变核（即 $\Theta(x, x')$ 仅依赖于 $i-j$），这使得其特征空间稳定。

6.4 长度外推的诊断指标 如果你想要通过计算经验 NTK 来诊断一个序列模型是否具有良好的长度外推能力，你会计算哪两个 NTK 矩阵，并比较它们之间的什么特性？

提示

需要计算训练数据的 Gram 矩阵 $K_{train} = \Theta(D_{short}, D_{short})$ 和交叉核矩阵 $K_{cross} = \Theta(D_{long}, D_{short})$。好的外推能力要求 $K_{cross}$ 的元素均值和方差不应比 $K_{train}$ 小太多，并且 $K_{cross}$ 应表现出与相对距离相关的结构。

挑战题与开放性思考

6.5 RoPE 稳定性的数学推演（高级） 假设一个单头注意力层，使用 RoPE，且忽略 FFN 和激活函数。请描述如何数学上证明，参数 $W_Q, W_K$ 对输出 $f(x)$ 的梯度内积 $\nabla_{W} f_i(x)^\top \nabla_{W} f_j(x')$ 依赖于相对距离 $|i-j|$，而非绝对位置 $i, j$。

提示

关注梯度流 $\nabla_W$. $W$ 的梯度与输入 $x$ 的特征有关。由于 RoPE 使得 Logits 仅依赖于 $i-j$，且 Softmax 在极限下可以近似为恒等映射（或集中），因此 Logits 的稳定结构会传递到梯度。需要利用 $R_i^\top R_j = R_{j-i}$ 的性质。

6.6 宽度与 OOD 泛化：NTK 的保守性 NTK 理论因其依赖于固定核 $\Theta_0$ 而被认为是“保守”的，即它无法捕捉到通过训练学习到的特征。请结合长度外推场景，论证为何有限宽度模型在理论上可能比无限宽模型更具优势。这种优势需要满足什么条件？

提示

有限宽模型允许 $\Theta(t)$ 随时间变化（特征学习）。如果训练能将 $x_{long}$ 的特征 $\phi(x_{long}, t)$ 旋转或伸到 $S_{train}$ 附近，模型就能泛化。这需要：1) 较大的学习率 $\eta$ 或较长的训练时间 $t$；2) 数据 $D_{train}$ 包含足够的结构信息来指导特征旋转。

6.7 外推失败与 NTK 特征值 假设我们计算了一个具有绝对 PE 的序列模型在 $D_{train}$ 上的 NTK Gram 矩阵 $K_{train}$。当我们将测试样本 $x_{long}$ 引入计算交叉核 $K_{cross}$ 时， $K_{cross}$ 的元素趋于零。请用特征分解 $K_{train} = V \Lambda V^\top$ 来解释 $x_{long}$ 是如何通过 $K_{train}$ 的特征函数进行表示的，并说明为什么这种表示在外推失败时失效。

提示

核回归解可以表示为 $f^*(x) = \sum_k c_k v_k(x)$，其中 $v_k(x)$ 是 $K_{train}$ 的特征函数。如果 $\Theta(x_{long}, x_i) \approx 0$，这意味着 $x_{long}$ 与所有特征函数 $v_k(x)$ 几乎正交，即 $\int v_k(x_{long}) dP(x_{long}) \approx 0$。因此，它不能被 $K_{train}$ 的特征基有效表。

常见陷阱与错误 (Gotchas)

  1. 混淆 NNGP 与 NTK 在长度外推中的角色 * 陷阱： 认为只要 NNGP 核是长度稳定的，NTK 就稳定。 * 纠正： NNGP 描述的是网络的先验（初始化时的协方差）。NTK 描述的是网络的训练动态。即使 NNGP 稳定，如果参数化或激活函数导致梯度在 OOD 区域不稳定（例如，梯度爆炸），NTK 也会迅速失效。在序列模型中，NTK 对位置编码的依赖性远高于 NNGP。

  2. 误认为 RoPE 实现了完全平移不变性 * 陷阱： RoPE 保证了 $\Theta$ 具有完美的平移不变性，因此可以无限外推。 * 纠正： RoPE 仅保证注意力 Logits 依赖于相对距离，这极大地稳定了 NTK。但 Transformer 结构中仍有 FFN、LayerNorm 等组件，它们的梯度流可能仍然受到绝对位置的轻微影响。此外，RoPE 依赖的三角函数性质在极大的相对距离上表现出训练未见的周期性导致外推能力并非无限，只是比绝对 PE 鲁棒得多。

  3. 忽略长度对特征维度的影响 * 陷阱： 认为 NTK 矩阵的维度只与训练样本数量 $N$ 有关。 * 纠正： 虽然 $K \in \mathbb{R}^{N \times N}$，但 $x$ 的特征图 $\phi(x) = \nabla_\theta f_\theta(x)$ 是 $P$ 维的，其中 $P$ 是参数数量。对于序列模型，如果 $L$ 变化，虽然 $P$ 不变，但 $\phi(x)$ 内部的结构（不同位置的梯度贡献）会改变。如果处理不同长度的序列，我们需要使用零填充等方法使输入标准化，这会引入新的、与填充相关的梯度项，影响经验 NTK 的计算。

  4. 忽略 LayerNorm/BatchNorm 对梯度的缩放 * 陷阱： 在分析深度 Transformer 的 NTK 时，忽略归一化层的影响。 * 纠正： LayerNorm 和 BatchNorm 涉及对输入进行均值和方差的归一化，这严重依赖于序列的长度 $L$ 和批次的大小 $B$。归一化层会动态地缩放梯度，尤其是 OOD 长度下，可能导致 NTK 的收敛速度和特征值分布发生变化，偏离理想的理论假设。