[論文レビュー] Parallel Sampling of Diffusion Models

Andy Y. Shih, Suneel Belkhale+3 more|arXiv (Cornell University)|May 25, 2023

Advanced Neuroimaging Techniques and Applications被引用数 9

ひとこと要約

ParaDiGMSは Picard反復法による並列化デノイジングを導入し、拡散モデルのサンプリングを高速化。ロボティクスと画像生成タスク全体で、品質の測定可能な低下なしに約2～4倍のスピードアップを達成する。

ABSTRACT

Diffusion models are powerful generative models but suffer from slow sampling, often taking 1000 sequential denoising steps for one sample. As a result, considerable efforts have been directed toward reducing the number of denoising steps, but these methods hurt sample quality. Instead of reducing the number of denoising steps (trading quality for speed), in this paper we explore an orthogonal approach: can we run the denoising steps in parallel (trading compute for speed)? In spite of the sequential nature of the denoising steps, we show that surprisingly it is possible to parallelize sampling via Picard iterations, by guessing the solution of future denoising steps and iteratively refining until convergence. With this insight, we present ParaDiGMS, a novel method to accelerate the sampling of pretrained diffusion models by denoising multiple steps in parallel. ParaDiGMS is the first diffusion sampling method that enables trading compute for speed and is even compatible with existing fast sampling techniques such as DDIM and DPMSolver. Using ParaDiGMS, we improve sampling speed by 2-4x across a range of robotics and image generation models, giving state-of-the-art sampling speeds of 0.2s on 100-step DiffusionPolicy and 14.6s on 1000-step StableDiffusion-v2 with no measurable degradation of task reward, FID score, or CLIP score.

研究の動機と目的

• 性能を速度と引換に計算資源を投じて、サンプル品質を犠牲にせずに拡散モデルのサンプリングを高速化する動機づけ。
• 事前学習済み拡散モデルや既存の高速サンプラーと互換性のある一般的な並列サンプリングフレームワークを開発する。
• 報酬、FID、CLIPなどの指標を保ちながら、ロボティクス方針と画像生成モデルの間で速度改善を実証する。
• 以前の高速サンプリング手法（DDIM、DPMSolver）および classifier-free ガイダンスとの互換性を示す。

提案手法

• 反転時拡散ODEを解くために Picard 繰り返し法を用い、デノイジングの全軌道の推定値を反復的に精緻化する。
• Picard更新を離散化して、時間ステップ全体でのドリフト s(x_i, t_i) の並列計算を可能にし、その後プレフィックス和の集約を行う。
• GPUメモリに適合させるスライディングウィンドウのミニバッチ戦略と、総変動距離を保証する停止基準を導入する。
• 逆SDEへ拡張する際に、固定前方ノイズ実現を決定論的ODEに吸収する。
• ParaDiGMS 内でオイラー法や高次解法へ置換し、ParaDDPM、ParaDDIM、ParaDPMSolver と統合することで互換性を示す。
• 許容基準停止基準を用いて、逐次ベースラインからの分布距離を上限化する。

実験結果

リサーチクエスチョン

• RQ1 サンプル品質を損なうことなく、デノイジングステップを並列化して拡散モデルのサンプリングを加速できるか？
• RQ2 Picard反復ベースの並列サンプリングは、速度と品質の点で既存の高速サンプラー（DDIM、DPMSolver）とどう比較されるか？
• RQ3 本手法は classifier-free ガイダンスや既存の高速化と互換性があり、画像とロボットポリシーの両方のタスクに一般化できるか？

主な発見

• ParaDiGMS は、ロボティクスおよび画像生成モデル全体で約2–4倍の速度アップを達成し、タスク報酬、FID、CLIPスコアに測定可能な低下を生じさせない。
• DiffusionPolicy の100ステップサンプリングでは、100ステップのアクション生成タスクにつきサンプルあたりの時間が0.74sから0.2sへ短縮（例を提示）。
• StableDiffusion-v2 の1000ステップサンプリングでは、サンプルあたりの時間が50.0sから14.6sへ短縮し、品質の低下なし。
• ParaDiGMSはDDPM、DDIM、DPMSolverと互換性があり、タスク全体で組み合わせた速度向上（ParaDDPM、ParaDDIM、ParaDPMSolver）を可能にする。
• StableDiffusion-v2の潜在空間拡散では、マルチGPU環境で総合的な速度アップが3倍超えとなり、バッチウィンドウのサイズとハードウェアに依存する。
• ParaDiGMSはCOCOキャプションのCLIPスコアで測定されるサンプル品質を、いくつかの設定に渡って維持する。