提出线性压缩注意力机制,让扩散Transformer生成图像更快更省资源。
EDiT: Efficient Diffusion Transformers with Linear Compressed Attention

- 用卷积网络调节查询,空间聚合键值,实现线性注意力。
- 在稳定扩散3.5中提速2.2倍,图像质量几乎不变。
- 适合移动端或高分辨率生成,兼顾速度与效果。
扩散Transformer(DiTs)已成为文本到图像合成的主流架构,能生成高质量、逼真的图像。然而,传统DiTs中注意力机制的二次方复杂度限制了其在高分辨率图像生成或资源受限设备上的应用。本文提出一种高效的扩散Transformer(EDiT),缓解这一效率瓶颈。首先,设计一种新型线性压缩注意力方法:通过多层卷积网络调节查询以引入局部信息,同时对键和值进行空间聚合。其次,针对多模态输入,提出混合注意力方案——图像间交互使用线性注意力,提示词交互仍采用标准缩放点积注意力。融合两种方法后,构建出表达能力强、时间复杂度线性的多模态高效扩散变压器(MM-EDiT)。将EDiT和MM-EDiT集成到PixArt-Sigma(常规DiT)和Stable Diffusion 3.5-Medium(MM-DiT)中,经蒸馏后最高实现2.2倍加速,且图像质量保持相当。
原文摘要 · Abstract (English)
Diffusion Transformers (DiTs) have emerged as a leading architecture for text-to-image synthesis, producing high-quality and photorealistic images. However, the quadratic scaling properties of the attention in DiTs hinder image generation with higher resolution or on devices with limited resources. This work introduces an efficient diffusion transformer (EDiT) to alleviate these efficiency bottlenecks in conventional DiTs and Multimodal DiTs (MM-DiTs). First, we present a novel linear compressed attention method that uses a multi-layer convolutional network to modulate queries with local information while keys and values are aggregated spatially. Second, we formulate a hybrid attention scheme for multimodal inputs that combines linear attention for image-to-image interactions and standard scaled dot-product attention for interactions involving prompts. Merging these two approaches leads to an expressive, linear-time Multimodal Efficient Diffusion Transformer (MM-EDiT). We demonstrate the effectiveness of the EDiT and MM-EDiT architectures by integrating them into PixArt-Sigma (conventional DiT) and Stable Diffusion 3.5-Medium (MM-DiT), achieving up to 2.2x speedup with comparable image quality after distillation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。