抛弃UNet，首个基于DiT的图像编辑框架！DiT4Edit：多尺寸编辑质量更优 | 北大港科大

文章链接：https://arxiv.org/pdf/2411.03286
工程链接：https://github.com/fkyyyy/DiT4Edit

亮点直击

• 基于基于Transformer的扩散模型在图像编辑中的优势，提出了DiT4Edit，第一个不需要调优的基于扩散Transformer（DiT）的图像编辑框架。

• 为了适应基于Transformer的去噪计算机制，首先提出了统一的注意力控制机制，以实现图像编辑。引入了DPM-Solver反演和图像块合并策略，以减少推理时间。

• 大量定性和定量的实验结果证明了DiT4Edit在对象编辑、风格编辑和形状感知编辑方面的优越性能，适用于各种图像尺寸，包括512 × 512、1024 × 1024和1024 × 2048。

总结速览

解决的问题：
现有的基于UNet的高分辨率图像编辑方法在实现形状感知的对象编辑方面仍存在不足。相比之下，Diffusion Transformers (DiT) 能更有效地捕捉图像块之间的长程依赖关系，但目前在图像编辑领域的应用仍较少。

提出的方案：
本文提出了基于Diffusion Transformers的首个图像编辑框架——DiT4Edit。DiT4Edit采用DPM-Solver反演算法来获取反演隐变量，与传统UNet框架常用的DDIM反演算法相比，显著减少了迭代步数。此外，设计了适用于Transformer计算流的统一注意力控制和图像块合并机制。

应用的技术：

1. Diffusion Transformers (DiT) 框架

2. DPM-Solver反演算法

3. 统一注意力控制和图像块合并机制

达到的效果：
DiT4Edit在编辑质量和速度上优于UNet框架，尤其在高分辨率和任意尺寸的图像编辑中表现突出。大量实验验证了该框架在各种编辑场景中的强大性能，展示了Diffusion Transformers在支持图像编辑方面的潜力。

方法

本文提出的框架旨在基于Diffusion Transformer实现各种尺寸图像的高质量编辑。本文的方法是基于预训练的文本到图像Transformer扩散模型（例如，）的首个编辑策略。通过我们的方法，用户可以通过提供目标提示，获得比现有的基于UNet的方法更好的编辑结果。本节首先介绍隐空间扩散模型和DPM反演。然后阐明基于Transformer的去噪在图像编辑任务中的优越性。最后讨论编辑方法的实现细节。

基础知识：隐空间扩散模型

隐空间扩散模型（Latent Diffusion Models，LDM） 提出了在隐空间Z中进行去噪处理的图像生成方法。具体而言，它使用编码器E将像素图像x编码为低分辨率隐空间变量z = E(x)。在训练过程中，模型通过去除人为噪声，优化去噪UNet ，该过程以文本提示embedding y和当前图像样本为条件，其中是步下的带噪样本：

经过训练后，它能够通过学习到的去噪过程将随机噪声转换为图像样本z。

DPM-Solver采样。在扩散概率模型（DPMs）的反演阶段，一个干净的图像逐渐被加入高斯噪声，转变为带噪样本。

其中， 是信噪比（SNR），它是t的严格递减函数。通过求解扩散常微分方程（ODE），DPM采样可以比其他方法更快。

其中，，并且，。在之前的工作中表明，使用指数积分器的ODE求解器在求解方程3时比传统求解器具有更快的收敛速度。通过设置的值，方程3的解xt可以通过以下方式计算：

其中， 是t的递减函数，具有反演函数，最近的研究表明，DPM-Solver可以在10–20步内生成逼真的图像。

扩散模型架构

与UNet结构相比，Diffusion Transformers (DiT)表现出更优的扩展性，能够生成更高质量的图像，并展现出更好的性能。是一个基于Transformer的文本到图像（T2I）扩散模型，由三个主要组件组成：跨注意力层（Cross-Attention layer）、AdaLN-single和重参数化（Re-parameterization）。研究人员通过三种复杂的设计训练了这个T2I扩散模型：分解训练策略、高效的T2I Transformer和高信息量的数据。大量实验结果表明，在图像质量、艺术性和语义控制方面表现更好。与先进的T2I SOTA模型相比，具有更快的训练速度、更低的推理成本和更出色的综合性能。本文将作为图像编辑方法的基准。

选择Transformer作为去噪模型的原因
与UNet结构相比，Transformer引入了全局注意力机制，使得模型能够关注图像中的更广泛区域。这种增强的可扩展性使得Transformer能够生成大尺寸（例如大于512×512）甚至任意尺寸的高质量图像。本文基于DiT的编辑框架在大尺寸图像上的编辑结果在下图1和图2中有所展示，这些是UNet框架之前未涉及的编辑任务。因此，采用了基于Transformer的去噪模型作为编辑框架，利用Transformer的能力来应对这些更复杂的编辑挑战。

基于扩散Transformer的图像编辑

如上图2所示，基于预训练的扩散Transformer图像编辑框架的流程。

DPM-Solver反演 如前所述，使用高阶的DPM-Solver（例如DPM-Solver++）可以有效提高采样速度。为了近似方程4中的积分项，给定时间时的，可以在处使用泰勒展开，DPM-Solver++能够在时间获得精确的解值：

特别是当时，方程5等价于DDIM采样器，如下所示：

在实际应用中，通常将，以实现快速推理并最小化离散化误差。这个DPM-Solver被称为DPM-Solver++ (2M)。

其中，2M意味着这个求解器是一个二阶多步求解器。然而，在高阶采样器（如DPM-Solver++(2M)）的反演阶段，为了获得当前时间步的反演结果，需要近似先前时间步{}中的值，以估计和解析计算方程5中的项。

最近的研究（Hong等，2024）提出了一种通过反向欧拉方法来获取方程8中高阶项近似的策略。

其中是去噪模型，是通过DDIM反演计算的值集合，用来估计方程8中的，并且。然后可以通过以下公式得到当前时间步的反演隐空间变量：

其中。在DiT4Edit中，利用DPM-Solver++反演策略从输入图像中获取反演隐空间变量用于编辑任务。此外，这一技术在之前基于UNet的图像编辑方法中未曾使用。即使不使用DDIM反演来计算的值，仍然可以获得一个良好的反演隐空间变量图。

统一控制注意力机制 在之前的研究Prompt to Prompt (P2P)中，研究人员展示了交叉注意力层包含来自提示文本的丰富语义信息。这一发现可以通过在扩散过程中替换源图像与目标图像之间的交叉注意力图来编辑图像。具体而言，常用的两种基于文本引导的交叉注意力控制策略是交叉注意力替换和交叉注意力细化。这两种方法确保了从目标提示到源提示的信息无缝流动，从而将隐空间变量图引导到期望的方向。

与交叉注意力不同，扩散Transformer中的自注意力机制用于引导图像布局的形成，这是交叉注意力机制无法实现的特性。如下图3所示，来自提示的对象和布局信息并未在Transformer的浅层查询向量中得到充分捕获，但在深层得到了很好的表示。此外，随着Transformer层数的增加，查询向量捕捉对象细节的能力变得更加清晰和具体。这表明，Transformer的全局注意力机制在捕捉长范围对象信息方面更加有效，使得DiT在大规模变形和广泛图像编辑中具有特别的优势。这一观察表明，通过控制自注意力机制可以实现图像的非刚性编辑。在MasaCtrl中，研究人员引入了互相自注意力控制机制。具体来说，在扩散的早期步中，编辑步中的特征、和将用于自注意力计算，以生成更接近目标提示的图像布局，而在后期步中，重建步中的特征和将用于引导生成更接近原始图像的目标图像布局。

然而，MasaCtrl可能仍然会遇到一些失败的情况，这可能是由于其在整个编辑过程中使用了，正如最近的研究（Xu等，2024）中提到的那样。为了解决这个问题，通过设置一个阈值来决定何时采用，这个阈值是基于步的数量来调整的。

Patches Merging: 为了提高推理速度，借鉴了Token Merging的方法，将patch合并技术嵌入到去噪模型中。这个方法的灵感来自于这样一个观察：在transformer架构中，涉及到的patch数量明显大于UNet。计算流程如下图4所示。对于一个特征图，我们首先计算每个patch之间的相似度，然后将最相似的patch合并，从而减少由注意力机制处理的patch数量。经过注意力计算后，我们会将patch拆分回来，以保持模型下一层的输入大小。通过将patch合并技术融入我们的框架中，我们旨在简化处理过程，提高整体效率，同时不改变每一层的基本操作。

实验

实现细节

对于512×512和更大尺寸（最高至1024×2048）的图像编辑任务，使用预训练模型−XL−512×512版本处理较小尺寸，−XL−1024×1024−MS版本处理较大尺寸（Chen等，2023）。对真实图像和生成图像进行编辑。对于真实图像输入，使用DPM-Solver反演来获取潜在噪声图。将DPM-Solver配置为30步，分类器无关引导设置为4.5，并且patch合并比率设置为0.8。所有实验都在NVIDIA Tesla A100 GPU上进行。

研究者们评估了DiT4Edit编辑框架与六个先前基准模型（包括Pix2Pix-Zero、PnPInversion、SDEdit、IP2P、MasaCtrl和InfEdit）的质量表现差异，所有方法均使用官方开源代码实现。

如下图5所示，在512×512和1024×1024图像上对比了本文的方案。图5的第一行展示了本文的框架在编辑真实512×512图像时，能够生成与原始内容一致的编辑图像，而现有方法通常会改变原始图像的背景或目标细节。此外，图5的第二行和第三行展示了在大规模图像和任意尺寸图像上的实验——这些任务是以前基于UNet的方法难以处理的。结果表明，本文提出的框架有效地处理了大图像中的风格和物体形状修改。相比之下，一些先进的基于UNet的方法，尽管能够进行编辑任务，但通常会导致背景和物体位置的显著变化和损坏。此外，由于UNet结构的局限性，这些方法通常只能生成512×512大小的目标图像。这些发现强调了基于transformer的扩散模型在大规模图像编辑中的巨大潜力。

定量比较

为了进行定量评估，使用了三个指标：Fréchet Inception Distance (FID)、峰值信噪比（PSNR）和CLIP，以评估我们模型与SOTA（最先进方法）在图像生成质量、背景保持和文本对齐方面的性能差异。本文比较了三种尺寸的图像：512×512、1024×1024 和1024×2048，结果详细列在下表1中。与Pix2Pix-Zero、PnPInversion、SDEdit、IP2P、MasaCtrl和InfEdit进行了性能比较。需要注意的是，由于此前没有基于DiT的编辑框架，所有对比基准都基于UNet架构。实验结果表明，DiT4Edit编辑策略在图像生成质量、背景保持和文本对齐方面优于SOTA方法。由于集成的transformer结构具备全局注意力能力，DiT4Edit框架在各种大小的编辑任务中表现出了强大的鲁棒性。生成的图像不仅显示出更高的质量，而且在背景和细节控制方面提供了更好的控制，保持了与原始图像的一致性。特别是在编辑大规模或任意尺寸的图像时，DiT4Edit相较于其他方法展示了显著的优势，充分展示了transformer架构强大的可扩展能力。同时，我们的编辑框架具有更短的推理时间，与无反演编辑方法（InfEdit）相当。

消融实验

本文进行了系列消融研究，旨在展示DPM-Solver反演和patch合并的有效性。关于patch合并的消融实验结果如前面图4和下表2所示。实现patch合并后，尽管保持了与未使用patch合并时相当的编辑质量，但大尺寸图像的编辑时间显著减少。这表明，patch合并可以显著提升图像编辑框架的整体性能。此外，DPM-Solver和DDIM的消融实验结果如图7所示。在相同推理步数（T = 30）下比较两种方法时，DPM-Solver在图像编辑质量方面始终优于DDIM。这表明本文采用的DPM-Solver反演策略能够生成更优的潜在图像，从而在较少的步中获得更好的编辑效果。

结论

本文提出了DiT4Edit，这是首个基于扩散transformer的图像编辑框架。与之前的UNet框架不同，DiT4Edit提供了更优的编辑质量，并支持多种尺寸的图像。通过利用DPM Solver反演、统一的注意力控制机制和patch合并，DiT4Edit在512×512和1024×1024尺寸图像的编辑任务中超越了UNet结构。特别是，DiT4Edit能够处理任意尺寸的图像，如1024×2048，展示了transformer在全局注意力和可扩展性方面的优势。本文的研究为基于DiT的图像编辑奠定了基础，并有助于进一步探索transformer结构在生成式AI中的潜力。

局限性
在实验中，观察到T5-tokenizer偶尔会遇到词汇分割问题，这可能导致最终编辑过程中的失败。此外，我们的模型与原始图像相比，可能会出现颜色不一致的情况。更多的编辑失败案例请参考补充材料。

参考文献

[1]DiT4Edit: Diffusion Transformer for Image Editing