压缩多模态扩散模型注意力机制，加速图像生成不降质
DiTFastAttnV2: Head-wise Attention Compression for Multi-Modality Diffusion Transformers
- 按注意力头粒度动态调整，减少冗余计算
- 降低68%注意力计算量，2K图生成功能提速1.5倍
- 适合需要高效推理的文本生成图像场景
文本到图像生成模型，尤其是多模态扩散变换器（MMDiT），在生成高质量图像方面取得了显著进展。然而，这些模型常面临显著的计算瓶颈，尤其在注意力机制上，限制了其可扩展性和效率。本文提出DiTFastAttnV2，一种针对MMDiT的后训练压缩方法，用于加速注意力计算。通过深入分析MMDiT的注意力模式，我们发现其与以往基于DiT的方法存在关键差异，并提出按头级的箭头注意力与缓存机制，以动态调节注意力头，有效弥合这一差距。同时设计了高效融合核以进一步加速。结合局部度量方法与优化技术，该方法将最优压缩方案搜索时间缩短至几分钟，且保持生成质量。此外，在定制化核支持下，DiTFastAttnV2在2K图像生成任务中实现68%的注意力FLOPs降低和1.5倍端到端速度提升，视觉保真度不受影响。
原文摘要 · Abstract (English)
Text-to-image generation models, especially Multimodal Diffusion Transformers (MMDiT), have shown remarkable progress in generating high-quality images. However, these models often face significant computational bottlenecks, particularly in attention mechanisms, which hinder their scalability and efficiency. In this paper, we introduce DiTFastAttnV2, a post-training compression method designed to accelerate attention in MMDiT. Through an in-depth analysis of MMDiT's attention patterns, we identify key differences from prior DiT-based methods and propose head-wise arrow attention and caching mechanisms to dynamically adjust attention heads, effectively bridging this gap. We also design an Efficient Fused Kernel for further acceleration. By leveraging local metric methods and optimization techniques, our approach significantly reduces the search time for optimal compression schemes to just minutes while maintaining generation quality. Furthermore, with the customized kernel, DiTFastAttnV2 achieves a 68% reduction in attention FLOPs and 1.5x end-to-end speedup on 2K image generation without compromising visual fidelity.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。