Abstract
一致性模型（CMs）是生成模型的一类新兴方法，它们提供了比传统扩散模型更快的采样速度。CMs要求采样轨迹上的所有点都映射到同一个初始点。但这一目标导致了资源密集型的训练：例如，截至2024年，在CIFAR-10上训练一个最先进的CM需要在一周内使用8个GPU。在这项工作中，作者提出了一个用于训练CMs的替代方案，大大提高了构建这类模型的效率。具体来说，通过用特定的微分方程表达CM轨迹，作者认为扩散模型可以被看作是具有特定离散化的CMs的特殊情况。因此，作者可以从预训练的扩散模型开始微调一致性模型，并在训练过程中逐步逼近完全一致性条件。作者提出的方法，称为简易一致性调整（ECT），在训练时间大幅缩短的同时确实比之前的方法质量更好：例如，ECT在单个A100 GPU上仅用1小时就在CIFAR10上达到了2步FID为2.73，与训练数百小时的一致性蒸馏相匹配。由于这种计算效率，作者研究了在ECT下CMs的扩展规律，结果显示它们似乎遵循经典幂律扩展，这暗示了它们能够在更大规模上提高效率和性能。
1 Introduction
扩散模型（DMs）[21; 63]，或基于分数的生成模型（SGMs）[65; 66]，极大地改变了视觉内容生成的领域，应用在图像[11; 17; 23; 55; 56; 57]、视频[2; 4; 6; 15; 22]和3D物体[1; 8; 36; 54; 70]上。DMs根据随机微分方程（SDE）[66]逐步将数据分布转换到已知的先验分布（例如高斯噪声），并训练模型去噪噪声观测。可以通过从噪声开始并使用训练后的模型逐步去噪的反向时间SDE生成样本。然而，由于扩散采样轨迹的曲率[26]，朴素地从DM中进行采样需要数百到数千次模型评估，这使得整个生成过程变得缓慢。已经提出了许多方法来解决这一问题，包括基于训练的技术，如蒸馏[13; 14; 44; 50; 58; 60; 74]，为主干模型设计的自适应计算架构[49; 69]，以及训练无关方法，如快速采样器[34; 42; 75; 80]或在采样过程中交错使用小和大主干模型[52]。然而，这些采样技术实现的加速通常以牺牲生成样本的质量为代价。
一致性模型（CMs）[64; 67]是一类新的生成模型，与扩散模型密切相关，作为更快的生成模型已显示出有希望的结果。这些模型学习噪声和数据之间的映射，使得采样轨迹的所有点都映射到同一个初始数据点。由于这个条件，这些模型能够在1-2次模型评估中生成高质量的样本。迄今为止，使用改进的一致性训练（iCT）[64]构建的最好的一致性模型，已将从头开始训练的1步CMs生成图像的质量提升到与使用数千步生成的SoTA DMs相媲美的水平。不幸的是，CMs在训练上仍然耗费时间并且实践上具有挑战性：迄今为止的最佳实践需要比类似质量的扩散模型长得多，并且在训练过程中涉及到复杂的超参数选择。总的来说，这大大限制了CMs在社区中的采用。
在这项工作中，作者提出了一个构建一致性模型更为高效的方案。具体来说，作者首先从动力系统的角度重写一致性条件，作为特定微分方程的有限差分近似，作者称之为“微分一致性条件”。直接使用这个条件作为损失函数训练一致性模型提出了几个挑战；作者提出了一些策略（包括选择特定的离散化课程和使用特定的加权函数）来解决它们。最重要的是，使用这个形式主义，作者观察到DMs可以被看作是具有宽松微分一致性条件的CMs的一个特例。这使得作者可以通过逐步加强一致性条件，平滑地从DM插值到CM，逐渐在训练过程中使时间离散化变得更小。将这些策略结合起来，作者得到了一个简化的训练方法，称为“简易一致性调整（ECT）”，它使作者能够从预训练的扩散模型开始构建一致性模型，从而“显著”降低训练成本。
ECT在CIFAR-10[35]和ImageNet 6464[10]上实现了比现有技术“更好”的两步样本质量，同时使用的训练FLOPs不到；即使考虑到相应扩散模型的预训练预算，ECT也只需要最先进iCT[64] 的时间。与预训练分数SDE/DMs[66]相比，ECT还显著降低了推理成本，为，同时保持了相当生成质量。
2 Preliminaries
扩散模型。令 表示数据分布。扩散模型（DMs）通过从 到 逐渐增加的独立同分布高斯噪声（标准差为 ）来扰动这个分布，使得 ，并且选择 使得 和 。这个过程由以下SDE描述[66]
其中 是标准维纳过程， 是漂移系数， 是扩散系数。可以通过从 到 求解反向时间SDE并采样 来生成样本。Song等人[66]表明这个SDE有一个相应的ODE，称为概率流ODE（PF-ODE），其轨迹具有与SDE相同的边缘概率密度。作者遵循Karras等人[26]的记法来描述ODE，即
其中 表示得分函数。先前的工作[26; 67]设置 得到
其中 表示一个去噪函数，它给定噪声图像 预测干净图像 。在本文的其余部分，作者将遵循这个参数化。注意，这种参数化下时间是等同于噪声水平的，作者将交替使用这两个术语。
一致性模型（CMs）建立在方程（3）中的PF-ODE之上，该方程在数据分布与噪声分布之间建立了一一对应的映射。CMs学习一个一致性函数，该函数将带噪声的图像映射回干净的图像
请注意，一致性函数需要满足在时的边界条件。先前的研究[26; 64; 67]通过将CM参数化为以下形式来施加这一边界条件：
其中是模型参数，是要训练的网络，和是与时间相关的缩放因子，使得，。这种参数化设计保证了边界条件。作者在附录C中讨论和的具体选择。
在训练期间，CMs首先将PF-ODE离散化为个子区间，边界由t_min=t_1<t_2<…<t_n=t
在这里， 是一个度量函数， 指的是一个可训练的神经网络，用于学习一致性函数， 指的是过去 值的指数移动平均（EMA），。此外，离散化课程 应该是自适应的，并在训练过程中调整以达到良好的性能。
在开创性的工作中，Song 等人[67]使用学习到的感知相似性得分（LPIPS）[76]作为度量函数，对于所有的 设置 ，并根据 Karras 等人[26]的采样调度程序采样 ：对于 和 ，。进一步，得分函数 可以从预训练的扩散模型中估计，这导致了一致性精馏（CD），或者可以使用无偏得分估计器估计
这导致了一致性训练（CT）。
后续工作iCT [64] 引入了几项改进，这些改进显著提高了CMs的训练效率以及性能。首先，用Pseudo-Huber度量替换了引入生成建模中不良偏见的LPIPS度量。其次，网络不再保持过去值的EMA。第三，iCT用权重方案替换了均匀权重方案。此外，仔细选择了噪声嵌入和dropout的缩放因子。第四，iCT在训练期间引入了一个复杂的离散化课程：
其中是当前的迭代次数，是总的迭代次数，和是训练的最大和最小噪声水平，和是超参数。最后，在训练期间，iCT从离散对数正态分布中采样，其中和。
3 Easy Consistency Tuning (ECT)
作者首先介绍构成作者方法核心的微分一致性条件，并基于此条件推导出损失目标。接下来，作者将分析这个损失目标，并强调使用它训练条件生成模型（CMs）的一些挑战。基于这一分析，作者提出了作者的方法——简易一致性调整（ECT）。ECT是一种简单、原则性的方法，可以高效地训练条件生成模型以满足（微分）一致性条件。由此产生的条件生成模型在1到2个采样步骤中即可生成高质量样本。
Differential Consistency Condition
正如在第2节中先前所述，CMs学习一个一致性函数，它将噪声图像映射回干净图像：。然而，作者认为CMs可以通过模型输出的时间导数来定义，由以下微分形式给出
由于干净图像与噪声水平无关，其时间导数为零。
然而，单独的微分形式不足以保证模型输出与干净图像相匹配，因为存在一些平凡解，其中模型将所有输入映射到一个常数，例如。为了消除这些坍缩的解，作者遵循Song等人[67]，Song和Dhariwal[64]的做法，对施加一个边界条件：
这个边界条件确保当噪声水平为零时，模型输出与干净图像相匹配。结合起来，方程(9)中的微分形式和边界条件定义了方程(10)中的_微分一致性条件_，简称_一致性条件_。
有限差分近似。为了学习一致性条件，作者使用有限差分近似将微分形式 离散化：
其中 ，，以及 表示 。对于给定的干净图像 ，作者使用相同的扰动方向 在两个噪声水平 和 上产生两个扰动图像 和 。具体来说，作者根据CMs的前向过程计算 和 ，即 和 。
为了满足一致性条件，作者最小化 和 之间的距离。对于边界条件（），作者优化 使其与干净图像 对齐。对于更高的噪声水平（），作者使用停止梯度算子 冻结 并优化 使其与 对齐。
损失函数。鉴于这种离散化，CMs的训练目标可以表述为：
作者最小化损失函数：
其中作者可以提取权重函数 ，，，使用共享的噪声方向 ，以及 是一个度量函数。在实践中，作者将 设置为平方的 度量。此外，作者在第3.3节讨论了Song和Dhariwal [64] 提出的伪胡伯度量与权重函数之间的有趣关系。为了正则化模型，辍学 Mask 也应该在 和 之间保持一致，通过设置共享的随机种子，确保噪声水平 和 是唯一变化的因素。
The "Curse of Consistency" and its Implications
等式（12）中的目标在时很难优化。这是因为每个离散化时间间隔的预测误差会累积，导致训练收敛缓慢，在极端情况下甚至发散。为了进一步说明，考虑一个大的噪声水平。作者首先遵循CT [67]的符号表示，并将噪声时间范围划分为个较小且连续的子区间。作者可以对CM的步预测误差进行界定：
以上观察的关键启示是，如果作者从零开始训练，严格遵循微分形式，使用非常小的，那么由此产生的模型可能会因为每个区间累积的一致性错误而在训练过程中陷入困境。作者通过使用不同数量的区间和相应的下的连续性条件训练一系列的CM，来验证这一假设，如图1所示。
Tackling the "Curse of Consistency"
第3.2节的讨论强调了当作者从零开始使用等式(12)中的损失目标训练CMs时，直接遵循可能会出现的训练不稳定问题。然而，通过适当的设计选择，这些问题是可以缓解的。在本节中，作者列举了几种解决前述训练不稳定问题的策略，并帮助提高CMs的效率。
以预训练的扩散模型开始ECT。受到iCT [64]的启发，作者从大的开始ECT，并逐渐缩小。在作者的问题设置中，由于，在时，作者有最大的，得到
用这个损失目标训练模型会导致去噪扩散模型[21]。这个观察提示了一个从DMs 平滑插值到CMs 的学习方案，通过在训练过程中逐渐缩小。基于这个推理，扩散预训练可以被认为是带有宽松离散化一致性条件的连贯训练的一个特例。在实践中，作者建议使用预训练的扩散模型来初始化ECT。这种初始化的另一个好处是，在训练过程中，尤其是在初始阶段，它确保了损失目标中有良好的目标。最后，作者注意到，即使作者从零开始使用等式(12)的目标训练CMs，作者的方法仍然有效（参见附录B中的消融研究的结果）。然而，使用预训练的扩散模型初始化ECT可以进一步提高效率，如第4节所指出的。
连续时间训练计划。作者研究了一个连续时间计划的设计原则，其“边界”条件导致扩散预训练，即，在开始时为所有构建的训练对。请注意，这与iCT中使用的训练计划不同（参见第2节和附录A）。作者考虑了一致性训练的重叠区间，这允许分解，并从噪声分布中连续采样无限的，例如，以及。
作者将称为映射函数。由于作者需要随着训练的进行缩小，作者增强了映射函数以依赖于训练迭代，，以控制。作者将映射函数参数化为
其中，作者取 ，其中 是sigmoid函数， 指的是训练迭代次数。通常，作者设置 ， 和 。由于 ，作者还将 限制以满足这一约束。在训练开始时，这个映射函数产生 ，这样就恢复了扩散预训练。作者在附录A中讨论这个函数的设计选择。
度量选择。如第2节所述，iCT使用了Pseudo-Huber度量[7]来减轻由LPIPS度量[76]引起的感知偏差。仔细观察这个度量，作者发现它提供了对梯度的自适应样本缩放，这减少了训练期间梯度的方差。具体来说，Pseudo-Huber损失的微分可以分解为两个项：自适应缩放因子 和平方 丢失的微分。作者在附录A中详细讨论这一点。因此，作者保留了DMs中使用的平方 度量，并探索变化的自适应加权项。算法1 简单一致性调整（ECT）
加权函数。加权函数通常会导致DMs的性能有显著差异，CMs也是如此。将方程(15)中的映射函数 代入方程(12)中的加权函数，作者得到 ，其中 是当 时损失的缩放因子，而 与DMs中的加权类似。这耦合了加权函数与映射函数。相反，作者考虑解耦的加权函数。受到Pseudo-Huber损失中出现的自适应缩放因子的启发（更多细节见附录A），作者将加权函数重写为
其中 。作者将 定义为常规权重，而 定义为自适应权重。自适应权重有助于使用 度量提高训练效率，因为当 （通常发生在 时），这种权重 可以放大梯度以避免在学习细粒度特征时出现的梯度消失问题，而 避免了在 时可能出现的数值问题。关于本文考虑的各种 和 的详细概述，请读者参考附录B中的表2和表3。总的来说，作者注意到条件模型（CMs）的生成能力在很大程度上得益于控制不同噪声水平下梯度方差的重权函数。
4 Experiments
作者测量了ECT在两个数据集上的效率和可扩展性：CIFAR-10 [35] 和 ImageNet [10]。作者使用Frechet Inception Distance（FID）[18]和DINOv2模型下的Frechet距离[51]（FD[68]）评估样本质量，并使用函数评估次数（NFEs）来衡量采样效率。作者还指出了这些方法的相对训练成本。详细的实现细节可以在附录C中找到。
Comparison of Training Schemes
作者比较了使用ECT训练的CM（记为ECM）与最先进的扩散模型、采用高级采样器的扩散模型、如一致性蒸馏（CD）[67]的蒸馏方法以及改进的一致性训练（iCT）[64]。关键结果在图2中总结。作者展示了四种训练方案的训练FLOPs、推理成本和生成性能。
扩散模型。作者将ECM与Score SDE [66]、EDM [26]以及采用DPM-Solver-v3 [79]的EDM进行了比较。经过仅k次迭代微调的步ECM，在FID方面与具有模型深度和 NFEs的Score SDE-deep相匹配。如图2所示，ECM仅需要的推理成本和延迟即可达到相同的样本质量。经过k次迭代微调的步ECM优于采用高级DPM-Solver-v3（NFE=10）的EDM。
扩散蒸馏。作者将ECT与一致性蒸馏（CD）[67]进行了比较，后者是一种将预训练的DM蒸馏到CM上的最先进方法。如表1所示，ECM在CIFAR-10和ImageNet 6464上的性能显著优于CD。作者注意到ECT摆脱了教师DM的错误，且不会产生运行教师DM的额外成本。在CIFAR-10上，步ECM在FID（对比）方面优于步CD（使用LPIPS [76]），同时使用大约的CD训练/微调计算。
从零开始的一致性训练。 改进的一致性训练（iCT）[64]是从零开始训练一致性模型而不推理扩散教师的最佳实践。与从零开始训练相比，ECT在图2和表1中显示，使用的训练计算量和的模型大小就能与iCT-deep相媲美。
ECT通过简单而原则性的方法解锁了最先进的少量步骤生成能力。在几乎可以忽略的调优成本下，ECT展示了强大的成果，同时从训练FLOPs的扩展中受益，以增强其少量步骤生成能力。
Scaling Laws of ECT
作者研究了ECT的扩展行为，重点关注三个方面：训练计算量、模型大小和模型FLOPs。当计算资源不是瓶颈时，ECT扩展良好，遵循经典幂律。
图3：在ImageNet 6464上，增加训练计算量和模型大小可提高样本质量。每组（从左至右）都有经过图像训练的ECM-S、经过图像训练的ECM-S以及经过图像训练的ECM-XL的步样本。
在不依赖冻结扩散教师的知识蒸馏或额外的对抗性监督的情况下，ECT使用的训练成本超过了一致性蒸馏（CD）[67]和从零开始训练的一致性模型（iCT）[64]。与扩散预训练（Score SDE/DMs）相比，ECT还显著降低了推理成本至，同时保持了相当的样本质量。
训练计算量。从EDM [26]的权重初始化，作者在CIFAR-10 [35]上针对六个计算规模对ECMs进行微调，并在图4（左）中绘制了与训练计算量的趋势。最大的计算量达到了扩散预训练预算的。当作者增加训练预算时，观察到的经典幂律衰减，表明在ECT中增加计算投资可显著提高生成性能。有趣的是，在增加训练计算量时，即使使用相同的计划，1步和2步生成之间的差距变得_更小_。作者进一步拟合幂律，其中是归一化的训练FLOPs。1步和2步生成与和之间的皮尔逊相关系数分别为和，都具有统计显著性（值分别为和）。
图4：左图：增加训练计算量会得到经典的幂律关系，即 与训练计算量之间的关系。右图：在相同的批大小和迭代次数下，增加模型大小和模型FLOPs与ImageNet 上的 改进强烈相关。直径与模型大小成比例。模型大小与FLOPs。从EDM2 [28]的权重初始化，作者训练了参数从到，模型FLOPs从到的ECM-S/M/L/XL模型。如图4所示，1步和2步生成能力相对于模型FLOPs和参数表现出对数线性扩展。这种扩展行为确认了ECT有效地利用增加的模型大小和计算能力来提高1步和2步生成性能。值得注意的是，在使用仅相当于iCT [64]训练预算，以及相当于EDM2 [28]预训练预算至（取决于模型大小）的训练图像时，ECT实现了比现有最佳CMs [64]更佳的2步生成性能。
5 Related Work
作者对ECMs使用1280万训练图像的预算（批量大小为128，10万次迭代）。 在ImageNet 6464上表示一个预算为102.4M训练图像（批量大小为1024，10万次迭代）。
一致性模型一致性模型[64, 67]是一类为高效生成而设计的新生成模型，无需对抗性训练即可在1或2个模型步骤中进行生成。CMs不依赖于预训练的扩散模型（DM）来生成训练目标，而是利用无偏的评分估计器。CMs已经被扩展到多步采样[29]，潜在空间模型[45]，ControlNet[71]，并结合额外的对抗性损失[29, 33]。尽管CMs在采样效率上有所提高，但它们通常更难训练，并且与扩散模型相比，需要更多的计算资源。作者的工作大大提高了CMs的训练效率，降低了未来在CMs上进行研究和部署的成本。扩散蒸馏受知识蒸馏[19]的启发，扩散蒸馏是迄今为止基于训练的最流行的加速扩散采样过程的方法。在扩散蒸馏中，一个预训练的扩散模型（DM），需要数百到数千次的模型评估来生成样本，充当教师模型。一个学生模型被训练以匹配教师模型的样本质量，使其能够在几个步骤中生成高质量的样本。
这个领域主要有两条主线。第一类涉及轨迹匹配，学生模型学习匹配教师模型采样轨迹上的点。这类方法包括离线蒸馏[13, 44, 78]，它们需要从预训练的DM中采样生成的离线合成数据集，以将教师模型提炼为几个步骤的学生模型；渐进式蒸馏[48, 58]，TRACT[3]，它们需要多次训练传递或离线数据集以实现同一目标；以及BOOT[14]，一致性蒸馏（CD）[67]，和Imagine-Flash[32]，它们最小化在采样轨迹上精心选择的点处学生预测之间的差异。
CD与作者的方法密切相关，因为它利用教师模型生成相邻点对，并强制这些点上的学生预测映射到初始数据点。然而，它限制了CMs的质量，使之受限于预训练扩散模型的质量，并且CD损失中使用的LPIPS-based度量[76]在评估FID等指标时引入了不希望有的偏差。
第二类方法力求减小数据分布与模型分布之间的概率分歧，即分布匹配[46, 54, 70, 74]。DreamFusion [54] 和 ProfileDreamer [70] 等方法也用于3D目标生成。这些方法[32, 37, 46, 50, 60, 72, 74] 使用得分蒸馏或对抗性损失将昂贵的教师模型提炼成高效的学生模型。然而，由于来自对抗性或得分蒸馏的交替更新方案，它们可能难以以稳定的方式进行训练。
基于训练方法的缺点是它们通常需要在预训练之后进行额外的训练过程，以提炼出高效的学生模型。关于扩散蒸馏的最新进展的详细讨论，作者参考[12]。
扩散模型的快速采样器。快速采样器通常是免训练的，并使用先进的求解器模拟扩散随机微分方程（SDE）或常微分方程（ODE），以减少采样步骤的数量。这些方法通过解析求解部分SDE或ODE [42, 43]，使用指数积分器和更高阶多项式对解进行更好的近似 [75]，使用更高阶数值方法 [26]，在采样过程中对噪声水平进行更好的近似 [34]，在每一步采样中修正预测 [77] 并确保ODE的解位于期望的流形上 [40]，从而减少采样过程中的离散化误差。另一种正交策略是通过并行化采样过程来实现加速 [53, 62]。这些快速采样器的一个缺点是，当采样步骤数量降至某个阈值以下，如10步时，样本的质量会急剧下降。## 6 限制
ECT的一个主要局限性在于，它需要一个数据集来将决策模块（DMs）调整到一致性模块（CMs），而近期的研究开发了无数据方法[14; 46; 74]用于扩散蒸馏。ECT与无数据方法之间的区别在于，ECT通过自教师在一个给定数据集上学习一致性条件，而无数据扩散蒸馏从冻结的扩散教师那里获取知识。ECT的这一点可能是一个潜在的局限性，因为某些模型的训练数据对公众不可用。然而，作者对使用与预训练不同的数据集来调整CMs持有乐观态度。一致性调整的数据组成和规模将是一个有价值的研究方向。
7 Conclusion
作者提出了一种简单而高效的训练一致性模型的方案。
首先，作者推导出微分一致性条件，并基于此条件制定了一个损失目标。
接下来，作者识别出学习CMs的“一致性诅咒”问题，并提出几种策略来减轻这些问题。
作者发现，扩散模型可以被视为一致性模型的一种特殊情况，其一致性条件具有宽松的离散化。
基于这些观察和策略，作者提出了一种简化的训练方法，称为简单一致性调优（ECT）。通过ECT获得的ECMs模型，在极小的调优成本下解锁了最先进的几步生成能力，并能从扩展中受益。
作者希望这项工作将有助于一致性模型在研究社区中的更广泛应用。
参考
[1].Consistency Models Made Easy.
