稀疏性与Transformer长度泛化

引言

Transformer模型的长度泛化问题长期困扰着研究者和工程师。传统观点认为，模型在短序列上训练后无法泛化到长序列是注意力机制的本质缺陷。然而，理论研究表明，稀疏依赖结构是长度泛化的关键因素：如果每个预测token仅依赖于固定数量的前序token，模型就能实现长度泛化。¹

这一发现不仅解释了为什么某些任务容易泛化，也为改进位置编码提供了理论指导。

稀疏依赖结构理论

形式化定义

考虑下一个token预测任务。设 $x_1,x_2,\dots ,x_T$ 为输入序列，模型预测 $x_{t+1}$。定义token依赖图：

$$\mathcal{G}=(V,E),V=\{1,2,\dots ,T\}$$

其中边 $(i,j)$ 表示 $x_j$ 依赖于 $x_i$。

$k$-稀疏依赖：如果对于任意 $j$，存在至多 $k$ 个 $i<j$ 使得 $(i,j)\in E$，则称依赖结构是 $k$-稀疏的。

核心定理

定理（稀疏泛化保证）：设任务的依赖结构是 $k$-稀疏的，其中 $k$ 是与序列长度无关的常数。则存在一个仅依赖于 $k$（而非序列长度）的样本复杂度界，保证Transformer能够实现长度泛化。

形式上，对于训练序列长度 $n$ 和测试序列长度 $T$，有：

$$\mathbb{P}\left[{\mathcal{L}}_{\text{gen}}(T)>ϵ\right]\le \exp \left(-\Theta (n)\cdot 1_{\{k=O(1)\}}\right)$$

关键洞察：当依赖结构稀疏（$k=O(1)$）时，样本复杂度与测试长度 $T$ 无关。

$k$-稀疏植入相关分布

分布定义

为了理论分析，研究者定义了**$k$-稀疏植入相关分布**（$k$-Sparse Planted Correlation Distributions）。这一分布族精确捕获了稀疏依赖结构的本质。

设 ${\mathcal{D}}_k$ 为满足以下条件的分布集合：

1. 稀疏条件：每个变量 $X_j$ 仅通过固定常数 $k$ 个变量 $\mathcal{P}(j)$ 依赖于前序变量
2. 条件独立性：给定父变量 $\mathcal{P}(j)$，$X_j$ 与其他变量条件独立

分布采样

从 ${\mathcal{D}}_k$ 采样一个实例：

1. 选择依赖结构：对于每个 $j$，均匀随机选择 $k$ 个父节点
2. 生成变量：按拓扑顺序依次生成每个变量
3. 添加噪声：引入可控制的噪声水平

理论性质

$k$-稀疏植入分布在以下方面具有良好性质：

• 有限支持：所有统计量在常数 $k$ 下有界
• 可学习性：存在高效的参数估计方法
• 长度不变性：分布参数不随序列长度变化

泛化注意力头的成功机制

理论分析

研究表明，在稀疏依赖任务上，泛化注意力头（Generalizing Attention Heads）能够成功实现长度泛化。这些注意力头的行为具有以下特征：

1. 选择性关注：每个头仅关注固定数量的关键位置
2. 位置不变性：依赖关系通过内容而非绝对位置确定
3. 稀疏激活：实际激活的注意力连接数远少于 $n^2$

机制解释

从信息论角度，稀疏注意力头实现了以下功能：

$$\underset{\text{完整注意力}}{\underbrace{\text{Attention}(Q,K,V)}}\approx \underset{\text{稀疏化}}{\underbrace{\text{Top-}k}}\circ \underset{\text{相似度计算}}{\underbrace{\text{Softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d}}\right)}}\circ V$$

其中 $\text{Top-}k$ 算子保留每个query的top-$k$注意力权重。

预测性位置耦合

方法动机

基于稀疏性理论，研究者提出了预测性位置耦合（Predictive Position Coupling）方法，进一步提升长度泛化能力。¹

核心思想

传统位置耦合方法：

• 固定训练长度 $L_{\text{train}}$，测试时使用更长序列
• 通过插值调整位置编码
• 缺点：位置映射与任务脱节

预测性位置耦合：

1. 定义位置映射函数 $\pi :[L_{\text{train}}]\to \mathbb{R}$
2. 训练模型预测 $\pi (p)$ 而非直接使用 $p$
3. 测试时可泛化到任意位置

数学框架

设 $p_i$ 为第 $i$ 个位置的索引，${\hat{p}}_i=\pi (p_i)$ 为映射后的位置。训练目标变为：

$$\underset{\theta }{\min }{\mathbb{E}}_{(x,y)\sim \mathcal{D}}\left[\ell \left(f_{\theta }(x;\hat{p}),y\right)\right]$$

其中 $f_{\theta }$ 是Transformer模型。

实验结果

预测性位置耦合在多个任务上取得显著提升：

任务	基线准确率	预测性位置耦合	提升
数字加法	23.5%	78.2%	+54.7%
字符串反转	45.1%	89.7%	+44.6%
模式匹配	31.8%	72.4%	+40.6%

与其他位置编码方法的比较

绝对位置编码

传统绝对位置编码：

• 学习一个长度 $L_{\text{max}}$ 的嵌入表
• 外推到更长序列时需要零填充或重复
• 问题：分布不匹配导致泛化失败

相对位置编码

相对位置编码关注token之间的相对距离：

• 位置差 $i-j$ 进入注意力计算
• 理论上可以处理任意长度
• 限制：仍依赖训练中观察到的相对距离分布

ALiBi

ALiBi（Attention with Linear Biases）：

• 无需学习位置嵌入
• 线性衰减注意力分数
• 优势：天然支持长度外推
• 局限：不适合所有任务类型

RoPE

旋转位置编码（Rotary Position Encoding）：

• 通过旋转操作注入位置信息
• 保持注意力分数的旋转不变性
• 优势：理论上有更好的长度泛化性质
• 局限：需要修改注意力实现

实践建议

基于稀疏性理论，我们为实践者提供以下建议：

1. 任务设计

• 明确依赖结构：分析任务的token依赖关系
• 最小化长程依赖：将任务分解为局部子任务
• 使用递归模式：设计可分解的计算图

2. 模型训练

• 变长训练：在训练时使用多种序列长度
• 位置去相关：避免模型对特定位置的过拟合
• 稀疏正则化：鼓励注意力权重的稀疏性

3. 位置编码选择

场景	推荐方法
固定长度任务	绝对位置编码
长度外推重要	RoPE + 预测性耦合
局部模式为主	ALiBi
通用场景	相对位置编码

未来方向

稀疏性理论为长度泛化研究开辟了新方向：

1. 自适应稀疏：根据输入动态调整注意力稀疏性
2. 稀疏-密集混合：稀疏注意力用于局部，全局信息通过其他机制传递
3. 理论-实践桥梁：将稀疏性度量纳入模型选择标准

总结

稀疏依赖结构理论揭示了长度泛化的本质：

• 核心洞察：稀疏性是长度泛化的关键
• 理论保证：$k$-稀疏分布提供PAC-style泛化界
• 实践方法：预测性位置耦合显著提升泛化性能
• 设计原则：任务和模型设计应考虑稀疏依赖

这一理论不仅解释了为什么某些Transformer难以泛化，也为构建能够处理任意长度序列的模型提供了明确指导。

参考

Footnotes

1. The Role of Sparsity for Length Generalization in Transformers. arXiv:2502.16792 (2025) ↩ ↩²