[论文解读] Real-Time Video Generation with Pyramid Attention Broadcast

Xinrong Zhao, Xiaolong Jin+2 more|arXiv (Cornell University)|Aug 22, 2024

Advanced Vision and Imaging被引用 4

一句话总结

本文提出金字塔注意力广播（PAB），一种无需训练的实时视频生成方法，适用于DiT架构模型。该方法通过基于方差的金字塔式策略，在扩散步骤间广播注意力输出，有效减少注意力冗余。PAB在Open-Sora、Open-Sora-Plan和Latte模型上实现了高达20.6 FPS的720p视频生成速度，推理速度提升达10.5倍，且无任何质量损失。

ABSTRACT

We present Pyramid Attention Broadcast (PAB), a real-time, high quality and training-free approach for DiT-based video generation. Our method is founded on the observation that attention difference in the diffusion process exhibits a U-shaped pattern, indicating significant redundancy. We mitigate this by broadcasting attention outputs to subsequent steps in a pyramid style. It applies different broadcast strategies to each attention based on their variance for best efficiency. We further introduce broadcast sequence parallel for more efficient distributed inference. PAB demonstrates up to 10.5x speedup across three models compared to baselines, achieving real-time generation for up to 720p videos. We anticipate that our simple yet effective method will serve as a robust baseline and facilitate future research and application for video generation.

研究动机与目标

• 解决基于DiT的视频扩散模型推理延迟过高的问题，尽管其输出质量高，但高延迟限制了实际部署。
• 克服扩散Transformer中注意力计算的低效性，尤其是在具有复杂时空依赖关系的视频生成任务中。
• 开发一种无需训练的加速方法，在不重新训练的前提下，保持多种基于DiT的视频模型的生成质量。
• 在不改变模型架构或分布的情况下，实现高分辨率（最高720p）视频生成的实时推理（≥20 FPS）。
• 设计一种可扩展、通用的加速框架，适用于未来基于DiT的视频模型，且无需重新训练。

提出的方法

• 提出金字塔注意力广播（PAB），一种无需训练的方法，通过利用注意力差异的U型模式，在扩散步骤间重用注意力输出。
• 根据注意力头的方差，对空间、时间及跨注意力头应用不同的广播范围——空间注意力广播范围最广，跨注意力最窄，形成金字塔策略。
• 引入广播序列并行机制，降低分布式推理中的通信开销，实现高效的多GPU部署。
• 使用连续步骤间注意力输出的均方误差（MSE）来量化并指导广播策略的选择。
• 将广播机制扩展至MLP层和扩散步骤，进一步提升整体效率。
• 在不修改模型架构或无需重新训练的前提下，将PAB无缝集成至现有基于DiT的视频模型（如Open-Sora、Latte）中。

实验结果

研究问题

• RQ1 能否有效利用基于DiT的视频扩散模型中的注意力冗余，在不损失质量的前提下加速推理？
• RQ2 扩散步骤间注意力差异的U型模式是否能够支持结构化、基于方差的广播策略，以实现最优效率？
• RQ3 针对不同注意力类型（空间、时间、跨注意力）定制的金字塔式广播策略，是否优于统一或固定范围的广播方法？
• RQ4 广播序列并行机制在降低视频生成分布式推理中的通信开销方面有多高效？
• RQ5 PAB能否在多个先进的基于DiT的视频模型上实现高分辨率（720p）视频生成的实时推理（≥20 FPS）？

主要发现

• PAB在三种领先的基于DiT的视频模型（Open-Sora、Open-Sora-Plan、Latte）上实现了最高达10.5倍的推理速度提升。
• 该方法在H100 GPU上实现了720p视频生成的20.6 FPS实时推理性能。
• 注意力差异的U型模式得到实证验证，中间70%的扩散步骤中变化最小。
• 基于注意力方差的广播策略提升了效率：空间注意力（方差最大）使用更广的广播范围，而跨注意力（最稳定）使用更窄的范围。
• 广播序列并行机制显著降低了分布式推理中的通信开销，支持可扩展且高效的多GPU部署。
• PAB在无需任何微调或重新训练的情况下保持了原始生成质量，证明了其在多种模型上的泛化能力与鲁棒性。