摘要：这一方法的核心创新在于：研究团队首先将隐藏空间中的 tokens 重塑为“视频”的结构，即按照时间、高度、宽度 [t, h, w] 的格式进行 reshape。随后对 reshaped 的 query 和 key 进行下采样，采用平均池化操作构建一个低分辨率的
近日，南京理工大学本科校友、美国西北大学沈轩博士生和所在团队提出一种全新的加速视频扩散模型的方法，名为 Draft Attention。
图 | 沈轩（来源：沈轩）
这一方法的核心创新在于：研究团队首先将隐藏空间中的 tokens 重塑为“视频”的结构，即按照时间、高度、宽度 [t, h, w] 的格式进行 reshape。随后对 reshaped 的 query 和 key 进行下采样，采用平均池化操作构建一个低分辨率的注意力图。这一草图注意力图能够高效捕捉视频中的关键区域，同时识别出存在冗余的部分。
接着，研究团队利用该草图注意力图生成稀疏掩码，来引导完整分辨率下的稀疏注意力计算，仅保留重要的注意力连接。通过这种方式他们在几乎不影响生成质量的情况下，显著降低了注意力模块的计算开销。
研究团队的理论分析也从数学上证明了：这种“草图”注意力图与原始全分辨率注意力图之间的误差是可控的，并且引入的稀疏性误差也是有界的。
实验结果进一步验证了方法的有效性——Draft Attention 在保持生成质量的同时，在 H100 和 A100 GPU 上分别实现了最高 1.75 倍和 2 倍的端到端加速效果。
简单来说，本次工作不仅提出了一种无需重新训练、即插即用的加速框架，也为高效视频生成提供了新的思路。
这项成果具有广泛的应用前景，尤其是在对生成速度和计算资源要求较高的场景中会非常有价值。比如，在未来一到两年内，研究团队可以预见它被用于提升视频生成平台的效率，降低内容创作者使用 AI 生成高质量视频的门槛。此外，它还可以用于多模态大模型中的视频生成模块，加速模型响应，提升用户交互体验。
往更长远看，随着边缘计算能力的提升，研究团队的方法有望被部署到本地终端设备，比如手机、XR 眼镜等，从而实现高质量视频的本地快速生成。这对于虚拟现实、游戏制作、数字人等领域都具有重要意义。
同时，草图注意力作为一种轻量、高效的稀疏计算机制，也可能被扩展到超清图像生成、视频编辑、三维建模等其他视觉生成任务中。总的来说，这项工作为大规模、高质量视觉生成的高效实现提供了一种新的思路。
据介绍，沈轩和朋友韩晨夏（https://cxhan.com/）在一次关于 AI 加速的讨论中，注意到当前的视频生成任务，尤其是高分辨率视频的生成，存在极高的时间成本。即便是在顶级 GPU 如 A100 和 H100 上，生成一个仅几秒钟的视频也要花费几十分钟。研究团队意识到，这一领域在生成效率方面仍有很大的优化空间，于是决定深入探索视频生成加速的可行性。
后来，他们关注到其他团队的一篇新论文 Sparse VideoGen (SVG)（https://arxiv.org/abs/2502.01776）。这项工作通过在空间和时间维度上分别构建稀疏注意力算子，在一定程度上缓解了视频生成中的冗余计算问题。
然而，研究团队在深入分析后发现了该方法的一些局限性：SVG 所采用的稀疏模式是预定义的，且只支持两种稀疏策略，这种固定的设计容易在高稀疏度下损失视频生成质量，限制了其适应性和效果。
随后，沈轩开始动手实践，基于腾讯开源的视频生成模型 HunyuanVideo（https://github.com/Tencent-Hunyuan/HunyuanVideo），尝试在其中引入 Max Pooling 机制，以 20% 的稀疏度生成视频。初步结果出人意料地好，这一实验结果让他们意识到：视频扩散过程中的注意力计算存在大量冗余，远比研究团队最初预想的要多。
为了进一步提升计算效率，研究团队尝试将草图注意力图引导下的稀疏模式对齐到模块化稀疏注意力结构中，并发现在满足 pool_h × pool_w = block_size 的条件下，可以更好地与现有的高效注意力框架兼容。
然而，新的挑战也随之而来：草图注意力生成的稀疏掩码在原始空间中是离散分布的，难以直接高效执行。为此，研究团队设计了一种重排策略，将这些稀疏块聚集成连续的内存布局，使得注意力计算可以在 GPU 上高效执行。
下图直观地展示了这一过程：草图注意力图（Draft Map）所对应的模块化稀疏注意力是离散的（下图左边的 Full Map），只有引入适当的重排（Reorder）才能够使得这些稀疏模块聚集在一起（下图右边 Reordered Full Map）。
在验证整体思路可行后，沈轩继续深入开展实验，进一步优化了设计细节。他们发现，在高稀疏度的设定下平均池化相比最大池化在生成质量上表现更佳，尤其是在保持背景一致性和主体细节方面更为稳定。
随后，沈轩采用了 80% 稀疏度生成了一系列视频，结果表明不仅生成速度显著提升，而且视频质量几乎没有明显下降。为了更全面地评估方法性能，他们引入了多项图像和视频评估指标，包括 PSNR、SSIM、LPIPS 以及视频生成领域的权威基准 VBench。
对比结果显示，在相同稀疏率下，研究团队的方法在生成质量上明显优于 SVG。此外，他们还对草图注意力图的计算开销进行了量化评估。结果表明，尽管他们在前期引入了 draft attention 的计算过程，但整体额外开销极小，相较于整个视频生成过程几乎可以忽略不计。这进一步证明了本次方法在加速效率和生成质量之间实现了良好平衡。
下一步，研究团队打算继续将量化（Quantization）引入到 Draft Attention 当中来进一步提高视频生成速度。
来源：科技列车