点击下方卡片，关注「集智书童」公众号
导读
扩散 Transformer （DiTs）已成为文本到图像合成的领先架构，能够生成高质量和逼真的图像。然而，DiT中的注意力机制在生成更高分辨率图像或在资源有限的设备上运行时，存在二次扩展性，这限制了效率。
本研究提出了一种高效的扩散 Transformer （EDiT），以缓解传统DiT和多模态DiT（MMDiTs）中的这些效率 Bottleneck 。首先，作者提出了一种新颖的线性压缩注意力方法，该方法使用多层卷积网络来调节 Query ，同时 Key和Value 在空间上聚合。其次，作者为多模态输入制定了一种混合注意力方案，该方案结合了线性注意力机制用于图像到图像的交互，以及标准缩放点积注意力机制用于涉及 Prompt 的交互。将这两种方法合并，得到了一种表达能力强、线性时间的多模态高效扩散 Transformer （MM-EDiT）。作者通过将EDiT和MM-EDiT架构集成到PixArt-≥（传统DiT）和Stable Diffusion 3.5-Medium（MM-DiT）中，实现了在蒸馏后图像质量相当的情况下，速度提升高达
。
- 引言
扩散 Transformer （DiTs）[16]迅速成为架构的首选，并在最先进的（SOTA）文本到图像生成模型中广泛使用，如SANA[22]、FLUX[11]、PixArt-
[2]、Stable Diffusion v3 [5]和Lumina-Next [24]。这些模型中的大多数都依赖于标准的二次缩放点积注意力[20]来捕捉长距离依赖和全局上下文，从而使它们能够生成高质量、逼真且美观的图像。DiTs的早期版本主要依赖于传统的 Transformer 架构，即图像 Token 通过自注意力进行处理，而 Prompt 条件是通过应用交叉注意力[2, 3]来实现的。最近，多模态DiTs（MM-DiTs）引入了联合注意力机制，以在单个自注意力层中处理图像和 Prompt Token [5]。虽然显著提高了图像质量，但联合注意力由于涉及 Token 数量的增加而扩展性较差。当生成高分辨率图像时，这一限制进一步加剧，主要是因为输入 Token 数量的增加，阻碍了SOTA文本到图像模型的效率部署，尤其是在资源受限的设备上。
现有研究主要探索了缓解DiT中这一计算 Bottleneck 的策略，主要通过关注线性化注意力机制[14, 22]或减少 Token 数[2]来实现。例如，Linfusion[14]依赖于多层非线性变换来映射输入 Token ，从而增强了线性注意力机制的表达能力。而SANA[22]将标准线性注意力[10]与 Transformer 前馈部分的卷积操作相结合，利用图像的固有结构，更有效地捕捉局部关系。最近的工作，例如PixArt-
[2]，减少了计算中涉及的 Key和Value 的数量，这是一种提高计算效率的有效策略。
本文提出了一种高效的扩散Transformer（EDiT）架构，采用了一种新颖的线性压缩注意力机制。EDiT采用以下方法：（i）多层卷积网络以有效地整合 Query Token 中的局部信息；（ii）使用深度卷积在 Key和Value 中对 Token 信息进行空间聚合。这些组件产生了一种表达能力强且计算效率高的注意力机制。作者通过蒸馏PixArt-
，一种传统的DiT模型，评估了EDiT，在获得可比图像质量的同时实现了显著的加速。
为了高效处理多模态输入，作者进一步引入了一种混合注意力机制，该机制结合了线性注意力和缩放点积注意力的优势，以处理基于MM-DiT的模型。具体来说，作者包括了EDiT的线性压缩注意力机制，用于处理计算密集型的交互，即计算图像到图像的依赖关系时；另一方面，使用缩放点积注意力来捕捉文本交互。这些修改引入了一种新型的多模态高效扩散 Transformer （MM-EDiT）架构，该架构在图像大小方面实现了线性时间复杂度，同时不牺牲质量。作者通过提炼基于MM-DiT的Stable Diffusion 3.5-Medium模型，展示了作者混合注意力方法的有效性，并表明MMEDiT优于现有的线性注意力替代方案。
作者的贡献可以总结如下：
-
- 作者引入了EDiT，它使用线性压缩注意力机制来构建扩散Transformer（第3.1节）。
-
- 作者提出了MM-EDiT，通过一种新颖的混合注意力公式将作者的线性压缩注意力集成到MM-DiTs中（第3.2节）。
-
- 最后，作者通过蒸馏（第4节）对PixArt-
[2] 和 StableDiffusion 3.5-Medium 进行适配，展示了这些新颖架构的有效性。作者的广泛实验表明，EDiT 和 MM-EDiT 可以实现高达
的加速，同时不牺牲图像质量。
- 背景
注意力机制，由 Vaswani 等人 [20] 提出，通过比较 Query
和键
来计算序列中值
的重要性。在自注意力中， Query 、 Key和Value 通常通过输入序列
的简单线性投影来计算。在这里，
和
分别表示 Token 数和模型维度。然后，使用缩放点积注意力 [20] 计算注意力块的输出。
虽然这一机制非常成功，但它需要显式计算
，其计算复杂度与 Token 数量呈二次方关系，这在使用缩放点积注意力的自注意力方法中构成了一个主要的计算 Bottleneck 。为了解决这个问题，已经提出了几种替代的注意力机制[7, 10]，这些机制的计算复杂度与 Token 数量呈线性关系，而不是二次方关系。
2.1 线性注意力
对公式（1）的进一步分析表明，注意力机制可以用任何相似度函数进行泛化，如下所示[10]：
，
其
中
。
attention在方程（1）中可以表示为核特征图的线性点积：
，通过应用矩阵乘法的结合律，可以进一步简化为
。在大多数线性注意力工作中，考虑不同的核函数来构建
，
和
。线性注意力结果的公式表示如下：
其中，
表示输出序列
中的一个单个 Token 。尽管方程（3）中线性注意力的公式在研究中保持一致，但文献中提供了多种定义核特征图
、
和
的方法。一些研究证明了，使用
或
等函数的基本公式来定义
和
，同时保持
线性，可以获得与缩放点积注意力相当的性能。
2.2 高效的扩散模型注意力机制
SANA [22] 在DiTs中采用了传统的线性注意力方法，通过线性投影值并使用
作为 Query 和键。然而，这种设置本身并不能产生令人满意的结果，这促使进一步修改transformer架构。具体来说，每个transformer块末尾的标准 FFN 被替换为一个名为MixFFN的复杂卷积模块，该模块包含多个卷积层以增强 Token 的局部信息。与使用缩放点积注意力的DiTs相比，这提高了延迟，同时实现了与SOTA相当的画面生成质量。
LinFusion [14] 探讨了基于 UNet 的扩散方法中使用的线性注意力机制，涉及 Stable Diffusion v1.5 [18] 和 Stable Diffusion XL [17]。为此，刘等人使用以下函数计算 Query 和键：
其中，LN表示层归一化[1]。第一层线性层将每个
映射到一个低维空间，而第二层线性层将其映射回原始空间。对于数值，他们使用线性投影。然后，这种线性注意力机制被整合到现有模型中，随后进行蒸馏。尽管这种方法对UNets有效，但[14]和本文中（见表1）呈现的结果显示，它对DiTs无效。
键值 Token 压缩。加速注意力的一种方法是键值（KV） Token 压缩，它减少了 Key和Value Token 的数量。这种技术在基于DiT的PixArt-
[2]中得到实现。Chen等人使用深度卷积来压缩 Key和Value 。在实践中，这个操作将
个 Token 块聚合为一个单一的压缩 Token ，从而将 Key和Value 的数量降低了
倍。因此，矩阵
的大小也相应地减小了
倍，从而提高了效率。在PixArt-
中，相同的压缩核用于 Key和Value 。由于这个核与传统的自注意力机制集成，计算复杂度仍然保持为二次方。
- 方法
首先，作者提出了EDiT，这是一种针对DiTs的新型线性压缩注意力机制，它结合了线性注意力和高效的 Token 压缩。随后，作者介绍了MM-EDiT，这是一种针对MM-DiT模型的混合注意力方法。
3.1. EDiT - 用于DiTs的线性注意力机制
EDiT使用多层卷积对 Query 进行线性压缩，并通过基于卷积的特征图在空间上压缩 Key和Value Token 。
处理 Query 。受LinFusion [14]和SANA [22]背后的直觉启发，作者提出了ConvFusion函数。此操作有效地将LinFusion和SANA的MixFFN结合成一个统一的从 Token 到 Query 的卷积映射。具体来说，作者使用
从输入 Token
处理 Query ：
其中，GN表示组归一化[21]。第一层卷积沿通道维度进行压缩，而第二层卷积将通道上采样回原始大小。与LinFusion[14]不同，作者提出的ConvFusion不是单独处理每个 Token ，而是利用 Token 中的局部信息。具体来说，作者将序列 Reshape 为原始潜在图像的形状，然后应用2D卷积。为了保持卷积的效率，压缩使用
的 Kernel 大小，上采样使用
的 Kernel 。作者在图1中说明了ConvFusion机制。
处理 Key和Value 。接下来，作者介绍空间压缩器，这是一种基于卷积的空间压缩技术，应用于 Key和Value （如图1所示）。该操作降低了方程3中
的复杂性，其公式如下：
其
中
，
在实践中，
具有一个
的深度可分离卷积核，步长为2，这使 Key和Value 的数量减少了4倍。尽管在PixArt-
[2]中已经探索了 Token 压缩技术，但它们的调查仅限于缩放点积注意力。相比之下，作者将空间压缩器集成到线性注意力机制中，在EDiT中提高了效率。现在，作者将ConvFusion中的 Query
与空间压缩器中的键
和值
插入到方程（3）中的线性注意力公式中。作者提出的用于高效DiTs的机制，称为EDiT，随后被用作DiTs中自注意力的替代。图1展示了EDiT概述。
3.2. MM-EDiT - 多模态多任务学习中的混合注意力机制
尽管DiTs为图像 Token 使用单独的自注意力，为文本 Prompt 使用交叉注意力，但多模态扩散 Transformer （MM-DiTs）通过联合注意力方法简化了流程，对拼接的 Token 应用自注意力。鉴于图像 Token 的数量通常远超过文本 Token 的数量，注意力的二次复杂度主要取决于图像 Token 的数量。对于高分辨率图像，这个问题更为严重，因为图像 Token 的数量显著更大。因此，作者设计了一种混合注意力机制，利用线性注意力来处理图像到图像的关系，同时利用更强大的缩放点积注意力来处理 Prompt Token 。将这种混合注意力与EDiT的线性压缩注意力相结合，形成了一种表达丰富且线性时间的多模态高效扩散 Transformer （MM-EDiT）。
更深入地探讨联合注意力。作者用
和
表示图像和 Prompt Token ，其中
和
分别是图像和 Prompt Token 的数量。作者将图像和 Prompt 的 Query 、 Key和Value 分别表示为
和
。最后，作者将
和
定义为对应图像和文本 Token 的拼接，这些拼接被插入到公式（1）中以计算注意力块的输出。然而，作者也可以将联合注意力机制分解为四个组件：（i）图像 Token 对图像 Token 的注意力，（ii）图像 Token 对 Prompt Token 的注意力，（iii） Prompt Token 对图像 Token 的注意力，以及（iv） Prompt Token 对自身的注意力。这种分解在图2(a)中进行了可视化，并导致
作者定义
和
。在这种分解下，与联合注意力相比，
中的softmax现在应用于不同的 Token 集，导致行和不为1。因此，需要添加一个归一化因子
，如公式（7）所示，以调整注意力矩阵的每个元素，使得输出保持不变：
在MM-EDiT中结合线性与缩放点积注意力。联合注意力中最计算密集的部分是
，其与图像分辨率呈二次方关系。为了优化它，作者应用EDiT来线性化图像-图像注意力段，从而显著减少计算时间。总体而言，EDiT将提出的EDiT线性注意力集成到图像-图像注意力块中，并将缩放点积注意力应用于其他部分。作者将这种方法称为MM-EDiT，并在图2(b)中进行了说明。
首先，作者使用ConvFusion（公式（4））计算
，并通过Spatial Compressor（公式（5））获得
和
。对于 Prompt Token ，
、
和
通过简单的线性投影得到，与原始模型相同。遵循MM-DiT方法，作者随后将图像和 Prompt Token 拼接。最后，作者将缩放点积联合注意力（公式（6））中的
替换为
，得到如下公式的混合联合注意力矩阵
：
此处，
。
使用这种公式，必须为每个块计算相应的归一化项。不幸的是，像FlashAttention这样的高效实现并没有提供对这些项的访问，这要求作者采用效率较低的、定制的缩放点积注意力实现。为了避免这种情况，作者可以通过 Token 数来近似
，即
Nr+N。在第4.2节中，作者表明这种近似不仅提高了计算速度，而且也略微提升了性能，因此作者认为它是MM-EDiT的标准。
- 实验
作者通过使用作者新颖的EDiT和MM-EDiT架构，对两个著名的扩散Transformer模型进行了评估。首先，作者将PixArtΣ[2]作为传统的DiT，其次，作者将StableDiffusion 3.5-Medium (SD-v3.5M) [5]作为MM-DiT实验的基础。数据集。在所有实验中，作者依赖
数据集进行训练。这个Laion-POP的子集包含大约480K张图片，每张图片有两个不同的标题。对于评估，作者使用PixartEval30K数据集
。EDiT的训练细节。作者使用了三个公开可用的PixArt-Σ预训练权重，这些预训练权重分别在
、
和
像素的分辨率下进行训练。作者使用
的PixArt-Σ预训练权重初始化EDiT，使用作者的线性压缩注意力方法修改自注意力层，并对模型进行30个epoch的蒸馏。具体来说，作者遵循[12, 14]，结合噪声预测损失、知识和特征蒸馏。对于特征 Level 的蒸馏，作者使用每个自注意力层的输出，并最小化教师和学生之间的差异。最后，作为知识蒸馏，作者最小化学生和教师预测之间的差异。
作者采用得到的模型，解冻EDiT的所有权重，并在
的分辨率下使用相应分辨率的教师预训练权重进行额外的10个epoch的训练。最后，作者在
的分辨率下重复相同的步骤，进行额外的8个epoch的训练。
MM-EDiT的训练细节。作者使用Stable Diffusion 3.5 Medium（SD-v3.5M）来评估作者MM-EDiT机制的效力。与PixArt-Σ不同，SD-v3.5M只有一个公开可用的预训练权重，可用于多个分辨率。作者采用这个预训练的SD-v3.5M预训练权重，并用MM-EDiT替换了多模态注意力层。作者首先在512×512图像上使用类似于EDiT的联合损失训练混合注意力层45个epoch。作者再次结合任务损失、知识和特征蒸馏，但这次，任务损失基于从SD-v3.5M中采用的修正流机制[15]。随后，作者解冻MM-EDiT模块的所有权重，并在1024×1024分辨率下再训练10个epoch。在这个阶段，联合损失表现出最佳性能，作者切换到纯知识蒸馏。
指标：作者报告了生成图像与文本之间的CLIP相似度得分，以评估 Prompt 的遵循程度；同时使用Inception-v3和CLIP编码器评估生成图像的质量，计算FID。
4.1. EDiT的评估
Baseline 与消融实验。作者比较了从PixArt-
模型蒸馏得到的EDiT版本与原始PixArt-
模型的性能。此外，作者还评估了EDiT与使用Linfusion的注意力机制[14]的Linfusion-DiT模型，以重现[14]的结果。作者还考虑了SANA-DiT，它使用SANA[22]中的线性注意力机制和MixFFN来增强PixArt-
架构，以及[2]中的原始KV Token压缩（KV Comp）方法。为了保持一致性，作者对所有架构都使用了作者的蒸馏方法。
最后，作者评估了EDiT在计算 Query 、 Key和Value 时的设计选择，通过三种设置进行评估：在 Query 和键上应用ConvFusion，仅在 Key和Value 上使用Spatial Compressor，以及在 Query 上应用ConvFusion而不压缩 Key和Value 。这有助于作者了解增强 Query 中 Token 交互的影响，并表明压缩不会损害性能。
定量结果。表1总结了基于Pixart-
的EDiT结果。在
和
分辨率下，EDiT的表现与教师模型（即Pixart-
Baseline 模型）相当——在较低分辨率下甚至超过了其性能。此外，对于
分辨率，EDiT分别比SANA-DiT和Linfusion的线性方法高2.6和8.8个Inception-v3 FID点，并且比KV Comp token压缩方法高2.5个FID点。作者在
分辨率下观察到EDiT具有相似的性能。这些结果表明，作者的EDiT模型优于当前的线性注意力机制和token压缩方法。将EDiT与其消融实验进行比较表明，作者的ConvFusion和Spatial Compressor混合方法在图像质量和运行时间之间取得了极佳的平衡。为了展示EDiT在生成高分辨率图像时的有效性，在图4中，作者展示了EDiT和PixArt-
在
分辨率下的结果，其中EDiT的表现相当，但所需时间大约减少了2.5倍。
4.2. MM-EDiT的评估
Baseline 与消融实验。除了将作者的基于MMEDiT的SD-v3.5M与原始模型进行比较外，作者还比较了MM-EDiT的性能与几个 Baseline 和消融实验的结果。作者使用线性注意力在MM-DiTs中与MM-EDiT进行了比较，包括以下三个 Baseline ：Linear MM-DiT-
，它像MM-EDiT一样计算 Query 、键和 Token ，但使用线性注意力处理所有关系；Linear MM-DiT-
，它结合了Linfusion中的
用于 Prompt Token ；以及SANA-MM-DiT，它使用
和MixFFN层。此外，作者还考虑了作者混合注意力机制的两种变体：MM-EDiT带有
（公式8），它使用理论上的正确调整项，以及不带
和
的MM-EDiT。这些比较突出了混合注意力和 Token 交互在提升性能中的重要性。
定量结果表2展示了基于
的MM-EDiT的性能，比较了混合和非混合注意力方法。EDiT优于SANA-MM-DiT、线性MMDiT-
和线性MM-DiT-
等非混合方法，突出了混合注意力在MM-DiTs中的重要性。请注意，作者的线性MMDiT-
方法在FID分数上取得了可比的结果，但定性结果劣于MM-EDiT（见附录A）。在混合技术中，使用真实的归一化因子
的性能略逊于作者的简单近似，并且处理速度较慢。此外，当作者用基本的线性函数替换
和
时，性能显著下降，突显了作者基于卷积函数的价值。
4.3 定性比较
为了展示EDiT和MM-EDiT生成的图像的视觉质量，图3展示了使用EDiT生成的示例图像。附录A提供了其他方法使用的 Prompt 和图像。这些图像展示了高质量的视觉效果，与基模型的视觉效果难以区分，并且高度遵循 Prompt 。
4.4 设备端运行时分析
在本节中，作者展示了EDiT在减少边缘设备运行时间方面的潜力。表5展示了MM-EDiT在Qualcomm SM8750-AB Snapdragon 8 Elite处理器上的运行时间分析——该处理器被用于最先进的智能手机。作者比较了作者的EDiT和MM-EDiT与基于SANA的 Baseline （SANA-DiT和SANA-MM-DiT）及其相应的基模型（Pixart-Σ和SD-v3.5M）的运行时间。对于DiTs，EDiT的延迟为896毫秒，分别比基模型和SANA-DiT模型快38%和26%。对于MM-DiTs，MMEDiT将SD-v3.5M的运行时间减少了2.2倍，尽管其运行时间略高于SANA-MMDiT。进一步对MM-EDiT的研究发现，对于图像-文本
和文本-图像
交互的缩放点积注意力操作在设备上计算成本高昂。这主要是因为没有对图像 Query
应用 Token 压缩，使得点积计算变得昂贵。尽管文本-图像交互
也资源密集，但由于对图像键
应用了 Token 压缩，其成本较低。相比之下，SANA的完全线性注意力机制可能导致了其运行时间略低于MM-EDiT，尽管这牺牲了图像质量。为了解决这个问题，作者鼓励未来的研究探索MM-DiTs中混合注意力内的替代交互策略。
- 相关工作
此外，SANA [22] 和 LinFusion [14]，作者在第2节中已讨论过，还有其他一些工作在DiTs中使用线性注意力或通过其他方式提高了其效率。SANA 1.5 [23] 将RMS归一化应用于 Query 和键，以实现线性注意力从1.6B到4.8B的有效缩放，同时保持效率和训练稳定性。ZIGMA [9] 利用一种称为Mamba [4] 的特殊因果线性注意力来开发用于文本和视频生成的有效扩散模型。然而，Mamba的因果性质在DiT框架中引入了复杂性，需要在每个块之后重新排列 Token 。DIMBA [6] 也将Mamba整合到DiT中，但与现有的自注意力块结合使用。虽然这种方法带来了一些性能提升，但也增加了计算需求，导致二次复杂性。值得注意的是，到目前为止，上述所有工作都没有讨论过MM-DiTs。
据作者所知，CLEAR [13] 是目前唯一现有的线性时间复杂度的 MM-DiTs 方法。然而，它们使用了Sparse邻域注意力 [8]，这需要复杂的内存访问模式，并且通常只有在处理非常高分辨率的图像（>2048×2048 像素）时才能实现实际效率的提升。此外，它依赖于专门的实现，这限制了其在不同硬件上的适用性。
更广泛地说，线性DiTs可以与 Mask 方法相结合以提高训练效率，正如在PatchMixer [19]中对vanilla DiTs所做的那样，这是一个有前景的研究方向。
- 结论
本文针对扩散 Transformer （DiTs）的计算 Bottleneck ，引入了高效扩散 Transformer （EDiT），该 Transformer 采用线性压缩注意力机制，将局部信息集成到 Query 中，并对 Key和Value 进行 Token 信息聚合。作者证明了使用EDiT蒸馏PixArt-Σ可以实现显著的加速，同时不牺牲图像质量。此外，作者将此方法扩展到多模态DiTs，提出了混合注意力机制，从而得到多模态高效扩散 Transformer （MM-EDiT），该 Transformer 结合了EDiT的线性注意力机制进行图像交互和标准注意力机制进行 Prompt 。使用MM-EDiT蒸馏Stable Diffusion 3.5-Medium在保持性能的同时实现了线性时间复杂度。局限性。作者的评估主要在PixArt-Σ和Stable Diffusion 3.5-Medium上进行。未来的工作应调查EDiT对更大模型（如Flux[11]）的泛化能力。此外，虽然使用了量化指标，但生成图像的定性评估仍然具有主观性，人类评估可以提供对图像质量的进一步见解。
参考
[1]. EDiT: Efficient Diffusion Transformers with Linear Compressed Attention
扫码加入👉「集智书童」交流群
（备注：方向+学校/公司+昵称）