EDiT: Efficient Diffusion Transformers with Linear Compressed Attention

📄 arXiv: 2503.16726v2 📥 PDF

作者: Philipp Becker, Abhinav Mehrotra, Ruchika Chavhan, Malcolm Chadwick, Luca Morreale, Mehdi Noroozi, Alberto Gil Ramos, Sourav Bhattacharya

分类: cs.CV, cs.LG

发布日期: 2025-03-20 (更新: 2025-08-11)

💡 一句话要点

提出EDiT：一种线性压缩注意力的高效扩散Transformer，加速高分辨率图像生成。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 扩散模型 Transformer 注意力机制 线性复杂度 图像生成 多模态学习 高效计算

📋 核心要点

  1. 传统扩散Transformer（DiT）的注意力机制计算复杂度高，限制了其在高分辨率图像生成和资源受限设备上的应用。
  2. EDiT通过线性压缩注意力机制，利用卷积网络提取局部信息调制查询，并空间聚合键和值，降低了计算复杂度。
  3. 实验表明，EDiT集成到现有DiT模型后，在保持图像质量的同时，实现了高达2.2倍的推理加速。

📝 摘要（中文）

扩散Transformer（DiT）已成为文本到图像合成的主流架构，能够生成高质量和逼真的图像。然而，DiT中注意力的二次方复杂度限制了其在高分辨率图像生成或资源受限设备上的应用。本文提出了一种高效扩散Transformer（EDiT），旨在缓解传统DiT和多模态DiT（MM-DiT）中的效率瓶颈。首先，我们提出了一种新颖的线性压缩注意力方法，该方法使用多层卷积网络来利用局部信息调制查询，同时在空间上聚合键和值。其次，我们为多模态输入设计了一种混合注意力机制，将线性注意力用于图像到图像的交互，并将标准缩放点积注意力用于涉及提示的交互。结合这两种方法，我们得到了一个富有表现力的线性时间多模态高效扩散Transformer（MM-EDiT）。通过将EDiT和MM-EDiT集成到PixArt-Sigma（传统DiT）和Stable Diffusion 3.5-Medium（MM-DiT）中，我们证明了其有效性，经过蒸馏后，实现了高达2.2倍的加速，同时保持了相当的图像质量。

🔬 方法详解

问题定义：现有扩散Transformer（DiT）在处理高分辨率图像时，由于注意力机制的二次方复杂度，计算成本显著增加，限制了其在资源受限设备上的应用。此外，多模态DiT在处理图像和文本提示的交互时，也面临着效率挑战。

核心思路：EDiT的核心思路是利用线性复杂度的注意力机制来替代传统的缩放点积注意力，从而降低计算成本。具体而言，通过卷积网络提取局部信息来调制查询，并对键和值进行空间聚合，从而在不损失过多信息的情况下，显著减少计算量。对于多模态输入，采用混合注意力机制，针对不同模态的交互选择合适的注意力方式。

技术框架：EDiT的整体架构与传统的DiT类似，主要区别在于注意力模块的实现。对于单模态图像生成，EDiT使用线性压缩注意力模块替换了原有的缩放点积注意力模块。对于多模态图像生成，MM-EDiT采用混合注意力机制，对图像-图像交互使用线性压缩注意力，对图像-文本提示交互使用传统的缩放点积注意力。

关键创新：EDiT最重要的技术创新点在于线性压缩注意力机制。与传统的缩放点积注意力相比，线性压缩注意力通过卷积网络提取局部信息并调制查询，同时对键和值进行空间聚合，从而将计算复杂度从O(N^2)降低到O(N)，其中N是序列长度。这种方法在保持图像质量的同时，显著提高了计算效率。

关键设计：线性压缩注意力模块的关键设计包括：1）使用多层卷积网络提取局部信息，卷积核大小和层数需要根据具体任务进行调整；2）对键和值进行空间聚合，可以使用平均池化或可学习的线性变换；3）混合注意力机制中，需要根据不同模态的特点选择合适的注意力方式，例如，对于文本提示，可以使用传统的缩放点积注意力，以更好地捕捉全局信息。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，将EDiT集成到PixArt-Sigma和Stable Diffusion 3.5-Medium后，在保持图像质量相当的情况下，推理速度分别提升了2.2倍。这表明EDiT能够有效降低扩散Transformer的计算成本，使其更适用于实际应用。

🎯 应用场景

EDiT可应用于各种图像生成任务，尤其是在需要处理高分辨率图像或在资源受限设备上进行推理的场景中。例如，它可以用于移动设备上的图像编辑、低带宽环境下的图像传输、以及需要快速生成大量图像的工业应用。此外，MM-EDiT可以应用于需要结合文本提示进行图像生成的场景，例如文本驱动的图像编辑和生成。

📄 摘要（原文）

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.