[论文解读] Learning to Efficiently Sample from Diffusion Probabilistic Models
提出一种动态规划方法，用于为预训练的 DDPMs 找到最优推理时间表，在仅需 32 次 refinement 步骤且无需再训练的情况下实现高质量采样。它在固定计算预算下通过优化 ELBO 选择时间步。
Denoising Diffusion Probabilistic Models (DDPMs) have emerged as a powerful family of generative models that can yield high-fidelity samples and competitive log-likelihoods across a range of domains, including image and speech synthesis. Key advantages of DDPMs include ease of training, in contrast to generative adversarial networks, and speed of generation, in contrast to autoregressive models. However, DDPMs typically require hundreds-to-thousands of steps to generate a high fidelity sample, making them prohibitively expensive for high dimensional problems. Fortunately, DDPMs allow trading generation speed for sample quality through adjusting the number of refinement steps as a post process. Prior work has been successful in improving generation speed through handcrafting the time schedule by trial and error. We instead view the selection of the inference time schedules as an optimization problem, and introduce an exact dynamic programming algorithm that finds the optimal discrete time schedules for any pre-trained DDPM. Our method exploits the fact that ELBO can be decomposed into separate KL terms, and given any computation budget, discovers the time schedule that maximizes the training ELBO exactly. Our method is efficient, has no hyper-parameters of its own, and can be applied to any pre-trained DDPM with no retraining. We discover inference time schedules requiring as few as 32 refinement steps, while sacrificing less than 0.1 bits per dimension compared to the default 4,000 steps used on ImageNet 64x64 [Ho et al., 2020; Nichol and Dhariwal, 2021].
研究动机与目标
- 激发在不重新训练的情况下降低 DDPM 采样的计算成本。
- 引入一种在给定精炼预算下选择最优推理时间步的精确动态规划方法。
- 利用 ELBO 可分解性以实现对推理路径的记忆化存储和精确优化。
提出的方法
- 将推理时间表的选择表述为时间步上的最短路径问题，利用 ELBO 的分解。
- 使用固定的预训练 DDPM 计算所有候选时间步的 KL 基 ELBO 项 L(t,s) 表。
- 应用动态规划算法，在恰好 K 次精炼步中找到精确的最优路径（0=t0<...<tK=1）。
- 通过记忆化利用，使填充 L(t,s) 项只需要 O(T) 次前向传播，其中 T 是网格时间步的数量。
- 通过构建从 0 开始、以连续时间步结束于 1 的有效 ELBO 路径，支持时间离散和时间连续的 DDPM。
- 可选地使用 Monte Carlo 采样估计 ELBO 项以降低计算量。
实验结果
研究问题
- RQ1在固定计算预算下，是否可以在训练后优化 DDPM 的推理时间表？
- RQ2在少步 regime 下，基于 DDPM 时间步的精确动态规划形式是否比手工设计的采样时间表获得更高的 ELBO（更低的负 ELBO）？
- RQ3需要多少次精炼步骤才能在显著降低计算量的同时尽量接近原始模型的对数似然？
- RQ4DP 推导的时间表是否可以在不同的预训练 DDPM 变体（时间离散和时间连续）之间无重训练地移植？
主要发现
| 模型 ∖ # 精炼步骤 | 8 | 16 | 32 | 64 | 128 | 256 | 全部 |
|---|---|---|---|---|---|---|---|
| DistAug Transformer (Jun et al., 2020) | – | – | – | – | – | – | 2.53 |
| DDPM++ (deep, sub-VP) (Song et al., 2021) | – | – | – | – | – | – | 2.99 |
| L_simple (Even stride) | 6.95 | 6.15 | 5.46 | 4.91 | 4.47 | 4.14 | 3.73 |
| L_simple (Quadratic stride) | 5.39 | 4.86 | 4.52 | 3.84 | 3.74 | 3.73 | – |
| L_simple (DP stride) | 4.59 | 3.99 | 3.79 | 3.74 | 3.73 | 3.72 | – |
| L_vlb (Even stride) | 6.20 | 5.48 | 4.89 | 4.42 | 4.03 | 3.73 | 2.94 |
| L_vlb (Quadratic stride) | 4.89 | 4.09 | 3.58 | 3.23 | 3.09 | 3.05 | – |
| L_vlb (DP stride) | 4.20 | 3.41 | 3.17 | 3.08 | 3.05 | 3.04 | – |
| L_hybrid (Even stride) | 6.14 | 5.39 | 4.77 | 4.29 | 3.92 | 3.66 | 3.17 |
| L_hybrid (Quadratic stride) | 4.91 | 4.15 | 3.71 | 3.42 | 3.30 | 3.26 | – |
| L_hybrid (DP stride) | 4.33 | 3.62 | 3.39 | 3.30 | 3.27 | 3.26 | – |
- 基于 DP 的方法能够为任意预算 K 找到最优推理路径，只需要 O(T) 次前向传播来计算 ELBO 项。
- 在 CIFAR-10 的 L_simple 和 ImageNet 64x64 的 L_hybrid 上，使用至多 32 次精炼步骤的时间表与原始 1000–4000 步模型在每维位给予 0.1 比特的差距内。
- DP 步长时间表在少步区间的对数似然（bits/dim）上超越了手工设计的等步或二次步长。
- DP 方法产生了较强的对数似然，但并不总是提升 FID 分数，突出表明对数似然与 FID 之间的已知不一致。
- 蒙特卡罗消融显示，使用多达 128 个样本来估计 ELBO 项仍能带来显著收益，特别是在 CIFAR-10 上；ImageNet 效益随样本增加而提升。
- 该方法无需训练，且可广泛应用于预训练的 DDPM。
更好的研究，从现在开始
从论文设计到论文写作，大幅缩短您的研究时间。
无需绑定信用卡
本解读由 AI 生成，并经人工编辑审核。