[論文レビュー] Multistep Consistency Models

Jonathan Heek, Emiel Hoogeboom+1 more|arXiv (Cornell University)|Mar 11, 2024

Complex Systems and Decision MakingDecision Sciences被引用数 3

ひとこと要約

本稿では、一時的整合性モデルとTRACTを統合し、1ステップの整合性から完全な拡散サンプリングまでの間を補間する、マルチステップ整合性モデルを提案する。共有パラメータを用いて2〜8ステップのセグメント単位で整合性モデルを訓練することで、8ステップでSOTAのFIDスコア（ImageNet64で1.4、ImageNet128で2.1）を達成し、標準的な拡散モデルの品質を維持しながら高いサンプリング速度を実現した。

ABSTRACT

Diffusion models are relatively easy to train but require many steps to generate samples. Consistency models are far more difficult to train, but generate samples in a single step. In this paper we propose Multistep Consistency Models: A unification between Consistency Models (Song et al., 2023) and TRACT (Berthelot et al., 2023) that can interpolate between a consistency model and a diffusion model: a trade-off between sampling speed and sampling quality. Specifically, a 1-step consistency model is a conventional consistency model whereas a $\infty$-step consistency model is a diffusion model. Multistep Consistency Models work really well in practice. By increasing the sample budget from a single step to 2-8 steps, we can train models more easily that generate higher quality samples, while retaining much of the sampling speed benefits. Notable results are 1.4 FID on Imagenet 64 in 8 step and 2.1 FID on Imagenet128 in 8 steps with consistency distillation, using simple losses without adversarial training. We also show that our method scales to a text-to-image diffusion model, generating samples that are close to the quality of the original model.

研究の動機と目的

• 標準的な拡散モデルと低ステップの整合性モデルの間の性能ギャップを埋めること。
• 複数ステップ推論を用いて、サンプリング速度と生成品質の滑らかなトレードオフを実現すること。
• TRACTと整合性蒸留の一般化により、整合性モデリングにおける訓練安定性とサンプル品質を向上させること。
• 統合誤差を補正する決定的サンプラー（aDDIM）を開発し、低ステップ生成におけるぼやけを低減すること。
• ImageNet128 やテキストから画像生成といった複雑なベンチマークで優れた性能を示すこと。

提案手法

• 本手法は拡散プロセスを事前に定義されたセグメントに分割し、各セグメントに対して個別の整合性モデルを訓練するが、すべてのセグメントでパラメータを共有する。
• 整合性学習（CT）と整合性蒸留（CD）を複数ステップに一般化することで、滑らかで学習可能なノイズ除去軌道を実現する。
• 整合性モデリングから得られる段階的スケジュールの緩和と同期ドロップアウトを組み込むことで、訓練安定性を向上させる。
• 統合誤差を補正する新しい決定的サンプラー、調整済みDDIM（aDDIM）を導入し、低ステップでのサンプリングにおけるぼやけを低減する。
• 1ステップの整合性モデルと完全な拡散モデル（無限ステップ）の間を補間可能なフレームワークを提供し、8ステップバージョンが標準的な拡散モデルの性能を再現した。
• 蒸留のため、事前に訓練済みの拡散モデルから、高品質な16ステップ整合性モデルをaDDIMを用いて訓練する。

実験結果

リサーチクエスチョン

• RQ1 整合性モデルは1ステップサンプリングを越えて、品質を向上させつつ速度優位性を維持できるか？
• RQ2 共有パラメータを用いた複数ステップ整合性学習は、1ステップ整合性モデルや標準的な拡散モデルよりも優れた性能を達成できるか？
• RQ3 aDDIMのような決定的サンプラーは、ImageNet128 などの複雑なデータセットで、少ないステップ数で競争力のあるFIDスコアを達成できるか？
• RQ4 段階的スケジュールの緩和は、低ステップマルチステップ整合性モデルの性能にどのように影響するか？
• RQ5 マルチステップ整合性モデルは、最小限のサンプリングステップ数で、テキストから画像生成において標準的な拡散モデルの性能を再現できるか？

主な発見

• 提案されたマルチステップ整合性モデルは、8ステップのサンプリングでImageNet64でFID 1.4、ImageNet128でFID 2.1というSOTAのスコアを達成した。
• 4ステップのサンプリングでも、ImageNet64でFID 1.6、ImageNet128でFID 2.3を達成し、低コストの推論でも優れた性能を示した。
• わずか8ステップで標準的な拡散モデル（例：100ステップのDDIM）の性能を再現し、品質ギャップを効果的に埋めた。
• 段階的スケジュールの緩和は、低ステップモデルの性能を顕著に向上させたが、ステップ数が増えると効果の逓減が見られた。
• aDDIMサンプラーにより、ImageNet128で競争力のあるFIDスコアが達成可能となり、かつて信頼できる決定的低ステップベースラインが欠如していたベンチマークをカバーした。
• テキストから画像生成のタスクでは、16ステップの蒸留モデルが、同じランダムシードを使用した場合、元の100ステップDDIMモデルとほぼ区別できないサンプルを生成した。