SwiftEdit：通过一步扩散实现闪电般快速的文本引导图像编辑

Paper是VinAI Research发表在Arxiv2024的工作

Paper Title:SwiftEdit: Lightning Fast Text-Guided Image Editing via One-Step Diffusion

Code地址

Abstract

文本引导的图像编辑方面的最新进展利用了基于多步骤扩散的文本到图像模型的大量先验，使用户能够通过简单的文本输入来执行图像编辑。然而，由于涉及昂贵的多步骤反演和采样过程，这些方法往往达不到现实世界和设备应用所需的速度要求。针对此问题，我们推出了 SwiftEdit，这是一款简单但高效的编辑工具，可实现即时文本引导的图像编辑（0.23 秒内）。SwiftEdit 的进步在于它的两个新颖的贡献：一步反演框架，可通过反演实现一步图像重建；以及采用我们提出的注意力重缩放机制的掩模引导编辑技术来执行局部图像编辑。我们进行了大量实验来证明 SwiftEdit 的有效性和效率。特别是，SwiftEdit 支持即时文本引导的图像编辑，这比以前的多步骤方法快得多（至少快 50 倍），同时保持了编辑结果的竞争力。

1. Introduction

最近的文本到图像扩散模型 [5, 24, 26, 27] 在生成与给定文本提示语义一致的高质量图像方面取得了显著成果。为了生成逼真的图像，它们中的大多数依赖于多步采样技术，这些技术将从随机噪声开始的扩散过程逆转为逼真的图像。为了克服这个耗时的采样过程，一些研究专注于通过蒸馏技术将采样步骤数减少到几个（4-8 步）[29] 甚至一步 [5, 20, 39, 40]，同时不影响结果。这些方法不仅可以加速图像生成，还可以为图像编辑等下游任务提供更快的推理。

对于文本引导的图像编辑，最近的方法 [11, 13, 19] 使用反转过程来确定源图像的初始噪声，从而允许 (1) 源图像重建和 (2) 与引导文本对齐的内容修改，同时保留其他细节。从这种反转噪声开始，在每个去噪步骤中应用额外的技术，例如注意力操纵和劫持 [3, 21, 35]，以在保留关键背景元素的同时逐渐注入编辑。然而，这种典型的方法是资源密集型的，需要两个漫长的多步骤过程：反转和编辑。为了解决这个问题，最近的研究 [6, 8, 33] 使用少步扩散模型，如 SD-Turbo [30]，来减少反转和编辑所需的采样步骤，并通过文本提示加入额外的指导以进行解开编辑。然而，这些方法仍然难以为设备应用程序实现足够快的文本引导图像编辑，同时保持与多步骤方法相媲美的性能。

在本研究中，我们采用了一种不同的方法，即在单步文本转图像模型的基础上构建图像编辑。我们推出了 SwiftEdit——第一个单步文本引导图像编辑工具——它的执行速度比以前的多步方法快至少 50 倍，同时保持了具有竞争力的编辑质量。值得注意的是，SwiftEdit 中的反转和编辑都是在一个步骤中完成的。

反转一步扩散模型具有挑战性，因为现有的技术如 DDIM 反转 [31] 和空文本反转 [19] 不适合我们的一步实时编辑目标。为了实现这一目标，我们设计了一个新颖的一步反转框架，该框架受到基于编码器的 GAN 反转方法 [36、37、41] 的启发。与需要特定领域网络和重新训练的 GAN 反转不同，我们的反转框架可以推广到任何输入图像。为此，我们利用 SwiftBrushv2 [5]，这是一种以速度、多样性和质量著称的最新一步文本到图像模型，将其用作一步图像生成器和一步反转网络的主干。然后，我们使用从 SwiftBrushv2 初始化的权重对其进行训练，以通过两阶段训练策略处理任何源输入，结合来自合成数据和真实数据的监督。

在一步反演之后，我们引入了一种高效的基于掩码的编辑技术。我们的方法可以接受输入的编辑掩码，也可以直接从训练后的反演网络和指导提示中推断出它。然后，该掩码将用于我们新颖的注意力重新缩放技术中，以混合和控制编辑强度，同时保留背景元素，从而实现高质量的编辑结果。

据我们所知，我们的工作是第一个探索基于扩散的一步反演，使用一步式文本到图像生成模型来立即执行文本引导的图像编辑（0.23 秒）。与其他多步和少步编辑方法相比，我们的方法速度明显更快，并且实现了具有竞争力的编辑结果，如图 2 所示。总之，我们的主要贡献包括：

• 我们提出了一种采用两阶段策略训练的新颖的一步反演框架。经过训练后，我们的框架可以在一个步骤中将任何输入图像反转为可编辑的潜在图像，而无需进一步重新训练或微调。
• 我们表明，我们训练有素的反演框架可以在单个批量前向传递中生成由源和目标文本提示引导的编辑蒙版。
• 我们提出了一种新颖的注意力重新缩放技术，用于基于蒙版的编辑，提供对编辑强度的灵活控制，同时保留关键背景信息。

图 2. 在背景保留 (PSNR)、编辑语义 (CLIP 分数) 和运行时间方面，将我们的单步 SwiftEdit 与少步和多步扩散编辑方法进行比较。我们的方法提供了闪电般的文本引导编辑，同时实现了具有竞争力的结果。

2. Related Work

2.1. Text-to-image Diffusion Models

---

基于扩散的文本到图像模型 [24、26、27] 通常依赖于计算成本高昂的迭代去噪，以从高斯噪声中生成逼真的图像。最近的进展 [16、18、28、32] 通过将多步骤教师模型中的知识蒸馏到几步学生网络中来缓解这一问题。值得注意的作品 [5、15、16、20、32、39、40] 表明，这些知识甚至可以蒸馏到一步学生模型中。具体来说，Instaflow [15] 使用整流流来训练一步网络，而 DMD [40] 应用分布匹配目标进行知识迁移。DMDv2 [39] 消除了昂贵的回归损失，实现了高效的几步采样。 SwiftBrush [20] 采用无图像蒸馏方法，具有文本到 3D 生成目标，SwiftBrushv2 [5] 集成了训练后模型合并和限制 CLIP 损失，超越了其教师模型，实现了最先进的一步式文本到图像性能。这些一步式模型提供了丰富的文本-图像对齐先验信息，并且速度极快，使其成为我们一步式基于文本的图像编辑方法的理想选择。

2.2. Text-based Image Editing

---

有几种方法利用文本到图像模型中图文关系的强先验，通过逆编辑方法执行文本引导的多步骤图像编辑。首先，它们将源图像反转为“信息丰富的”噪声。DDIM 反转 [31] 等方法使用噪声预测的线性近似，而 Nulltext 反转 [19] 通过昂贵的每步优化来提高重建质量。直接反转 [11] 通过解开源和目标生成分支来绕过这些问题。其次，诸如 [3、10、21、22、35] 之类的编辑方法操纵注意力图来嵌入编辑，同时保留背景内容。然而，它们的多步扩散过程对于实际应用来说仍然太慢了。

为了解决这个问题，一些研究 [6, 8, 33] 使用快速生成模型 [29] 实现了几步图像编辑。 ICD [33] 使用一致性蒸馏框架在 3-4 步内实现精确反演，然后进行文本引导编辑。ReNoise [8] 在每个步骤中使用迭代重噪技术细化采样过程。TurboEdit [6] 使用移位噪声时间表将反演噪声与 SDXL Turbo [29] 等快速模型中的预期时间表对齐。 虽然这些方法减少了推理时间，但它们达不到快速应用所需的即时基于文本的图像编辑。我们的一步反演和一步局部编辑方法大大提高了时间效率，同时在编辑性能上超越了几步方法。

2.3. GAN Inversion

---

GAN 逆向 [2, 4, 14, 17, 23, 36, 41] 将源图像映射到预先训练的 GAN 的潜在空间中，从而使生成器能够重新创建图像，这对于图像编辑等任务非常有价值。有效的编辑需要一个既可以重建图像又可以通过潜在代码的变化支持逼真的编辑的潜在空间。方法分为三类：基于编码器的 [23, 41, 42]、基于优化的 [4, 14, 17] 和混合 [1, 2, 41]。基于编码器的方法学习从图像到潜在代码的映射以快速重建。基于优化的方法通过迭代优化来完善代码，而混合方法则将两者结合起来，使用编码器的输出作为进一步优化的初始化。受基于编码器的速度的启发，我们开发了一个一步式逆向网络，但我们没有使用 GAN，而是利用了一步式文本到图像扩散模型。这使我们能够跨不同领域实现基​​于文本的图像编辑，而不像基于 GAN 的方法那样局限于特定领域。

3. Preliminaries

多步扩散模型（Multi-step diffusion model）：文本到图像扩散模型 ${\mathbf{ϵ}}_{\varphi }$ 尝试通过 $T$ 次迭代去噪步骤生成目标提示嵌入 ${\mathbf{c}}_y$（从 CLIP 文本编码器中提取的文本提示 $y$ 对应的嵌入），最终生成图像 $\hat{\mathbf{x}}$。其过程从高斯噪声 ${\mathbf{z}}_T=\mathbf{ϵ}\sim \mathcal{N}(0,I)$ 开始：

$${\mathbf{z}}_{t-1}=\frac{{\mathbf{z}}_t-{\sigma }_t{ϵ}_{\varphi }\left({\mathbf{z}}_t,t,{\mathbf{c}}_y\right)}{{\alpha }_t}+{\delta }_t{\mathbf{ϵ}}_t,{\mathbf{ϵ}}_t\sim \mathcal{N}(0,I)$$

其中，$t$ 表示时间步，${\sigma }_t,{\alpha }_t,{\delta }_t$ 是三个系数。最终的潜变量 $\mathbf{z}={\mathbf{z}}_0$ 被输入到一个 VAE 解码器 $\mathcal{D}$ 中，以生成图像 $\hat{\mathbf{x}}=\mathcal{D}(\mathbf{z})$。

单步扩散模型（One-step diffusion model）：传统扩散模型的采样过程需要多次迭代，因而耗时较长。为了解决这一问题，单步文本到图像扩散模型（如 InstaFlow [15]、DMD [40]、DMD2 [39]、SwiftBrush [20] 和 SwiftBrushv2 [5]）被开发出来，将采样步骤减少到单步。具体而言，单步文本到图像扩散模型 $\mathbf{G}$ 试图在给定文本提示嵌入 ${\mathbf{c}}_y$ 的情况下，从噪声输入 $\mathbf{ϵ}\sim \mathcal{N}(0,1)$ 直接生成图像潜变量 $\hat{\mathbf{z}}$，无需迭代去噪步骤，即：

$$\hat{\mathbf{z}}=\mathbf{G}\left(\mathbf{ϵ},{\mathbf{c}}_y\right)$$

SwiftBrushv2（SBv2）在单步图像生成中表现出色，能够快速生成高质量且多样化的输出，成为我们方法的基础。在其前身的基础上，SBv2 集成了关键改进：

• 使用 SD-Turbo 初始化以增强输出质量；
• 引入夹紧的 CLIP 损失以强化视觉-文本对齐；
• 结合模型融合与后处理增强技术。

这些改进共同提升了性能与视觉保真度。

Score Distillation Sampling (SDS)：SDS 是一种流行的目标函数，通过利用 2D 扩散模型中学到的强先验来优化目标数据点 $\mathbf{z}$，其梯度的计算公式如下：

$${\nabla }_{\theta }{\mathcal{L}}_{\mathrm{SDS}}\triangleq {\mathbb{E}}_{t,\mathbf{ϵ}}\left[w(t)\left({\mathbf{ϵ}}_{\varphi }\left({\mathbf{z}}_t,t,{\mathbf{c}}_y\right)-\mathbf{ϵ}\right)\frac{\partial \mathbf{z}}{\partial \theta }\right]$$

其中：

• $\mathbf{z}=g(\theta )$ 是通过一个可微分的图像生成器 $g$ 渲染得到的，$g$ 由参数 $\theta$ 参数化；
• ${\mathbf{z}}_t$ 表示 $\mathbf{z}$ 被随机噪声 $\mathbf{ϵ}$ 扰动后的版本；
• $w(t)$ 是与时间步 $t$ 对应的缩放函数。

SDS 损失的目标是通过扩散模型的评分函数 ${ϵ}_{\varphi }\left({\mathbf{z}}_t,t,{\mathbf{c}}_y\right)$ 提供一个更新方向，将 $\mathbf{z}$ 移动到数据流形的高密度区域。

值得注意的是，SDS 梯度省略了扩散模型骨干（如 U-Net）的雅可比项，从而避免了反向传播整个扩散模型时的高昂计算成本。

Image-Prompt via Decoupled Cross-Attention：IPAdapter [38] 引入了一种图像提示条件 $\mathbf{x}$，可以无缝集成到预训练的文本到图像生成模型中。它通过分离文本和图像特征的条件效应的解耦交叉注意力机制（Decoupled Cross-Attention）实现了这一目标。具体来说，IPAdapter 为原始 U-Net 中的每个交叉注意力层添加了额外的交叉注意力层。给定以下特征：

• 图像特征 ${\mathbf{c}}_{\mathbf{x}}$（由 CLIP 图像编码器从 $\mathbf{x}$ 中提取）；
• 文本特征 ${\mathbf{c}}_y$（通过 CLIP 文本编码器从文本提示 $y$ 中提取）；
• 查询特征 ${\mathbf{Z}}_l$（来自前一层 U-Net 层 $l-1$），

解耦交叉注意力的输出 ${\mathbf{h}}_l$ 计算公式为：

$${\mathbf{h}}_l=\mathrm{Attn}\left(Q_l,K_y,V_y\right)+s_{\mathbf{x}}\mathrm{Attn}\left(Q_l,K_{\mathbf{x}},V_{\mathbf{x}}\right)$$

其中：

• $\mathrm{Attn}(.)$ 表示注意力操作；
• 缩放因子 $s_{\mathbf{x}}$ 用于控制图像特征 ${\mathbf{c}}_{\mathbf{x}}$ 对生成输出的影响；
• $Q_l=W^Q{\mathbf{Z}}_l$ 是通过权重矩阵 $W^Q$ 投影得到的查询矩阵；
• 文本特征 ${\mathbf{c}}_y$ 的键和值矩阵分别为 $K_y=W_y^K{\mathbf{c}}_y$ 和 $V_y=W_y^V{\mathbf{c}}_y$；
• 图像特征 ${\mathbf{c}}_{\mathbf{x}}$ 的投影键和值矩阵分别为 $K_{\mathbf{x}}=W_{\mathbf{x}}^K{\mathbf{c}}_{\mathbf{x}}$ 和 $V_{\mathbf{x}}=W_{\mathbf{x}}^V{\mathbf{c}}_{\mathbf{x}}$。

值得注意的是：

• 只有两个权重矩阵 $W_{\mathrm{x}}^K$ 和 $W_{\mathrm{x}}^V$ 是可训练的；
• 其余权重保持冻结状态，以保留预训练扩散模型的原始行为。

4. Proposed Method

我们的目标是使用单步文本转图像模型 SBv2 实现即时图像编辑。在第 4.1 节中，我们开发了一个单步反转网络，该网络可预测反转噪声，以便在通过 SBv2 时重建源图像。 我们为该反转网络引入了一种两阶段训练策略，无需进一步重新训练即可单步重建任何输入图像。概览如图 3 所示。在推理过程中，如第 4.2 节所述，我们使用自导编辑掩码来定位已编辑区域。 然后，我们的注意力重新缩放技术利用掩码实现解缠编辑并控制编辑强度，同时保留背景。

图 3. 我们提出的一步反演框架的两阶段训练方案。在第 1 阶段，我们利用 SwiftBrushv2 生成的合成数据来预热反演网络。在第 2 阶段，我们将重点转移到真实图像，使我们的反演框架能够立即反演任何输入图像，而无需额外的微调或重新训练。

4.1. Inversion Network and Two-stage Training

---

给定一个可能为合成（由 SBv2 等模型生成）或真实的输入图像，我们的首要目标是使用 SBv2 模型对其进行逆向操作和重构。为此，我们开发了一个单步逆向网络 ${\mathbf{F}}_{\theta }$，用于将图像潜变量 $\mathbf{z}$ 转换为逆向噪声 $\widehat{\mathbf{ϵ}}={\mathbf{F}}_{\theta }\left(\mathbf{z},{\mathbf{c}}_y\right)$，然后将其反馈到 SBv2 中以计算重构的潜变量：

$$\hat{\mathbf{z}}=\mathbf{G}\left(\widehat{\mathbf{ϵ}},{\mathbf{c}}_y\right)=\mathbf{G}\left({\mathbf{F}}_{\theta }\left(\mathbf{z},{\mathbf{c}}_y\right),{\mathbf{c}}_y\right)$$

合成图像

对于合成图像，训练 ${\mathbf{F}}_{\theta }$ 较为直接，因为我们可以获得噪声与潜变量对 $(\mathbf{ϵ},\mathbf{z})$，其中 $\mathbf{ϵ}$ 是生成 $\mathbf{z}$ 时使用的噪声。这允许我们直接回归 $\widehat{\mathbf{ϵ}}$ 到 $\mathbf{ϵ}$，从而使逆向噪声对齐 SBv2 的输入噪声分布。

真实图像

然而，对于真实图像，领域差异（domain gap）带来了挑战，因为原始噪声 $\mathbf{ϵ}$ 不可用，无法计算回归目标，可能导致 $\widehat{\mathbf{ϵ}}$ 偏离目标分布。

在接下来的部分中，我们将讨论逆向网络的架构及其设计的两阶段训练策略，以有效应对这些挑战。

---

逆向网络 ${\mathbf{F}}_{\theta }$

我们的逆向网络 ${\mathbf{F}}_{\theta }$ 使用单步扩散模型 $\mathbf{G}$ 的架构，并初始化为 $\mathbf{G}$ 的权重。然而，我们发现这种方法存在以下问题：

• ${\mathbf{F}}_{\theta }$ 预测的逆向噪声 $\widehat{\mathbf{ϵ}}$ 会尝试完美重构输入图像。
• 这导致过拟合到输入的特定模式，使得逆向噪声过度依赖于输入特征。

结果：
这种依赖限制了编辑的灵活性，因为噪声被过度调整以适配特定输入，削弱了生成的多样性和鲁棒性。

为了解决上述问题，我们在单步生成器 $G$ 中引入了一个辅助的图像条件分支，类似于 IP-Adapter [38]，并将其命名为 $G^{\mathbb{IP}}$。该分支将输入图像 $\mathbf{x}$ 编码的图像特征与文本提示 $y$ 一起集成，从而辅助重构并减少逆向网络 ${\mathbf{F}}_{\theta }$ 嵌入输入图像的视觉细节的负担。此方法有效地降低了对 $\widehat{\mathbf{ϵ}}$ 的依赖，增强了重构能力和编辑能力。逆向噪声 $\widehat{\mathbf{ϵ}}$ 和重构的图像潜变量 $\hat{\mathbf{z}}$ 计算如下：

$$\widehat{\mathbf{ϵ}}={\mathbf{F}}_{\theta }\left(\mathbf{z},c_y\right),\hat{\mathbf{z}}={\mathbf{G}}^{\mathrm{IP}}\left(\widehat{\mathbf{ϵ}},{\mathbf{c}}_y,{\mathbf{c}}_{\mathbf{x}}\right)$$

---

阶段 1：使用合成图像进行训练

如上所述，此阶段旨在使用从文本到图像扩散网络 $G$（即 SBv2）中采样的合成训练数据预训练逆向网络 ${\mathbf{F}}_{\theta }$。在图 3 中，我们用橙色展示了阶段 1 的训练流程。训练样本对 $(\mathbf{ϵ},\mathbf{z})$ 的生成方式如下：

$$\mathbf{ϵ}\sim \mathcal{N}(0,1),\mathbf{z}=\mathbf{G}\left(\mathbf{ϵ},{\mathbf{c}}_y\right)$$

我们结合重构损失 ${\mathcal{L}}_{\text{rec }}^{\text{stage }1}$ 和回归损失 ${\mathcal{L}}_{\text{regr }}^{\text{stage }1}$ 来训练逆向网络 ${\mathbf{F}}_{\theta }$ 以及部分 IP-Adapter 分支（包括图像条件的线性映射和交叉注意力层）。其中：

• 回归损失 ${\mathcal{L}}_{\text{regr }}^{\text{stage }1}$：鼓励 ${\mathbf{F}}_{\theta }$ 生成的逆向噪声 $\widehat{\mathbf{ϵ}}$ 能够紧密贴合 SBv2 的输入噪声分布，通过回归 $\widehat{\mathbf{ϵ}}$ 到 $\mathbf{ϵ}$ 实现。这确保了逆向噪声接近多元正态分布，这对于实现高效编辑能力至关重要（如之前的研究 [19] 所示）。

• 重构损失 ${\mathcal{L}}_{\text{rec }}^{\text{stage }1}$：强制重构的潜变量 $\hat{\mathbf{z}}$ 与原始源潜变量 $\mathbf{z}$ 对齐，从而保留输入图像的细节。

训练目标总结

阶段 1 的训练目标如下：

$$\begin{array}{c}{\mathcal{L}}_{\text{rec }}^{\text{stage }}=\Vert \mathbf{z}-\hat{\mathbf{z}}{\Vert }_2^2,{\mathcal{L}}_{\text{regr }}^{\text{stage }1}=\Vert \mathbf{ϵ}-\widehat{\mathbf{ϵ}}{\Vert }_2^2\\ {\mathcal{L}}^{\text{stage }1}={\mathcal{L}}_{\text{rec }}^{\text{stage }1}+{\lambda }^{\text{stage }1}\cdot {\mathcal{L}}_{\text{regr }}^{\text{stage }1}\end{array}$$

其中，在训练过程中设置 ${\lambda }^{\text{stage }1}=1$。完成该阶段后，我们的逆向框架能够重构由 SBv2 模型生成的源输入图像。然而，由于领域差异（domain gap），其无法很好地适用于真实图像，这促使我们继续进行阶段 2 的训练。

阶段 2：使用真实图像进行训练

在阶段 2 中，我们用基于感知的损失替换了阶段 1 的重构损失，采用 Deep Image Structure and Texture Similarity (DISTS) 度量 [7] 来定义感知损失：

$${\mathcal{L}}_{\text{perceptual }}^{\text{stage2 }}=\mathrm{DISTS}(\mathbf{x},\hat{\mathbf{x}})$$

其中，$\hat{\mathbf{x}}=\mathcal{D}(\hat{\mathbf{z}})$，而 $\hat{\mathbf{z}}$ 由以下公式生成：
$$\hat{\mathbf{z}}={\mathbf{G}}^{\mathrm{IP}}\left(\widehat{\mathbf{ϵ}},{\mathbf{c}}_y,{\mathbf{c}}_{\mathbf{x}}\right)$$

DISTS 在真实图像上进行训练，捕捉结构和纹理的感知细节，比阶段 1 中的像素级重构损失更鲁棒。

---

面临的挑战

在此阶段，由于无法获得用于在 SBv2 中重构 $\mathbf{z}$ 的原始噪声 $\mathbf{ϵ}$，我们无法直接应用阶段 1 的回归目标。若仅使用感知损失 ${\mathcal{L}}_{\text{perceptual }}^{\text{stage }2}$ 进行训练，可能会导致逆向噪声 $\widehat{\mathbf{ϵ}}$ 偏离理想的噪声分布 $\mathcal{N}(0,I)$。感知损失鼓励 $\widehat{\mathbf{ϵ}}$ 捕捉源图像的模式，从而有助于重构，但会限制未来的编辑灵活性（见图 4，第 2 列）。

为了解决这一问题，我们引入了一个新的正则化项 ${\mathcal{L}}_{\mathrm{regu}}^{\text{stage }2}$，灵感来源于 Score Distillation Sampling (SDS)（参见公式 (2)）。SDS 梯度将优化潜变量引导至数据流形的高密度区域。由于真实图像的潜变量 $\mathbf{z}=\mathcal{E}(\mathbf{x})$ 已位于高密度区域，我们将优化重点转向噪声项 $\widehat{\mathbf{ϵ}}$，将其视为添加到 $\mathbf{z}$ 的噪声。具体损失梯度计算如下：

$$\begin{array}{c}\widehat{\mathbf{ϵ}}={\mathbf{F}}_{\theta }\left(\mathbf{z},{\mathbf{c}}_y\right),{\mathbf{z}}_t={\alpha }_{\mathrm{t}}\mathbf{z}+{\sigma }_{\mathrm{t}}\widehat{\mathbf{ϵ}},\\ {\nabla }_{\theta }{\mathcal{L}}_{\mathrm{regu}}^{\text{stage }2}\triangleq {\mathbb{E}}_{t,\widehat{\mathbf{ϵ}}}\left[w(t)\left(\widehat{\mathbf{ϵ}}-{\mathbf{ϵ}}_{\varphi }\left({\mathbf{z}}_t,t,{\mathbf{c}}_y\right)\right)\frac{\partial \widehat{\mathbf{ϵ}}}{\partial \theta }\right]\end{array}$$

---

正则化梯度的作用

上述正则化梯度的符号与公式 (2) 相反，因为我们优化的是 $\widehat{\mathbf{ϵ}}$ 而非 $\mathbf{z}$（推导细节见附录）。

1. 初始条件：

   • 从阶段 1 初始化后，$\widehat{\mathbf{ϵ}}$ 类似于高斯噪声 $\mathcal{N}(0,1)$。
   • 这使得带噪潜变量 ${\mathbf{z}}_t$ 与多步模型的训练数据兼容，从而使教师模型准确预测 ${ϵ}_{\varphi }\left({\mathbf{z}}_t,t,{\mathrm{c}}_y\right)$，达到 ${ϵ}_{\varphi }\left({\mathbf{z}}_t,t,{\mathrm{c}}_y\right)-\widehat{\mathbf{ϵ}}\approx 0$。

2. 防止偏离原始分布：

   • 随着时间推移，重构损失推动 ${\mathbf{F}}_{\theta }$ 生成逆向噪声 $\widehat{\mathbf{ϵ}}$，以便于重构。
   • 但 $\widehat{\mathbf{ϵ}}$ 会逐渐偏离 $\mathcal{N}(0,1)$，从而生成不熟悉的 ${\mathbf{z}}_t$。
   • 正则化梯度阻止 $\widehat{\mathbf{ϵ}}$ 过度偏离原始分布，强化了阶段 1 的稳定性（如图 4 第 3 列所示）。

---

训练策略

和阶段 1 类似，我们结合感知损失 ${\mathcal{L}}_{\text{perceptual }}^{\text{stage }2}$ 和正则化损失 ${\mathcal{L}}_{\text{regu }}^{\text{stage }2}$，训练目标为：

$${\mathcal{L}}^{\text{stage2 }}={\mathcal{L}}_{\text{perceptual }}^{\text{stage }2}+{\lambda }^{\text{stage }2}\cdot {\mathcal{L}}_{\text{regu }}^{\text{stage }2}$$

其中，${\lambda }^{\text{stage }2}=1$。在训练过程中：

• 冻结 IP-Adapter 分支和解耦交叉注意力层，以保留阶段 1 学到的图像先验特征；
• 仅训练逆向网络 ${\mathbf{F}}_{\theta }$。

图 3 中以青绿色展示了阶段 2 的训练流程。

4.2. Attention Rescaling for Mask-aware Editing (ARaM)

---

在推理过程中，给定一个源图像 ${\mathbf{x}}^{\text{source}}$、源提示 $y^{\text{source}}$ 和编辑提示 $y^{\text{edit}}$，我们的目标是生成一个遵循编辑提示的编辑图像 ${\mathbf{x}}^{\text{edit}}$，同时保持与编辑无关的背景元素不被修改。经过两阶段的训练后，我们获得了一个训练良好的逆向网络 ${\mathbf{F}}_{\theta }$，能够将源图像潜变量 ${\mathbf{z}}^{\text{source}}=\mathcal{E}({\mathbf{x}}^{\text{source}})$ 转换为逆向噪声 $\widehat{\mathbf{ϵ}}$。直观地，我们可以使用单步生成器 ${\mathbf{G}}^{\mathrm{IP}}$ 再生图像，但用编辑提示嵌入 ${\mathbf{c}}_y^{\text{edit}}$ 作为引导提示。编辑图像的潜变量通过以下公式计算：

$${\mathbf{z}}^{\text{edit}}={\mathbf{G}}^{\mathrm{IP}}\left(\widehat{\mathbf{ϵ}},{\mathbf{c}}_y^{\text{edit}},{\mathbf{c}}_{\mathbf{x}}\right)$$

如 4.1 节所述，源图像条件 ${\mathbf{c}}_{\mathbf{x}}$ 对重构至关重要，其影响由 Eq. (3) 中的缩放因子 $s_{\mathbf{x}}$ 调节。为说明这一点，我们在生成编辑图像 ${\mathbf{x}}^{\text{edit}}=\mathcal{D}({\mathbf{z}}^{\text{edit}})$ 时改变 $s_{\mathbf{x}}$ 的值（见图 5b 的橙色模块）。如图所示，较高的 $s_{\mathbf{x}}$ 值加强了对源图像的保真度，但由于 ${\mathbf{c}}_{\mathbf{x}}$ 的严格控制，限制了编辑灵活性；相反，较低的 $s_{\mathbf{x}}$ 值允许更灵活的编辑，但会降低重构质量。

---

引入 Attention Rescaling for Mask-aware editing (ARaM)

基于上述观察，我们在 ${\mathbf{G}}^{\mathrm{IP}}$ 中引入了注意力重缩放机制，用于基于掩码的编辑（ARaM）。核心思想是在未编辑区域放大 ${\mathbf{c}}_{\mathbf{x}}$ 的影响以更好地保留原图，而在编辑区域减少其影响以提高编辑灵活性。为此，我们对 ${\mathbf{G}}^{\mathrm{IP}}$ 中 Eq. (3) 的计算进行重新公式化，移除全局缩放因子 $s_{\mathbf{x}}$，并引入区域特定的缩放因子，如下：

$$\begin{array}{rl}{\mathbf{h}}_l=& s_y\cdot M\cdot \mathrm{Attn}\left(Q_l,K_y,V_y\right)\\ & +s_{\text{edit}}\cdot M\cdot \mathrm{Attn}\left(Q_l,K_{\mathbf{x}},V_{\mathbf{x}}\right)\\ & +s_{\text{non-edit}}\cdot (1-M)\cdot \mathrm{Attn}\left(Q_l,K_{\mathbf{x}},V_{\mathbf{x}}\right)\end{array}$$

与 Eq. (3) 的区别

该解耦交叉注意力引入了三个缩放因子：

• $s_y$：应用于编辑区域 $M$ 的提示注意力，控制编辑提示的对齐强度；
• $s_{\text{edit}}$ 和 $s_{\text{non-edit}}$：分别控制图像条件 ${\mathbf{c}}_{\mathbf{x}}$ 在编辑和非编辑区域的影响。

如图 5b 的紫色模块所示，与使用相同的 $s_{\mathbf{x}}$ 相比，该方法生成的编辑图像能够更好地遵循编辑提示语义，同时实现良好的背景保留。此外，$s_y$ 的引入使我们能够在编辑区域 $M$ 内减弱或加强编辑提示对齐效果，从而控制编辑强度（见图 5c）。

---

编辑掩码 $M$ 的生成

上述的编辑掩码 $M$ 可以由用户提供，也可以由我们的逆向网络 ${\mathbf{F}}_{\theta }$ 自动生成。

自引导编辑掩码提取

我们观察到，经过训练的 ${\mathbf{F}}_{\theta }$ 在条件化不同的文本提示时，能够辨别逆向噪声图中的空间语义差异。具体步骤如下（见图 5a）：

1. 将源图像潜变量 ${\mathbf{z}}^{\text{source}}$ 输入 ${\mathbf{F}}_{\theta }$，分别使用源提示 ${\mathbf{c}}_y^{\text{source}}$ 和编辑提示 ${\mathbf{c}}_y^{\text{edit}}$ 生成逆向噪声 ${\widehat{ϵ}}^{\text{source}}$ 和 ${\widehat{ϵ}}^{\text{edit}}$；
2. 计算噪声差分图 ${\widehat{ϵ}}^{\text{source}}-{\widehat{ϵ}}^{\text{edit}}$；
3. 对差分图进行归一化，生成编辑掩码 $M$，高亮编辑区域。

这种方法能够有效突出需要编辑的区域，同时保留与背景相关的非编辑区域。

图 5. 蒙版感知编辑 (ARaM) 的注意力重新缩放说明。我们将注意力重新缩放与自导式编辑蒙版结合使用，以实现局部图像编辑并实现编辑强度控制。