[论文解读] Optimizing DDPM Sampling with Shortcut Fine-Tuning

Ying Fan, Kangwook Lee|arXiv (Cornell University)|Jan 31, 2023

Machine Learning in Healthcare被引用 4

一句话总结

本文提出快捷微调（SFT），一种受强化学习启发的方法，通过直接最小化积分概率度量（IPM）来优化DDPM采样，绕过反向扩散过程。通过将采样建模为策略梯度问题，SFT仅用10步即实现最先进样本质量——在CIFAR-10和CelebA上FID分数分别达到2.28和2.01，优于完整1000步的DDPM。

ABSTRACT

In this study, we propose Shortcut Fine-Tuning (SFT), a new approach for addressing the challenge of fast sampling of pretrained Denoising Diffusion Probabilistic Models (DDPMs). SFT advocates for the fine-tuning of DDPM samplers through the direct minimization of Integral Probability Metrics (IPM), instead of learning the backward diffusion process. This enables samplers to discover an alternative and more efficient sampling shortcut, deviating from the backward diffusion process. Inspired by a control perspective, we propose a new algorithm SFT-PG: Shortcut Fine-Tuning with Policy Gradient, and prove that under certain assumptions, gradient descent of diffusion models with respect to IPM is equivalent to performing policy gradient. To our best knowledge, this is the first attempt to utilize reinforcement learning (RL) methods to train diffusion models. Through empirical evaluation, we demonstrate that our fine-tuning method can further enhance existing fast DDPM samplers, resulting in sample quality comparable to or even surpassing that of the full-step model across various datasets.

研究动机与目标

• 为解决标准DDPM采样效率低下的问题，其高质量生成通常需要大量步数（T ≈ 1000）。
• 探究是否可通过偏离已学习的反向扩散过程，来改进采样质量。
• 开发一种直接通过IPM最小化优化样本质量的方法，而非模仿反向过程。
• 建立DDPM上IPM的梯度下降与策略梯度之间的理论联系，从而实现基于强化学习的扩散模型训练。

提出的方法

• 提出SFT-PG，一种新颖算法，将DDPM采样建模为强化学习问题，其中生成器（DDPM采样器）通过策略梯度进行优化。
• 证明DDPM采样器关于IPM的梯度下降在数学上等价于随机策略梯度，从而支持基于强化学习的训练方式。
• 引入IPM的代理函数，为训练过程中单调改进的条件提供理论洞察。
• 采用判别器-生成器交替训练机制，结合梯度惩罚与基线正则化，以稳定策略梯度训练。
• 使用FastDPM中的固定方差调度和子采样调度作为初始化，仅微调均值预测头。
• 应用梯度裁剪和每次判别器更新时的多步生成器更新，以提升训练稳定性和性能。

Figure 1 : Image denoising is similar to a closed-loop control system: finding paths from pure noise to natural images.

实验结果

研究问题

• RQ1 我们能否通过不遵循反向扩散过程，而是发现替代采样路径，来提升DDPM采样质量？
• RQ2 DDPM采样器关于IPM的梯度下降是否等价于策略梯度，从而支持基于强化学习的扩散模型训练？
• RQ3 IPM的代理函数能否在微调过程中引导稳定且单调的性能提升？
• RQ4 与基于模仿的快速采样方法（如FastDPM和Analytic-DPM）相比，直接IPM最小化是否能在更少步数下实现更高的样本质量？
• RQ5 使用SFT-PG微调能否仅用T′=10步采样，实现与完整1000步DDPM（T=1000）相当或更优的性能？

主要发现

• SFT-PG在10步采样下，CIFAR-10上FID达2.28，CelebA上达2.01，优于完整1000步DDPM（FID分别为3.03和3.26）。
• 在瑞士卷数据集上，SFT-PG将FID从100步时的2.36降低至10步时的0.64，表明极少数步数下亦能实现显著性能提升。
• 该方法优于现有快速采样器：SFT-PG在CIFAR-10上T′=10时FID优于FastDPM（29.43）、Analytic-DPM（22.94）和SN-DPM（16.33）。
• 在MNIST上，使用5步生成器更新和γ=1.0时，FID最低（0.82），表明梯度裁剪与步数设置显著影响性能。
• SFT-PG结合基线正则化（SFT-PG B）在玩具数据集上表现略优于SFT与SFT-PG，表明其训练更稳定。
• 理论分析证实，基于IPM的DDPM梯度下降等价于策略梯度，这是首次将强化学习方法用于训练扩散模型。

Figure 2 : A visual illustration of the key idea of Shortcut Fine-Tuning (SFT). DDPMs aim at learning the backward diffusion model, but this approach is limited to a small number of steps. We propose the idea of not following the backward process and exploring other unexplored paths that can lead to