当前位置：首页 > news >正文

【AI论文】使用滑动磁贴注意力实现快速视频生成

news 来源：原创 2025/5/18 10:23:04

摘要：扩散变换器（DiTs）凭借3D全局注意力机制在视频生成领域达到了最先进水平，但其计算成本高昂——生成一段仅5秒的720P视频时，仅注意力计算就占用了总推理时间的945秒中的800秒。本文引入了滑动磁贴注意力（STA）机制来应对这一挑战。STA基于这样一个观察结果：在预训练的视频扩散模型中，注意力分数主要集中在局部的3D窗口内。通过滑动并在局部时空区域内进行注意力计算，STA消除了全局注意力的冗余。与传统基于标记的滑动窗口注意力（SWA）不同，STA采用了一种新颖的硬件感知滑动窗口设计，以瓦片为单位进行操作，既保持了表达能力，又提高了硬件效率。通过细致的核级优化，STA提供了首个高效的2D/3D滑动窗口式注意力实现，达到了58.79%的多功能单元（MFU）利用率。具体来说，与FlashAttention-2（FA2）相比，STA将注意力计算速度提高了2.8至17倍，与FlashAttention-3（FA3）相比则提高了1.6至10倍。在领先的视频DiT模型HunyuanVideo上，STA在不降低质量的情况下，将端到端延迟从945秒（使用FA3）缩短至685秒，且无需重新训练。启用微调后，延迟可进一步降低至268秒，同时在VBench上的性能仅下降了0.09%。Huggingface链接：Paper page，论文链接：2502.04507

1. 引言

背景与动机：

扩散变换器（DiTs）：作为高分辨率视频生成的主流架构，DiTs通过3D注意力机制对视频帧进行空间和时间依赖性的建模，实现了长时间、视觉连贯的视频输出。然而，这种设计引入了显著的计算开销，尤其是在视频分辨率和持续时间增加时，训练和推理速度变得极为缓慢。
计算成本问题：生成一段仅5秒的720P视频时，注意力计算就占用了总推理时间的绝大部分（例如，在HunyuanVideo模型中，注意力计算占了945秒中的800秒）。这一瓶颈严重限制了DiTs在实际部署中的应用。
本文贡献：针对上述问题，本文引入了滑动磁贴注意力（STA）机制，旨在通过消除全局注意力的冗余来加速视频生成过程。

2. 问题分析

全局注意力的局限性：

计算复杂度：全局注意力机制的计算复杂度为O(N^2)，其中N是视频帧中标记（token）的数量。随着视频分辨率和帧数的增加，N迅速增长，导致计算成本急剧上升。
注意力分数分布：通过对预训练视频扩散模型的注意力分数进行可视化，发现注意力分数主要集中在局部的3D窗口内。这表明，全局注意力机制中存在大量的冗余计算。

滑动窗口注意力（SWA）的局限性：

现有SWA实现的问题：尽管SWA机制通过限制注意力计算到局部窗口来减少计算成本，但现有的2D或3D SWA实现（如NATTEN和CLEAR）未能将浮点运算减少量转化为相应的时钟加速比。这主要是因为它们的计算模式对GPU不友好，导致硬件利用率低。
混合块问题：在SWA中，由于每个查询标记（query token）关注的键标记（key token）组不同，导致注意力图中存在大量的混合块（mixed blocks）。这些混合块不仅不减少浮点运算量，还引入了额外的掩码评估开销。

3. 方法

滑动磁贴注意力（STA）：

基本思想：STA通过以磁贴（tile）为单位进行滑动窗口注意力计算，消除了混合块，提高了计算效率。具体来说，STA将视频帧划分为多个非重叠的磁贴，每个瓦片内的标记具有连续的序列索引。然后，对每个瓦片内的查询标记和键标记进行注意力计算。
磁贴大小与窗口大小：STA设置磁贴大小（T）等于FlashAttention的块大小（B）的立方根（即B=T^3），并确保视频大小和窗口大小都是瓦片大小的整数倍。这样，每个磁贴都可以完整地滑入注意力窗口内，形成密集的注意力块。
核级优化：STA基于ThunderKittens和FlashAttention3实现了高效的注意力核。通过拆分线程块为计算warp组和数据warp组，并将块间掩码逻辑与计算核解耦，STA进一步提高了计算效率。

应用于视频扩散模型：

无需训练的应用：由于视频扩散模型表现出显著的3D局部性和头部专业化模式，STA可以通过在少量提示上搜索每个头部的最优窗口大小，然后直接替换预训练视频DiTs中的3D注意力，而无需重新训练。
微调：为了进一步降低延迟，可以对应用STA后的模型进行微调。微调过程中采用了注意力蒸馏损失、最终层损失和数据损失来确保模型保持高质量的视频生成能力。

4. 实验

实验设置：

数据集与模型：实验在HunyuanVideo模型上进行，该模型是一个领先的开源视频DiT模型。此外，还在图像扩散模型FLUX上应用了STA，以展示其在2D场景下的有效性。
评估指标：效率方面，使用MFU（多功能单元利用率）和延迟来衡量注意力核的效率；质量方面，除了报告自动化指标（如VBench、SSIM、PSNR和CD-FVD）外，还进行了人类评估。

效率评估：

与现有方法的比较：与CLEAR、NATTEN、Swin和∆-DiT等现有方法相比，STA在保持高表达能力的同时，实现了更高的计算效率。具体来说，STA的MFU达到了58.79%，比现有方法高出许多倍。与FlashAttention-3相比，STA将注意力计算速度提高了1.6至10倍。
在HunyuanVideo上的应用：在不降低质量的情况下，STA将HunyuanVideo的端到端延迟从945秒缩短至685秒，且无需重新训练。启用微调后，延迟可进一步降低至268秒。

质量评估：

人类评估：在人类评估中，STA生成的视频在质量上与原始HunyuanVideo相当，且明显优于∆-DiT生成的视频。这表明STA在加速视频生成的同时，没有牺牲视频质量。
自动化指标：在自动化指标上，STA也表现出了与原始模型相当或更优的性能。例如，在VBench上，STA的得分与原始模型相差无几，甚至在某些配置下略有提高。

5. 相关工作

线性注意力方法：

线性注意力方法通过分解二次注意力操作来实现线性复杂度。然而，这些方法在视频DiTs中尚未取得成功。

扩散模型加速技术：

除了注意力加速外，还有许多技术被提出用于加速扩散模型的采样过程，如DDIM、更快的ODE和PDE求解器等。STA与这些技术是互补的，可以结合使用以实现更高的加速比。

6. 结论与未来工作

结论：

本文提出的滑动磁贴注意力（STA）机制通过消除全局注意力的冗余，显著加速了视频扩散模型的推理过程。实验结果表明，STA在不降低视频质量的情况下，实现了高达3.53倍的端到端加速比。

未来工作：

STA的概念与其他加速技术（如缓存和一致性蒸馏）是正交的，未来可以探索它们的组合使用以实现更高的加速比。此外，还可以将STA应用于其他类型的扩散模型（如图像生成模型）中，以验证其通用性。

7. 影响声明

本文的工作通过引入高效的注意力核来减少视频生成时间，同时保持输出质量。这一改进使得视频生成对于计算资源有限的研究人员和开发者来说更加实用，有望促进AI驱动的视频应用在创意产业、教育等领域的发展。尽管更快的视频生成可能会带来滥用风险，但现有的内容检测和水印技术可以帮助缓解这些问题。总体而言，本文的工作对推动可访问的视频AI系统的发展具有重要意义。

8. 详细实现与优化

磁贴与窗口的设计：

STA通过精心设计的磁贴大小和窗口大小来确保每个磁贴都可以完整地滑入注意力窗口内，从而形成密集的注意力块。这种设计不仅消除了混合块，还提高了GPU的硬件利用率。

核级优化技术：

STA采用了多种核级优化技术来提高计算效率，包括拆分线程块为计算warp组和数据warp组、将块间掩码逻辑与计算核解耦等。这些优化技术使得STA能够在保持高表达能力的同时，实现高效的注意力计算。

微调策略：

为了进一步降低延迟并保持高质量的视频生成能力，本文对应用STA后的模型进行了微调。微调过程中采用了多种损失函数来确保模型在加速的同时不损失性能。这些损失函数包括注意力蒸馏损失、最终层损失和数据损失等。

9. 实验细节与结果分析

实验设置与基准方法：

实验在HunyuanVideo模型上进行，并与CLEAR、NATTEN、Swin和∆-DiT等现有方法进行了比较。实验设置包括生成720P 5秒的视频、使用不同的注意力核实现等。

效率与质量权衡：

实验结果表明，STA在保持高质量视频生成的同时，实现了显著的加速比。与现有方法相比，STA在MFU、延迟和计算效率等方面均表现出色。此外，人类评估和自动化指标也验证了STA在视频质量方面的优越性。

不同配置下的性能：

本文还探索了不同配置下STA的性能表现。例如，在不同的稀疏度设置下，STA仍然能够保持高效的计算性能和高质量的视频生成能力。此外，STA还可以应用于不同的视频分辨率和帧数设置下，以适应不同的应用场景需求。

10. 未来研究方向

与其他加速技术的结合：

如前所述，STA的概念与其他加速技术（如缓存和一致性蒸馏）是正交的。未来可以探索这些技术的组合使用以实现更高的加速比和更好的性能表现。

应用于其他扩散模型：

除了视频生成模型外，STA还可以应用于其他类型的扩散模型（如图像生成模型）中。通过验证STA在不同类型扩散模型中的通用性，可以进一步推动其在实际应用中的发展。

硬件加速与优化：

随着硬件技术的不断发展，未来可以探索针对STA的硬件加速和优化技术。例如，可以设计专门的硬件加速器来支持STA的高效计算，或者优化现有的GPU架构以更好地支持STA的执行。这些硬件加速和优化技术有望进一步提高STA的计算效率和性能表现。