[論文レビュー] Optimizing DDPM Sampling with Shortcut Fine-Tuning

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

Machine Learning in Healthcare被引用数 4

ひとこと要約

本稿では、強化学習にインspiredされた手法であるShortcut Fine-Tuning (SFT)を提案する。SFTは、逆拡散プロセスを回避するため、DDPMのサンプリングを直接最適化するためのIntegral Probability Metrics (IPM)を最小化する。サンプリングをポリシー勾配問題として定式化することで、わずか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に関する勾配降下法とポリシー勾配の理論的関連を確立し、拡散モデルの強化学習ベースの訓練を可能にすること。

提案手法

• DDPMのサンプリングを強化学習問題として定式化する新規アルゴリズムSFT-PGを提案。生成器（DDPMサンプラー）をポリシー勾配により最適化する。
• DDPMサンプラーに対するIPMに関する勾配降下法が、確率的ポリシー勾配と数学的に同等であることを示し、RL風の訓練を可能にする。
• IPMの代理関数を導入し、訓練中の単調な改善が達成される条件について理論的洞察を提供する。
• 勾配ペナルティとベースライン正則化を用いた、クリティック・ジェネレータの交互訓練スキームを採用し、ポリシー勾配訓練の安定性を向上させる。
• FastDPMからの固定分散スケジュールとサブサンプリングスケジュールを初期化として採用し、平均予測ヘッドのみをファインチューニングする。
• 勾配クリッピングと、1回のクリティック更新あたり複数回のジェネレータステップを適用することで、訓練の安定性と性能を向上させる。

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

実験結果

リサーチクエスチョン

• RQ1 逆拡散プロセスに従わないことで、代替のサンプリングパスを発見することで、DDPMのサンプリング品質を向上させられるか？
• RQ2 IPMに関するDDPMサンプラーに対する勾配降下法は、ポリシー勾配と同等であるか？これにより、拡散モデルの強化学習ベースの訓練が可能になるか？
• RQ3 IPMの代理関数は、ファインチューニング中の安定的かつ単調な改善を導けるか？
• RQ4 直接的なIPM最小化は、FastDPM や Analytic-DPM といったイミテーションベースの高速サンプリング手法に比べ、少ないステップ数で優れたサンプル品質を達成できるか？
• RQ5 SFT-PGによるファインチューニングは、T′=10のサンプリングステップで、T=1000のフルステップDDPMと同等またはそれ以上の性能を達成できるか？

主な発見

• SFT-PGは10ステップのサンプリングで、CIFAR-10ではFID 2.28、CelebAではFID 2.01を達成。これは、それぞれFID 3.03および3.26を記録したフルステップDDPMを上回る性能である。
• スイスロールデータセットでは、SFT-PGによりFIDがT=100の2.36からT′=10の0.64に低下し、極めて少ないステップ数で顕著な改善が確認された。
• 既存の高速サンプラーの性能も向上：SFT-PGはCIFAR-10において、T′=10でFastDPM（29.43）、Analytic-DPM（22.94）、SN-DPM（16.33）を上回るFIDを達成した。
• MNISTでは、5ステップのジェネレータステップとγ=1.0の組み合わせが最良のFID（0.82）を達成し、勾配クリッピングとステップ数が性能に影響を与えることが示された。
• ベースライン正則化を施したSFT-PG B（SFT-PG B）は、トゥイティーデータセットにおいてSFTおよびSFT-PGよりもわずかに優れた最終性能を達成しており、安定性の向上が示された。
• 理論的分析により、IPMに基づくDDPMに対する勾配降下法がポリシー勾配と同等であることが確認され、拡散モデルの訓練にRL手法を初めて適用したとされる。

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