SpargeAttention2：基于混合Top-k+Top-p掩码和蒸馏微调的可训练稀疏注意力机制 SpargeAttention2: Trainable Sparse Attention via Hybrid Top-k+Top-p Masking and Distillation Fine-Tuning
混合Top-k/Top-p掩码+速度蒸馏，实现95%稀疏率且不损质量
前置知识
稀疏注意力（Sparse Attention）
标准自注意力机制的计算复杂度为 $O(N^2)$，其中 $N$ 是序列长度。稀疏注意力通过引入二元掩码 $M \in \{0, 1\}^{N \times N}$，仅保留注意力权重矩阵 $P$ 中最重要的元素，将对应位置的计算跳过。具体来说，对注意力权重施加逐元素乘法 $P \leftarrow P \odot M$，当 $M_{ij} = 0$ 时跳过对应的点积计算 $Q_i K_j^\top$ 和值更新 $P_{ij} V_j$。这种方法可以显著减少计算量，但需要设计合理的掩码策略来决定保留哪些位置。
本文的核心就是改进稀疏注意力的掩码设计和训练方法，理解稀疏注意力的基本原理是理解全文的基础。
Top-k 和 Top-p 掩码
Top-k 掩码对注意力权重矩阵的每一行，固定保留概率最大的 $k$ 个位置。Top-p 掩码（也称 nucleus sampling）则保留概率累积和达到阈值 $p\%$ 的最小位置集合。两种方法各有优劣：Top-k 在概率分布均匀时可能遗漏重要 token，Top-p 在概率分布高度集中时可能被注意力沉降（attention sink）主导而丢弃其他有用信息。在块稀疏实现中，这两种掩码作用于池化后的块级注意力权重 $\bar{P}$。
本文的核心创新之一就是发现并解决这两种掩码在高稀疏率下的失效问题，提出混合策略。
块稀疏注意力（Block Sparse Attention）
由于细粒度的逐元素稀疏在 GPU 上映射效果差，实际实现通常采用块结构。将查询、键、值张量划分为大小为 $b_q$ 和 $b_{kv}$ 的块，每个块对 $(i, j)$ 的掩码 $M_{ij}$ 是全 1 或全 0 的块，与 GPU 的 tiling 对齐。FlashAttention 等高效内核可以利用这种块结构跳过被掩码的矩阵乘法和 softmax 计算，从而实现实际的加速。
论文的高效实现建立在块稀疏注意力之上，理解块结构对于理解算法 1 的内核实现至关重要。
流匹配（Flow Matching）
流匹配是扩散模型的一种连续时间建模框架，用速度场而非离散去噪步骤定义生成过程。给定干净样本 $x_1$、噪声样本 $x_0 \sim \mathcal{N}(0, I)$ 和时间步 $t \in [0, 1]$，构造中间潜变量 $x_t = t x_1 + (1-t) x_0$。真实速度定义为 $v_t = x_1 - x_0$，模型 $\theta$ 训练来预测这个速度。训练目标是 MSE 损失 $\mathcal{L} = \mathbb{E}\|u(x_t, c_{\text{text}}, t; \theta) - v_t\|^2$。Wan 视频模型采用了这种框架。
论文的速度蒸馏损失建立在流匹配框架之上，理解速度预测的含义是理解蒸馏目标的关键。
知识蒸馏（Knowledge Distillation）
知识蒸馏是一种模型压缩技术，由 Hinton 等人在 2015 年提出。核心思想是用一个较大的教师模型的输出作为监督信号来训练一个较小的学生模型。学生模型不仅学习数据的真实标签，还学习教师模型的输出分布，因为教师的输出包含了更丰富的暗知识。本文中，原始的全注意力模型作为冻结的教师，稀疏注意力模型作为学生，通过速度预测对齐来传递知识。
论文提出的速度蒸馏损失是核心创新之一，用教师模型的速度预测替代标准扩散损失来微调稀疏注意力模型。
注意力沉降（Attention Sink）
注意力沉降是指在 Transformer 模型中，某些特定 token（通常是序列开头的 token）会吸引不成比例的大量注意力权重，即使这些 token 的语义内容并不重要。这种现象在 LLM 和扩散模型中都有观察到（Xiao et al., 2024; Gu et al., 2024）。在稀疏注意力中，注意力沉降会导致 Top-p 掩码的问题：当某行的概率分布高度倾斜时，累积概率阈值可能仅由注意力沉降 token 就能满足，导致其他有信息量的 token 被丢弃。
理解注意力沉降现象是理解 Top-p 掩码在倾斜分布下失效的关键原因。
研究动机
当前的可训练稀疏注意力方法存在两个关键限制。第一个限制（L1）是，在非常高的注意力稀疏率（如 >90%）下，Top-k 和 Top-p 两种常见的掩码策略都会失效。具体来说，对于概率分布相对均匀的行，Top-k 保持固定数量的 token 可能只捕获了总概率的一小部分，遗漏有用的上下文信息。例如，如果一行包含十个概率为 0.1 的 token，Top-20% 只保留其中两个。对于概率分布高度倾斜的行，Top-p 可能仅由几个注意力沉降 token 就满足累积概率阈值，导致丢弃其他有信息量的 token。例如，对于 [0.6 (sink), 0.2, 0.1, ...] 这样的行，Top-p(60%) 只选择沉降 token 的概率。第二个限制（L2）是，现有方法使用标准扩散损失微调视频扩散模型时，当微调数据分布与预训练数据不匹配时会导致性能显著下降。以 Wan2.1 模型为例，即使保持全注意力，使用扩散损失微调后，1.3B 模型的美学质量从 0.6441 降到 0.6183，视觉奖励从 0.1084 降到 0.0936；14B 模型也出现了类似的退化。
本文的目标是本文的具体目标是开发一种可训练稀疏注意力方法，在实现极高注意力稀疏率（目标 95%）的同时，不降低视频生成质量。具体来说，需要同时解决三个问题：设计在高稀疏率下稳健的掩码策略、实现高效的可训练稀疏注意力内核、以及找到在数据分布不匹配情况下仍能保持生成质量的微调目标。最终目标是实现注意力计算的大幅加速（预期 16 倍以上）和端到端视频生成的显著提速。
与已有工作不同的是，本文的独特切入角度在于三个方面。首先，通过系统分析 Top-k 和 Top-p 在不同注意力权重分布（均匀 vs 倾斜）下的失败模式，提出了混合策略而非单独改进某一种掩码。其次，认识到数据分布不匹配是微调退化的根本原因，而这个问题与稀疏注意力本身无关，因此不能仅通过修改注意力结构来解决。最后，受蒸馏启发，提出用速度级别的蒸馏损失替代标准扩散损失，让稀疏注意力模型学习匹配冻结的全注意力教师模型的行为，而不是拟合低质量的微调数据。这种设计从根本上改变了微调的目标：从拟合数据变为模仿教师。
核心方法
SpargeAttention2 的整体思路可以分为三个相互配合的组件。直觉上，要让稀疏注意力在极高稀疏率下工作，需要同时解决选对 token 和学对行为两个问题。首先，对于掩码设计，作者观察到 Top-k 和 Top-p 各自在不同分布下有优势，因此将两者结合为统一的混合掩码器，取两者并集。其次，对于训练目标，作者发现标准扩散损失会让模型拟合低质量微调数据，因此改用蒸馏方式，让稀疏注意力模型模仿全注意力模型的速度预测输出。技术路线上，先通过块池化计算块级注意力权重 $\bar{P}$，然后同时应用 Top-k 和 Top-p 掩码取并集得到块稀疏掩码 $\bar{M}$，最后用高效的 CUDA 内核实现块稀疏注意力的前向和反向传播。训练时，冻结原始全注意力模型作为教师，用速度蒸馏损失对齐学生模型的输出。
核心创新点有两个。第一个是混合 Top-k+Top-p 掩码策略，与已有方法的本质区别在于：已有方法要么用 Top-k 要么用 Top-p，而本文证明两者互补。对于均匀分布的行，Top-p 优于 Top-k（Top-p 的 L1 误差 0.3726 vs Top-k 的 0.4150）；对于倾斜分布的行，Top-k 优于 Top-p（Top-k 的 L1 误差 0.1664 vs Top-p 的 0.2160）。混合策略取并集后，在两种分布下都能达到接近最优的效果。第二个创新是速度蒸馏损失 $\mathcal{L}_{VD} = \mathbb{E}\|u_{\text{sparse}}(x_t, c_{\text{text}}, t) - u_{\text{full}}(x_t, c_{\text{text}}, t)\|^2$，用全注意力教师的速度预测作为监督信号。这与标准扩散损失的本质区别在于：扩散损失驱动模型拟合微调数据分布，而蒸馏损失驱动模型模仿教师的行为，避免了数据分布不匹配导致的退化。
方法步骤详情
SpargeAttention2 的方法包含以下步骤。步骤 1：块划分，将查询 $Q$ 划分为 $T_m = N/b_q$ 个块 $\{Q_i\}$，键 $K$ 和值 $V$ 划分为 $T_n = N/b_{kv}$ 个块 $\{K_j\}, \{V_j\}$。步骤 2：计算块级注意力权重，通过对池化后的查询和键计算 $\bar{P} = \text{softmax}(\text{pool}(Q)\text{pool}(K)^\top / \sqrt{d})$。步骤 3：构建混合掩码，分别计算 $M_1 = \text{Top-k}(\bar{P}, k\%)$ 和 $M_2 = \text{Top-p}(\bar{P}, p\%)$，取并集 $\bar{M} = M_1 \cup M_2$。步骤 4：块稀疏注意力计算，对于每个被保留的块对 $(i, j)$，计算注意力分数 $S_{ij} = Q_i K_j^\top / \sqrt{d}$，应用在线 softmax 得到注意力权重 $P_{ij}$，并累积输出 $O_{ij}$。步骤 5：模型适配，将预训练模型的所有注意力层替换为 SpargeAttention2，冻结全注意力教师模型。步骤 6：速度蒸馏训练，采样数据 $(x_1, c_{\text{text}})$、噪声 $x_0$ 和时间步 $t$，构造含噪潜变量 $x_t$，分别计算教师速度 $u_{\text{full}}$ 和学生速度 $u_{\text{sparse}}$，最小化蒸馏损失 $\mathcal{L}_{VD}$。
技术新颖性
SpargeAttention2 的技术新颖性体现在多个层面。首先，混合掩码策略的形式化定义 $\bar{M}_{ij} = 1 \text{ if } j \in \text{Top-k}(\bar{P}_{i,:}, k\%) \cup \text{Top-p}(\bar{P}_{i,:}, p\%)$ 虽然简单，但基于对两种掩码失效模式的深入理论分析，特别是错误分解 $e = (p \odot (1-m) + (1-1/\tau)(p \odot m))V$ 揭示了丢弃误差和归一化误差两个分量。其次，速度蒸馏损失的提出是基于对微调退化原因的深刻洞察：退化源于数据分布不匹配而非稀疏注意力本身，因此需要改变微调目标而非注意力结构。这种模仿而非拟合的思路在可训练稀疏注意力领域是全新的。此外，微调后的注意力分布变得更集中（稀疏率从 51.3% 提升到 56.9%），这与误差分解理论一致，证明了训练确实让模型学会了更稀疏的注意力模式。
实验结果
SpargeAttention2 在多个实验设置下都取得了优异的结果。在 Wan2.1-1.3B 模型（480p 分辨率）上，SpargeAttention2 达到 95% 注意力稀疏率，注意力延迟从全注意力的 97 秒降低到 6 秒，实现 16.2 倍加速；端到端生成时间从 159 秒降到 68 秒，实现 2.3 倍加速。更重要的是，生成质量不仅没有下降，反而有所提升：图像质量从 63.67 提升到 67.68，美学质量从 64.41 提升到 65.05，VQA-a 从 81.28 提升到 83.86。在 Wan2.1-14B 模型（720p 分辨率）上，同样达到 95% 稀疏率，注意力延迟从 2550 秒降到 157 秒（16.2 倍加速），端到端时间从 3043 秒降到 650 秒（4.7 倍加速），图像质量从 68.01 提升到 69.08。与其他方法对比，SpargeAttention2 在 95% 稀疏率下仍保持高质量，而 SLA 在相同稀疏率下出现明显退化，VSA 和 VMoBA 在 90% 稀疏率下也表现出质量下降。消融实验进一步验证了每个组件的贡献：混合掩码优于单独的 Top-k 或 Top-p，训练优于不训练（训练-free 变体的 VQA 仅为 20.40），速度蒸馏优于标准扩散损失。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 视频生成质量（Wan2.1-1.3B-480p） | 图像质量 IQ | 67.68 | 63.67（全注意力） | +6.3%（质量提升同时实现 16.2× 注意力加速） |
| 视频生成质量（Wan2.1-1.3B-480p） | VQA-a 准确率 | 83.86 | 81.28（全注意力）/ 33.35（VSA@90%） | +3.2% vs 全注意力，+151% vs VSA |
| 视频生成质量（Wan2.1-14B-720p） | 图像质量 IQ | 69.08 | 68.01（全注意力） | +1.6%（质量提升同时实现 4.7× 端到端加速） |
| 注意力计算效率（Wan2.1-1.3B-480p） | 注意力延迟（秒） | 6s | 97s（全注意力）/ 11s（SLA@95%） | 16.2× vs 全注意力，1.8× vs SLA |
| 端到端生成效率（Wan2.1-14B-720p） | 端到端时间（秒） | 650s | 3043s（全注意力）/ 778s（SLA@95%） | 4.7× vs 全注意力，1.2× vs SLA |
| 消融实验：掩码策略（1.3B-480p） | VQA 综合分数 | 86.73（混合 Top-k+Top-p） | 86.90（Top-k only）/ 62.57（Top-p only） | Top-p only 严重退化，混合策略在 14B 上更优（88.22 vs 84.25） |
| 消融实验：训练目标（1.3B-480p） | VQA 综合分数 | 86.73（速度蒸馏） | 85.05（扩散损失）/ 20.40（无训练） | 蒸馏比扩散损失 +2.0%，无训练几乎不可用 |
局限与改进
论文存在几个值得关注的局限性。首先，实验仅在视频扩散模型（Wan2.1）上验证，未在图像扩散模型或 LLM 上测试，方法的通用性尚不确定。其次，训练数据使用了 3000 个私有视频，虽然来源是公开的，但数据集本身未公开发布，这影响了结果的可复现性。第三，速度蒸馏需要同时运行教师模型和学生模型，训练成本约为普通微调的两倍，尽管教师模型不需要梯度计算。第四，论文报告的加速数据基于 RTX 5090 GPU，不同硬件上的加速比可能有差异。第五，消融实验在 14B 模型上仅训练 100 步，而主结果使用 500 步，消融结果的代表性可能不完全充分。此外，论文未讨论混合掩码中 Top-k 和 Top-p 比例的选择策略，目前是手动调参。
独立分析的弱点
本文存在几个可以改进的弱点。第一，混合掩码的 $k\%$ 和 $p\%$ 超参数需要手动校准以达到目标稀疏率，缺乏自适应机制。改进方向可以是开发基于注意力分布统计量的自动参数调整方法，例如根据每行注意力权重的熵来动态调整 Top-k 和 Top-p 的比例。第二，速度蒸馏损失仅在速度预测层面对齐，可能忽略了更高层次的语义对齐。可以探索多尺度蒸馏，同时对齐中间层特征和最终输出。第三，块大小 $b_q = 128$ 和 $b_{kv} = 64$ 是固定的，不同层可能需要不同的粒度。可以研究自适应块大小策略。第四，论文未探索混合掩码的软组合方式（如加权并集），目前是硬取并集，可能损失了一些精细信息。第五，蒸馏训练的收敛性和最优训练步数未充分分析，500 步的选择缺乏理论依据。
未来方向
基于本文成果，可以延伸出多个研究方向。首先，将 SpargeAttention2 扩展到其他模态和模型架构，包括图像扩散模型、文本到图像生成、以及 LLM 的长上下文推理。其次，探索训练时的渐进式稀疏率调度，从低稀疏率逐步提升到高稀疏率，可能有助于更稳定的训练。第三，研究混合掩码的自动学习机制，让模型自己学会在不同层和不同输入上选择最优的掩码策略。第四，将速度蒸馏与其他压缩技术（如量化、剪枝）结合，实现更极致的效率优化。第五，探索在线蒸馏方案，即在推理时动态调整稀疏率，对简单输入用更高稀疏率，对复杂输入用更低稀疏率。作者也提到了将方法推广到其他扩散模型架构和任务的计划。
复现评估
复现评估方面，本文的可复现性存在一定挑战。代码和模型权重：论文未明确说明是否会开源代码和预训练权重，但作者来自清华和 UC Berkeley，有开源的传统。数据：训练使用了 3000 个私有视频，虽然来源是公开的，但数据集未公开，复现者需要自行收集类似数据。算力需求：1.3B 模型的训练相对可行，但 14B 模型的训练需要大量 GPU 资源（论文使用了 RTX 5090），且蒸馏训练需要同时加载教师和学生模型。技术难度：CUDA 内核实现基于 FlashAttention，需要较强的系统编程能力。评估基准：使用了 VBench 等公开基准，评估部分可复现。总体来说，核心算法思想清晰，但完整复现需要相当的工程投入和计算资源。
论文图表
展示了两种典型的注意力权重矩阵 P 的分布情况：(a) 均匀分布，每行的概率相对平均分散；(b) 倾斜分布，每行的概率高度集中在少数位置。两种情况下都展示了保持 60% 累积概率的掩码效果，直观说明了 Top-k 和 Top-p 在不同分布下的行为差异。
这张图帮助读者理解为什么单一掩码策略会在某些分布下失败，是理解混合掩码动机的关键视觉辅助。
在两个模型规模上进行了系统消融：掩码设计（混合 vs Top-k only vs Top-p only）、训练效果（可训练 vs 训练-free）、训练目标（速度蒸馏 vs 扩散损失）。结果显示混合掩码在 14B 上优势明显（88.22 vs 84.25），训练-free 几乎不可用（VQA 20.40/45.85），蒸馏优于扩散损失（86.73 vs 85.05 / 88.22 vs 85.05）。
这个表格通过消融实验验证了每个组件的贡献，是理解方法设计合理性的关键证据。