Sparse Transformers(稀疏 Transformer)深度解读报告
论文:Generating Long Sequences with Sparse Transformers(arXiv:1904.10509v1,2019-04-23 提交,后发表于 NeurIPS 2019)
作者:Rewon Child、Scott Gray、Alec Radford、Ilya Sutskever(全部 OpenAI)
篇幅:10 页;关键词:Machine Learning, self-attention, attention, generative modeling
数据依据声明:本报告所有数字唯一来源于同目录《00-源数据提取.md》(对应原论文 PDF 的元数据、章节树、核心主张表、方法细节与实验要点)。凡属本报告的解读、外推或背景联想,一律标注「(推断)」。不引用任何其他文件。
摘要
Transformer 的自注意力对序列长度 n 的平方复杂度 O(n²),使其无法处理长序列。本文提出因子化自注意力(factorized self-attention):用 p 个注意力头分别关注位置的一个稀疏子集,每个子集大小 ∝ ⁿ√n,且保证任意两个位置在 p 步注意力内连通,从而将复杂度降至 O(n√n)。配合 pre-activation 残差与按深度 1/√(2N) 缩放的初始化使数百层网络可稳定训练、梯度检查点重计算节省显存、定制 block-sparse GPU kernel 与混合精度加速后,模型可处理数万 timestep、数百层深。实验上:CIFAR-10 达 2.80 bits/dim(超越前 SOTA PixelSNAIL 的 2.85);Enwik8 达 0.99 bits/dim(95M 参数追平参数量翻倍以上的 Transformer-XL 277M);ImageNet 64×64 达 3.44(超越前 SOTA SPN 的 3.52);原始音频古典音乐在 65,536 长度下达 1.97 bits/byte。反直觉的是,稀疏模式不仅更快,在被比较的数据集上误差还更低。
一、问题定位
1.1 核心问题:O(n²) 的三重挤压
自注意力机制中每个位置都要与所有历史位置计算相关性,计算量与显存随序列长度 n 平方增长(Abstract、§1)。这带来三重挤压:
- 计算量:n² 的注意力矩阵乘法在 n 上万时变得不可承受;
- 显存:n×n 的注意力权重矩阵必须在反向传播时保留(或在重计算策略下至少在前向时可重建),直接限制序列长度与网络深度的乘积;
- 建模需求与能力的错位:许多真实数据——高分辨率图像逐像素展开后的数千字节序列、长文本、连续音频波形——天然具有远超当时常规上下文(通常数百至一千 token 量级)的时间跨度,而稠密注意力模型既算不动也训不稳这样的规模。(推断:源文件记载 Enwik8 实验"上下文 12,288,长于以往方法",侧面印证此前方法的上下文普遍显著短于此。)
1.2 论文对问题的精确切分
本文没有把问题笼统地表述为"让注意力变快",而是拆成四个相互咬合的子问题,并逐一给出工程化答案:
| 子问题 | 论文应对方案 | 出处 |
|---|---|---|
| --- | --- | --- |
| 注意力连接数过多 | 因子化自注意力,每头只看 ⁿ√n 个位置,p=2 头保证连通 | §4.2–4.3 |
| 数百层网络训不动 | pre-activation 残差 + 1/√(2N) 初始化缩放 | §5.2 |
| n² 显存装不下 | 反向传播时重算 attention/FFN(梯度检查点) | §5.4 |
| 稀疏计算跑不快 | 定制 block-sparse kernels + FP16 混合精度 | §5.5–5.6 |
(推断:这种"算法结构 × 数值稳定性 × 显存管理 × 硬件 kernel"四位一体的打法,是本文能真正把理论复杂度落地为实测收益的关键;只做第一点的论文往往停留在理论层面。)
1.3 一个关键的先导观察
方法设计并非凭空设定稀疏模式,而是先做了诊断实验(§4.1):可视化一个 128 层全注意力网络在 CIFAR-10 上学到的注意力模式(Figure 2),发现大多数层呈现稀疏模式——早期层局部连接(类似卷积),第 19–20 层自发分解为行注意力 + 列注意力,部分层呈全局数据依赖访问,64–128 层高度稀疏。据此作者假设:引入预定稀疏性不会显著损害性能。这一"先观察模型自己想学什么,再照着设计归纳偏置"的研究路径,是全文方法论上最值得注意的起点。(推断:这也解释了为何固定的稀疏模式最终能在多个任务上打平甚至超过稠密注意力——它并非强行约束,而是对已学行为的形式化。)
二、概念内核
2.1 因子化自注意力的形式化定义
注意力层由连接模式 S={S₁,…,Sₙ} 参数化,其中每个位置 i 只与集合内的位置发生注意力交互。因子化版本使用 p 个注意力头,第 m 头定义子集 Aᵢ⁽ᵐ⁾ ⊂ {j : j ≤ i},要求满足两条性质(§4.2):
- (a) 稀疏性:每个 |Aᵢ⁽ᵐ⁾| ∝ ⁿ√n —— 单头连接数从 n 降到 ⁿ√n 量级;
- (b) 连通性:跨 p 步注意力,所有位置保持连通——任意 j ≤ i 都存在一条长度 ≤ p+1 的路径。
本文取 p=2。这两条性质合起来给出了复杂度 O(n√n) 的来源(Abstract、§1、§4.2):总连接数约为 p·n·ⁿ√n,而连通性保证信息两跳内可达任意历史位置,不以牺牲全局感受野为代价。(推断:性质 (b) 是设计的关键约束——若只有稀疏而无连通保证,深层网络的信息流会被切断;p=2 意味着任意位置的 query 至多经一次"中转位置"即可触达全部历史。)
2.2 两种二维分解模式(§4.3)
Strided attention(跨步注意力,Figure 3b):
- 第 1 头:Aᵢ⁽¹⁾ = {t, …, i},其中 t = max(0, i−l)——看最近 l 个位置(局部带宽);
- 第 2 头:Aᵢ⁽²⁾ = {j : (i−j) mod l = 0}——沿 stride l 每隔 l 个位置取一点(跨步汇总);
- stride l 取接近 √n。适合图像、音乐等有周期结构的数据。
Fixed attention(固定注意力,Figure 3c):
- 第 1 头:块内连接 Aᵢ⁽¹⁾ = {j : ⌊j/l⌋ = ⌊i/l⌋}——只看同一块内的位置;
- 第 2 头:列连接 Aᵢ⁽²⁾ = {j : j mod l ∈ {t, …, l}},t = l − c——每个块内最后 c 个位置作为"汇总单元"向所有后续位置传播;
- 论文给出的具体示例:l=128、c=8 时,所有 >128 的位置可以看到 120–128,>256 的位置看到 248–256,依此类推。
- 调参经验:c∈{8,16,32}、l∈{128,256} 表现良好;c=1 会显著限制表达力;fixed 比 strided 多约 c 倍计算。多头时,让各头注意同一 l 块内不同的 c 子块更优。
两种模式的选择准则(§4.3 末尾):文本等无周期结构数据用 strided 会失败,因为元素的空间坐标与其未来相关位置不对应(空间坐标与相关性不对齐),此时改用 fixed。这一失效后来被 Table 2 的数字直接证实(Enwik8 上 strided 仅 1.13 vs dense 1.00,见第四节)。
2.3 概念内核小结
整篇论文的概念贡献可以压缩为一句话:把"每个位置看所有位置"替换为"每个位置看 ⁿ√n 个精心选定的位置,且两步之内全网可达"。strided 用"局部带 + 等距骨架"实现这一点,fixed 用"块内稠密 + 列状汇总枢纽"实现这一点;前者契合周期性数据,后者契合无周期数据。(推断:fixed 模式的列汇总单元本质上引入了一种层级式信息汇聚结构,可视为后续 hierarchical/slab 结构稀疏注意力的雏形。)
三、方法拆解
3.1 因子化头的三种整合方式(§5.1,Eq.6–8)
有了两个稀疏头,如何嵌入网络?论文给出三种方式:
- 逐残差块交替(Eq.6)——相邻残差块轮流使用两种注意力模式;
- merged head 合并模式(Eq.7)——在同一层内合并两种模式;
- 多头并行(Eq.8)——标准多头机制下各头分别执行不同稀疏模式。
经验结论:通常多头效果好;但当序列极长、注意力计算占主导时,逐个串行(交替)更划算。(推断:这是因为多头并行会重复支付注意力开销,而串行复用了同一份表示。)
3.2 百层深度的可训练性改造(§5.2)
要把网络堆到上百层(CIFAR-10 实验用到 128 层),标准 post-norm Transformer 无法稳定收敛。本文采用:
- He et al. 2016 pre-activation 残差块:Hₖ = Hₖ₋₁ + resblock(Hₖ₋₁)(Eq.9–14),即恒等主干 + 分支计算;
- 深度尺度不变的初始化:W₂ 与 Wₚ 按 1/√(2N) 缩放初始化,保持激活方差跨层数不变(N 为层数)(推断:这一缩放抵消了残差累加导致的方差增长,是"数百层"得以成立的核心数值技巧);
- GELU 高效近似:f(X) = X ⊙ sigmoid(1.702X),避免调用 erf;
- FFN 输出维度为 4d(保持标准比例)。
3.3 多数据类型的统一嵌入(§5.3,Eq.15)
为了让同一个稀疏注意力框架同时吃图像、文本、音频:
- 图像:d_data = 3 维数据嵌入(行、列、通道);
- 文本与音频:d_attn = 2 维注意力嵌入,按 stride 宽度矩阵的行列索引。(推断:即把一维序列折叠成宽度为 l 的二维网格,使 strided/fixed 的行、列注意力与数据的二维组织方式一一对应——这正是 strided 对图像有效、对纯文本失效的结构根源。)
3.4 显存优化:重计算(§5.4)
采用 Chen et al. 2016 的 gradient checkpointing:前向时不保存注意力权重,反向传播时重算 attention 与 FFN 块。论文明确声称:仅凭重计算即可在长度 16,384 上训练数百层的稠密注意力网络——这在现代硬件上 otherwise 不可行。配套细节:为简化实现,注意力块内部不加 dropout,只在每次残差相加末尾施加。(推断:以约 +33% 的额外前向计算换取 O(n²)→O(n) 量级的激活显存降幅,属于典型的时空权衡。)
3.5 定制 block-sparse GPU kernels(§5.5)
理论稀疏 ≠ 实际加速,必须绕过非结构化稀疏在 GPU 上的低效:
- 从 Q/K/V 切出子块,分块(blocked)计算;
- softmax 融合进单一 kernel,中间结果驻留寄存器避免反复加载显存;
- 从不计算注意力矩阵的上三角——因果语言模型只需下三角,此举免去负偏置项(用于屏蔽上三角的 mask 加法)并将运算量减半。(推断:这个细节说明作者的优化直达硬件指令层面,"省掉 mask 加法"这类微优化在大矩阵场景有可观的实际收益。)
3.6 混合精度训练(§5.6)
- 权重以 FP32 存储,激活与梯度用 FP16(利用 V100 Tensor Core 加速);
- 动态 loss scaling 防止 FP16 下溢;
- 采样时将 query-key 乘积转回 FP32 防溢出。(推断:采样阶段输入可能更长、数值极值风险更高,单独回退精度是稳健做法。)
3.7 训练配方(§6)
Adam 优化器;5000 迭代线性 warmup;梯度裁剪阈值 1.0;weight decay 0.01;cosine 学习率衰减;嵌入维度 d ∈ {256, 512, 1024};权重初始化 N(0, 0.125/√d_in);logits 输出权重初始化为 0;默认硬件 8×V100。作者特别指出:warmup 与梯度裁剪对稳定性至关重要。(推断:结合 §5.2 的深度缩放初始化,本文实际上构建了一套完整的深稀疏网络训练稳定化配方,三者缺一不可。)
四、证据链
4.1 主结果(Table 1):四个模态全面刷新或追平 SOTA
| 数据集 | 模型配置 | 结果 | 对比基线 |
|---|---|---|---|
| --- | --- | --- | --- |
| CIFAR-10 | 59M strided | 2.80 bits/dim(2.798±0.004 over seeds 1,2,3) | 前 SOTA PixelSNAIL 2.85 |
| Enwik8 | 95M fixed | 0.99 bits/dim(0.992±0.001 over seeds 1,2,3) | 同规模 Transformer-XL 88M 为 1.03;Transformer-XL 277M 为 0.99 |
| ImageNet 64×64 | 152M strided | 3.44 bits/dim(3.437,单次运行) | 前 SOTA SPN 150M 为 3.52 |
| 经典音乐(原始音频,12kHz) | 152M | 1.97 bits/byte @ 65,536 序列长 | 见 4.4 的长度-容量规律 |
三点观察:
- Enwik8 的含金量最高:95M 参数追平 277M(近 3 倍参数)的 Transformer-XL,且上下文 12,288 长于以往方法——参数效率优势明确。
- 多 seed 报告:CIFAR-10 与 Enwik8 给出 seeds 1,2,3 的均值±标准差,而 ImageNet 64 仅单次运行(3.437)——后者结论强度稍弱。(推断:64×V100×7 天的成本使得多种子重复不现实,这也是大规模实验的通病。)
- 模态覆盖:图像(2D 周期)、文本(非周期)、音频(1D 长序列)三类数据同时验证,说明方法不是单一任务的 trick。
4.2 效率与精度的双重优势(Table 2)
| 数据集(上下文) | 模式 | bits/dim | 相对时延 |
|---|---|---|---|
| --- | --- | --- | --- |
| Enwik8(12,288 ctx) | Dense | 1.00 | 1.31 |
| Enwik8(12,288 ctx) | Fixed | 0.99 | 0.55 |
| Enwik8(12,288 ctx) | Strided | 1.13 | 0.35 |
| CIFAR-10(3,072 ctx) | Dense | 2.82 | 0.54 |
| CIFAR-10(3,072 ctx) | Fixed | 2.85 | 0.47 |
| CIFAR-10(3,072 ctx) | Strided | 2.80 | 0.38 |
这张表是全文最有力的证据:
- 图像上 strided 全胜:2.80 vs dense 2.82,时延却只有 0.38——又快又准;
- 文本上 strided 如预言般失败:1.13 明显差于 dense 1.00,验证了 §4.3 "无周期数据 strided 失效"的设计判断;fixed 则以 0.55 时延做到 0.99,仍优于 dense 的精度与时延;
- 总体结论(§8):在被比较的数据集上,稀疏模式不仅更快,还收敛到更低误差。作者推测原因是稀疏模式的归纳偏置或稠密注意力存在潜在优化问题——注意这只是推测,论文未做消融验证归因。(推断:无论哪种解释成立,都意味着"全连接"并非注意力容量的必要条件,甚至可能是一种过参数化的负担。)
4.3 上下文长度的单调收益(Table 3)
Enwik8 上评估上下文越长性能越好(训练上下文 12,288):
| min context | bits/dim |
|---|---|
| --- | --- |
| 6,144 | 0.9952 |
| 12,096 | 0.9922 |
| 12,160 | 0.9908 |
三个点严格单调改善,证明长上下文本身携带真实信息增益,而非收益饱和——这为"值得投入长序列建模"提供了直接的经验动机。(推断:该消融同时暗示 12,288 的训练上下文尚未触及收益天花板,更长序列仍有空间,呼应了音频实验推向百万长度的尝试。)
4.4 序列长度与模型容量的权衡规律(Table 4,音乐)
原始音频古典音乐实验揭示了一条经验规律:序列长 ×4 所需容量约 ÷8(⁴√4 = 8,即容量按序列长的四次方根反比缩减):
| 序列长度 | 参数量 | bits/byte |
|---|---|---|
| --- | --- | --- |
| 65,536(约 5 秒 @12kHz) | 152M | 1.97 |
| 262,144 | 25M | 2.17 |
| 1,048,576 | 3M | 2.99 |
65,536 长度样本展现全局连贯性与多样演奏风格;但更长序列因容量被迫缩减,质量快速退化。(推断:这条规律暴露了当时硬件下的真实约束——稀疏注意力解决了"算得动",但没解决"给定算力预算下长度与容量的分配问题";质量退化源于容量不足而非注意力机制本身失效。)
4.5 成本与训练配置证据(§6–7)
- 默认配置 8×V100;ImageNet 64 的 152M 模型(48 层 / 16 heads / d=512 / stride 128 / dropout 0.01 / 70 epochs)用了 64×V100 训练 7 天;
- CIFAR-10 细节:3072 字节序列,2 头 / 128 层 / d=256 / half-size FFN 与 QK 投影 / 120 epochs / lr 3.5e-4 / dropout 0.25,48,000 训练 + 2,000 验证;strided 收敛快于稠密注意力;
- Enwik8 细节:30 层 fixed / 8 heads / d=512 / dropout 0.40 / stride 128 / c=32 / merged heads / 80 epochs,90M tokens 训练;
- 另有 300M 参数的 ImageNet 64 无条件采样模型,视觉上无稀疏伪影且具长期结构。(推断:这是对"稀疏会不会产生可见 artifact"的定性回应,虽非量化指标,但对生成质量的疑虑是必要的补充证据。)
五、批判性分析
5.1 方法论亮点
- 观察驱动的归纳偏置设计:先可视化 128 层稠密网络的习得模式(Figure 2),再形式化为因子化注意力——设计有经验依据而非拍脑袋,这使得"稀疏不降精度"的结果显得合理而非侥幸。
- 理论与系统协同设计:O(n√n) 的算法主张由 block-sparse kernel、重计算、混合精度三条工程线共同兑现,Table 2 的相对时延(0.35–0.55)是系统级证据而非纸面推导。
- 诚实的负结果:Table 2 公开报告了 strided 在 Enwik8 上的失败(1.13 vs 1.00),并与 §4.3 的理论预判互相印证,增强了整体可信度。
5.2 归因未验证:为什么稀疏反而更准?
作者给出两个候选解释——稀疏模式的归纳偏置、或稠密注意力的潜在优化问题(§8)——但均为推测措辞,论文未设计消融来区分二者。(推断:若是归纳偏置,则应观察到不同稀疏模式间的系统性差异与数据结构相关;若是优化问题,则应在稠密基线上做同等训练稳定化处理后重测。缺此一步,"稀疏更准"的普适性存疑——尤其 Table 2 中 Fixed 在 CIFAR-10 上 2.85 其实差于 Dense 2.82,说明"更准"并非在所有〈数据集×模式〉组合上都成立,而是依赖模式-数据匹配。)
5.3 固定模式的表达能力天花板
预定稀疏模式无法学习 Figure 2 中出现的全局/数据依赖型注意力(§6 局限 2);c=1 时网络表达力被严重限制(§4.3)。也就是说,本文的稀疏是对"平均注意力行为"的一阶近似,对那些真正需要动态路由信息的样本(如长程指代)可能系统性吃亏。(推断:Enwik8 上 fixed 0.99 与 dense 1.00 的差距极小,可能恰恰因为字符级语言模型的多数依赖是局部的;换到强语义依赖的任务,差距未必如此乐观。)
5.4 超长序列的容量墙
音乐实验显示序列增至 1,048,576 时参数被迫压到 3M,bits/byte 从 1.97 恶化到 2.99(Table 4)。稀疏注意力把注意力的墙推远了,但没有拆除——墙外是模型容量与序列长度的硬权衡。(推断:这说明 2019 年的瓶颈已从"注意力算不动"部分转移到"整体算力预算不够",预示了后续工作需要在效率之外同时解决扩容问题。)
5.5 工程门槛与可复现性
要复现本文结果需要:定制 block-sparse GPU kernel、FP16 动态 loss scaling、梯度检查点、精细的初始化/warmup/裁剪配方(§5.4–5.6、§6)。任一环节缺失都可能导致不收敛或无加速。(推断:这套系统工程复杂度在当时远超学术实验室的平均能力,客观上提高了独立复现与公平对比的难度;ImageNet 64 的 64×V100×7 天成本也使第三方重复实验近乎不可能。)
5.6 实验设计的其他保留项
- ImageNet 64 的 3.437 为单次运行,无 seed 方差;
- bits/dim 与 bits/byte 跨数据集不可直接比较,横向结论只能在各数据集内部成立;
- 音频"质量快速退化"以定性描述为主,缺少感知类指标。(以上均基于源文件所载数据范围的保守评估。)
六、启示与影响
6.1 直接启示:四条可迁移的工程原则
- 先看模型在学什么,再设计约束:Figure 2 的可视化→因子化设计路径,适用于一切"给模型加结构先验"的研究;
- 连通性是稀疏化的底线约束:p 步可达(p+1 长度路径)保证信息全网流通,任何稀疏化方案都应显式论证这一点;
- 深稀疏网络的稳定化三角:pre-activation 残差 + 1/√(2N) 深度缩放初始化 + warmup/裁剪,三位一体缺一不可;
- 理论加速必须过 kernel 这一关:块稀疏分块计算、softmax kernel 融合、跳过上三角减半运算、寄存器驻留——算法创新只有落到硬件友好的实现上才产生真实收益。
6.2 学术谱系定位(推断)
本文是"高效注意力"研究线的奠基性节点之一:(推断:其后涌现的 Longformer、BigBird、Routing Transformer 等工作均可视为对本路线的扩展——或引入局部+全局混合窗口,或将固定稀疏升级为可学习/数据依赖稀疏,或给出稀疏模式表达能力的理论刻画。Fixed 模式的"列汇总枢纽"结构与后来的全局 token 设计思路相通;而其"strided 适配周期数据、fixed 适配非周期数据"的二分法,成为后续稀疏模式选型的基本参照。)此外,本文作者团队出自 OpenAI,(推断:同年前后该团队正推进大规模语言模型的序列建模,长上下文注意力工程在此后的大模型时代成为标配能力,本文的重计算、块稀疏 kernel、混合精度等组件均被后续基础设施吸收演化。)
6.3 对当下长上下文建模的镜鉴
- "更长上下文 ⇒ 更好"已被本文定量预演(Table 3 的单调改善),今天的长上下文竞赛本质上延续着同一逻辑;
- 长度-容量权衡仍未消失:Table 4 的 ÷8 规律提醒我们,拉长序列必然挤占单位长度的模型容量,现代长上下文模型同样面临此约束,只是表现形式变为 KV cache 与算力预算之争;(推断)
- 数据依赖型稀疏仍是开放问题:本文承认固定模式学不出全局数据依赖路由(局限 2),如何以接近固定模式的成本获得接近数据依赖的效果,至今仍是高效注意力研究的核心张力;(推断)
- 归纳偏置的双刃剑:稀疏带来速度与可能的正则化收益,但也锁死了信息通路——选择哪种稀疏,本质是在押注数据的结构假设(周期 or 层级),这个决策框架在本文之后并未过时。(推断)
*报告完。除标注「(推断)」处外,全部事实与数字以《00-源数据提取.md》为唯一依据。*