[論文レビュー] Consistent Diffusion Models: Mitigating Sampling Drift by Learning to be Consistent

Giannis Daras, Yuval Dagan+2 more|arXiv (Cornell University)|Feb 17, 2023

Model Reduction and Neural Networks被引用数 7

ひとこと要約

本稿では、サンプリングドリフトを軽減するための新しいトレーニング目的である一貫性のある拡散モデル（CDM）を提案する。この手法は、モデルの予測が自身が生成したサンプル上で時間ステップにわたって一貫性を持つように制約することで、サンプリングドリフトを緩和する。この一貫性損失を弱いノイズ除去スコアマッチング目的と組み合わせることで、ドリフトしていないデータでのみトレーニングされた場合でも、CIFAR-10で5.45のSOTA FIDスコアを達成し、AFHQおよびFFHQでも性能向上を達成した。

ABSTRACT

Imperfect score-matching leads to a shift between the training and the sampling distribution of diffusion models. Due to the recursive nature of the generation process, errors in previous steps yield sampling iterates that drift away from the training distribution. Yet, the standard training objective via Denoising Score Matching (DSM) is only designed to optimize over non-drifted data. To train on drifted data, we propose to enforce a \emph{consistency} property which states that predictions of the model on its own generated data are consistent across time. Theoretically, we show that if the score is learned perfectly on some non-drifted points (via DSM) and if the consistency property is enforced everywhere, then the score is learned accurately everywhere. Empirically we show that our novel training objective yields state-of-the-art results for conditional and unconditional generation in CIFAR-10 and baseline improvements in AFHQ and FFHQ. We open-source our code and models: https://github.com/giannisdaras/cdm

研究の動機と目的

• 逆方向サンプリング中の誤差累積と不完全なスコア関数学習に起因する、拡散モデルにおけるサンプリングドリフトを是正すること。
• トレーニング中に観測されなかった分布外の点（例：ドリフトしたサンプル）へのスコア関数の一般化を向上させること。
• トレーニング中に真のスコア関数やターゲット分布にアクセスする必要がないトレーニング目的を開発すること。
• 理論的および実験的に、時間ステップにわたる一貫性の強制が、より正確で安定したサンプリングをもたらすことを示すこと。
• サンプリングプロセスの変更なしに、画像生成ベンチマークでSOTAの性能を達成すること。

提案手法

• 一貫性特性を導入：時間tにおける、時間tに生成されたサンプルに対するモデルのノイズ除去予測が、同じサンプルについて時間tにおける予測と一致すること。
• モデルの自身が生成した軌跡における時間ステップ間の予測差をペナルティとする一貫性損失を定式化する。
• 一貫性損失を、ドリフトしていないデータ（ノイズで汚れたクリアデータ）でのみトレーニングされる弱いノイズ除去スコアマッチング（DSM）目的と組み合わせる。
• DSM損失と一貫性損失の両方を用いて、エンドツーエンドでモデルをトレーニングし、推論時にドリフトしたサンプルへの一般化を可能にする。
• 一貫性損失を用いて、明示的に最適化されていない時間ステップに対しても、生成軌跡全体にわたるスコア関数の正則化を間接的に実現する。
• アブレーションとして、サンプリング中に早期停止を実装し、モデルが明示的にトレーニングされた時間ステップを超えて一般化できることを検証する。

(a) Consistency Property Testing on AFHQ. The plot illustrates how the Consistency Loss, $L^{2}_{t,t^{\prime},x_{t}}$ , behaves for $t^{\prime}=0$ , as $t$ changes.

実験結果

リサーチクエスチョン

• RQ1 時間ステップにわたるモデル予測の一貫性を強制することで、拡散モデルにおけるサンプリングドリフトを軽減できるか？
• RQ2 弱いDSM監視と組み合わせた一貫性ベースの正則化は、ドリフトしたデータへのスコア関数の一般化を向上させるか？
• RQ3 ドリフトしたデータへの明示的トレーニングやサンプリングプロセスの変更なしに、モデルがSOTAの性能を達成できるか？
• RQ4 一貫性特性と学習済みスコア関数の正しさの間には、理論的にどのような関係があるか？
• RQ5 一貫性損失は、明示的にトレーニングされていない時間ステップへの一般化能力にどのように影響を与えるか？

主な発見

• 提案手法のCDMは、CIFAR-10でSOTAのFIDスコア5.45を達成し、ベースラインのEDMモデル（FID 5.81）を上回った。
• アブレーションスタディでは、特定の時間ステップにおけるDSM損失を除去すると性能が著しく低下（FID 6.59）し、サンプリングの早期停止も悪影響を及ぼした（FID 14.52）。
• 一貫性損失により、DSMで明示的に最適化されていない時間ステップに対しても、モデルが一般化できていることが実証された。特に、DSMで使用された時間ステップを超えてサンプリングを継続しても良好な性能が得られた。
• 理論的分析により、ある時間ステップでスコアが完全に学習されており、かつ一貫性特性がすべての時間ステップで成り立つならば、すべての時間ステップでスコアが完全に学習されていることが示された。
• 実験的結果では、AFHQおよびFFHQでも一貫した性能向上が確認され、CIFAR-10以外の多様なデータセットに対しても一般化可能であることが示された。
• トレーニング時間は約1.5倍に増加したため、推論効率の最適化の余地があることが示唆された。

(b) Consistency Property Testing on AFHQ. The plot illustrates how the Consistency Loss, $L^{2}_{t,t^{\prime},x_{t}}$ , behaves for $t=0$ , as $t^{\prime}$ changes.