论文题目EDiT: Efficient Diffusion Transformers with Linear Compressed Attention基于线性压缩注意力的高效扩散 Transformer会议ICCV 2025。摘要扩散变压器(dit)已经成为文本到图像合成的领先架构可以产生高质量和逼真的图像。然而dit中注意力的二次缩放特性阻碍了在资源有限的设备上生成更高分辨率的图像。本文介绍了一种有效的扩散变压器(EDiT)来缓解传统dit和多模态dit (mmdit)中的这些效率瓶颈。首先我们提出了一种新颖的线性压缩注意力方法该方法使用多层卷积网络来调制具有局部信息的查询同时键和值在空间上聚合。其次我们为多模态输入制定了一个混合注意方案该方案结合了图像到图像交互的线性注意和涉及提示的交互的标准缩放点积注意。将这两种方法结合在一起就形成了一种具有表现力的线性时间多模态高效扩散变压器(MM-EDiT)。我们通过将EDiT和MM-EDiT架构集成到PixArt-Σ(传统DiT)和Stable Diffusion 3.5-Medium (MM-DiT)中证明了它们的有效性在蒸馏后的图像质量相当的情况下实现了高达2.2倍的加速。博主认为的核心本质就是在图像端做了一个卷积性的线性注意力然后保持文本端的不变说白了MM-EDiT 并没有把整个多模态 Attention 都线性化而是只对计算量最大的图像—图像 Attention 下手用卷积增强的 Linear Attention 替掉它涉及文本的 Attention 继续保留原来的标准 Attention。一、研究背景与核心问题1.1 DiT 真正贵在哪里DiT 中最直接的性能瓶颈来自标准 Scaled Dot-Product Attention假设图像被转换为 (N) 个 token那么 (QK^T) 需要构造一个 () 的 Attention Matrix因此计算复杂度会随 token 数量呈 () 增长。图像分辨率越高视觉 token 越多这个问题就越严重。这也是为什么 Attention 在普通分辨率下可能还可以接受但当生成分辨率来到 2K、4K 甚至 8K 时计算成本会迅速爆炸。MM-DiT 的问题更加明显。传统 DiT 往往分别进行 image self-attention 和 text cross-attention而 Stable Diffusion 3 一类 MM-DiT 会把图像 token 和文本 token 放进一个 Joint Attention 中统一建模。问题在于图像 token 通常远远多于文本 token因此整个 Joint Attention 最昂贵的部分实际上仍然是 image-to-image interaction。所以这篇论文真正想解决的并不是简单的“如何把 Attention 做快”而是两个更具体的问题第一普通 Linear Attention 虽然把复杂度降下来了但直接用于 DiT 时表达能力往往不足生成质量容易下降第二在 MM-DiT 中把所有 Attention 一刀切地线性化未必合理因为视觉 token 与文本 prompt 的作用并不相同。1.2 为什么已有高效 Attention 还不够已有工作大致有两条路线。一类是直接使用 Linear Attention。例如 SANA 使用线性 Attention同时在 FFN 部分加入卷积结构 MixFFN 来补充局部信息LinFusion 则设计了更复杂的非线性特征映射但论文指出 LinFusion 在 UNet 上有效迁移到 DiT 后效果明显下降。另一类是减少 Key/Value 的 token 数量。PixArt-Σ 就利用 depthwise convolution 对 K/V 做压缩。这样虽然能够缩小 (QK^T) 的矩阵但核心运算仍然是标准 Attention因此复杂度本质上仍然是二次的。EDiT 的关键洞察就在这里Linear Attention 本身不是问题的全部答案。真正重要的是如何在线性化以后仍然保留图像的局部空间结构同时继续压缩计算量。二、方法整体框架EDiT 与 MM-EDiT论文其实包含两个递进的模型。EDiT面向传统 DiT用 Linear Compressed Attention 替换原来的图像 Self-AttentionMM-EDiT则面向 MM-DiT把 EDiT 扩展成 Hybrid Attention只线性化 image-to-image interaction而文本相关的 interaction 仍然使用 Scaled Dot-Product Attention。论文 Figure 1EDiT 的 ConvFusion、Spatial Compressor 与 Linear Compressed AttentionFigure 1 是理解 EDiT 最重要的一张图。左侧分别给出两个核心组件Query 走ConvFusionKey 和 Value 走Spatial Compressor右侧再把它们放进 Linear Attention。换句话说作者并不是让 Q、K、V 使用完全相同的映射而是针对它们在 Linear Attention 中承担的不同角色分别设计。整个逻辑可以概括成一句话Query 负责“看懂局部结构”Key/Value 负责“减少需要被看的 token”最后再利用矩阵乘法结合律消除 () Attention Matrix。对于 MM-DiT作者则进一步观察到没有必要把四种 token interaction 全部线性化。最昂贵的 () 使用 EDiT文本相关部分仍然使用标准 Attention从而在计算效率和 Prompt 建模能力之间做折中。三、EDiTLinear Attention 为什么还要加入卷积3.1 Linear Attention 的关键变化标准 Attention 的计算顺序可以理解为它首先计算所有 Query 与所有 Key 之间的两两关系因此出现 () 矩阵。Linear Attention 利用矩阵乘法结合律把顺序改成论文采用的形式可以写成这样不再显式计算完整的 (QK^T)计算量就可以随 token 数量线性增长。但问题也随之出现图像并不是普通的一维 token 序列。相邻位置往往对应相邻像素或 latent patch而普通线性映射会弱化这种天然的二维局部结构。这也是为什么作者没有停留在“把 Attention 改成 Linear Attention”这一步。3.2 ConvFusion让 Query 重新看到二维邻域作者首先针对 Query 设计了ConvFusion其中 (X) 是输入 tokenGN 表示 Group Normalization。实现时作者先把 token sequence 重新恢复成 latent image 的二维形状然后进行 2D convolution第一层采用 () 卷积压缩 channel第二层用 () 卷积恢复 channel 数。这里最重要的不是公式有多复杂而是Query 不再独立处理每一个 token。普通 Linear Attention 中一个 Query 的特征映射通常主要来自它自己ConvFusion 则提前把邻域信息融合进 Query。直观来看它相当于先告诉 Attention“你在进行全局匹配之前先看看这个 token 周围长什么样。”这样既保留了卷积对二维局部结构的归纳偏置又没有重新引入完整的二次 Attention。3.3 Spatial Compressor减少真正参与聚合的 K/V第二个核心模块是Spatial Compressor。论文定义其中作者使用 () depthwise convolution并设置 stride2因此在二维空间的两个方向上都减半最终Key 和 Value 的 token 数量降低到原来的约 (1/4)。这一步和 PixArt-Σ 的 KV compression 看起来相似但两者作用的位置不同。PixArt-Σ 是“压缩 K/V 标准 Attention”所以 (QK^T) 依然存在EDiT 则是“压缩 K/V Linear Attention”不仅减少 (K^TV) 的计算量还从根本上避免构造完整的 image-to-image Attention Matrix。因此 EDiT 的设计实际上有两层加速Linear Attention 解决复杂度形式Spatial Compressor 进一步减少线性 Attention 内部真正需要处理的 token。四、MM-EDiT不是所有 Attention 都应该被线性化4.1 把 Joint Attention 拆开看MM-DiT 将 image tokens (X_I) 与 prompt tokens (X_P) 一起进行 Joint Attention。如果把这个 Attention Matrix 拆开可以得到四种关系image → image、image → prompt、prompt → image以及 prompt → prompt。其中真正随高分辨率迅速爆炸的是因为 image token 数量 () 很大这一块矩阵的尺寸为 ()。论文 Figure 2(a) 正是在强调这一点。论文 Figure 2标准 Joint Attention 与 MM-EDiT Hybrid Attention 对比这张图重点看左右两部分的差异标准 MM-DiT 中左上角 image-to-image block 直接计算 ()MM-EDiT 则把这一块改写成 Linear Attention 的 () 再与 () 相乘而另外三个包含文本的 interaction 仍然保留标准 Attention。论文因此把相对于图像 token 数量的 Attention 复杂度降低到 (O(N_I))。这实际上是 MM-EDiT 最值得记住的设计。作者并没有追求“100% Linear Attention”而是只线性化真正造成计算爆炸的部分。4.2 Hybrid Attention 怎么算MM-EDiT 中图像 Query (Q_I) 继续通过 ConvFusion 获得图像 Key/Value (K_I,V_I) 使用 Spatial Compressor文本侧 (Q_P,K_P,V_P) 则继续使用原模型中的线性投影。其核心可以理解为第一行对应 image queryimage-image 使用 Linear Attentionimage-text 仍然使用标准 Attention第二行对应 text query其 interaction 继续保留标准 Attention。这里还有一个容易被忽略的工程细节。将 Joint Attention 拆成几个独立 Attention 后需要额外的归一化系数否则不同 block 单独 softmax 后的权重无法直接还原原来的 Joint Attention。理论上可以严格计算 ()但这样需要访问 Attention 内部的归一化统计而 FlashAttention 一类高效实现并不会方便地暴露这些量。因此作者最终直接根据 image/text token 数量近似这一系数有意思的是后面的消融实验表明这个更简单的近似不仅计算更快实验结果反而还略好因此成为 MM-EDiT 的默认方案。这也是论文一个很工程化的选择理论上更精确并不意味着实际系统中一定更好。五、实验结果与消融分析5.1 实验设置作者分别选择两个模型进行验证传统 DiT 使用PixArt-ΣMM-DiT 使用Stable Diffusion 3.5-Medium。训练数据采用 YE-POP大约包含 48 万张图像每张图像对应两种 caption评估使用 PixArtEval30K。指标方面CLIP similarity 衡量文本和生成图像的一致性同时使用 Inception-v3 FID 和 CLIP FID 衡量生成质量。这里还有一点很重要EDiT 并不是把预训练模型中的 Attention 直接替换后零训练使用而是通过Knowledge Distillation Feature Distillation恢复能力。PixArt-Σ 版本先在 () 上蒸馏 30 个 epoch再在 1024 分辨率训练 10 个 epoch最后在 2048 分辨率继续 8 个 epochMM-EDiT 则先在 512 分辨率训练 Hybrid Attention 45 个 epoch再在 1024 上训练 10 个 epoch。因此EDiT 更准确的定位是把已经训练好的强 DiT 蒸馏成更高效的 Attention 架构。5.2 Table 1EDiT 能否保住 PixArt-Σ 的生成质量Table 1 是 EDiT 最重要的主实验。在 () 下PixArt-Σ 的 CLIP / Inception-FID / CLIP-FID 分别为0.285 / 7.57 / 2.50EDiT 为0.283 / 7.06 / 2.57。EDiT 的 CLIP 与 CLIP-FID 略逊但 Inception-FID 从 7.57 降到了 7.06。相比之下SANA-DiT 的 Inception-FID 是8.43LinFusion-DiT 达到15.87单纯使用 KV Compression 则为10.69。因此至少从这组实验来看EDiT 确实比几种直接的线性化或 token compression 方案更好地维持了教师模型的生成能力。到了 (10241024)PixArt-Σ 的 CLIP / Inception-FID / CLIP-FID 为0.285 / 7.09 / 2.53EDiT 则为0.290 / 7.82 / 2.64。这里需要注意论文结果的边界EDiT 并不是所有指标都超过教师模型。它在 CLIP 上更高但两个 FID 略差更准确的说法应该是“整体保持在接近教师模型的水平”而不是全面优于 PixArt-Σ。5.3 Table 1 的消融真正负责“保质量”的是谁Table 1 下半部分其实比主结果更能解释 EDiT 为什么有效。如果只给 K/V 使用 Spatial Compressor而没有使用 ConvFusion 增强 Query那么 () 的 Inception-FID 会从完整 EDiT 的7.06恶化到14.531024 下更是达到26.59。相反只在 Query 上使用 ConvFusion、不压缩 K/V 时FID 分别为7.06和7.70已经非常接近甚至略好于完整模型的7.06 / 7.82。这组消融把两个模块的分工解释得非常清楚ConvFusion 主要负责弥补 Linear Attention 的表达能力损失而 Spatial Compressor 主要负责进一步降低计算成本。单独压缩 K/V 会明显破坏质量但是 Query 已经通过 ConvFusion 获得足够的局部结构信息后再压缩 K/V精度损失就很小。换句话说EDiT 的关键不是“压缩”而是先增强 Query再放心压缩 K/V。5.4 Figure 3分辨率越高EDiT 的价值越明显论文 Figure 3PixArt-Σ 与 EDiT 的 2048×2048 对比在 (2048\times2048) 下PixArt-Σ 的 CLIP 为0.274EDiT 为0.288Inception-FID 分别为7.48和8.29CLIP-FID 为2.59和2.71。也就是说EDiT 在 CLIP 上更高但 FID 略差整体仍然处于比较接近的生成质量区间。真正拉开差距的是速度论文报告这一分辨率下 EDiT 大约实现了2.5×的加速。这正好体现了 EDiT 的目标它并不是为了刷新一个纯生成质量指标而是希望在质量基本可接受的情况下把高分辨率 DiT 的计算成本压下来。5.5 Table 2为什么 MM-DiT 不能全部改成 Linear AttentionTable 2 用 Stable Diffusion 3.5-Medium 验证 Hybrid Attention。在 () 下原始 SD-v3.5M 的 CLIP / Inception-FID / CLIP-FID 为0.283 / 10.49 / 3.86MM-EDiT 为0.285 / 11.60 / 3.91在 () 下二者分别为0.286 / 8.83 / 2.89和0.287 / 9.73 / 2.95。因此 MM-EDiT 的 FID 有一定退化但 CLIP 基本持平甚至略高整体仍然维持在教师模型附近。更关键的是其他 Linear MM-DiT baseline。在 512 分辨率下SANA-MM-DiT、Linear MM-DiT-α 和 Linear MM-DiT-β 的 Inception-FID 分别是14.94、13.59、13.53均明显高于 MM-EDiT 的11.60。这说明把 image-image、image-text、text-image、text-text 全部线性化虽然更彻底但生成质量也更难保住。这组结果实际上验证了论文最重要的 MM-DiT 假设视觉内部关系非常昂贵所以应该线性化文本相关 Attention token 数较少却承担关键的 Prompt Conditioning因此值得保留更强的标准 Attention。Table 2 还验证了两个设计细节。首先去掉 ConvFusion 和 Spatial Compressor 后512 分辨率 Inception-FID 直接恶化到20.98说明简单的 Linear Attention 无法替代作者设计的卷积特征映射。其次使用理论上的 () 时 FID 为13.05反而不如默认近似方案的11.60支持作者最终采用 token 数量近似归一化系数。5.6 Figure 4定性结果是否还能保持论文 Figure 4EDiT、MM-EDiT 与对应 Base Model 的定性结果Figure 4 给出了 (10241024) 的生成样例。上半部分比较 PixArt-Σ、EDiT 和 SANA-DiT下半部分比较 SD-v3.5M、MM-EDiT 和 SANA-MM-DiT。主文展示了风景、北极熊、人物、马、海滩、宇航员和花卉等多种 Prompt。作者认为 EDiT/MM-EDiT 在整体视觉质量以及 Prompt Adherence 上与对应 Base Model 接近。这张图的重要性在于FID 和 CLIP 只能告诉我们平均统计结果而图像生成模型还必须确认加速之后有没有明显丢掉主体、构图或 Prompt 语义。至少从主文给出的这些案例来看没有看到因为 Attention 线性化而产生非常明显的结构崩坏。5.7 Table 3真正的优势出现在高分辨率如果只看 512 或 1024 的 FIDEDiT 的优势可能还不够直观但 Table 3 把分辨率一路拉到 8K 后线性复杂度与二次复杂度之间的差距开始迅速放大。在 NVIDIA A100 上单次 Transformer forward 的 latency 如下PixArt-Σ 在 1024、2048、4096、8192 分辨率下分别需要0.047、0.387、4.770、72.96 秒而 EDiT 只需要0.034、0.121、0.461、1.693 秒。到了 8192EDiT 相比 PixArt-Σ 已经快约43×相比 SANA-DiT 的 21.76 秒也快约13×。MM-EDiT 的趋势也一样。在 4096 分辨率下SD-v3.5M 需要5.603 秒SANA-MM-DiT 需要0.861 秒MM-EDiT 为0.623 秒相较原始 SD-v3.5M 大约是9× 加速。所以 Table 3 才是 EDiT “Linear-Time” 最有说服力的证据。分辨率越高它相对于普通 Attention 的优势越大。5.8 Figure 5手机端还能不能跑得更快论文 Figure 5Samsung S25 Ultra 端侧 latency作者还在搭载 Snapdragon 8 Elite 的 Samsung S25 Ultra 上测试了 4-bit quantized 模型测量生成 (10241024) 图像时 Transformer 单次 forward 的 latency。Figure 5 中PixArt-Σ 为1.243 sEDiT 降到0.896 sSANA-DiT 为1.206 sSD-v3.5M 为3.115 sMM-EDiT 则降到1.423 s约为原模型的2.2× 加速。SANA-MM-DiT 的 1.277 s 仍然略快于 MM-EDiT。这个结果也暴露了 Hybrid Attention 的剩余瓶颈。作者发现MM-EDiT 中保留下来的 image-to-text () 和 text-to-image () 仍然会产生较高端侧成本其中 image query () 没有经过 token compression使 image-text dot product 依然比较昂贵。SANA-MM-DiT 因为完全使用 Linear Attention所以速度可以更快但论文认为这种优势伴随着生成质量上的代价。这也给出了一个非常自然的后续方向如何在不破坏 Prompt Conditioning 的前提下把剩余的跨模态 Attention 进一步高效化。六、总结与思考如果把 EDiT 压缩成一句话它真正做的事情其实是不要粗暴地把 DiT 的 Attention 全部换成 Linear Attention而是利用图像的二维结构重新设计 Query/K/V并且只在线性化真正昂贵的 interaction。EDiT 中最值得记住的是ConvFusion Spatial Compressor的分工。ConvFusion 负责向 Query 注入局部空间信息解决 Linear Attention 表达能力不足的问题Spatial Compressor 再压缩 K/V进一步减少计算量。Table 1 的消融尤其说明ConvFusion 是保持生成质量的关键而 Spatial Compressor 更像是在已有表达能力的基础上换取额外效率。MM-EDiT 中更值得借鉴的是Hybrid Attention的设计思想。作者没有把“Linear Attention”当成目标本身而是先拆解计算结构既然 image-image interaction 才是高分辨率下真正的二次复杂度来源那么就只把这一块线性化文本相关 Attention 数量较小、但关系到 Prompt Conditioning就继续保留标准 Attention。Figure 2 和 Table 2 基本形成了从机制到实验的一条完整证据链。从实验边界来看EDiT/MM-EDiT 并不是在所有质量指标上都优于教师模型因此把它理解成“质量更强的新生成模型”并不准确。它真正追求的是quality-efficiency trade-off允许极小的 FID 波动换取随着分辨率升高越来越明显的速度优势。尤其是 Table 3 中EDiT 从 1024 的 0.034 秒扩展到 8192 仍只有 1.693 秒而 PixArt-Σ 已经增长到 72.96 秒这才是这篇论文最核心的实验结果。从工程角度看还有一点需要记住EDiT 不是一个完全 training-free 的 Attention 替换方案而是依赖预训练教师模型和后续蒸馏来恢复能力。因此它更适合被理解成一种“预训练 DiT → 高效 DiT”的模型转换路线。而在 MM-DiT 上剩余的跨模态标准 Attention 仍然构成端侧瓶颈这也是作者明确提出值得进一步研究的方向。最终这篇论文留下的思路可能比具体结构更值得记住高效 Attention 的关键不只是把复杂度从 (O(N^2)) 写成 (O(N))而是先判断哪些 token interaction 真正需要强表达能力、哪些 interaction 才是真正的计算瓶颈再针对不同部分使用不同计算机制。这也是 EDiT 从普通 Linear Attention 工作中最明显的区别。