截断一致性模型

一致性模型（Consistency Models）近期被提出，旨在通过直接从初始噪声预测概率流常微分方程（PF ODE）的解（即数据）来加速扩散模型的采样过程。然而，一致性模型的训练需要学习将PF ODE轨迹上的所有中间点映射到其对应的终点，这一任务远比最终的一步生成目标更具挑战性——后者仅关注PF ODE中从噪声到数据的映射关系。我们通过实验发现，这种训练范式限制了一致性模型在一步生成任务中的性能表现。为解决该问题，我们提出将一致性训练推广至截断的时间范围（truncated time range），使模型能够忽略早期时间步的去噪任务，从而将网络容量集中于生成能力的提升。为此，我们引入了一种新的一致性函数参数化方式，并设计了一种两阶段训练流程，有效防止截断时间训练陷入平凡解（trivial solution）的困境。在CIFAR-10和ImageNet $64\times64$数据集上的实验表明，我们的方法在一步和两步生成任务中均取得了优于当前最优一致性模型（如iCT-deep）的FID指标，且所使用的网络规模比其小超过2倍。项目主页：https://truncated-cm.github.io/