论文状态:已完成

Consistency Models

发表:2023/03/03
原文链接PDF 下载
价格:0.100000
价格:0.100000
已有 2 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

本文提出一致性模型,以解决扩散模型生成慢的问题。这类模型通过将噪声直接映射到数据,支持极速的一步生成,并允许多步采样以平衡计算和样本质量。此外,一致性模型具备零样本数据编辑能力,无需针对特定任务训练,实验表明其在CIFAR-10等数据集上的性能优于现有技术。

摘要

Diffusion models have significantly advanced the fields of image, audio, and video generation, but they depend on an iterative sampling process that causes slow generation. To overcome this limitation, we propose consistency models, a new family of models that generate high quality samples by directly mapping noise to data. They support fast one-step generation by design, while still allowing multistep sampling to trade compute for sample quality. They also support zero-shot data editing, such as image inpainting, colorization, and super-resolution, without requiring explicit training on these tasks. Consistency models can be trained either by distilling pre-trained diffusion models, or as standalone generative models altogether. Through extensive experiments, we demonstrate that they outperform existing distillation techniques for diffusion models in one- and few-step sampling, achieving the new state-of-the-art FID of 3.55 on CIFAR-10 and 6.20 on ImageNet 64x64 for one-step generation. When trained in isolation, consistency models become a new family of generative models that can outperform existing one-step, non-adversarial generative models on standard benchmarks such as CIFAR-10, ImageNet 64x64 and LSUN 256x256.

思维导图

论文精读

中文精读

1. 论文基本信息

1.1. 标题

Consistency Models

1.2. 作者

Yang Song, Prafulla Dhariwal, Mark Chen, Ilya Sutskever

1.3. 隶属机构

OpenAI

1.4. 发表年份

2023

1.5. 摘要

扩散模型(Diffusion Models)极大地推动了图像、音频和视频生成领域的发展,但其迭代采样过程导致生成速度缓慢。为了克服这一限制,本文提出了一种新型模型家族——一致性模型(Consistency Models)。这类模型通过直接将噪声映射到数据来生成高质量样本。它们天生支持快速一步(one-step)生成,同时允许多步(multistep)采样以权衡计算成本和样本质量。一致性模型还支持零样本(zero-shot)数据编辑,例如图像修复(inpainting)、着色(colorization)和超分辨率(super-resolution),而无需针对这些任务进行显式训练。一致性模型可以通过蒸馏(distilling)预训练的扩散模型进行训练,也可以作为独立的生成模型进行训练。通过广泛的实验,本文证明一致性模型在一步和少步(few-step)采样方面优于现有扩散模型蒸馏技术,在 CIFAR-10 上实现了一步生成的最先进 FID 3.55,在 ImageNet 64x64 上达到了 6.20。当独立训练时,一致性模型成为一个全新的生成模型家族,在 CIFAR-10、ImageNet 64x64 和 LSUN 256x256 等标准基准测试中,其性能可以超越现有的一步、非对抗性生成模型。

1.6. 原文链接

2. 整体概括

2.1. 研究背景与动机

2.1.1. 扩散模型的成功与局限

扩散模型(Diffusion Models),又称基于分数的生成模型(Score-based Generative Models),在图像、音频、视频生成等领域取得了前所未有的成功。它们通过逐步向数据添加噪声,然后学习如何逆转这个过程来生成样本。这种迭代去噪(iterative denoising)过程提供了计算成本和样本质量之间的灵活权衡,并且是许多零样本数据编辑能力(如图像修复、着色、笔触引导图像编辑等)的基础。

然而,扩散模型也存在一个显著的局限性:生成速度缓慢。与生成对抗网络(GANs)、变分自编码器(VAEs)或归一化流(Normalizing Flows)等单步生成模型(single-step generative models)相比,扩散模型的迭代生成过程通常需要 10-2000 倍的计算量来生成样本,这导致推理速度慢,限制了其实时应用。

2.1.2. 现有解决方案的不足

为了加速扩散模型的采样过程,现有方法主要包括:

  • 更快的数值常微分方程(ODE)求解器: 这些方法旨在减少迭代步数,但通常仍需要超过 10 步才能生成有竞争力的样本。
  • 蒸馏技术(Distillation Techniques): 大多数蒸馏方法(如 Knowledge DistillationDFNO)需要先从扩散模型中生成大量样本来构建一个大型合成数据集,这个过程本身计算成本很高。虽然 Progressive Distillation (PD) 不需要构建合成数据,但在少步生成方面仍有提升空间。

2.1.3. 本文的切入点与创新思路

本文的目标是创建一种生成模型,既能实现高效的单步生成,又能保留迭代采样带来的重要优势,例如:

  • 在需要时权衡计算成本和样本质量的能力。

  • 执行零样本数据编辑任务的能力。

    为了实现这一目标,本文基于连续时间扩散模型中的概率流常微分方程(PF-ODE)提出了一致性模型(Consistency Models)。核心思想是学习一个模型,将 PF-ODE 轨迹上任意时间步的任何点直接映射到轨迹的起始点(原始数据)。这种模型的关键特性是自一致性(self-consistency):同一轨迹上的不同点映射到相同的初始点。

2.2. 核心贡献/主要发现

本文的主要贡献和发现可以总结如下:

  • 提出一致性模型: 引入了一种新型的生成模型家族,其核心设计理念是支持单步和少步高效生成。
  • 两种训练方法:
    • 一致性蒸馏 (Consistency Distillation, CD): 通过蒸馏预训练的扩散模型来训练一致性模型,显著提高了少步采样时的样本质量,并支持零样本图像编辑。
    • 独立训练 (Isolation Training, CT): 无需预训练的扩散模型,一致性模型可以作为独立的生成模型进行训练,成为一个新的生成模型类别。
  • 性能突破:
    • 蒸馏性能超越 SOTA: 在一步和少步采样中,Consistency Models 优于现有扩散模型蒸馏技术。在 CIFAR-10 上,一步生成实现了新的最先进 FID 3.55,两步生成达到了 2.93。在 ImageNet 64x64 上,一步生成实现了 6.20,两步生成实现了 4.70。这些结果均优于 Progressive Distillation (PD) 等现有方法,甚至超越了需要构建合成数据集的蒸馏方法。
    • 独立训练性能优越: 当独立训练时,一致性模型可以匹敌或超越 Progressive Distillation 的一步样本质量,并且在 CIFAR-10、ImageNet 64x64 和 LSUN 256x256 等基准测试上优于许多非对抗性、单步生成模型(如 VAEs 和 Normalizing Flows)。
  • 零样本数据编辑能力: Consistency Models 继承了扩散模型的零样本数据编辑能力,无需显式训练即可执行图像去噪、插值、修复、着色、超分辨率和笔触引导图像编辑(如 SDEdit)等任务。
  • 理论支撑: 提供了理论分析,证明了在零一致性蒸馏损失下,模型可以任意准确地估计一致性函数,并且连续时间目标与离散时间目标之间的关系。
  • 与强化学习和对比学习的联系: 发现 Consistency Models 与深度 Q-学习和基于动量的对比学习技术存在相似之处,预示着跨领域思想和方法交叉融合的潜力。

3. 预备知识与相关工作

3.1. 基础概念

3.1.1. 扩散模型 (Diffusion Models)

扩散模型(Diffusion Models),或称基于分数的生成模型(Score-based Generative Models),是一种生成模型,通过模拟一个逐步向数据添加噪声的扩散过程,并学习如何逆转这个过程来生成新的数据样本。

其核心思想包括两个阶段:

  1. 前向扩散过程 (Forward Diffusion Process): 将数据样本 x0pdata(x)\mathbf{x}_0 \sim p_{\text{data}}(\mathbf{x}) 逐步转换为纯高斯噪声。这个过程通常通过一个随机微分方程(Stochastic Differential Equation, SDE)或马尔可夫链来定义,其特点是每一步都向数据添加少量高斯噪声。随着时间的推移,数据分布逐渐扩散为简单的可追踪的噪声分布。
  2. 逆向去噪过程 (Reverse Denoising Process): 从纯噪声开始,通过学习一个模型来逐步逆转扩散过程,即从噪声中恢复出数据。这个逆向过程通常由一个参数化的模型来估计每个时间步的梯度(score function),然后利用数值求解器来迭代地从噪声中去噪,最终生成数据样本。

3.1.2. 概率流常微分方程 (Probability Flow Ordinary Differential Equation, PF-ODE)

在连续时间扩散模型中,与随机微分方程(SDE)对应的,存在一个确定性的常微分方程(ODE),称为概率流常微分方程(PF-ODE)。PF-ODE 的解轨迹在每个时间步 tt 的分布与 SDE 在该时间步的分布 pt(x)p_t(\mathbf{x}) 相同。这意味着,我们可以通过求解 PF-ODE 来生成样本,这与 SDE 相比具有一些优势,例如确定性采样和更易于控制。

PF-ODE 的形式如下: dxt=[μ(xt,t)12σ(t)2logpt(xt)]dt \mathrm{d}\mathbf{x}_t = \left[ \pmb{\mu}(\mathbf{x}_t, t) - \frac{1}{2}\sigma(t)^2 \nabla \log p_t(\mathbf{x}_t) \right] \mathrm{d}t 其中:

  • xt\mathbf{x}_t: 在时间 tt 的数据状态。

  • μ(xt,t)\pmb{\mu}(\mathbf{x}_t, t): 漂移系数(drift coefficient),描述了数据在没有噪声情况下的平均变化趋势。

  • σ(t)\sigma(t): 扩散系数(diffusion coefficient),描述了噪声的强度。

  • logpt(xt)\nabla \log p_t(\mathbf{x}_t): 评分函数(score function),是时间 tt 时数据分布 pt(xt)p_t(\mathbf{x}_t) 的对数梯度的负值,它指出了如何从噪声中去除数据。

  • dt\mathrm{d}t: 无穷小时间步。

    在本文使用的设置中(Karras et al., 2022),μ(x,t)=0\pmb{\mu}(\mathbf{x}, t) = \mathbf{0}σ(t)=2t\sigma(t) = \sqrt{2t}。在这种情况下,经验 PF-ODE 简化为: dxtdt=tsϕ(xt,t) \frac{\mathrm{d}\mathbf{x}_t}{\mathrm{d}t} = - t \pmb{s}_{\phi}(\mathbf{x}_t, t) 其中 sϕ(xt,t)\pmb{s}_{\phi}(\mathbf{x}_t, t) 是学习到的评分模型,近似于 logpt(xt)\nabla \log p_t(\mathbf{x}_t)

3.1.3. 评分函数 (Score Function)

评分函数 logpt(x)\nabla \log p_t(\mathbf{x}) 是扩散模型中的核心概念。它表示在给定数据点 x\mathbf{x} 和时间步 tt 时,数据分布 pt(x)p_t(\mathbf{x}) 的对数似然梯度。这个梯度可以被解释为“去噪方向”,因为它指出了如何将当前嘈杂的数据点推向更“真实”的数据区域。扩散模型通过训练一个神经网络来估计这个评分函数,从而实现从噪声中生成数据的能力。

3.1.4. FID (Fréchet Inception Distance)

FID(Fréchet Inception Distance)是衡量生成模型图像质量和多样性的一个常用指标。它通过比较真实图像分布和生成图像分布在 Inception V3 模型特征空间中的统计量来计算。 概念定义: FID 旨在量化生成样本的真实性和多样性。它假设 Inception V3 模型提取的特征能够很好地表征图像的语义信息。FID 值越低,表示生成图像的质量越高,与真实图像的分布越接近。 数学公式: FID(X,G)=μXμG22+Tr(ΣX+ΣG2(ΣXΣG)1/2) \mathrm{FID}(X, G) = ||\mu_X - \mu_G||_2^2 + \mathrm{Tr}(\Sigma_X + \Sigma_G - 2(\Sigma_X\Sigma_G)^{1/2}) 符号解释:

  • XX: 真实图像集的特征向量集合。
  • GG: 生成图像集的特征向量集合。
  • μX\mu_X: 真实图像集在 Inception V3 特征空间中的均值向量。
  • μG\mu_G: 生成图像集在 Inception V3 特征空间中的均值向量。
  • ΣX\Sigma_X: 真实图像集在 Inception V3 特征空间中的协方差矩阵。
  • ΣG\Sigma_G: 生成图像集在 Inception V3 特征空间中的协方差矩阵。
  • 22||\cdot||_2^2: L2 范数的平方。
  • Tr()\mathrm{Tr}(\cdot): 矩阵的迹(trace)。
  • ()1/2(\cdot)^{1/2}: 矩阵的平方根。

3.1.5. LPIPS (Learned Perceptual Image Patch Similarity)

LPIPS(Learned Perceptual Image Patch Similarity)是一种衡量两张图像之间感知相似度的指标。它通过计算两张图像经过预训练深度特征提取网络(如 AlexNet、VGG、Inception)后的特征向量之间的距离来评估。 概念定义: LPIPS 旨在捕捉人类对图像相似度的感知。传统的像素级距离(如 MSE)往往与人类感知不符,而 LPIPS 利用深度学习模型提取的特征,能够更好地反映人眼对图像差异的判断。LPIPS 值越低,表示两张图像在感知上越相似。 数学公式: LPIPS(x,y)=l1HlWlh,wwl(ϕl(x)h,wϕl(y)h,w)22 \mathrm{LPIPS}(\mathbf{x}, \mathbf{y}) = \sum_{l} \frac{1}{H_l W_l} \sum_{h,w} ||w_l \odot (\phi_l(\mathbf{x})_{h,w} - \phi_l(\mathbf{y})_{h,w})||_2^2 符号解释:

  • x,y\mathbf{x}, \mathbf{y}: 两张待比较的图像。
  • ϕl\phi_l: 深度特征提取网络(如 AlexNet)中第 ll 层的激活。
  • wlw_l: 缩放因子,用于在不同层之间进行加权。
  • Hl,WlH_l, W_l: 第 ll 层特征图的高度和宽度。
  • \odot: 逐元素乘积(element-wise product)。
  • 22||\cdot||_2^2: L2 范数的平方。

3.2. 前人工作与技术演进

3.2.1. 扩散模型的演进

  • Sohl-Dickstein et al. (2015): 首次提出了扩散模型(Diffusion Models)的基本思想,通过引入一个扩散过程来学习数据分布。
  • Song & Ermon (2019, 2020), Ho et al. (2020), Song et al. (2021): 推动了扩散模型的发展,特别是基于分数的生成模型(Score-based Generative Models)和去噪扩散概率模型(Denoising Diffusion Probabilistic Models, DDPM)的提出,使其在图像生成领域取得了突破性进展。
  • Karras et al. (2022): 提出了 Elucidating the Design Space of Diffusion-based Generative Models (EDM),系统性地阐明了扩散模型的设计空间,并提出了改进的架构和训练技术,显著提升了样本质量和训练效率。本文的工作即基于 EDM 模型。

3.2.2. 扩散模型加速采样技术

  • 更快的 ODE 求解器:
    • Song et al. (2020, 2021): 提出了 DDIM (Denoising Diffusion Implicit Models),通过使用非马尔可夫的逆向过程,允许更快的采样步数。
    • Zhang & Chen (2022), Lu et al. (2022), Dockhorn et al. (2022): 专注于开发更高效的数值 ODE 求解器(如 DPM-solver3-DEIS),以减少扩散模型生成样本所需的评估步数。
  • 蒸馏技术 (Distillation Techniques):
    • Luhman & Luhman (2021), Zheng et al. (2022): 提出了 Knowledge DistillationDFNO 等方法,通过训练一个更小的模型来模仿预训练扩散模型的行为,以加速采样。然而,这些方法通常需要生成大量合成数据进行蒸馏,本身计算成本较高。
    • Salimans & Ho (2022): 提出了 Progressive Distillation (PD),这是一种不需要生成合成数据即可蒸馏扩散模型的方法。它通过将 NN 步采样过程蒸馏到 N/2N/2 步,然后重复进行,从而实现少步生成。本文将 Consistency ModelsPD 进行了广泛比较。

3.2.3. 其他单步生成模型

  • GANs (Goodfellow et al., 2014): 生成对抗网络,通过判别器和生成器的对抗训练来生成高质量样本,通常是单步生成。
  • VAEs (Kingma & Welling, 2014; Rezende et al., 2014): 变分自编码器,通过学习数据的潜在表示来生成样本。
  • Normalizing Flows (Dinh et al., 2015, 2017; Kingma & Dhariwal, 2018): 归一化流,通过可逆变换将简单分布映射到复杂数据分布。

3.3. 差异化分析

Consistency Models 的核心创新在于直接学习从噪声到数据的一步映射,同时保持了扩散模型的一些关键优势。

特性/模型类型 扩散模型 (Diffusion Models) Progressive Distillation (PD) Consistency Models (CM)
采样过程 迭代式,需 10-2000 步 迭代式,经蒸馏后步数减少(如 1-10 步) 支持一步生成,也可多步采样权衡质量
生成速度 中等,快于原始扩散模型 ,一步生成最快
训练方式 学习去噪评分函数 蒸馏预训练扩散模型,将 NN 步蒸馏为 N/2N/2 蒸馏预训练扩散模型 (CD) 或独立训练 (CT)
对预训练模型依赖 自身无需,但依赖 score-based 训练范式 强依赖于预训练扩散模型 CD 依赖预训练扩散模型,CT 不依赖
零样本数据编辑能力 强大,通过迭代采样实现 继承自扩散模型 强大,通过多步采样过程实现
对合成数据集依赖
自一致性属性 隐式存在于 ODE 轨迹中 间接体现 显式设计并强制学习
核心创新点 逐步去噪生成高质量样本 迭代蒸馏以加速采样 直接映射噪声到数据,确保轨迹一致性
GANs/VAEs 等对比 迭代式,质量更高,但速度慢 速度提升,但仍迭代 可实现单步生成,非对抗性,质量优于一般 GANs/VAEs

Consistency Models 的核心差异化在于其自一致性(self-consistency)定义和训练目标。它不只是简单地模仿扩散模型的输出,而是主动学习一个函数,这个函数能够将 PF-ODE 轨迹上的任何点映射到该轨迹的唯一起始数据点。这种设计使得模型能够自然地支持一步生成,同时通过链式操作实现多步采样,并拥有强大的零样本数据编辑能力。这使其成为一个独立于扩散模型蒸馏之外,甚至可以独立训练的全新生成模型家族。

4. 方法论

本文提出了一致性模型(Consistency Models),旨在实现高效的单步生成,同时保留扩散模型在样本质量和零样本数据编辑方面的优势。一致性模型的核心思想是学习一个函数,该函数能够将概率流常微分方程(PF-ODE)轨迹上的任何点映射到该轨迹的起始数据点。

4.1. 方法原理

4.1.1. 一致性函数的定义

给定连续时间扩散模型中的 PF-ODE 轨迹 {xt}t[ϵ,T]\{ \mathbf{x}_t \}_{t \in [\epsilon, T]},一致性模型定义了一个一致性函数 f:(xt,t)xϵf: (\mathbf{x}_t, t) \mapsto \mathbf{x}_{\epsilon}。这个函数的目标是将轨迹上任意时间 tt 的状态 xt\mathbf{x}_t 映射回原始数据点 xϵ\mathbf{x}_{\epsilon}

自一致性(Self-consistency)属性: 一致性函数具有自一致性属性:对于属于同一 PF-ODE 轨迹的任意一对 (xt,t)(\mathbf{x}_t, t)(xt,t)(\mathbf{x}_{t'}, t'),其输出是相同的,即 f(xt,t)=f(xt,t)f(\mathbf{x}_t, t) = f(\mathbf{x}_{t'}, t'). 这意味着无论从轨迹的哪个点开始,通过一致性函数映射回的原始数据点都应该是一样的。

边界条件(Boundary Condition): 一个关键的约束是边界条件:f(xϵ,ϵ)=xϵf(\mathbf{x}_{\epsilon}, \epsilon) = \mathbf{x}_{\epsilon}。这意味着当时间步 tt 达到最小时间 ϵ\epsilon 时,一致性函数应该直接返回输入本身。这个条件对于成功训练一致性模型至关重要。

4.1.2. 一致性模型的参数化

一致性模型 fθ(x,t)f_{\theta}(\mathbf{x}, t) 通常使用深度神经网络进行参数化。为了满足边界条件 f(xϵ,ϵ)=xϵf(\mathbf{x}_{\epsilon}, \epsilon) = \mathbf{x}_{\epsilon},本文采用了带有跳跃连接(skip connections)的参数化方法,这与许多成功的扩散模型(如 EDM)类似,方便借用其强大的架构。

参数化形式如下: fθ(x,t)=cskip(t)x+cout(t)Fθ(x,t) f_{\theta}(\mathbf{x}, t) = c_{\mathrm{skip}}(t) \mathbf{x} + c_{\mathrm{out}}(t) F_{\theta}(\mathbf{x}, t) 其中:

  • Fθ(x,t)F_{\theta}(\mathbf{x}, t): 是一个自由形式的深度神经网络,其输出维度与 x\mathbf{x} 相同。

  • cskip(t)c_{\mathrm{skip}}(t)cout(t)c_{\mathrm{out}}(t): 是可微函数,需要满足边界条件:cskip(ϵ)=1c_{\mathrm{skip}}(\epsilon) = 1cout(ϵ)=0c_{\mathrm{out}}(\epsilon) = 0

    EDM (Karras et al., 2022) 中,原本的 cskip(t)c_{\mathrm{skip}}(t)cout(t)c_{\mathrm{out}}(t) 定义为: cskip(t)=σdata2t2+σdata2,cout(t)=σdatatσdata2+t2 c_{\mathrm{skip}}(t) = \frac{\sigma_{\mathrm{data}}^2}{t^2 + \sigma_{\mathrm{data}}^2}, \quad c_{\mathrm{out}}(t) = \frac{\sigma_{\mathrm{data}} t}{\sqrt{\sigma_{\mathrm{data}}^2 + t^2}} 为了满足 ϵ0\epsilon \neq 0 时的边界条件,本文对其进行了修改: cskip(t)=σdata2(tϵ)2+σdata2,cout(t)=σdata(tϵ)σdata2+t2 c_{\mathrm{skip}}(t) = \frac{\sigma_{\mathrm{data}}^2}{(t - \epsilon)^2 + \sigma_{\mathrm{data}}^2}, \quad c_{\mathrm{out}}(t) = \frac{\sigma_{\mathrm{data}} (t - \epsilon)}{\sqrt{\sigma_{\mathrm{data}}^2 + t^2}} 其中 σdata=0.5\sigma_{\mathrm{data}} = 0.5。这种修改确保了当 t=ϵt=\epsilon 时,cskip(ϵ)=1c_{\mathrm{skip}}(\epsilon)=1cout(ϵ)=0c_{\mathrm{out}}(\epsilon)=0,从而使 fθ(x,ϵ)=xf_{\theta}(\mathbf{x}, \epsilon) = \mathbf{x}

4.1.3. 采样过程

一旦训练好一致性模型 fθ(,)f_{\theta}(\cdot, \cdot),可以采用两种方式进行采样:

  1. 单步生成 (One-step Generation):

    • 从初始分布中采样噪声:x^TN(0,T2I)\hat{\mathbf{x}}_T \sim \mathcal{N}(\mathbf{0}, T^2 I)
    • 通过一致性模型一步生成样本:x^ϵ=fθ(x^T,T)\hat{\mathbf{x}}_{\epsilon} = f_{\theta}(\hat{\mathbf{x}}_T, T)。 这只需要一次网络前向传播,因此速度极快。
  2. 多步采样 (Multistep Sampling): 为了在样本质量和计算成本之间进行权衡,也可以通过交替去噪和注入噪声步骤来多次评估一致性模型。这个过程在 Algorithm 1 中进行了总结。

Algorithm 1 Multistep Consistency Sampling

Input: Consistency model fo(., ·), sequence of time points τ1 > τ2 > · . · > TN−1, initial noise XT
1: x ← fθ(xT, T )
2: for n = 1 to N − 1 do
3:    Sample z ~ N (0, I)
4:    xτn ← x + √τ n − 2Z
5:    x ← fθ(xτn, τn)
6: end for
Output: x

符号解释:

  • fθ(,)f_{\theta}(\cdot, \cdot): 训练好的一致性模型。
  • τ1>τ2>>τN1\tau_1 > \tau_2 > \cdots > \tau_{N-1}: 预定义的递减时间点序列。
  • xT\mathbf{x}_T: 初始噪声向量,通常从标准高斯分布 N(0,T2I)\mathcal{N}(\mathbf{0}, T^2I) 中采样。
  • x\mathbf{x}: 当前样本(在循环中会被更新)。
  • zN(0,I)\mathbf{z} \sim \mathcal{N}(0, I): 每次循环中采样的高斯噪声。
  • xτn\mathbf{x}_{\tau_n}: 在时间步 τn\tau_n 的噪声注入后的样本。
  • xoutput\mathbf{x}_{\text{output}}: 最终生成的样本。

算法流程:

  1. 初始化: 使用一致性模型将初始噪声 xT\mathbf{x}_T 在时间 TT 处一步映射到初始估计 x\mathbf{x}

  2. 迭代细化: 对于 nn1N-1: a. 采样一个新的高斯噪声 z\mathbf{z}。 b. 根据当前样本 x\mathbf{x} 和时间步 τn\tau_n,注入噪声以模拟在时间步 τn\tau_n 处的数据点 xτn\mathbf{x}_{\tau_n}。这里 τn2Z\sqrt{\tau_n - 2Z} 的具体形式可能需要根据上下文(如扩散过程的逆过程)进一步明确,但基本思想是根据当前去噪的样本,在较小的时间步上重新引入一些噪声。 c. 使用一致性模型 fθf_{\theta} 将带噪声的 xτn\mathbf{x}_{\tau_n} 重新映射回其原始数据估计 x\mathbf{x}

  3. 输出: 返回最终的样本 x\mathbf{x}

    通过这种多步采样,模型可以在每一步都修正和细化样本,从而提高生成质量,同时仍然比传统扩散模型的迭代步数少得多。

4.2. 训练一致性模型

本文提供了两种训练一致性模型的方法:一致性蒸馏(Consistency Distillation, CD)独立一致性训练(Consistency Training, CT)

4.2.1. 一致性蒸馏 (Consistency Distillation, CD)

这种方法通过蒸馏预训练的评分模型(score modelsϕ(x,t)s_{\phi}(\mathbf{x}, t) 来训练一致性模型。它利用数值 ODE 求解器生成 PF-ODE 轨迹上相邻点对,然后最小化一致性模型在这些点对上的输出差异。

训练目标构建:

  1. 离散化时间步: 将时间范围 [ϵ,T][\epsilon, T] 离散化为 N-1 个子区间,得到时间点 t1=ϵ<t2<<tN=Tt_1 = \epsilon < t_2 < \cdots < t_N = T。这些时间点通常使用 Karras et al. (2022) 提出的公式 ti=(ϵ1/ρ+(i1)/(N1)(T1/ρϵ1/ρ))ρt_i = (\epsilon^{1/\rho} + (i-1)/(N-1)(T^{1/\rho} - \epsilon^{1/\rho}))^{\rho} 来确定,其中 ρ=7\rho=7

  2. 生成相邻点对:

    • 从数据分布中采样真实数据点 xpdata\mathbf{x} \sim p_{\text{data}}
    • SDE 的转移密度中采样 xtn+1N(x,tn+12I)\mathbf{x}_{t_{n+1}} \sim \mathcal{N}(\mathbf{x}, t_{n+1}^2 I)
    • 使用数值 ODE 求解器(如 Euler 求解器)从 xtn+1\mathbf{x}_{t_{n+1}} 估计前一个时间点 tnt_n 的状态 x^tnϕ\hat{\mathbf{x}}_{t_n}^{\phi}。 若使用 Euler 求解器,更新规则为: x^tnϕ:=xtn+1+(tntn+1)Φ(xtn+1,tn+1;ϕ) \hat{\mathbf{x}}_{t_n}^{\phi} := \mathbf{x}_{t_{n+1}} + (t_n - t_{n+1}) \Phi(\mathbf{x}_{t_{n+1}}, t_{n+1}; \phi) 其中 Φ(x,t;ϕ)=tsϕ(x,t)\Phi(\mathbf{x}, t; \phi) = -t s_{\phi}(\mathbf{x}, t)。 所以: x^tnϕ=xtn+1(tntn+1)tn+1sϕ(xtn+1,tn+1) \hat{\mathbf{x}}_{t_n}^{\phi} = \mathbf{x}_{t_{n+1}} - (t_n - t_{n+1}) t_{n+1} s_{\phi}(\mathbf{x}_{t_{n+1}}, t_{n+1}) 这里的 x^tnϕ\hat{\mathbf{x}}_{t_n}^{\phi} 代表通过预训练评分模型 sϕs_{\phi}ODE 求解器从 xtn+1\mathbf{x}_{t_{n+1}} 逆向一步得到的 tnt_n 估计。
  3. 一致性蒸馏损失 (Consistency Distillation Loss): 本文定义了一致性蒸馏损失 LCDN\mathcal{L}_{CD}^N 如下: LCDN(θ,θ;ϕ):=E[λ(tn)d(fθ(xtn+1,tn+1),fθ(x^tnϕ,tn))] \mathcal{L}_{CD}^N (\pmb{\theta}, \pmb{\theta}^-; \phi) := \mathbb{E} [ \lambda(t_n) d( \pmb{f_{\theta}}(\mathbf{x}_{t_{n+1}}, t_{n+1}), \pmb{f_{\theta^-}}(\hat{\mathbf{x}}_{t_n}^{\phi}, t_n) ) ] 符号解释:

    • E[]\mathbb{E}[\cdot]: 对所有随机变量求期望。

    • xpdata\mathbf{x} \sim p_{\text{data}}: 从数据分布中采样真实数据 x\mathbf{x}

    • nU[[1,N1]]n \sim \mathcal{U}[[1, N-1]]: 从离散时间步索引中均匀采样 nn

    • xtn+1N(x;tn+12I)\mathbf{x}_{t_{n+1}} \sim \mathcal{N}(\mathbf{x}; t_{n+1}^2 I): 从以 x\mathbf{x} 为均值、方差为 tn+12It_{n+1}^2 I 的高斯分布中采样 xtn+1\mathbf{x}_{t_{n+1}}。这实际上是在 x\mathbf{x} 上添加噪声,得到 tn+1t_{n+1} 时刻的噪声数据。

    • λ(tn)R+\lambda(t_n) \in \mathbb{R}^+: 一个正的权重函数,通常设为 1。

    • d(,)d(\cdot, \cdot): 一个度量函数(如 L2 距离或 LPIPS),衡量两个向量之间的距离。

    • fθ(,)f_{\theta}(\cdot, \cdot): 当前正在训练的在线一致性模型(online network)。

    • fθ(,)f_{\theta^-}(\cdot, \cdot): 目标一致性模型(target network),其参数 θ\theta^-θ\theta 的指数移动平均(EMA)。

    • x^tnϕ\hat{\mathbf{x}}_{t_n}^{\phi}: 通过预训练评分模型 sϕs_{\phi}ODE 求解器从 xtn+1\mathbf{x}_{t_{n+1}} 逆向一步得到的 tnt_n 估计。

      这个损失函数的核心思想是,online networktn+1t_{n+1} 处对 xtn+1\mathbf{x}_{t_{n+1}} 的预测,应该与 target networktnt_n 处对 x^tnϕ\hat{\mathbf{x}}_{t_n}^{\phi} 的预测保持一致。

Algorithm 2 Consistency Distillation (CD) Input:Pretrainedscoremodelsφ(.,),consistencymodelfθ(.,),EMAparametersµ,T,ε,ρ,NInitializeθθforeachtrainingiterationdoSamplex pdataSamplen U[[1,N1]]Samplextn+1 N(x;tn+12I)//CalculateˆxtnφusingtheODEsolver(e.g.,Euler)//ForEuler:ˆxtnφxtn+1(tntn+1)tn+1sφ(xtn+1,tn+1)ˆxtnφxtn+1+(tntn+1)Φ(xtn+1,tn+1;φ)L(θ,θ;φ)λ(tn)d(fθ(xtn+1,tn+1),fθ(ˆxtnφ,tn))θθηθL(θ,θ;φ)θstopgrad(µθ+(1µ)θ)endforOutput:Trainedfθ(.,) Input: Pre-trained score model sφ(., ·), consistency model fθ(., ·), EMA parameters µ, T, ε, ρ, N Initialize θ⁻ ← θ for each training iteration do Sample x ~ pdata Sample n ~ U[[1, N-1]] Sample x_{tn+1} ~ N(x; t_{n+1}^2 I) // Calculate ˆx_{tn}^φ using the ODE solver (e.g., Euler) // For Euler: ˆx_{tn}^φ ← x_{tn+1} - (tn - tn+1) * tn+1 * sφ(x_{tn+1}, tn+1) ˆx_{tn}^φ ← x_{tn+1} + (tn - tn+1) * Φ(x_{tn+1}, tn+1; φ) L(θ, θ⁻; φ) ← λ(tn) * d(fθ(x_{tn+1}, tn+1), fθ⁻(ˆx_{tn}^φ, tn)) θ ← θ - η * ∇θ L(θ, θ⁻; φ) θ⁻ ← stopgrad(µ * θ⁻ + (1 - µ) * θ) end for Output: Trained fθ(., ·) 算法流程:

  1. 初始化: 目标网络参数 θ\theta^- 用在线网络参数 θ\theta 初始化。
  2. 迭代训练: 在每次训练迭代中: a. 采样真实数据 x\mathbf{x},采样时间步索引 nn。 b. 通过在 x\mathbf{x} 上添加噪声得到 xtn+1\mathbf{x}_{t_{n+1}}。 c. 利用预训练评分模型 sϕs_{\phi}ODE 求解器(如 Euler)从 xtn+1\mathbf{x}_{t_{n+1}} 计算前一步的估计 x^tnϕ\hat{\mathbf{x}}_{t_n}^{\phi}。 d. 计算一致性蒸馏损失 L(θ,θ;ϕ)\mathcal{L}(\theta, \theta^-; \phi)。 e. 使用随机梯度下降(SGD)更新在线网络参数 θ\theta。 f. 使用指数移动平均(EMA)更新目标网络参数 θ\theta^-,并应用 stopgrad 操作以避免梯度回传到 θ\theta^-

EMA 更新: θstopgrad(μθ+(1μ)θ)\pmb{\theta}^- \gets \mathrm{stopgrad}(\mu \pmb{\theta}^- + (1 - \mu) \pmb{\theta}) 其中 μ\mu 是衰减率。这种 EMA 更新和 stopgrad 操作类似于深度强化学习(DQN 中的 target network)和基于动量的对比学习,有助于稳定训练过程并提升性能。

理论支撑: 定理 1: 假设 fθf_{\theta} 满足 Lipschitz 条件,ODE 求解器具有局部误差界限 O((tn+1tn)p+1)O((t_{n+1} - t_n)^{p+1}),且 LCDN(θ,θ;ϕ)=0\mathcal{L}_{CD}^N (\pmb{\theta}, \pmb{\theta}; \phi) = 0(即在线网络和目标网络参数相同且损失为零),则估计的一致性模型与真实一致性函数之间的误差为 O((Δt)p)O((\Delta t)^p)。这意味着,只要时间步长 Δt\Delta t 足够小,一致性模型可以任意精确。

4.2.2. 独立一致性训练 (Consistency Training, CT)

Consistency Models 也可以在没有预训练扩散模型的情况下进行训练,使其成为一个独立的生成模型家族。CT 的关键在于,它利用了评分函数的一个无偏估计,从而避免了对预训练评分模型的依赖。

评分函数的无偏估计: 引理 1:xpdata(x)\mathbf{x} \sim p_{\text{data}}(\mathbf{x}), xtN(x;t2I)\mathbf{x}_t \sim \mathcal{N}(\mathbf{x}; t^2 \pmb{I}),且 pt(xt)=pdata(x)N(0,t2I)p_t(\mathbf{x}_t) = p_{\text{data}}(\mathbf{x}) \otimes \mathcal{N}(\mathbf{0}, t^2 I)。则评分函数可以无偏估计为: logpt(xt)=E[xtxt2 | xt] \nabla \log p_t(\mathbf{x}_t) = - \mathbb{E} \left[ \frac{\mathbf{x}_t - \mathbf{x}}{t^2} \ \middle| \ \mathbf{x}_t \right] 这意味着,给定真实数据 x\mathbf{x} 和带噪声的数据 xt\mathbf{x}_t,可以通过 xtx\mathbf{x}_t - \mathbf{x} 的比例来估计评分函数。

一致性训练损失 (Consistency Training Loss): 将引理 1 中的无偏估计代入一致性蒸馏损失,并假设使用 Euler 求解器,在 NN \to \infty 的极限情况下,可以得到一致性训练(CT)损失 LCTN\mathcal{L}_{CT}^NLCTN(θ,θ):=E[λ(tn)d(fθ(x+tn+1z,tn+1),fθ(x+tnz,tn))] \mathcal{L}_{CT}^N (\pmb{\theta}, \pmb{\theta}^-) := \mathbb{E} [ \lambda(t_n) d( f_{\theta}(\mathbf{x} + t_{n+1}\mathbf{z}, t_{n+1}), f_{\theta^-}(\mathbf{x} + t_n\mathbf{z}, t_n) ) ] 其中 zN(0,I)\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})。 这个损失函数与一致性蒸馏损失 LCDN\mathcal{L}_{CD}^N 非常相似,但不再显式依赖预训练的评分模型 sϕs_{\phi}。相反,它直接使用 x+tn+1z\mathbf{x} + t_{n+1}\mathbf{z} 作为 xtn+1\mathbf{x}_{t_{n+1}} 的噪声版本,和 x+tnz\mathbf{x} + t_n\mathbf{z} 作为 xtn\mathbf{x}_{t_n} 的噪声版本。

Algorithm 3 Consistency Training (CT) Input:Consistencymodelfθ(.,),EMAparametersµ(k),T,ε,ρ,N(k)Initializeθθ,k0foreachtrainingiterationkdoSamplex pdataSamplen U[[1,N1]]Samplez N(0,I)//xtn+1isimplicitlyx+tn+1z//xtnisimplicitlyx+tnzL(θ,θ)λ(tn)d(fθ(x+tn+1z,tn+1),fθ(x+tnz,tn))θθηθL(θ,θ)θstopgrad(µ(k)θ+(1µ(k))θ)kk+1endforOutput:Trainedfθ(.,) Input: Consistency model fθ(., ·), EMA parameters µ(k), T, ε, ρ, N(k) Initialize θ⁻ ← θ, k ← 0 for each training iteration k do Sample x ~ pdata Sample n ~ U[[1, N-1]] Sample z ~ N(0, I) // x_{tn+1} is implicitly x + t_{n+1}*z // x_{tn} is implicitly x + tn*z L(θ, θ⁻) ← λ(tn) * d(fθ(x + t_{n+1}*z, t_{n+1}), fθ⁻(x + tn*z, tn)) θ ← θ - η * ∇θ L(θ, θ⁻) θ⁻ ← stopgrad(µ(k) * θ⁻ + (1 - µ(k)) * θ) k ← k + 1 end for Output: Trained fθ(., ·) 算法流程:

  1. 初始化: 目标网络参数 θ\theta^- 用在线网络参数 θ\theta 初始化,迭代计数器 kk 归零。
  2. 迭代训练: 在每次训练迭代中: a. 采样真实数据 x\mathbf{x},采样时间步索引 nn,采样高斯噪声 z\mathbf{z}。 b. 计算一致性训练损失 L(θ,θ)\mathcal{L}(\theta, \theta^-)。 c. 使用随机梯度下降(SGD)更新在线网络参数 θ\theta。 d. 使用动态的指数移动平均(EMA)衰减率 μ(k)\mu(k) 更新目标网络参数 θ\theta^-,并应用 stopgrad 操作。 e. 递增迭代计数器 kk

渐进式 NNμ\mu 调度: 为了平衡收敛速度和样本质量,CT 采用渐进式增加 NN(离散步数)和调整 μ\muEMA 衰减率)的策略。

  • NN (大 Δt\Delta t): 损失函数的“方差”较小但“偏差”较大,有助于训练初期快速收敛。
  • NN (小 Δt\Delta t): 损失函数的“方差”较大但“偏差”较小,有助于训练后期获得更高质量的样本。 通过调度函数 N(k)μ(k)\mu(k),可以在训练过程中动态调整这些参数,以达到最佳性能。

调度函数的形式为: N(k)=kK((s1+1)2s02)+s021+1μ(k)=exp(s0logμ0N(k)) N(k) = \left\lceil \sqrt{\frac{k}{K} ((s_1 + 1)^2 - s_0^2) + s_0^2 - 1} \right\rceil + 1 \\ \mu(k) = \exp \left( \frac{s_0 \log \mu_0}{N(k)} \right) 符号解释:

  • kk: 当前训练迭代次数。
  • KK: 总训练迭代次数。
  • s0s_0: 初始离散步数。
  • s1>s0s_1 > s_0: 训练结束时的目标离散步数。
  • μ0>0\mu_0 > 0: 模型训练开始时的 EMA 衰减率。

4.2.3. 连续时间扩展

本文还探讨了 CDCT 损失在连续时间(即 NN \to \infty)下的形式。这些连续时间损失函数不需要指定离散步数 NN 或时间步 {ti}\{t_i\},但通常需要使用前向模式自动微分(forward-mode automatic differentiation)才能高效实现。理论分析表明,在连续时间极限下,CT 损失没有离散时间损失相关的“偏差”,因为它有效地将 Δt\Delta t 驱动到 0。

5. 实验设置

5.1. 数据集

本文在以下图像数据集上进行了实验:

  • CIFAR-10 (Krizhevsky et al., 2009):
    • 特点: 包含 10 个类别的 32x32 彩色图像,每个类别有 6000 张图像(5000 张训练,1000 张测试)。

    • 样本示例:

      该图像是一个展示生成模型效果的图表,包含两个部分:上半部分展示了EDM模型生成的图像,FID为2.04。图中包含多种类别的图像,主要以动物和车辆为主,展示了该模型的生成能力与多样性。 该图像是一个展示生成模型效果的图表,包含两个部分:上半部分展示了EDM模型生成的图像,FID为2.04。图中包含多种类别的图像,主要以动物和车辆为主,展示了该模型的生成能力与多样性。

      CIFAR-10 数据集中的图像样本,涵盖了动物、车辆等类别。

  • ImageNet 64x64 (Deng et al., 2009):
    • 特点: ImageNet 数据集的一个子集,包含 64x64 像素的彩色图像。

    • 样本示例:

      该图像是一个示意图,展示了EDM模型在图像生成中的表现,FID值为2.44。上半部分是一个生成结果的例子,与下半部分对比,展现了模型在图像生成质量上的提升。 该图像是一个示意图,展示了EDM模型在图像生成中的表现,FID值为2.44。上半部分是一个生成结果的例子,与下半部分对比,展现了模型在图像生成质量上的提升。

      ImageNet 64x64 数据集中的图像样本,包含各种物体和场景。

  • LSUN Bedroom 256x256 (Yu et al., 2015):
    • 特点: 包含卧室图像,分辨率为 256x256 像素。

    • 样本示例:

      该图像是一个图表,展示了基于EDM生成的多个房间图像样例,FID值为3.57。图中包含了不同风格的卧室,显示了模型在生成高质量样本方面的能力。 该图像是一个图表,展示了基于EDM生成的多个房间图像样例,FID值为3.57。图中包含了不同风格的卧室,显示了模型在生成高质量样本方面的能力。

      LSUN Bedroom 256x256 数据集中的图像样本,展示了多种卧室场景。

  • LSUN Cat 256x256 (Yu et al., 2015):
    • 特点: 包含猫的图像,分辨率为 256x256 像素。

    • 样本示例:

      该图像是一个展示猫咪图像生成效果的图表,显示了EDM(FID=6.69)生成的猫咪图像。多个图像在同一张图中并排展示,以比较图像生成技术的效果。 该图像是一个展示猫咪图像生成效果的图表,显示了EDM(FID=6.69)生成的猫咪图像。多个图像在同一张图中并排展示,以比较图像生成技术的效果。

      LSUN Cat 256x256 数据集中的图像样本,展示了多种猫的姿态。

选择这些数据集是为了在不同分辨率、不同复杂度和不同类别数量的图像上评估模型性能,从小型分类数据集 CIFAR-10 到中等分辨率的 ImageNet 64x64,再到高分辨率的特定类别数据集 LSUN Bedroom/Cat 256x256。这些数据集能有效地验证模型在生成多样化和高质量图像方面的能力。

5.2. 评估指标

5.2.1. FID (Fréchet Inception Distance)

  • 概念定义: FID 旨在量化生成样本的真实性和多样性。它通过比较真实图像分布和生成图像分布在 Inception V3 模型特征空间中的统计量来计算。FID 值越低,表示生成图像的质量越高,与真实图像的分布越接近。
  • 数学公式: FID(X,G)=μXμG22+Tr(ΣX+ΣG2(ΣXΣG)1/2) \mathrm{FID}(X, G) = ||\mu_X - \mu_G||_2^2 + \mathrm{Tr}(\Sigma_X + \Sigma_G - 2(\Sigma_X\Sigma_G)^{1/2})
  • 符号解释:
    • XX: 真实图像集的特征向量集合。
    • GG: 生成图像集的特征向量集合。
    • μX\mu_X: 真实图像集在 Inception V3 特征空间中的均值向量。
    • μG\mu_G: 生成图像集在 Inception V3 特征空间中的均值向量。
    • ΣX\Sigma_X: 真实图像集在 Inception V3 特征空间中的协方差矩阵。
    • ΣG\Sigma_G: 生成图像集在 Inception V3 特征空间中的协方差矩阵。
    • 22||\cdot||_2^2: L2 范数的平方。
    • Tr()\mathrm{Tr}(\cdot): 矩阵的迹(trace)。
    • ()1/2(\cdot)^{1/2}: 矩阵的平方根。

5.2.2. IS (Inception Score)

  • 概念定义: IS 用于评估生成图像的质量和多样性。高 IS 值通常表示生成图像不仅语义清晰(由 softmax 输出的低熵表示),而且涵盖了多种模式(由多个类别上的高熵表示)。
  • 数学公式: IS(G)=exp(ExG[DKL(p(yx)p(y))]) \mathrm{IS}(G) = \exp(\mathbb{E}_{\mathbf{x} \sim G} [D_{\mathrm{KL}}(p(y|\mathbf{x}) || p(y))])
  • 符号解释:
    • GG: 生成图像的分布。
    • x\mathbf{x}: 从生成器 GG 中采样得到的图像。
    • p(yx)p(y|\mathbf{x}): Inception 模型对生成图像 x\mathbf{x} 的预测类别分布。
    • p(y): Inception 模型对所有生成图像的平均预测类别分布。
    • DKL()D_{\mathrm{KL}}(\cdot || \cdot): Kullback-Leibler (KL) 散度,衡量两个分布之间的差异。
    • ExG[]\mathbb{E}_{\mathbf{x} \sim G}[\cdot]: 对生成图像的期望。

5.2.3. Precision (Prec.)

  • 概念定义: Precision 衡量生成图像的真实性,即生成图像中有多少是“可信的”或“真实的”。高 Precision 值表示生成的样本很少是离群点或不真实的。
  • 数学公式: 论文中未直接给出 Precision 的数学公式,但通常在生成模型评估中,Precision 是通过计算生成样本与真实样本的特征空间中的近邻来定义的。例如,它可以是生成样本在其最近的 kk 个真实样本邻居中的比例。

5.2.4. Recall (Rec.)

  • 概念定义: Recall 衡量生成模型的多样性,即模型能覆盖多少真实数据分布的模式。高 Recall 值表示模型能够生成各种各样的真实图像类型。
  • 数学公式: 论文中未直接给出 Recall 的数学公式,但通常在生成模型评估中,Recall 是通过计算真实样本与生成样本的特征空间中的近邻来定义的。例如,它可以是真实样本在其最近的 kk 个生成样本邻居中的比例。

5.3. 对比基线

本文将 Consistency Models 与以下几类模型进行了比较:

5.3.1. 扩散模型及其快速采样器

  • DDIM (Song et al., 2020): Denoising Diffusion Implicit Models,通过修改逆向过程,允许用更少的步数进行采样。
  • DPM-solver-2 (Lu et al., 2022): 一种用于扩散概率模型的快速 ODE 求解器。
  • DPM-solver-fast (Lu et al., 2022): DPM-solver 的加速版本。
  • 3-DEIS (Zhang & Chen, 2022): 另一种高效的 ODE 求解器,用于加速扩散模型的采样。
  • EDM (Karras et al., 2022): Elucidating the Design Space of Diffusion-based Generative Models,本文蒸馏的预训练模型。
  • Score SDE (Song et al., 2021): 基于随机微分方程的评分生成模型。
  • DDPM (Ho et al., 2020): Denoising Diffusion Probabilistic Models,扩散模型的基础工作之一。
  • LSGM (Vahdat et al., 2021): 潜在空间中的评分生成模型。
  • PFGM (Xu et al., 2022): Poisson Flow Generative Models

5.3.2. 扩散模型蒸馏技术

  • Knowledge Distillation (Luhman & Luhman, 2021):* 一种需要构建大量合成数据集进行蒸馏的方法。
  • DFNO (Zheng et al., 2022):* 另一种需要构建合成数据集的蒸馏方法。
  • Rectified Flow (+distill) (Liu et al., 2022):* 基于 Rectified Flow 的蒸馏方法,也需要合成数据集。
  • Progressive Distillation (PD, Salimans & Ho, 2022): 一种不需要合成数据集的蒸馏方法,本文与之进行了广泛比较。

5.3.3. 其他直接生成模型 (单步,非对抗性)

  • GANs (Generative Adversarial Networks):
    • BigGAN (Brock et al., 2019): 大规模 GAN 训练,实现高保真图像合成。
    • Diffusion GAN (Xiao et al., 2022): 结合了扩散模型和 GAN 的思想。
    • AutoGAN (Gong et al., 2019): 利用神经架构搜索来优化 GAN 结构。
    • E2GAN (Tian et al., 2020): 一种高效 GAN 架构搜索方法。
    • ViTGAN (Lee et al., 2021): 使用 Vision Transformers 训练 GAN
    • TransGAN (Jiang et al., 2021): 使用纯 Transformer 结构构建 GAN
    • StyleGAN2-ADA (Karras et al., 2020): StyleGAN2 的自适应判别器增强版本。
    • StyleGAN-XL (Sauer et al., 2022):StyleGAN 扩展到大型多样化数据集。
    • BigGAN-deep (Brock et al., 2019): BigGAN 的深度版本。
    • PGGAN (Karras et al., 2018): 渐进式增长 GAN,以提高质量和稳定性。
    • PG-SWGAN (Wu et al., 2019): 基于 Sliced Wasserstein 距离的渐进式 GAN
    • TDPM (GAN) (Zheng et al., 2023): 截断扩散概率模型与 GAN 结合。
    • StyleGAN2 (Karras et al., 2020): StyleGAN 的改进版本。
  • VAEs & Normalizing Flows:
    • 1-Rectified Flow (Liu et al., 2022): 一种流模型。

    • Glow (Kingma & Dhariwal, 2018): 基于可逆 1x1 卷积的生成流模型。

    • Residual Flow (Chen et al., 2019): 用于可逆生成建模的残差流。

    • GLFlow (Xiao et al., 2019): 生成潜在流模型。

    • DenseFlow (Grci et al., 2021): 密集连接的归一化流。

    • DC-VAE (Parmar et al., 2021): 对比生成自编码器。

      这些基线模型涵盖了扩散模型、扩散模型加速技术以及其他主流生成模型,能够全面评估 Consistency Models 在不同场景下的性能表现。

5.4. 训练细节

模型架构:

  • CIFAR-10 实验: 使用 Song et al. (2021) 中的 NCSN++NCSN++ 架构。
  • ImageNet 64x64, LSUN Bedroom 256x256, LSUN Cat 256x256 实验: 使用 Dhariwal & Nichol (2021) 中的网络架构。
  • 一致性模型架构与 EDM 模型相同,但修改了 cskip(t)c_{\mathrm{skip}}(t)cout(t)c_{\mathrm{out}}(t) 函数以满足边界条件(见方法论部分)。

预训练 EDM 模型:

  • 本文使用的 EDM 模型是内部预训练的,以便与 Progressive DistillationConsistency Distillation 进行公平比较。
  • Karras et al. (2022)EDM 论文中未提供 LSUN Bedroom 256x256LSUN Cat 256x256 的超参数,因此本文主要沿用了 ImageNet 64x64 的超参数。

超参数配置: 以下是 CDCT 模型训练所使用的超参数(原文 Table 3):

Hyperparameter CIFAR-10 ImageNet 64 × 64 LSUN 256 × 256
CD CT CD
Learning rate 4e-4 4e-4 8e-6
Batch size 512 512 2048
μ\mu 0 0.95
μ0\mu_0 0.9
s0s_0 2
s1s_1 150
NN 18 40
ODE solver Heun Heun
EMA decay rate 0.9999 0.9999 0.999943
Training iterations 800k 800k 600k
Mixed-Precision (FP16) No No Yes
Dropout probability 0.0 0.0 0.0
Number of GPUs 8 8 64

特殊说明:

  • EMA 衰减率: LSUN 256x256 数据集的 EMA 衰减率与 ImageNet 64x64 相同,但 LSUN Bedroom 256x256 上的 CD 训练发现零 EMA 效果更好。
  • Progressive Distillation (PD): 对于 PD,沿用了 Salimans & Ho (2022) 论文中 CIFAR-10ImageNet 64x64 的设置。对于 LSUN 256x256 数据集,虽然原始论文没有测试,但本文使用 ImageNet 64x64 的设置也取得了良好效果。
  • LPIPS 指标计算:
    • CIFAR-10ImageNet 64x64: 图像被双线性上采样到 224x224 分辨率,然后输入 LPIPS 网络。
    • LSUN 256x256: 直接使用原始分辨率评估 LPIPS

初始化: 所有一致性模型均随机初始化,然后使用 EDM 预训练模型中的 EMA 权重对其进行指数移动平均更新。

6. 实验结果与分析

6.1. 核心结果分析

6.1.1. 训练一致性模型的影响因素

本文首先在 CIFAR-10 上进行了一系列实验,以了解不同超参数对 CDCT 性能的影响。

1. 度量函数 d(.,.) 的影响:Figure 3a 中可以看出,LPIPS 作为度量函数在所有训练迭代中都明显优于 1\ell_12\ell_2 距离。这符合预期,因为一致性模型的输出是图像,而 LPIPS 专门设计用于衡量自然图像之间的感知相似度。

2. ODE 求解器和离散步数 NN 的影响 (针对 CD):

  • ODE 求解器: Figure 3b 显示,在相同的 NN 下,Heun 二阶 ODE 求解器始终优于 Euler 一阶求解器。这与定理 1(高阶 ODE 求解器训练的一致性模型具有更小的估计误差)的理论预测一致。

  • 离散步数 NN: Figure 3c 表明,当 NN 足够大时,CD 的性能对 NN 不再敏感。在 CIFAR-10 上,N=18N=18 是最佳选择。这些发现与 Karras et al. (2022) 的建议一致。

    综合以上发现,对于 CD,本文后续实验主要采用 LPIPS 度量和 Heun ODE 求解器。

3. 调度函数 N(k) 和 $$\mu(k)$ 对 CT 的影响: Figure 3d 显示,CT 的收敛速度对 NN 高度敏感。

  • NN: 导致更快的收敛速度,但最终样本质量较差。

  • NN: 导致收敛速度较慢,但收敛后样本质量更好。 这验证了第 5 节中的分析,即 CT 损失在小 NN 时“方差”较小但“偏差”较大,大 NN 时“方差”较大但“偏差”较小。因此,为了平衡收敛速度和样本质量,本文采用了渐进式增长 NNμ\mu 的自适应调度函数,这显著提高了 CT 的收敛速度和样本质量。

    该图像是一个图表,展示了不同训练迭代中CD(Consistency Distillation)和CT(Consistency Training)方法下的FID指标变化。图中包含多条曲线,分别代表了不同的方法和参数设置,直观呈现了不同模型训练效果的对比。 该图像是一个图表,展示了不同训练迭代中CD(Consistency Distillation)和CT(Consistency Training)方法下的FID指标变化。图中包含多条曲线,分别代表了不同的方法和参数设置,直观呈现了不同模型训练效果的对比。

图 3: CD 和 CT 训练性能随超参数的变化。 (a) 度量函数对 CD 训练的影响。(b) ODE 求解器对 CD 训练的影响。(c) N 值对 CD 训练的影响。(d) N 和 μ\mu 调度函数对 CT 训练的影响。

6.1.2. 少步图像生成结果

6.1.2.1. 蒸馏方法对比 (CD vs. PD)

本文将 Consistency Distillation (CD) 与目前无需构建合成数据的最直接可比方法 Progressive Distillation (PD) 进行了全面比较。

以下是原文 Figure 4 的结果,展示了 CDPD 在不同数据集和采样步数下的 FID 性能:

该图像是示意图,展示了不同采样步骤下的FID衡量结果,包括CIFAR-10、ImageNet 64×64、Bedroom 256×256和Cat 256×256四个子图。结果表明,随着采样步骤的增加,某些方法的FID值逐渐降低,从而反映出生成模型在生成质量上的改进。 该图像是示意图,展示了不同采样步骤下的FID衡量结果,包括CIFAR-10、ImageNet 64×64、Bedroom 256×256和Cat 256×256四个子图。结果表明,随着采样步骤的增加,某些方法的FID值逐渐降低,从而反映出生成模型在生成质量上的改进。

图 4: 在一致性蒸馏 (CD) 中使用 LPIPS 度量在所有数据集和采样步骤上都优于渐进式蒸馏 (PD),唯一的例外是 Bedroom 256x256 上的单步生成。

主要观察:

  • LPIPS 优于 2\ell_2: 使用 LPIPS 度量普遍改进了 PD 的性能,这表明感知距离在图像生成任务中的重要性。
  • 步数增加带来提升: 随着采样步数的增加,PDCD 的性能(FID 值降低)都有所提升,这表明多步采样可以提升样本质量。
  • CD 普遍优于 PD: 在所有数据集、采样步数和度量函数下,CD 几乎都优于 PD。唯一的例外是 Bedroom 256x256 上的单步生成,CD (使用 2\ell_2 度量) 略低于 PD (使用 2\ell_2 度量)。
  • CD 达到 SOTA: CDCIFAR-10 上达到了单步生成 FID 3.55 和两步生成 FID 2.93 的新最先进结果。在 ImageNet 64x64 上,单步生成 FID 6.20 和两步生成 FID 4.70 同样创下新纪录。

6.1.2.2. 与其他蒸馏技术的对比

以下是原文 Table 1 的结果,展示了 CIFAR-10 上各种扩散模型、蒸馏技术和直接生成模型的样本质量对比:

METHODNFE (↓)FID (↓)IS (↑)
Diffusion + Samplers
DDIM (Song et al., 2020)504.67
DDIM (Song et al., 2020)206.84
DDIM (Song et al., 2020)108.23
DPM-solver-2 (Lu et al., 2022)105.94
DPM-solver-fast (Lu et al., 2022)104.70
3-DEIS (Zhang & Chen, 2022)104.17
Diffusion + Distillation
Knowledge Distillation* (Luhman & Luhman, 2021)19.36
DFNO* (Zheng et al., 2022)14.12
1-Rectified Flow (+distill)* (Liu et al., 2022)16.189.08
2-Rectified Flow (+distill)* (Liu et al., 2022)14.859.01
3-Rectified Flow (+distill)* (Liu et al., 2022)15.218.79
PD (Salimans & Ho, 2022)18.348.69
CD13.559.48
PD (Salimans & Ho, 2022)25.589.05
CD22.939.75
Direct Generation
BigGAN (Brock et al., 2019)114.79.22
Diffusion GAN (Xiao et al., 2022)114.68.93
AutoGAN (Gong et al., 2019)112.48.55
E2GAN (Tian et al., 2020)111.38.51
ViTGAN (Lee et al., 2021)16.669.30
TransGAN (Jiang et al., 2021)19.269.05
StyleGAN2-ADA (Karras et al., 2020)12.929.83
StyleGAN-XL (Sauer et al., 2022)11.85
Score SDE (Song et al., 2021)20002.209.89
DDPM (Ho et al., 2020)10003.179.46
LSGM (Vahdat et al., 2021)1472.10
PFGM (Xu et al., 2022)1102.359.68
EDM (Karras et al., 2022)352.049.84
1-Rectified Flow (Liu et al., 2022)13781.13
Glow (Kingma & Dhariwal, 2018)148.93.92
Residual Flow (Chen et al., 2019)146.4
GLFlow (Xiao et al., 2019)144.6
DenseFlow (Grci et al., 2021)134.9
DC-VAE (Parmar et al., 2021)117.98.20
CT18.708.49
CT25.838.85

主要观察:

  • CD 优于现有蒸馏方法: CDCIFAR-10 上的一步生成 (FID 3.55) 和两步生成 (FID 2.93) 显著优于所有其他蒸馏方法,包括需要合成数据的方法(如 Knowledge Distillation*DFNODFNO*Rectified Flow*)以及 PD
  • 与最佳扩散模型相当: CD 的两步生成 FID 2.93 甚至可以与一些高性能扩散模型(如 DDPMFID 3.17)在更少步数下相媲美,甚至接近 Score SDE (2000步, FID 2.20) 和 EDM (35步, FID 2.04) 的表现。

6.1.2.3. 直接生成对比 (CT vs. 其他单步生成模型)

以下是原文 Table 1 和 Table 2 的结果,展示了 CT 作为独立生成模型与其他直接生成模型的样本质量对比:

CIFAR-10 (Table 1 节选):

METHODNFE (↓)FID (↓)IS (↑)
Direct Generation
BigGAN (Brock et al., 2019)114.79.22
Diffusion GAN (Xiao et al., 2022)114.68.93
AutoGAN (Gong et al., 2019)112.48.55
E2GAN (Tian et al., 2020)111.38.51
ViTGAN (Lee et al., 2021)16.669.30
TransGAN (Jiang et al., 2021)19.269.05
StyleGAN2-ADA (Karras et al., 2020)12.929.83
StyleGAN-XL (Sauer et al., 2022)11.85
1-Rectified Flow (Liu et al., 2022)13781.13
Glow (Kingma & Dhariwal, 2018)148.93.92
Residual Flow (Chen et al., 2019)146.4
GLFlow (Xiao et al., 2019)144.6
DenseFlow (Grci et al., 2021)134.9
DC-VAE (Parmar et al., 2021)117.98.20
CT18.708.49
CT25.838.85

ImageNet 64x64, LSUN Bedroom 256x256LSUN Cat 256x256 (Table 2):

METHODNFE (↓)FID (↓)Prec. ()Rec. (↑)
ImageNet 64 × 64
PD† (Salimans & Ho, 2022)115.390.590.62
DFNO† (Zheng et al., 2022)18.35
CD†16.200.680.63
PD† (Salimans & Ho, 2022)28.950.630.65
CD24.700.690.64
ADM (Dhariwal & Nichol, 2021)2502.070.740.63
EDM (Karras et al., 2022)792.440.710.67
BigGAN-deep (Brock et al., 2019)14.060.790.48
CT113.00.710.47
CT211.10.690.56
LSUN Bedroom 256 × 256
PD† (Salimans & Ho, 2022)116.920.470.27
PD† (Salimans & Ho, 2022)28.470.560.39
CD†17.800.660.34
CD†25.220.680.39
DDPM (Ho et al., 2020)10004.890.600.45
ADM (Dhariwal & Nichol, 2021)10001.900.660.51
EDM (Karras et al., 2022)793.570.660.45
PGGAN (Karras et al., 2018)18.34
PG-SWGAN (Wu et al., 2019)18.0
TDPM (GAN) (Zheng et al., 2023)15.24
StyleGAN2 (Karras et al., 2020)12.350.590.48
CT116.00.600.17
CT27.850.680.33
LSUN Cat 256 × 256
PD† (Salimans & Ho, 2022)129.60.510.25
PD† (Salimans & Ho, 2022)215.50.590.36
CD†111.00.650.36
CD†28.840.660.40
DDPM (Ho et al., 2020)100017.10.530.48
ADM (Dhariwal & Nichol, 2021)10005.570.630.52
EDM (Karras et al., 2022)796.690.700.43
PGGAN (Karras et al., 2018)137.5
StyleGAN2 (Karras et al., 2020)17.250.580.43
CT120.70.560.23
CT211.70.630.36

主要观察:

  • CT 优于非对抗性单步生成模型: CT 作为独立生成模型,在 CIFAR-10 上的一步生成 (FID 8.70) 和两步生成 (FID 5.83) 明显优于现有的 VAEsNormalizing Flows(如 GlowResidual Flow 等),这些模型的 FID 值通常非常高。

  • CTPD (单步) 质量相当: CT 在没有预训练扩散模型的情况下,其一步生成质量可以匹敌 PD 的一步生成(例如,CIFAR-10CT FID 8.70 vs. PD FID 8.34)。

  • GANs 对比: 虽然 CT 无法超越最顶尖的 GANs(如 StyleGAN-XLStyleGAN2-ADA),但它在非对抗性单步生成模型中表现优异。例如,在 CIFAR-10 上,CT 的 FID 8.70 优于 BigGAN (FID 14.7) 等一些 GAN 模型。

  • 多步 CT 的提升: CT 的两步生成在所有数据集上都显著提高了样本质量(FID 值降低),例如 CIFAR-10 从 8.70 降至 5.83,LSUN Bedroom 从 16.0 降至 7.85。

  • 样本视觉效果: Figure 5 展示了 EDM、单步 CT 和两步 CT 生成的样本。尽管 CTEDM 是独立训练的,但从相同初始噪声生成的样本表现出显著的结构相似性,这表明 CT 不太可能出现模式崩溃问题。

    Figure 5: Samples generated by EDM (top), \(\\mathrm { C T + }\) single-step generation (middle), and \(\\mathrm { C T } + 2\) -step generation (Bottom). Al corresponding images are generated from the same initial noise. 该图像是多个图像样本的展示,左侧为不同动物和场景的生成图像,右侧为室内环境和猫的生成图像,这些样本体现了模型在不同主题下的生成能力。

图 5: 由 EDM(上),CT+单步生成(中)和 CT+两步生成(下)生成的样本。所有对应图像都由相同的初始噪声生成。

6.2. 零样本数据编辑能力

Consistency Models 继承了扩散模型的优势,可以通过修改多步采样过程(如 Algorithm 1 中描述的)实现各种零样本数据编辑和操作任务,而无需针对特定任务进行显式训练。

Algorithm 4 总结了零样本图像编辑的一般过程,它结合了 Consistency Models 的多步采样能力与特定任务的线性变换 A\pmb{A} 和二值掩码 Ω\pmb{\Omega}

Algorithm 4 Zero-Shot Image Editing

Input: Consistency model fθ(., ·), sequence of time points t1 > t2 > · · · > tN, reference image y, invertible linear transformation A, and binary image mask Ω
1: y ← A⁻¹ [ (A y) ⊙ (1 - Ω) + 0 ⊙ Ω ]
2: Sample x ~ N(y, t1^2 I)
3: x ← fθ(x, t1)
4: x ← A⁻¹ [ (A y) ⊙ (1 - Ω) + (A x) ⊙ Ω ]
5: for n = 2 to N do
6:    Sample x ~ N(x, (tn^2 - ε^2) I)
7:    x ← fθ(x, tn)
8:    x ← A⁻¹ [ (A y) ⊙ (1 - Ω) + (A x) ⊙ Ω ]
9: end for
10: Output: x

符号解释:

  • fθ(,)f_{\theta}(\cdot, \cdot): 训练好的一致性模型。
  • t1>t2>>tNt_1 > t_2 > \cdots > t_N: 递减的时间点序列。
  • y\mathbf{y}: 参考图像,用于引导样本生成。
  • A\pmb{A}: 可逆线性变换,将图像映射到潜在空间,其中条件信息通过 y\mathbf{y} 注入。
  • Ω\pmb{\Omega}: 二值图像掩码,指示需要编辑的区域。
  • \odot: 逐元素乘积。
  • x\mathbf{x}: 当前样本(在循环中会被更新)。
  • zN(0,I)\mathbf{z} \sim \mathcal{N}(0, I): 每次循环中采样的高斯噪声。

算法流程:

  1. 初始化参考图像 y\mathbf{y}: 根据任务需求,通过 A\pmb{A}Ω\pmb{\Omega} 对参考图像进行预处理。例如,对于着色,将参考图像的颜色信息通过掩码去除。

  2. 初始噪声采样: 从以 y\mathbf{y} 为均值、方差为 t12It_1^2 I 的高斯分布中采样初始噪声 x\mathbf{x}

  3. 第一次去噪: 使用一致性模型 fθf_{\theta}x\mathbf{x} 在时间 t1t_1 处去噪。

  4. 第一次条件注入: 根据参考图像 y\mathbf{y} 和掩码 Ω\pmb{\Omega},将条件信息注入到去噪后的 x\mathbf{x} 中。

  5. 迭代细化: 对于 nn2NN: a. 在当前样本 x\mathbf{x} 上添加噪声,模拟在时间步 tnt_n 处的数据点。这里方差为 (tn2ϵ2)I(t_n^2 - \epsilon^2)I,表示添加的噪声与当前时间步的尺度相关。 b. 使用一致性模型 fθf_{\theta} 将带噪声的 x\mathbf{x} 在时间 tnt_n 处去噪。 c. 再次将条件信息注入到去噪后的 x\mathbf{x} 中。

  6. 输出: 返回最终的编辑图像 x\mathbf{x}

    以下是 Consistency ModelsLSUN Bedroom 256x256 数据集上进行零样本图像编辑的示例。

    该图像是示意图,展示了图像着色过程的不同阶段。第一行左侧为灰度图,中间为着色图,右侧为真实图像;第二行左侧为下采样图(32×32),中间为全分辨率图(256×256),右侧为真实图像(256×256)。 该图像是示意图,展示了图像着色过程的不同阶段。第一行左侧为灰度图,中间为着色图,右侧为真实图像;第二行左侧为下采样图(32×32),中间为全分辨率图(256×256),右侧为真实图像(256×256)。

图 6: 使用一致性蒸馏在 LSUN Bedroom 256x256 上训练的一致性模型进行零样本图像编辑。左图:用户提供的笔触输入。右图:笔触引导的图像生成。

具体的零样本任务展示:

  1. 图像着色 (Image Colorization):

    • 方法:y\mathbf{y} 为灰度图像,Ω\pmb{\Omega} 为颜色通道的掩码(例如,只保留灰度通道信息),A\pmb{A} 为将 RGB 转换为 luminance-chrominance 空间的变换。

    • 效果: Figure 8 展示了模型可以成功地为灰度卧室图像添加颜色,即使它没有专门为着色任务训练过。

      Figure 8: Gray-scale images (left, colorized images by a consistency model (middle), and ground truth (right). 该图像是图表,展示了灰度图像(左侧)、通过一致性模型生成的上色图像(中间)和真实图像(右侧)的对比。此图表显示了一致性模型在图像上色任务中的效果,体现了其在零样本数据编辑方面的能力。

    图 8: 灰度图像(左),一致性模型着色图像(中),真实图像(右)。

  2. 超分辨率 (Super-resolution):

    • 方法:y\mathbf{y} 为低分辨率图像,Ω\pmb{\Omega} 为高频信息的掩码,A\pmb{A} 为下采样操作。

    • 效果: Figure 9 展示了模型可以从低分辨率输入生成高分辨率图像。

      Figure 9: Downsampled images of resolution \(3 2 \\times 3 2\) (left), full resolution \(2 5 6 \\times 2 5 6 )\) images generated by a consistency model (middle), and ground truth images of resolution \(2 5 6 \\times 2 5 6\) (right). 该图像是示意图,展示了使用一致性模型生成的不同分辨率的室内图像。左侧为分辨率32×3232 \times 32的低分辨率图像,中间为一致性模型生成的全分辨率256×256256 \times 256图像,右侧为同样分辨率的真实图像。该图像展示了模型生成图像在细节和真实感方面的对比。

    图 9: 32x32 分辨率的下采样图像(左),一致性模型生成的 256x256 全分辨率图像(中),256x256 分辨率的真实图像(右)。

  3. 笔触引导图像生成 (Stroke-guided Image Generation):

    • 方法: 类似于 SDEdit (Meng et al., 2021)。令 y\mathbf{y} 为用户绘制的笔触图像,A\pmb{A} 为恒等变换,Ω\pmb{\Omega} 为笔触区域的掩码。

    • 效果: Figure 6Figure 13 展示了模型可以根据用户提供的笔触输入生成图像。

      该图像是生成模型在图像编辑任务(如图像修复和上色)中的结果展示,包括多个样本生成的视觉效果。图像展示了不同场景和细节上生成质量的多样性,体现出一致性模型在快速生成能力方面的优势。 该图像是生成模型在图像编辑任务(如图像修复和上色)中的结果展示,包括多个样本生成的视觉效果。图像展示了不同场景和细节上生成质量的多样性,体现出一致性模型在快速生成能力方面的优势。

    图 13: 笔触引导图像生成(SDEdit)的结果。

  4. 图像修复 (Image Inpainting):

    • 方法:y\mathbf{y} 为部分遮盖的图像,Ω\pmb{\Omega} 为遮盖区域的二值掩码,A\pmb{A} 为恒等变换。

    • 效果: Figure 10 展示了模型可以填充图像中的缺失区域。

      Figure 10: Masked images (left), imputed images by a consistency model (middle), and ground truth (right). 该图像是多个房间的图像展示,左侧为原始图像,中央为通过一致性模型生成的修复图像,右侧为真实图像。展示了模型在图像修复方面的能力,帮助可视化其生成效果。

    图 10: 遮盖图像(左),一致性模型修复图像(中),真实图像(右)。

  5. 图像去噪 (Image Denoising):

    • 方法: 给定带噪声的输入图像 x\mathbf{x},其噪声水平为 σ[ϵ,T]\sigma \in [\epsilon, T],直接评估 fθ(x,σ)f_{\theta}(\mathbf{x}, \sigma) 即可生成去噪后的图像。

    • 效果: Figure 12 展示了模型可以有效地去除不同噪声水平下的图像噪声。

      该图像是示意图,展示了使用一致性模型生成的高质量样本的对比,包括多个不同环境的室内场景,体现了噪声到数据映射的效果。图中显示了不同步骤生成效果的变化,突显了一步生成和多步生成的能力。 该图像是示意图,展示了使用一致性模型生成的高质量样本的对比,包括多个不同环境的室内场景,体现了噪声到数据映射的效果。图中显示了不同步骤生成效果的变化,突显了一步生成和多步生成的能力。

    图 12: 顶部一行显示不同噪声水平的噪声图像,底部一行显示去噪后的图像。

  6. 图像插值 (Image Interpolation):

    • 方法: 通过在潜在噪声空间中进行球形线性插值(spherical linear interpolation),然后将插值后的噪声输入一致性模型生成图像。

    • 效果: Figure 11 展示了模型可以平滑地在两个图像之间插值。

      该图像是一个示意图,展示了多种室内卧室设计风格的组合。这些图像展示了不同床铺、装饰和色彩搭配,体现了多样化的室内设计理念。 该图像是一个示意图,展示了多种室内卧室设计风格的组合。这些图像展示了不同床铺、装饰和色彩搭配,体现了多样化的室内设计理念。

    图 11: 通过在潜在空间中对噪声向量进行球形线性插值生成的一致性模型样本。

这些实验结果充分证明了 Consistency Models 强大的零样本数据编辑能力,使其在多种应用场景下都具有灵活性和实用性。

6.3. 连续时间扩展的实验验证

Figure 7 比较了离散一致性蒸馏/训练算法与其连续时间对应物在 CIFAR-10 上的表现。

Figure 7: Comparing discrete consistency distillation/training algorithms with continuous counterparts. 该图像是一个图表,展示了离散一致性蒸馏和训练算法相对于连续算法在FID指标上的效果。图(a)显示了一致性蒸馏过程中的FID变化,图(b)则展示了一致性训练过程中的FID变化,横轴为训练迭代次数,纵轴为FID值。

图 7: 比较离散一致性蒸馏/训练算法与连续对应物。

对于蒸馏 (CD):

  • CD (2\ell_2) 和 CD (LPIPS): 离散版本 (N=18) 的 CDLPIPS2\ell_2 度量下都表现出色。
  • CD (\infty, 2\ell_2) 和 CD (\infty, 1\ell_1): 连续时间的 CD 损失,特别是那些没有 stopgrad 操作的,在实验中表现不佳。
  • CD (\infty, stopgrad, 2\ell_2) 和 CD (\infty, stopgrad, LPIPS): 带有 stopgrad 的连续时间 CD 损失(如 Theorem 5 所述的“伪目标”)表现更好,其中 LPIPS 度量再次优于 2\ell_2。这表明 stopgrad 对于稳定连续时间训练是重要的。

对于训练 (CT):

  • CT (LPIPS): 离散版本 (N=120) 的 CT 表现良好。
  • CT (\infty, 2\ell_2) 和 CT (\infty, LPIPS): 连续时间的 CT 损失在随机初始化下表现不佳。这与离散 CT 形成对比,后者在随机初始化下也能很好地工作。

结论:

  • 离散时间目标 (尤其是使用 LPIPSHeun 求解器) 表现出强大的性能。
  • 连续时间目标在实践中更具挑战性,尤其是对于 CT 的随机初始化。
  • stopgrad 操作对于稳定连续时间蒸馏至关重要。
  • 离散 CT 损失在小 NN 时“偏差”较大,但随着 NN 的增加而减小,而连续 CT 没有这种“偏差”。

6.4. 额外样本

附录 EE 中提供了更多由 CDCTCIFAR-10ImageNet 64x64LSUN Bedroom 256x256LSUN Cat 256x256 上生成的样本图。这些图进一步展示了模型在不同数据集上的生成能力和视觉效果。

例如,CIFAR-10 的样本:

该图像是一个展示生成模型效果的图表,包含两个部分:上半部分展示了EDM模型生成的图像,FID为2.04。图中包含多种类别的图像,主要以动物和车辆为主,展示了该模型的生成能力与多样性。 该图像是一个展示生成模型效果的图表,包含两个部分:上半部分展示了EDM模型生成的图像,FID为2.04。图中包含多种类别的图像,主要以动物和车辆为主,展示了该模型的生成能力与多样性。

图 14: 来自 CIFAR-10 32x32 的未筛选样本。所有对应样本都使用相同的初始噪声。

ImageNet 64x64 的样本:

该图像是一个示意图,展示了EDM模型在图像生成中的表现,FID值为2.44。上半部分是一个生成结果的例子,与下半部分对比,展现了模型在图像生成质量上的提升。 该图像是一个示意图,展示了EDM模型在图像生成中的表现,FID值为2.44。上半部分是一个生成结果的例子,与下半部分对比,展现了模型在图像生成质量上的提升。

图 15: 来自 ImageNet 64x64 的未筛选样本。所有对应样本都使用相同的初始噪声。

LSUN Bedroom 256x256 的样本:

该图像是一个图表,展示了基于EDM生成的多个房间图像样例,FID值为3.57。图中包含了不同风格的卧室,显示了模型在生成高质量样本方面的能力。 该图像是一个图表,展示了基于EDM生成的多个房间图像样例,FID值为3.57。图中包含了不同风格的卧室,显示了模型在生成高质量样本方面的能力。

图 16: 来自 LSUN Bedroom 256x256 的未筛选样本。所有对应样本都使用相同的初始噪声。

LSUN Cat 256x256 的样本:

该图像是一个展示猫咪图像生成效果的图表,显示了EDM(FID=6.69)生成的猫咪图像。多个图像在同一张图中并排展示,以比较图像生成技术的效果。 该图像是一个展示猫咪图像生成效果的图表,显示了EDM(FID=6.69)生成的猫咪图像。多个图像在同一张图中并排展示,以比较图像生成技术的效果。

图 17: 来自 LSUN Cat 256x256 的未筛选样本。所有对应样本都使用相同的初始噪声。

这些视觉样本进一步印证了 Consistency Models 在高质量图像生成方面的有效性。

7. 总结与思考

7.1. 结论总结

本文引入了一致性模型(Consistency Models),这是一种旨在支持单步和少步高效生成的新型生成模型。通过广泛的实验,本文证明了其有效性:

  1. 蒸馏性能卓越: 作为一种蒸馏方法(Consistency Distillation, CD),一致性模型在多个图像基准测试(CIFAR-10, ImageNet 64x64, LSUN 256x256)的单步和少步采样中,显著优于现有的扩散模型蒸馏技术,取得了新的最先进(state-of-the-art)FID 分数。
  2. 独立生成能力强大: 作为独立的生成模型(Consistency Training, CT),一致性模型可以超越现有的单步、非对抗性生成模型(如 VAEs 和 Normalizing Flows),并且可以匹敌甚至在某些情况下优于一些 GANs。
  3. 零样本数据编辑: 一致性模型保留了扩散模型的零样本数据编辑能力,能够执行图像修复、着色、超分辨率、去噪、插值和笔触引导图像生成等多种任务,无需专门训练。
  4. 理论与实践结合: 论文提供了理论依据,并通过实验验证了其核心思想和训练范式的有效性,包括 EMA 更新和动态调度函数对于训练稳定性和性能的重要性。

7.2. 局限性与未来工作

论文中虽然没有明确列出“局限性”的章节,但从实验和方法论的讨论中可以推断出一些潜在的局限和未来工作方向:

  1. 连续时间损失的实现挑战: 连续时间(continuous-time)的 CDCT 损失虽然在理论上更优雅,但其实际实现需要前向模式自动微分(forward-mode automatic differentiation),这在某些深度学习框架中可能支持不佳,并且在随机初始化下,连续时间 CT 表现不佳,需要进一步研究方差降低技术。
  2. 多步采样的时间点优化: 论文提到多步采样的离散时间点序列 {τ1,τ2,,τN1}\left\{ \tau _ { 1 } , \tau _ { 2 } , \cdot \cdot \cdot , \tau _ { N - 1 } \right\} 是通过贪婪算法和三元搜索来优化的,并假设 FID 是一个单峰函数。探索更好的策略来优化这些时间点是一个未来研究方向。
  3. 更高阶 ODE 求解器的泛化: 虽然本文框架可以扩展到多步 ODE 求解器,但目前主要考虑了单步 ODE 求解器。将框架泛化到更高阶的 ODE 求解器是一个未来的工作。
  4. GANs 的性能差距: 尽管 CT 在非对抗性模型中表现优异,但与最先进的 GANs(如 StyleGAN-XL)相比仍有性能差距。未来可以探索如何进一步提升一致性模型的生成质量,缩小这一差距。

7.3. 个人启发与批判

7.3.1. 个人启发

  1. “一致性”的优雅: 论文提出的“自一致性”概念非常优雅。将 PF-ODE 轨迹上所有点映射到同一原点,并以此为训练目标,简洁而强大。这提供了一种全新的视角来处理扩散模型的采样加速问题,超越了单纯的数值求解器优化或黑盒蒸馏。
  2. 单步生成与多步权衡: 能够同时支持单步快速生成和多步质量细化,这一点非常实用。它使得模型可以根据不同的应用场景(例如,实时应用需要速度,离线内容创作需要质量)进行灵活调整,极大地扩展了其应用范围。
  3. 零样本能力的重要性: 继承并强化了扩散模型的零样本数据编辑能力,使得模型无需针对特定任务进行再训练即可应对多种图像处理需求,极大地提高了模型的通用性和经济性。
  4. 跨领域联系的潜力: 论文提及一致性模型与深度 Q-学习和基于动量的对比学习存在相似之处,这暗示了这些看似不同的领域之间可能存在更深层次的理论连接。这种交叉融合的思考模式对于推动人工智能的整体发展至关重要,可能会带来新的训练范式或理论突破。
  5. 独立训练的价值: Consistency Training (CT) 使得模型可以独立于预训练的扩散模型进行训练,这使其成为一个全新的生成模型家族,而不是仅仅作为扩散模型的“附属品”。这为生成模型研究开辟了新的路径。

7.3.2. 批判

  1. 理论与实践的差距: 论文虽然提供了连续时间损失的理论推导,但实验结果显示连续时间 CT 在随机初始化下表现不佳,这表明理论上的优雅与实际训练的稳定性之间仍存在差距。如何弥合这一差距,使连续时间训练在实践中更具可用性,是一个值得深入研究的问题。

  2. 调度函数的依赖性: CT 在实践中对 N(k)μ(k)\mu(k) 调度函数的高度依赖,意味着其训练过程可能比 CD 更为复杂和需要精心调参。虽然论文提供了这些函数的具体形式,但其通用性和最佳参数选择仍需进一步探索。

  3. LPIPS 的依赖: 实验结果强烈表明 LPIPS 度量在图像生成任务中的优越性。这在一定程度上说明了模型在学习感知相似度方面的能力,但也可能引入 LPIPS 本身的一些潜在偏差或局限性。如何在不依赖特定感知度量的情况下,提高模型在客观指标上的表现,也值得思考。

  4. 计算资源的消耗: 尽管 Consistency Models 实现了快速生成,但其训练过程,尤其是蒸馏大型扩散模型或在大型数据集上独立训练时,仍然需要大量的计算资源(如论文中提到使用 64 块 GPU 进行训练)。如何降低训练成本,使其更易于普及和应用,也是一个实际挑战。

  5. 特定扩散模型依赖: Consistency Distillation (CD) 仍然依赖于一个预训练好的扩散模型(如 EDM)。这意味着 CD 的性能上限受限于源扩散模型的质量,并且需要先训练好一个高质量的扩散模型。

    总体而言,Consistency Models 是生成模型领域的一项重要进展,它在提高生成效率和保持高质量方面取得了显著成就,并为零样本数据编辑提供了强大的工具。同时,它也提出了一些有趣的理论和实践挑战,为未来的研究提供了肥沃的土壤。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。