[论文解读] Consistent Diffusion Models: Mitigating Sampling Drift by Learning to be Consistent

Giannis Daras, Yuval Dagan+2 more|arXiv (Cornell University)|Feb 17, 2023

Model Reduction and Neural Networks被引用 7

一句话总结

本文提出了一种新型训练目标——一致扩散模型（CDM），通过强制执行一致性属性来缓解扩散模型中的采样漂移问题：即模型在其自身生成的样本上，于不同时间步的预测结果保持一致。通过将该一致性损失与弱去噪分数匹配目标相结合，该方法在仅使用非漂移数据进行训练的情况下，实现了CIFAR-10上的SOTA FID分数5.45，以及在AFHQ和FFHQ上的性能提升。

ABSTRACT

Imperfect score-matching leads to a shift between the training and the sampling distribution of diffusion models. Due to the recursive nature of the generation process, errors in previous steps yield sampling iterates that drift away from the training distribution. Yet, the standard training objective via Denoising Score Matching (DSM) is only designed to optimize over non-drifted data. To train on drifted data, we propose to enforce a \emph{consistency} property which states that predictions of the model on its own generated data are consistent across time. Theoretically, we show that if the score is learned perfectly on some non-drifted points (via DSM) and if the consistency property is enforced everywhere, then the score is learned accurately everywhere. Empirically we show that our novel training objective yields state-of-the-art results for conditional and unconditional generation in CIFAR-10 and baseline improvements in AFHQ and FFHQ. We open-source our code and models: https://github.com/giannisdaras/cdm

研究动机与目标

• 为解决扩散模型中的采样漂移问题，该问题源于分数函数学习不完善以及反向采样过程中误差的累积。
• 提升所学习分数函数对分布外点（例如，训练期间未见的漂移样本）的泛化能力。
• 开发一种在训练过程中无需访问真实分数函数或目标数据分布的训练目标。
• 从理论上和实证上证明，在时间步之间强制执行一致性可带来更准确、更稳定的采样。
• 在不修改采样过程的前提下，实现出图像生成基准上的SOTA性能。

提出的方法

• 引入一致性属性：在时间t对时间t生成的样本进行去噪预测时，其预测结果必须与在时间t对该样本的去噪预测结果一致。
• 制定一致性损失，以惩罚模型在其自身生成轨迹上不同时间步的预测差异。
• 将一致性损失与弱去噪分数匹配（DSM）目标相结合，仅使用非漂移数据（即，被噪声污染的干净数据）进行训练。
• 端到端地联合使用DSM损失和一致性损失进行模型训练，使其在推理阶段能够泛化到漂移样本。
• 利用一致性损失隐式地对整个生成轨迹上的分数函数进行正则化，即使在未显式优化的时间步也有效。
• 在采样过程中引入早停作为消融实验，以验证模型是否能泛化到显式训练时间步之外的区域。

(a) Consistency Property Testing on AFHQ. The plot illustrates how the Consistency Loss, $L^{2}_{t,t^{\prime},x_{t}}$ , behaves for $t^{\prime}=0$ , as $t$ changes.

实验结果

研究问题

• RQ1 在时间步之间强制执行模型预测的一致性，是否能缓解扩散模型中的采样漂移？
• RQ2 基于一致性的正则化与弱DSM监督相结合，是否能提升分数函数对漂移数据的泛化能力？
• RQ3 模型是否能在未显式训练于漂移数据或不修改采样过程的情况下，实现SOTA性能？
• RQ4 一致性属性与所学习分数函数正确性之间的理论关系是什么？
• RQ5 一致性损失如何影响模型对未显式训练时间步的泛化能力？

主要发现

• 所提出的CDM方法在CIFAR-10上实现了SOTA的FID分数5.45，优于基线EDM模型（FID 5.81）。
• 消融实验表明，若在某些时间步移除DSM损失，性能显著下降（FID 6.59）；而采样早停则导致性能极差（FID 14.52）。
• 一致性损失使模型能够泛化到未显式训练的时间步，表现为即使在DSM训练时间步之外继续采样，性能依然强劲。
• 理论分析表明，若在单个时间步完美学习到分数函数，且一致性属性在整个时间轴上成立，则分数函数在所有时间步均被完美学习。
• 实证结果表明，CDM在AFHQ和FFHQ上也表现出一致的性能提升，表明该方法在CIFAR-10之外的多样化数据集上具有良好的泛化能力。
• 该方法使训练时间增加约1.5倍，提示在推理效率方面仍有优化空间。

(b) Consistency Property Testing on AFHQ. The plot illustrates how the Consistency Loss, $L^{2}_{t,t^{\prime},x_{t}}$ , behaves for $t=0$ , as $t^{\prime}$ changes.