质量差距有三个可分离来源
稀疏注意力是降低扩散 Transformer 在长视频和高分辨率 3D 序列上二次开销的常用手段,但高稀疏度通常会损害生成质量。这篇论文要问的是:这一质量差距究竟来自哪里,而不是默认它只是选择器精度问题。
作者在相同注意力密度下进行受控 oracle 对比,把朴素块稀疏注意力(BSA)与三种 oracle 相比较,这些 oracle 都用精确的稠密注意力概率来最大化保留的注意力质量:块 oracle 为每个查询块选择 KV 块;token oracle 选择单个 KV token,但在每个查询块内共享同一套选择;逐查询 oracle 则让每个查询各自选择 KV token。
对比这些 oracle 可分离出差距的三个来源:结构绑定误差,源于把 KV token 绑定成块并迫使同组查询共享一套选择;选择误差,源于均值池化等近似打分;被丢弃尾部误差,源于被省略的 token 仍带有非零注意力质量。这一分解,而非某个单一精度数字,才是论文的诊断性贡献。
均值池化块打分漏掉高权重 token
朴素 BSA 用均值池化后的查询与键表示的点积为每个查询块/KV 块对打分,然后在给定注意力密度下保留得分最高的 KV 块。论文表明这种池化打分并不能保持平均注意力概率:池化查询对键做 softmax 的结果,并不等于该组内各查询真实注意力权重的均值。
这种不匹配之所以重要,是因为池化打分可能漏掉对该组具有高累积注意力质量的 KV token。oracle 对比证实两种约束都会造成损失:从块 oracle 转到共享查询块的 token oracle 可消除 KV 绑定,再进一步转到逐查询 oracle 可消除查询绑定,每一步都带来精度提升。在相同实际密度下,更细的 KV 粒度和把更相似的查询分组,都能提升保真度。
token 级选择需要瓦片对齐的查询分组
MC-Sparse 通过跨块边界为每个查询组选择单个 KV token 来消除 KV 绑定,按索引在线收集被选中的键和值,而不把它们重新打包成块。由于多头注意力使用各头专属的键和值,高效瓦片必须在每个头内批量处理查询,因此该方法把 C 个相似查询组成一组共享同一套 KV 选择,C 与内核的查询瓦片大小相匹配。
最小化组内查询方差等价于最大化组间方差,但 k-means 等标准聚类不保证组大小相等,而强制均衡约束会增加分配开销并使 GPU 并行复杂化。作者采用主方向分裂划分的中位数分裂版本:递归地把每组投影到其主方向,并在中位数处切分有序查询,直到每组恰好包含 C 个查询,从而在促进组内相似性的同时保持瓦片对齐。
精确的 token 级选择需要注意力分数,而 FlashAttention 并不暴露这些分数,因此作者实现了用于精确 KV 选择的两遍内核,以及用 CuTeDSL 编写的 token 稀疏注意力内核。正是这一工程代价让更细的粒度变得可行,而不只是准确。
复用精确选择优于重算均值池化
在没有训练好的索引器时,每个去噪步都重算精确注意力分数代价过高,因此 MC-Sparse 利用时间稳定性:只在锚点步执行精确选择和查询分组,并在随后的复用步中复用得到的查询组、KV 索引和稠密-稀疏输出残差。当前的查询、键和值仍在每一步重新计算。
论文报告,来自较早步骤的精确选择优于每一步重算的均值池化选择,且注意力模式在相邻去噪步之间保持稳定。这是把分组和选择都摊销掉、而非在每一步近似它们的经验依据。
锚点步是稠密预热或刷新段的最后一步。所报告的配置对视频模型使用锚点步 {10, 26},对 3D 模型使用 {1},并按模型给出预热层和预热步,例如 Minimax-H3-Base 为 1/52 预热层和 10/49 预热步。
残差补偿找回被丢弃的尾部
被丢弃的 token 通常注意力权重很小,但仍做出非零贡献,即使采用逐查询 oracle 选择,也会与稠密注意力之间留下输出差距。块稀疏方法用块均值表示或块级统计来近似这些贡献,但 token 级选择会不规则地切过这些块,无法直接复用这些摘要。
MC-Sparse 转而利用时间维度:稠密-稀疏残差在去噪步之间的变化小于注意力输出本身,而相比均值池化选择,精确选择进一步减小了这种变化。在锚点步,该方法记录稠密与稀疏注意力输出之间的残差;在复用步,它把缓存的残差加到稀疏注意力输出上。残差与缓存的查询组和 KV 索引遵循相同的刷新计划,且构造它只需在每个锚点额外做一次稀疏注意力评估和一次逐元素相减。
加速伴随可忽略的实测质量损失
在视频和 3D 生成模型上,MC-Sparse 相比所评估的稀疏注意力基线,对稠密注意力输出有更高的保真度,并取得更大的去噪加速,且没有可见的质量下降。相对稠密注意力,它在 Minimax-H3-Base 上取得 1.80 倍去噪加速,在 3D 资产生成上取得 2.32 倍加速,两者质量损失均可忽略。
这些是相对稠密注意力的去噪加速,而非端到端生成加速,并且与论文列出的具体模型配置和预热计划绑定。保真度主张表述为对稠密注意力输出的保真度,这与感知质量或任务精度是不同的度量,因此两者不应互相替代。
论文还报告,当 GPU 显存不足时,可将不活跃的缓存条目卸载到 CPU 内存,并在使用前预取,使传输与计算重叠以隐藏延迟。这是缓存方案的部署细节,而非对精度论证的改变。
oracle 分析界定了稀疏能挽回多少
三类误差分解来自 oracle 对比,这些 oracle 在相同密度下用精确稠密注意力概率最大化保留的注意力质量。这使它成为一种诊断性的上界分析:它显示差距中有多少可归因于结构绑定、选择误差和被丢弃尾部质量,但 oracle 本身并不是可部署的选择器。
MC-Sparse 免训练,并在去噪步之间复用缓存元数据,因此其行为取决于锚点步与复用步之间的注意力稳定性。所报告的锚点与预热配置是模型专属的,加速数字与这些设置绑定,而非与任意稀疏度或序列长度绑定。
评估覆盖视频和 3D 资产生成模型。论文并未证明同样的误差分解或同样的复用计划可迁移到其他扩散 Transformer 任务、其他稀疏度目标或训练期使用。
在相同实际密度下比较稀疏方法
最有用的结论是一条比较规则,而非某个单一数字:稀疏注意力方法应在相同实际注意力密度下比较,其质量差距应分别归因于选择粒度、打分精度和被丢弃尾部质量。MC-Sparse 的贡献在于表明这三者都重要,且 token 级选择加残差补偿能解决块级选择无法解决的部分。
对于判断该方法是否适用于自身场景的读者,论文自身的变量就是决策输入:工作负载是长序列视频还是高分辨率 3D 生成,模型能否容忍一段稠密预热和周期性锚点步,以及目标稀疏度是否高到使被丢弃尾部质量成为约束瓶颈。所报告的 1.80 倍和 2.32 倍去噪加速是这些特定配置的参考点,而非普遍保证。
本文引用的文献
MC-Sparse: Deconstructing and Closing the Dense-Sparse Attention Gap in Diffusion Transformers
MC-Sparse 是一个免训练的稀疏注意力框架,它选择单个 KV token,把相似查询组织成与瓦片对齐的分组,并在去噪步之间复用缓存的选点结果与稠密-稀疏残差,相对稠密注意力在 Minimax-H3-Base 上取得 1.80 倍去噪加速、在 3D 资产生成上取得 2.32 倍加速,且质量损失可忽略。
