[论文解读] Soft Truncation: A Universal Training Technique of Score-based Diffusion Model for High Precision Score Estimation

Dong‐Jun Kim, Seungjae Shin+3 more|arXiv (Cornell University)|Jun 10, 2021

Advanced Neuroimaging Techniques and Applications被引用 10

一句话总结

本文提出 Soft Truncation，一种适用于基于分数的扩散模型的通用训练技术，该技术用一个随机变量替代固定的截断超参数，从而在所有扩散时间范围内实现改进的分数估计。通过在每个训练步骤中采样一个随机的最小扩散时间，该方法在 CIFAR-10、CelebA、CelebA-HQ 和 STL-10 上实现了负对数似然（NLL）和 Fréchet Inception 距离（FID）的最先进性能，解决了长期存在的密度估计与生成样本质量之间的权衡问题。

ABSTRACT

Recent advances in diffusion models bring state-of-the-art performance on image generation tasks. However, empirical results from previous research in diffusion models imply an inverse correlation between density estimation and sample generation performances. This paper investigates with sufficient empirical evidence that such inverse correlation happens because density estimation is significantly contributed by small diffusion time, whereas sample generation mainly depends on large diffusion time. However, training a score network well across the entire diffusion time is demanding because the loss scale is significantly imbalanced at each diffusion time. For successful training, therefore, we introduce Soft Truncation, a universally applicable training technique for diffusion models, that softens the fixed and static truncation hyperparameter into a random variable. In experiments, Soft Truncation achieves state-of-the-art performance on CIFAR-10, CelebA, CelebA-HQ 256x256, and STL-10 datasets.

研究动机与目标

• 解决基于分数的扩散模型中密度估计（NLL）与样本生成质量（FID）之间的反向相关性问题。
• 改进分数网络在整个扩散时间范围内的训练，特别是在 NLL 最为敏感的低噪声区域。
• 通过引入统一的训练技术，消除对分别针对 FID 或 NLL 优化的独立模型配置的需求。
• 提供一种理论基础坚实、普遍适用的方法，在不改变网络架构的前提下，同时提升似然估计与生成性能。

提出的方法

• 引入一个随机截断变量 τ 代替固定的 ε，其中 τ 在每个小批量训练中重新采样，以定义分数估计的最小扩散时间。
• 仅在区间 [τ, T] 上训练分数网络，忽略低于 τ 的分数，从而通过随机采样确保整个时间范围的覆盖。
• 该方法隐式地在扩散时间上诱导出一个加权函数，等价于损失中的广义期望，这解释了其在平衡 NLL 与 FID 方面的成功。
• Soft Truncation 与任何基于分数的扩散模型和 SDE 构型兼容，无需对模型架构或采样过程进行修改。
• 将训练目标表述为最大扰动似然估计（MPLE），这是专用于扩散模型的 MLE 广义形式。
• 该技术在多个数据集和 SDE 上进行了实证验证，包括 VE、VP 和基于 ODE 的构型。

实验结果

研究问题

• RQ1 是否存在一种单一的训练技术，能够同时提升扩散模型中的密度估计（NLL）与样本质量（FID）？
• RQ2 为何在实践中，固定的截断超参数 ε 会导致 NLL 与 FID 之间的权衡？
• RQ3 在训练期间随机化最小扩散时间是否能实现对整个时间范围的更好分数估计？
• RQ4 Soft Truncation 是否在理论上等价于加权似然目标？如果是，这种加权是如何产生的？
• RQ5 Soft Truncation 是否在不同模型架构、SDE 构型和数据集上具有泛化能力？

主要发现

• Soft Truncation 在 CIFAR-10（VP SDE，DDPM++）上实现了 3.45 的最先进 FID，与注重 FID 的模型相当，同时保持了 3.03 的 NLL，与注重 NLL 的模型相当。
• 在 CelebA-HQ 256 上，Soft Truncation 实现了 FID 3.96 和 NLL 3.01，优于所有注重 FID 或 NLL 的基线模型。
• 消融实验表明，与固定 ε=10⁻⁵ 相比，Soft Truncation 使用 τ ∼ P₁ 在 CIFAR-10（RVE，UNCSN++）上将 FID 降低了 2.74 个点，同时保持了较低的 NLL。
• 与固定截断相比，该方法在 CIFAR-10 上将 FID 最多降低 2.74 个点，在 CelebA 上降低 1.26 个点，且 NLL 无下降。
• 使用 P₁（在 [0.1, 1.0] 上均匀分布）的 Soft Truncation 实现了最佳平衡，在 CIFAR-10（VP，DDPM++）上将 FID 从 6.70 降低至 3.96，同时保持 NLL 稳定。
• 该技术在不同架构上表现稳健：在 CIFAR-10 上将 FID 提升 1.26–2.74 个点，在 CelebA 上提升 0.7–1.2 个点，且 NLL 表现一致。