一致性模型:单步生成高质量图像,破解扩散模型速度瓶颈

这篇名为《一致性模型》的论文提出了一类新的生成模型,能够一步到位地生成高质量图像 ,解决了扩散模型的主要瓶颈------采样过程慢、需要多次迭代。它的核心贡献在于证明了:在保持扩散模型强大表达能力和训练稳定性的同时,实现快速的一步生成和灵活的多步优化是可行的 。通过理论推导和大量实验,作者表明一致性模型在一步和少步图像生成任务上达到了顶尖水平,这标志着高效、多功能生成建模向前迈出了重要一步。

背景

扩散模型通过逆转一个逐渐添加噪声的过程 来生成数据,这个过程可以用随机微分方程来描述。每个数据点 x₀ 会随时间 t 演变成一个带噪声的状态 xₜ,这个过程也对应着一个概率流常微分方程

在实践中,通常需要使用一个 "分数网络" 来数值求解这个ODE,该网络负责估计 ∇xₜ log pₜ(xₜ)。这种迭代式的采样能产生高质量图像,但需要成百上千次计算,非常耗时

这就促使研究者们去寻找一个新框架,既能保留扩散模型的表达能力和稳定性 ,又能实现快速的一步生成。理想情况下,这种模型应该保留扩散模型的解释性和编辑灵活性,同时摆脱对迭代采样的依赖。当前方法中存在的这个空白,正是本篇论文提出 "一致性模型" 框架的基础。

方法

  • 模型

一致性模型的核心思想是,绕过扩散模型缓慢的迭代过程,通过学习一个直接映射,将任何沿着同一概率流ODE轨迹的噪声样本 xₜ 直接映射回其干净的原始数据 x₀。其关键直觉在于:同一条轨迹上的所有状态在本质上是确定性地关联的;因此,一个能在这些状态上保持输出一致的模型,就能通过单次前向传播生成 x₀。这个特性构成了一致性函数的基础。

  • 定义

假设 xₜ 是通过概率流ODE从一个干净样本 x₀ 演化得到的、在扩散时间 t 上的噪声样本。一个一致性函数 f(xₜ, t) 将任何 xₜ 映射回其原始数据点 x₀,并且要求同一条轨迹上的所有点都产生相同的输出:

这种自一致性确保了无论噪声水平如何,模型在轨迹中间点的预测结果都是一致的。一致性模型就是一个用于近似 f 的神经网络 fθ。当成功训练后,fθ 可以通过以下形式实现一步生成:

  • 参数化与边界条件

为了保证当 t 趋近于一个极小值 ε(接近真实数据分布)时模型行为正确,模型必须满足边界条件。论文提出了一种稳定的跳跃连接参数化方法

这里,Fθ(x, t) 是一个U-Net类型的骨干网络,负责学习将噪声输入向数据分布修正的残差 。这种设计保持了在边界处的可微性 ,并使得同一个网络能够无缝地处理所有扩散时间

  • 采样过程

训练完成后,一致性模型支持一步采样和少步采样。

  • 一步生成: 直接从一个随机噪声潜在表示开始,通过一次前向传播解码出图像:

这单次前向传播就能产生与扩散模型多步输出相媲美的近似样本。

  • 少步优化: 模型在预先选定的离散时间点 tᵢ 上,交替进行"去噪"和"受控的再加噪"操作,仅用少量几次计算就能显著提升样本质量。
  • 零样本数据编辑

同样的自一致性特性使得模型无需针对特定任务训练,就能实现零样本编辑。通过在将噪声输入 xₜ 送入模型 fθ(xₜ, t) 之前,部分地替换其中已知区域或加入引导信号,模型就能完成图像修复、上色、超分辨率、插值或笔触引导生成等任务。

  • 通过蒸馏训练一致性模型

第一种训练方法叫做一致性蒸馏 。这种方法通过迫使一致性模型去模仿一个预训练好的扩散模型的行为来训练它,具体做法是强化模型在概率流ODE轨迹上邻近点之间预测的一致性。核心思想是:既然扩散模型定义了一个从数据到噪声再返回的连续映射,那么一致性模型可以直接学习这个映射,而无需反复求解ODE。

形式上,考虑一个由概率流ODE产生的扩散轨迹 {xₜ}。CD目标要求这条轨迹上的两个点,xₜ 和一个轻微扰动后的点 x̂ₜ₋₁,在模型映射下应该指向同一个干净样本:

其中 θ' 是模型参数的指数移动平均值,x̂ₜ₋₁ 是使用数值求解器沿ODE向前一步得到的状态。这个一致性损失惩罚了模型在相邻时间步输出上的差异,有效地鼓励模型在整个扩散过程中保持自一致性。

为了稳定训练,模型预测之间的差异使用一种感知度量(例如LPIPS距离)来衡量,而不是简单的 ℓ₁ 或 ℓ₂ 范数。总的CD损失是:

这个过程在概念上类似于扩散模型中的渐进式蒸馏,但它引入了一个更通用、更原则化的一致性约束,作用于整个轨迹,而不仅仅是离散的端点。最终,被蒸馏的模型会收敛到基础概率流ODE的真实解。

通过CD,一个预训练的扩散模型可以被压缩成一个单步生成器,同时保留其分布特性,从而实现快速采样并保持扩散级别的图像质量。

  • 独立训练一致性模型

第二种方法叫做一致性训练 ,它完全摒弃了对预训练扩散模型的依赖。CT不是通过蒸馏获取知识,而是通过强制模型在不同噪声水平的扩散过程下保持输出稳定,直接从零开始学习自一致性。

在CT中,模型使用一个随机噪声调度来模拟自己的扩散轨迹,从一个数据点 x₀ 生成成对的噪声样本 (xₜ, xₜ₋₁)。目标是确保模型对这两个噪声版本进行预测时,结果是一致的:

这个损失在结构上与CD目标类似,但不依赖于教师模型。相反,噪声过程本身定义了监督信号

为了保持训练稳定性,CT采用了自适应的调度函数来控制时间步数 N(t) 和EMA衰减率 μ(t),随着训练的进行,逐渐增强模型在时间上的一致性。作者证明,在极限情况下,当概率流ODE求解器变得精确且EMA更新趋近于恒等映射时,CT和CD在理论上是等价的------两者都收敛于一个完美的自一致性函数 f*,该函数代表了真实的数据生成过程。

这个公式表明,一致性模型可以自成一体,成为一个独立的生成模型家族,而不仅仅是蒸馏的目标。CD利用了现有扩散网络的能力,而CT则证明了自一致性本身就是一个强大的生成建模学习原则,无需显式的似然训练或分数监督也能产生逼真的图像。

  • 两种框架的意义

一致性蒸馏和一致性训练共同提供了两条互补的训练路径。

  • CD 利用了预训练的扩散模型,提供了一种高效的方法将其转换为快速的一步生成器,同时保留其学到的数据分布。
  • CT 则从零开始学习相同的特性,证明了自一致性本身就是一个强大的生成原则。

两种方法都强制了同一个约束------扩散轨迹上邻近点的输出必须相同------这使得模型能够内化数据分布的全局结构。这种统一的视角连接了基于扩散的生成建模和直接生成建模,为这两种范式之间架起了一座新的理论和实践桥梁。

实验

实验研究评估了一致性模型在图像生成和零样本图像编辑任务上的表现,检验了其在一步和少步采样下的性能。作者旨在证明CMs能够达到扩散级别的样本质量,同时大幅减少计算次数。实验分为三部分分析:训练行为和超参数选择、少步图像生成质量、以及零样本编辑能力。

  • 数据集

模型在几个广泛使用的图像生成基准上进行了训练和测试:

  • CIFAR-10: 用于评估低分辨率生成的标准数据集,便于与扩散基线直接比较。
  • ImageNet 64×64: 中等规模基准,测试类别条件生成和多样性。
  • LSUN Bedroom & Cat 256×256: 大规模高分辨率数据集,用于评估少步优化和视觉保真度。

这些数据集覆盖了不同的分辨率和复杂度,使得作者能够验证CM的可扩展性和泛化能力。

  • 模型设置

所有模型都采用与标准扩散架构兼容的U-Net骨干网络,以确保公平比较。

  • 训练机制
  • 一致性蒸馏: 从预训练的扩散模型初始化,并使用Heun ODE求解器进行蒸馏。
  • 一致性训练: 使用随机噪声调度和EMA更新,从头开始训练。
  • 超参数
  • 求解器: 海恩法,以获得更高精度。
  • 距离度量: LPIPS,发现在保持感知细节上优于 ℓ₁ 和 ℓ₂ 度量。
  • 步数 N: 通常在10-40之间,在不同的消融实验中会变化。
  • EMA衰减 μ: 自适应调度,随时间推移增加稳定性。
  • 硬件: 实验在8张A100 GPU上使用混合精度进行。
  • 对比基线

提出的CMs与几个强大的基线进行了比较:

  • 扩散模型: 原始基于分数的模型,需要数百次采样步。
  • 渐进式蒸馏: 一种先前的加速方法,通过分阶段的师生训练,逐步将扩散模型的步数减半。
  • GANs和流模型: 代表性的单步生成模型,如StyleGAN-2和Glow,作为样本效率和质量的参考。
  • DDIM / DPM-Solver: 快速采样扩散变体,在可比步数下进行评估。
  • 评估指标

使用以下指标衡量生成和编辑质量:

  • FID: 衡量生成图像与真实图像之间的统计距离(越低越好)。
  • IS: 评估生成图像的多样性和可识别性(越高越好)。
  • 精确度和召回率: 量化样本保真度和多样性之间的权衡。
  • LPIPS: 在训练中用作内部感知距离度量,并作为重建一致性的补充衡量标准。

结果

  • 训练一致性模型

作者首先分析了训练过程中设计选择的影响。总结来说,使用海恩求解器训练的模型 consistently 优于使用欧拉更新的模型,尤其是在一步采样下。同样,基于LPIPS的损失函数比 ℓ₁/ℓ₂ 目标能带来更平滑的收敛和更好的视觉质量。ODE步数 N 也起着关键作用:步数太少会导致不稳定,而适中的值(10-20)能提供最佳性能。对于CT,N(t) 和 EMA衰减 μ(t) 的自适应调度显著增强了稳定性,证实了逐步的时间对齐有助于端到端的自一致性学习。

  • 少步图像生成

定量结果显示,一致性模型在多个数据集上的一步和两步图像生成都达到了顶尖水平。

在CIFAR-10上,CD达到了FID 2.82(1步)和2.20(2步),优于渐进式蒸馏,并接近完整扩散模型。

在ImageNet 64×64上,CD获得了FID 4.17(1步)和3.55(2步),再次超过了所有蒸馏基线。

CT变体虽然无需教师模型训练,在两步后也达到了FID 3.90,验证了仅靠自一致性就足以获得有竞争力的质量。

定性地看,一步生成的图像展示了锐利的纹理和自然的色彩分布,与多步扩散样本相媲美。即使是少量的优化步数也能产生接近扩散模型真实感的图像,证明了CMs在灵活平衡速度与质量方面的能力。

  • 零样本图像编辑

CMs的一个定义性特征是它们无需显式微调就能执行多样的图像编辑任务。利用相同的少步生成过程,作者演示了去噪、修复、上色、超分、笔触引导合成和插值。例如,被遮罩的修复区域能够被无缝地填充上合理的纹理,稀疏的颜色提示或笔触能够连贯地传播到整个图像。这些结果表明,CMs保留了扩散模型的可控性和语义连贯性,同时提供了数量级更快的推理速度。

总而言之,实验证实了一致性模型将单步生成器的效率与扩散模型的灵活性和质量结合了起来,确立了它们作为高效、可控生成建模的新范式。

结论

这项研究通过引入一个能保留扩散模型优点、同时消除其对缓慢迭代采样依赖的框架,向快速灵活的生成建模迈出了重要一步。通过在扩散轨迹上强制执行自一致性 ,所提出的模型能够通过单次前向传播生成高质量图像,并且只需少量几步就能进一步优化。两种互补的训练策略------一致性蒸馏一致性训练------证明了自一致性既可以从预训练的扩散模型中学习,也可以从零开始学习,为高效生成提供了一个统一的原则。大量实验表明,一致性模型不仅匹配了多步扩散方法的视觉保真度,还实现了如修复和超分辨率等零样本编辑能力。

评述

这篇论文为扩散模型的核心瓶颈之一------推理效率低下------提供了一个优雅的解决方案。通过将生成重新定义为一种一致性约束,而非序列去噪过程,它在保留扩散模型表征能力的同时,将整个过程压缩成了一个单步变换。该框架概念清晰、有坚实的理论奠基、CD和CT都表现出强大的实证效果,使其成为一项杰出的贡献。所展示的速度-质量权衡和零样本编辑能力表明,一致性模型不仅仅是加速工具,更是一个独特且多功能的生成模型类别。

尽管前景广阔,但一致性模型主要还是在图像数据和中等分辨率下进行评估。该方法在高分辨率或跨模态生成任务上的有效性仍有待测试。此外,虽然一步生成很快,但由于需要跨多个时间对计算一致性损失,训练过程------尤其是CT------可能计算量很大。未来的研究可以专注于扩展该框架并将其应用于条件生成和多模态领域。尽管如此,这项工作确立了一个强大且可推广的原则------以一致性作为生成建模的基础------这很可能将影响下一代快速可控的扩散架构的发展。

相关推荐
音视频牛哥3 小时前
AI智能体从系统智能到生态智能:SmartMediaKit 如何成为智能体时代的视频神经系统
人工智能·计算机视觉·音视频·大牛直播sdk·多智能体协同·rtsp播放器rtmp播放器·视频感知低延迟音视频
JMzz3 小时前
Rust 中的数据结构选择与性能影响:从算法复杂度到硬件特性 [特殊字符]
开发语言·数据结构·后端·算法·性能优化·rust
CoovallyAIHub4 小时前
搞定边缘AI部署:开源神器RamaLama,让视觉语言模型无处不在
深度学习·算法·计算机视觉
CyberSoma4 小时前
机器人模仿学习运动基元数学编码方法还有用吗?
人工智能·算法·计算机视觉·机器人
CoovallyAIHub4 小时前
英伟达再出「神作」!黄仁勋华盛顿GTC宣布Vera Rubin超级芯片,联手诺基亚进军6G,市值直逼5万亿美元
深度学习·算法·计算机视觉
黑菜钟4 小时前
代码随想录第50天 | 图论 基础介绍(新篇章
算法·深度优先·图论
mit6.8245 小时前
[nanoGPT] 检查点 | `ckpt.pt`记忆 | 预训练模型加载`from_pretrained`
人工智能·深度学习·机器学习
草莓熊Lotso5 小时前
《算法闯关指南:优选算法--前缀和》--27.寻找数组的中心下标,28.除自身以外数组的乘积
开发语言·c++·算法·rpc
CLubiy5 小时前
【研究生随笔】Pytorch中的卷积神经网络(2)
人工智能·pytorch·python·深度学习·cnn·卷积神经网络·池化