E2EDiff：增强扩散模型的噪声到数据直接映射

paper是复旦大学发表在Arxiv 2024的工作

paper title：E2EDiff: Direct Mapping from Noise to Data for Enhanced Diffusion Models

Abstract

扩散模型已成为生成建模的强大框架，在各种任务中都取得了最先进的性能。然而，它们面临着几个固有的局限性，包括训练采样差距、渐进噪声过程中的信息泄漏，以及无法在训练期间纳入感知和对抗损失等高级损失函数。为了应对这些挑战，我们提出了一个创新的端到端训练框架，通过直接优化最终的重建输出来协调训练和采样过程。我们的方法消除了训练采样差距，通过将训练过程视为从纯噪声到目标数据分布的直接映射来减轻信息泄漏，并将感知和对抗损失整合到目标中。在 COCO30K 和 HW30K 等基准上进行的大量实验表明，我们的方法始终优于传统的扩散模型，即使在减少采样步骤的情况下，也能在 Frechet 初始距离 (FID) 和 ´ CLIP 分数方面取得优异成绩。这些发现强调了端到端训练的潜力，可以使基于扩散的生成模型朝着更为稳健和高效的解决方案发展。

1. Introduction

生成模型已成为现代机器学习的基石，为合成复杂数据分布提供了强大的框架。该领域的早期突破包括变分自动编码器 (VAE)[6]，它利用概率潜变量模型来学习数据的结构化表示并促进生成。随后出现了生成对抗网络 (GAN)[3]，它引入了对抗性训练策略，能够生成高质量且逼真的样本。最近，扩散模型作为一种前沿的生成方法而备受瞩目，它利用迭代去噪过程来实现输出的出色保真度和多样性 [5]。同时，一致性模型[17] 因其在整个生成过程中对齐中间表示的能力而备受关注，从而能够以更高的稳定性和可靠性实现更快的采样。这些进步凸显了生成建模领域中模型质量、采样效率和一致性之间的持续权衡。

在生成建模方法中，扩散模型在各种任务中取得了显著成功，建立了最先进的性能 [5, 18]。然而，尽管取得了令人瞩目的成就，扩散模型也面临着一些固有的局限性。这些挑战可以分为三个主要问题：1.训练采样差距：扩散模型的训练和测试过程存在根本性的不匹配。 在训练期间，模型经过优化，可在单步去噪任务中针对随机采样的时间步骤预测噪声。相反，测试阶段涉及多步迭代去噪过程。这种差异使模型无法有效捕获和减轻采样过程中出现的复合误差，尤其是在使用较少采样步骤时。因此，训练目标无法与测试条件很好地保持一致，从而限制了生成过程的效率和稳健性。2.$x_T$中的信息泄漏：训练中使用的渐进式噪声过程出现了一个关键问题。理想情况下，最终的噪声状态$x_T$应该近似于纯高斯噪声，以确保与反向去噪过程的理论假设一致 [5]。然而，由于前向过程中的细微不一致，$x_T$往往会偏离真正的高斯噪声，从而导致信息泄露。这种偏差损害了模型准确重建数据的能力，引入了偏差，从而降低了生成输出的质量。3. 有限的损失函数灵活性：扩散模型的训练范式涉及中间步骤的随机采样，这对将高级损失函数（例如感知损失或 GAN 损失）直接结合到中间输出上提出了挑战。人们普遍认为，使用 GAN 损失进行训练可以有效增强生成图像的细节丰富度，使其看起来更逼真。 这些损失函数对于提高生成结果的感知质量和语义一致性至关重要。 然而，扩散模型的固有结构限制了它们充分利用此类损失函数的能力，从而限制了它们实现卓越视觉保真度和高质量合成的潜力，特别是对于需要细粒度细节生成的任务。总的来说，这些限制阻碍了扩散模型的性能、一致性和灵活性，凸显了创新的必要性，这些创新既能解决这些挑战，又能保持其生成优势。

为了解决这些限制，我们提出了一个用于扩散模型的端到端训练框架，该框架在解决关键挑战的同时协调训练和采样过程。与专注于随机采样中间步骤的噪声预测的传统方法不同，我们的方法直接优化最终重建，通过统一的过程将纯高斯噪声转换为目标分布。这种设计通过使模型能够学习处理所有采样步骤中的复合误差来弥补训练采样差距。此外，它通过将训练过程视为从纯噪声到数据分布的直接映射来消除$x_T$中的信息泄露。统一目标还支持感知和对抗损失的集成，增强生成输出的保真度和语义一致性，同时保持理论一致性和稳健性。

在 COCO30K[8] 和 HW30K[2] 等广受认可的基准上进行的大量实验验证了我们提出的方法的有效性。我们的方法在关键性能指标（包括 Frechet Inception Distance (FID)[4] 和 CLIP ´ 分数[12]）方面始终超越现有的最先进扩散模型，同时以更少的采样步骤展示了强大的性能和优势。这些发现强调了端到端训练在推进基于扩散的生成模型方面的变革潜力。通过直接解决现有方法的核心限制，我们的框架为生成建模的效率、质量和一致性树立了新标准。

2. Related Work

2.1. Generative models.

---

从变分自编码器 (VAE)[6] 和生成对抗网络 (GAN)[3] 开始，生成模型经历了重大的演变。VAE 利用概率隐变量模型来生成结构化数据，但由于依赖于显式似然优化，因此通常难以生成高保真样本。另一方面，GAN 使用对抗训练来合成视觉上吸引人的数据，但容易出现不稳定和模式崩溃，从而限制了其可扩展性。现代的进步，例如扩散模型和潜在生成方法，通过引入迭代细化和高效的潜在表示解决了其中一些问题。然而，这些方法通常涉及复杂的训练动态，并且训练和推理阶段之间缺乏一致性[5, 15, 13]。与这些方法不同，我们的端到端方法直接优化了生成轨迹，确保了训练和采样之间的一致性，同时降低了复杂性。

2.2. Diffusion Models

---

扩散模型，例如去噪扩散概率模型 (DDPM) [5] 及其扩展，例如随机微分方程 (SDE) [18] 和常微分方程 (ODE)，已经为各种任务中的生成建模设定了基准。这些模型通过多步去噪过程逐步将高斯噪声转换为目标数据分布，利用学习到的逆变换序列。尽管扩散模型取得了显著的成功，但它们面临着一个重大限制，这个限制源于根本的训练-采样不匹配。在训练期间，模型经过优化，可在单步去噪任务中针对随机采样的时间步骤预测噪声。然而，在采样阶段，模型必须跨多个步骤迭代细化数据 [15]。这种差异引入了训练采样差距，导致推理过程中的复合误差和低效率。我们提出的方法通过端到端优化框架统一训练和采样过程，解决了这些限制。我们不再关注中间噪声预测，而是直接优化最终重建，使模型的训练目标与采样过程保持一致。这种方法不仅弥补了训练采样差距，还减轻了误差积累并提高了推理效率。通过直接瞄准生成输出的质量，我们的方法实现了卓越的性能，为更强大的基于扩散的生成模型铺平了道路。

2.3. Accelerating DMs.

---

为加速扩散模型，研究人员引入了几种创新技术。潜在扩散模型 (LDM)[13] 通过在压缩的潜在空间中运行来减少计算量，而渐进式蒸馏通过分阶段的知识转移最大限度地减少了采样步骤的数量[14, 7]。一致性模型对齐中间表示以促进快速采样[17, 16]，潜在一致性模型将此概念扩展到潜在空间以进一步提高效率[10]。然而，这些方法通常涉及额外的假设，例如特定的网络架构或辅助损失，并且可能仍然依赖于多步采样程序。我们的方法通过以端到端的方式学习整个轨迹来简化生成过程，消除了对逐步细化的依赖，并以更少的约束实现了稳健的性能。

3. Method

3.1. Background

---

扩散模型 [5] 作为基础框架，其定义了一个正向过程，通过 $T$ 个时间步逐步向数据样本 ${\mathbf{x}}_0\sim q\left({\mathbf{x}}_0\right)$ 添加高斯噪声，该过程建模为：

$$q\left({\mathbf{x}}_t\mid {\mathbf{x}}_{t-1}\right)=\mathcal{N}\left({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I}\right)$$

其中 ${\beta }_t\in (0,1)$ 是噪声方差的调度参数。反向过程旨在重构干净数据，参数化为：

$$p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)=\mathcal{N}\left({\mathbf{x}}_{t-1};{\mu }_{\theta }\left({\mathbf{x}}_t,t\right),{\Sigma }_{\theta }\left({\mathbf{x}}_t,t\right)\right)$$

在训练过程中，模型学习预测在每个时间步添加到数据 ${\mathbf{x}}_t$ 的噪声 $ϵ$，通过优化以下目标函数完成：

$$L(\theta )={\mathbb{E}}_{q\left({\mathbf{x}}_0,{\mathbf{x}}_t,ϵ\right)}\left[{\Vert ϵ-{ϵ}_{\theta }\left({\mathbf{x}}_t,t\right)\Vert }^2\right]$$

关键在于，这个损失函数将模型优化为单步去噪，假设各时间步之间相互独立。然而，在采样过程中，模型执行多步去噪，每一步都依赖于前一步的准确性，从而在时间步之间形成耦合。单步训练与多步采样之间的不匹配通常会限制生成样本的质量，因为误差可能在多个步骤中累积。

潜在扩散模型（Latent Diffusion Model, LDM）[13] 通过引入潜在表示 ${\mathbf{z}}_0=\mathcal{E}\left({\mathbf{x}}_0\right)$ 解决了在高维数据空间中操作的计算效率问题。该潜在表示由编码器 $\mathcal{E}$ 获得，正向和反向扩散过程随后在潜在空间中进行：

$$\begin{array}{rl}q\left({\mathbf{z}}_t\mid {\mathbf{z}}_{t-1}\right)& =\mathcal{N}\left({\mathbf{z}}_t;\sqrt{1-{\beta }_t}{\mathbf{z}}_{t-1},{\beta }_t\mathbf{I}\right),\\ p_{\theta }\left({\mathbf{z}}_{t-1}\mid {\mathbf{z}}_t\right)& =\mathcal{N}\left({\mathbf{z}}_{t-1};{\mu }_{\theta }\left({\mathbf{z}}_t,t\right),{\Sigma }_{\theta }\left({\mathbf{z}}_t,t\right)\right)\end{array}$$

训练目标保持相似形式：

$$L(\theta )={\mathbb{E}}_{q\left({\mathbf{z}}_0,{\mathbf{z}}_t,ϵ\right)}\left[{\Vert ϵ-{ϵ}_{\theta }\left({\mathbf{z}}_t,t\right)\Vert }^2\right]$$

尽管 LDM 通过在低维空间中操作降低了计算成本，但它继承了扩散模型的相同局限性：训练过程优化的是单步去噪，而采样过程则涉及多步的逐步细化。由于误差会在采样步骤中传播，最终结果可能因此劣化。

3.2. E2EDIFF

---

为了解决扩散模型的上述局限性，特别是单步训练与多步采样之间的不匹配问题，我们提出了一种端到端的训练方法。该方法直接优化从纯高斯噪声 ${\mathbf{z}}_T$ 到重构潜在表示 ${\mathbf{z}}_0$ 的生成过程，并以文本描述作为条件输入。此方法简化了训练目标，使得模型能够在潜在空间内的整个采样轨迹上进行联合优化，从而有效对齐训练目标和采样目标。

在传统扩散模型中，训练目标将损失分布在所有 $T$ 个时间步上，模型被训练以预测每个中间时间步 ${\mathbf{z}}_t$ 上添加的噪声 $ϵ$。这涉及 $T$ 个独立的训练损失，而最终的生成性能则受到采样过程中中间时间步耦合的间接影响。相比之下，我们的方法专注于 ${\mathbf{z}}_0$ 的最终重建质量，通过单一损失函数端到端训练模型。具体而言，我们通过对重建的潜在结果 ${\hat{\mathbf{z}}}_0$ 直接应用损失来重新定义训练目标，如下所示：

$$L_{\text{recon }}(\theta )={\mathbb{E}}_{q\left({\mathbf{z}}_T,{\mathbf{z}}_0\right)}\left[d\left({\mathbf{z}}_0,{\hat{\mathbf{z}}}_0\right)\right]$$

其中，${\hat{\mathbf{z}}}_0$ 是从纯高斯噪声 ${\mathbf{z}}_T$ 开始的生成过程的输出，并以文本嵌入作为条件输入。距离度量函数 $d(\cdot ,\cdot )$ 用于量化真实潜在表示 ${\mathbf{z}}_0$ 与重建输出 ${\hat{\mathbf{z}}}_0$ 之间的距离。通过在潜在空间而非像素空间中操作，我们的方法利用了预训练的图像编码器，实现了高效且语义上有意义的表示。

训练过程（详见算法 1）旨在学习从带噪潜在输入 ${\mathbf{z}}_T$ 到干净潜在表示 ${\mathbf{z}}_0$ 的映射，该映射以来自预训练文本编码器的文本嵌入 ${\mathbf{e}}_{\mathbf{c}}$ 作为条件输入。为了确保稳定性并提升模型性能，训练中采用指数移动平均（EMA）方法对模型参数进行逐步更新。

训练的关键在于定义度量函数 $d(\cdot ,\cdot )$，用于衡量预测的潜在表示 ${\hat{\mathbf{z}}}_0$ 与目标 ${\mathbf{z}}_0$ 之间的相似性。根据任务目标的不同，有多种常用的选择：

1. L1 损失：定义为 $d_{L1}\left({\mathbf{z}}_0,{\hat{\mathbf{z}}}_0\right)={\Vert {\mathbf{z}}_0-{\hat{\mathbf{z}}}_0\Vert }_1$，计算平均绝对误差（MAE），以其对异常值的鲁棒性著称。
2. L2 损失：定义为 $d_{L2}\left({\mathbf{z}}_0,{\hat{\mathbf{z}}}_0\right)={\Vert {\mathbf{z}}_0-{\hat{\mathbf{z}}}_0\Vert }_2^2$，计算均方误差（MSE）。这种损失对较大的偏差惩罚更为严格，有助于加速收敛，但对异常值更加敏感。
3. LPIPS 损失：定义为 $d_{\text{LPIPS}}\left({\mathbf{z}}_0,{\hat{\mathbf{z}}}_0\right)=\mathrm{LPIPS}\left({\mathbf{z}}_0,{\hat{\mathbf{z}}}_0\right)$。此损失利用预训练神经网络的特征来更好地对齐与人类感知相符的视觉相似性，在提高图像生成任务的输出质量方面尤为有效。
4. 混合损失：可以结合上述方法，通过加权公式 $d\left({\mathbf{z}}_0,{\hat{\mathbf{z}}}_0\right)={\lambda }_{L1}{d}_{L1}+{\lambda }_{L2}{d}_{L2}+{\lambda }_{\text{LPIPS}}{d}_{\text{LPIPS}}$ 进行组合。其中权重 ${\lambda }_{L1},{\lambda }_{L2},{\lambda }_{\text{LPIPS}}$ 用于平衡各个损失的贡献，使模型能够利用这些损失函数的互补优势。

训练过程中，每次迭代从时间步 t∼{1,…,T}t \sim\{1, \ldots, T\}t∼{1,…,T} 中随机采样，用于构造带噪样本 ${\mathbf{z}}_t$。与逐时间步优化的传统方法不同，我们的方法仅对最终输出 ${\hat{\mathbf{z}}}_0$ 应用重建损失，确保在整个生成轨迹上进行端到端优化。该方法将训练目标与采样过程对齐，减轻了传统扩散模型中误差传播的问题。

此外，重建损失 $L_{\text{recon}}$ 灵活支持多种度量函数 $d(\cdot ,\cdot )$，能够适应多样化的目标。例如，通过引入 LPIPS 损失可以优先提高感知质量，而使用 L1 或 L2 损失可以增强结构准确性。在本文的实验部分对这些选择及其影响进行了详细评估。

另一优势在于，该方法可以灵活地引入辅助目标，例如对抗性训练。通过添加基于 GAN 的损失 $L_{\mathrm{GAN}}$，模型可以进一步提高生成数据的感知质量。综合训练目标为：

$$L_{\text{total}}(\theta )=L_{\mathrm{recon}}(\theta )+{\lambda }_{\mathrm{GAN}}{L}_{\mathrm{GAN}}(\theta )$$

其中，${\lambda }_{\mathrm{GAN}}$ 控制对抗性损失的权重。重建损失确保数据的真实还原，而对抗性损失鼓励生成高质量和真实感强的输出。这种方法解决了现有扩散模型的核心局限性，提供了一个高效且灵活的框架，用于高质量生成任务。

图 1. 两种扩散模型训练方法的比较。 (a) 展示了单步训练方法，其中模型被训练为在随机采样的时间步骤中通过单个去噪步骤预测噪声。这种方法引入了训练采样差距，因为训练侧重于单步去噪，而测试需要迭代多步去噪。此外，前向噪声过程可能导致信息泄露，其中最终状态 xT 偏离理想的高斯噪声，从而损害重建质量。 (b) 展示了端到端训练方法，该方法直接优化了整个采样轨迹。从纯高斯噪声开始，该模型通过多步采样生成图像，无缝地对齐训练和测试过程。这种方法有效地消除了信息泄露，并实现了高级损失函数（如感知和 GAN 损失）的集成，从而提高了生成图像的保真度、语义一致性和整体质量。