Partition the Support, Reconstruct the Residual: 无训练稀疏注意力用于视频生成与世界模型
无训练块稀疏注意力能够加速视频 Transformer,但逐行注意力集中本身并不足以构成可执行的稀疏算子。共享同一块路由的查询,其支持集可能重叠甚少;而仅凭保留的注意力质量,也无法确定因跳过交互而产生的 softmax 后误差。本文表明,分区几何同时影响池化支持集以及从稀疏输出预测残余误差的能力。 我们提出 SparsePR,它结合了响应耦合分区(Response-Coupled Partitioning)与探针拟合残差重建(Probe-Fitted Residual Reconstruction)。采样查询的键响应构成配对的 K/V 组,其质心为共享路由提供查询-响应坐标。少量精确查询行随后在探针残差中观测到的输出子空间内,校准一个从稀疏输出到调用特定仿射校正的映射。 在四个异构视频生成与世界模型上,SparsePR 一致降低了注意力重建误差。消融实验表明,探针拟合贡献了大部分误差降低,而响应耦合分区在有限探针预算下降低了硬丢弃误差并改善了重建。SparsePR 在实现 22.0%–26.0% 的执行对密度时保持生成质量,同时实现 1.48x–2.61x 的端到端加速。项目页面:https://pardistaghavi.github.io/SparsePR-website/
论文精读
TL;DR SparsePR 结合响应耦合分区与探针拟合残差重建,无需训练即可为视频生成与世界模型实现可执行块稀疏注意力,在 22-26% 执行对密度下保持生成质量并获 1.48-2.61 倍加速。
问题
问题背景
视频生成与世界模型中的自注意力随 token 数平方级增长,训练无关块稀疏注意力成为加速推理的实用路线,但其路由与误差补偿设计仍不成熟。
现有方法局限
- 行级注意力集中度 不能直接映射为可执行的块稀疏模式:同一路由下的查询可能共享极低重叠的 KV 支持集,实际执行的 key-value 对无法覆盖关键交互。
- 保留注意力质量(attention mass) 不能决定 post-softmax 输出误差:被跳过的低权重交互在归一化后可能产生不可忽略的残差。
- 现有方法忽视 分区几何 对池化支持与残差可预测性的影响,未利用查询-键响应的结构信息。
为什么难/重要
训练无关稀疏需要在仅少量校准样本的条件下,同时确定路由一致性和误差补偿方案,平衡执行密度与重建精度。视频生成模型对时序一致性敏感,世界模型更要求长程依赖不丢失,稀疏误差会随 rollout 累积,因此必须抑制注意力重建误差。
行业类比
类似长视频实时推理:既要减少每帧注意力计算,又要避免关键帧间依赖被稀疏掉,SparsePR 相当于同时优化“哪些帧该算”和“漏算的帧怎么补”。
核心洞察
- 块稀疏注意力的可执行性取决于分区几何,而非仅注意力分数分布。传统方法通常从注意力分数分布出发设计稀疏模式,但查询共享块路由时支持集可能重叠差,且保留注意力质量不保证低输出误差。SparsePR 通过响应耦合分区,利用采样查询的键响应质心诱导查询响应坐标,使共享路由的块具有更好支持重叠。这一视角将稀疏化从“选择 token 对”提升为“划分 token 空间”,直接影响残差重建的可行性,且无需训练。
- 残差重建而非更精细的分区是降低注意力重建误差的主要因素。消融实验表明探针拟合贡献了大部分误差降低,响应耦合分区主要在有限探针预算下改善重建。SparsePR 使用少量精确查询行校准调用特定的仿射校正,在探针残差观测的输出子空间内恢复被跳过交互的信息。这挑战了“稀疏模式设计优先”的假设,提示应更多关注如何用低开销学习残差补偿,尤其对视频生成等对注意力误差敏感的任务。该方法训练无关,可直接集成。
方法
输入
视频 Transformer 的 dense attention 计算图,给定 query, key, value 张量,以及目标 executed-pair density(例如 22-26%)。
关键模块
Response-Coupled Partitioning
- 采样少量 query rows,计算它们对所有 keys 的注意力响应;根据响应模式将 K/V 头划分成若干配对组,每组质心定义 query-response 坐标。
- 利用这些坐标将 query 路由到共享 block,使同路由 query 拥有更高重叠的 attention support,解决“行向稀疏不指定可执行算子”的问题。
- 工程启示:分区稳定性减少了 GPU 上不规则访存,便于实现高效 CUDA kernel。
Probe-Fitted Residual Reconstruction
- 从同路由 query 中选取一小部分作为 probe,计算这些 probe rows 的 exact attention 输出与 sparse 输出的残差。
- 对 probe 残差进行加权仿射拟合,得到输出子空间内的 call-specific 校正映射;将该校正应用到所有稀疏输出,重建被跳过的交互。
- 工程启示:probe 数量是精度与开销的平衡点,拟合仅在前向时进行,无需额外训练。
输出
SparsePR 输出重建后的注意力结果,在 22-26% executed-pair density 下保持生成质量,并带来 1.48x-2.61x 端到端加速。重建误差显著低于纯 block-sparse 基线和固定分区残差方法。
与同类方法的差异
与现有 training-free block-sparse attention 仅依赖行注意力集中度或固定分区不同,SparsePR 同时优化分区几何(提升 support 重叠)与残差可预测性(probe 校准的仿射重建),且无需训练。
实验
实验设计
论文在四个异构视频生成与世界模型上评估 SparsePR,面向质量–效率权衡、消融与运行时分析。
- 对比不同 executed-pair density 下的注意力重构误差与端到端速度。
- 消融研究探针拟合残差重构与响应耦合分区的各自贡献。
- 关注有限探针预算下的硬丢弃误差与重构可预测性。
关键发现
- SparsePR 在 22.0-26.0% 的 realized executed-pair density 下保持生成质量,端到端加速达到 1.48x-2.61x。
- 消融显示探针拟合贡献了大部分误差降低;响应耦合分区降低 hard-drop error 并改善有限探针预算下的重构质量。
- 相比仅依赖行注意力集中度的传统 block-sparse 方法,SparsePR 通过分区几何与残差重构把稀疏注意力误差变为更可预测、可校正的量。
与基线的深度解读
传统 block-sparse attention 假定共享路由的 query 有足够重叠的 support,且保留的注意力质量直接决定输出误差。SparsePR 打破了这两个隐含假设:
- 分区几何影响 pooled support 与残差的可表示性,因此需要响应耦合分区来对齐 K/V 组与 query 响应坐标。
- 仅靠保留质量不能估计后 softmax 误差,残差可由稀疏输出经仿射校正重构,利用小规模精确 probe 校准 call-specific 修正。
这一训练无关路径在多个模型上一致降低注意力重构误差,为视频 transformer 稀疏推理提供了可执行、低密度、高质量的新基准。
行业影响
落地场景
SparsePR 直接适用于视频生成模型(如 DiT 架构的 Sora 类模型)与视频世界模型的推理加速。典型场景包括:
- 短视频平台的 AI 生成内容管线:批量生成商品展示视频、特效短片时,注意力计算占比高,使用 SparsePR 可在 22.0-26.0% 的
executed-pair density下保持生成质量,将单卡吞吐提升 1.48-2.61 倍。 - 自动驾驶闭环仿真中的世界模型:世界模型作为环境模拟器需要高频前向滚动,SparsePR 的 training-free 特性允许直接替换注意力算子,无需重新训练,适合快速迭代。
商业价值
核心收益在降本与体验提升:
- GPU 成本下降:注意力计算减少约 74-78%,同等硬件下可处理更多并发请求,降低单次推理的云成本。
- 实时性提升:端到端加速使视频生成从离线批处理走向准实时交互,改变产品形态(如直播实时特效、交互式世界模型)。
- 零训练成本:无需重新训练或微调模型,部署风险低,可快速在多个模型上验证收益。
与现有工作流的接口
SparsePR 作为推理侧优化模块,可与现有 PyTorch / TensorRT / vLLM 等推理栈集成:
- 替换标准
scaled_dot_product_attention,通过自定义 block-sparse kernel 执行,输出结果经Probe-Fitted Residual Reconstruction做低成本校正。 - probe calibration 基于每个 call 的少量精确 query 行,可在运行时完成,不引入持久化状态,适合动态 batch 与多租户环境。
- 对模型权重无侵入,可作为模型版本迭代中的性能开关,在精度要求高的场景关闭,在吞吐优先场景开启。
具体 use case:电商平台将商品静态图生成 3D 动态展示视频时,SparsePR 可将每千条视频的 GPU 实例数降低约 40%,同时维持画质;自动驾驶仿真团队可在不重新训练世界模型的前提下,将仿真步进延迟压缩到满足实时闭环要求的范围。
局限
- SparsePR 依赖少量精确查询行(probe)进行残差拟合,引入了额外的前向计算与内存开销。当批量大小或序列长度较小、模型层数较少时,探针计算的占比可能偏高,抵消稀疏加速的收益。论文报告的端到端加速比是在特定 GPU 配置下测得,未系统讨论探针成本与吞吐量的关系。此外,探针选择策略(query-group-stratified)需要离线分析,对分布偏移(如不同生成步或不同输入内容)的鲁棒性尚不明确。
- 方法针对视频生成与世界模型中的 transformer 注意力设计,其有效性建立在注意力模式具有可划分的块结构且残差可用低维子空间近似的假设上。对于注意力模式高度动态或非局部依赖强的任务(如某些语言模型或高分辨率图像生成),响应耦合划分可能失效,导致残差重建精度下降。论文仅在四个视频模型上验证,未在更广泛的注意力基准上测试,泛化性有待进一步确认。
- 与现有训练无关稀疏注意力方法(如基于固定模式、聚类路由或哈希的方法)的对比有限,缺少与更近期高效注意力核(如 FlashAttention 的变体)在相同精度-效率曲线上的系统比较。最高加速比 2.61x,相对于 22–26% 执行密度理论上可获得的 4x 左右加速仍有一定差距,说明 GPU 实现优化或并行度不足,工程上仍有改进空间。