[论文解读] Restart Sampling for Improving Generative Processes

Yilun Xu, Mingyang Deng+4 more|arXiv (Cornell University)|Jun 26, 2023

Generative Adversarial Networks and Image Synthesis被引用 9

一句话总结

本文提出 Restart 采样，其将大噪声前向和后向 ODE 交错结合，以将随机性带来的收缩与较低离散化误差相结合，从而实现比现有 ODE 或 SDE 采样器更快且更高质量的生成采样。

ABSTRACT

Generative processes that involve solving differential equations, such as diffusion models, frequently necessitate balancing speed and quality. ODE-based samplers are fast but plateau in performance while SDE-based samplers deliver higher sample quality at the cost of increased sampling time. We attribute this difference to sampling errors: ODE-samplers involve smaller discretization errors while stochasticity in SDE contracts accumulated errors. Based on these findings, we propose a novel sampling algorithm called Restart in order to better balance discretization errors and contraction. The sampling method alternates between adding substantial noise in additional forward steps and strictly following a backward ODE. Empirically, Restart sampler surpasses previous SDE and ODE samplers in both speed and accuracy. Restart not only outperforms the previous best SDE results, but also accelerates the sampling speed by 10-fold / 2-fold on CIFAR-10 / ImageNet $64 imes 64$. In addition, it attains significantly better sample quality than ODE samplers within comparable sampling times. Moreover, Restart better balances text-image alignment/visual quality versus diversity than previous samplers in the large-scale text-to-image Stable Diffusion model pre-trained on LAION $512 imes 512$. Code is available at https://github.com/Newbeeer/diffusion_restart_sampling

研究动机与目标

• 在基于微分方程的生成模型（ODE 与 SDE 采样）中，激发速度与质量之间的权衡。
• 从理论上分析离散化误差和收缩在 ODE 与 SDE 采样中的交互。
• 提出一个 Restart 采样算法，在随机性与确定性的 ODE 步骤之间混合，以提高性能。
• 给出理论界限和实证证据，表明 Restart 在速度和质量上优于先前求解器。

提出的方法

• 表征 ODE 和 SDE 采样器并推导对数据分布的一阶 Wasserstein 距离的上界。
• 引入 Restart 采样，在一个时间区间内重复前向噪声与后向 ODE 循环，以放大初始误差的收敛。
• 在每次 Restart 迭代中，通过执行大噪声前向步骤和确定性的后向 ODE，将随机性与漂移分离。
• 提供关于在不同模型和数据集间平衡收敛性与额外采样误差来选择 Restart 间隔和 Restart 迭代次数（K）的指南。
• 表明 Restart 允许在获得随机收缩的同时获得更小的离散化误差，从而实现快速采样。

实验结果

研究问题

• RQ1 离散化误差和随机收缩如何影响 ODE 和 SDE 采样器的采样质量？
• RQ2 Restart 策略是否能利用较大前向噪声与确定性后向动力学，超越 ODE 和 SDE 采样器？
• RQ3 有哪些理论保障（Wasserstein 上界）支持 Restart 相对于现有方法的有效性？
• RQ4 针对不同模型和数据集，选择 Restart 间隔和 K 的实际指南是什么？
• RQ5 Restart 是否在标准基准和大规模文本到图像模型上提升速度和质量？

主要发现

模型	NFE	FID
EDM-VP [13] ODE (Heun)	63	1.97
EDM-VP [13] ODE (Heun)	35	1.97
EDM-VP [13] Improved SDE	63	2.27
EDM-VP [13] Improved SDE	35	2.45
CIFAR-10 Restart	43	1.90
PFGM++ [28] ODE (Heun)	63	1.91
PFGM++ [28] ODE (Heun)	35	1.91
PFGM++ [28] Restart	43	1.88

• Restart 在样本质量和速度方面在 NFEs、数据集和预训练模型上始终超越之前的 ODE 和 SDE 采样器。
• 在 CIFAR-10 的 VP 条件下，Restart 在相似 FID 下实现比以往最佳 SDE 结果快 10 倍或 2 倍，且优于快速 ODE 求解器。
• Restart 在不额外训练的情况下达到 UNet 架构的新 SOTA FID（如在 EDM 的 64×64 class-conditional ImageNet 上 1.36，在 PFGM++ 的无条件 CIFAR-10 上 1.88）。
• Restart 在 Stable Diffusion 的 LAION-5B 上改进文本到图像的权衡，优于先前采样器在文本-图像对齐/视觉质量和多样性之间的平衡。
• 理论分析表明 Restart 产生更小的 Wasserstein 上界，原因是更强的收缩以及类似 ODE 的离散化误差，且收缩随 Restart 迭代次数 K 指数增长（在受其他因素影响的一个区间内）。
• 实证结果包括 Pareto 前沿分析，显示 Restart 在给定额外采样误差下实现更低的收缩误差，并在低-NFE 与高-NFE 情况下均具优势。