[論文レビュー] SA-Solver: Stochastic Adams Solver for Fast Sampling of Diffusion Models

Shuchen Xue, Mingyang Yi+5 more|arXiv (Cornell University)|Sep 10, 2023

Bayesian Methods and Mixture Models被引用数 6

ひとこと要約

本稿では、拡散確率的モデルにおける高速かつ高品質なサンプリングのための新しい確率的アダムス法、SA-Solverを提案する。分散を制御する拡散スドルーディング方程式（SDE）の族を導入し、効率的な線形多段階SDEソルバーを活用することで、関数評価回数（NFE）を減らしながら、複数のベンチマークで最先端のFIDスコアを達成した。これは、決定的および既存の確率的サンプラーを凌駕しており、特にNFEが少ない設定や部分的に訓練されたモデルにおいて顕著な優位性を示した。

ABSTRACT

Diffusion Probabilistic Models (DPMs) have achieved considerable success in generation tasks. As sampling from DPMs is equivalent to solving diffusion SDE or ODE which is time-consuming, numerous fast sampling methods built upon improved differential equation solvers are proposed. The majority of such techniques consider solving the diffusion ODE due to its superior efficiency. However, stochastic sampling could offer additional advantages in generating diverse and high-quality data. In this work, we engage in a comprehensive analysis of stochastic sampling from two aspects: variance-controlled diffusion SDE and linear multi-step SDE solver. Based on our analysis, we propose extit{SA-Solver}, which is an improved efficient stochastic Adams method for solving diffusion SDE to generate data with high quality. Our experiments show that extit{SA-Solver} achieves: 1) improved or comparable performance compared with the existing state-of-the-art (SOTA) sampling methods for few-step sampling; 2) SOTA FID on substantial benchmark datasets under a suitable number of function evaluations (NFEs). Code is available at https://github.com/scxue/SA-Solver.

研究の動機と目的

• 拡散モデルのトレーニングフリーなサンプリング手法における非効率性と品質制限を解消すること。
• 決定的ソルバーと比較して、より高品質で多様なデータを生成できる可能性を有する確率的サンプリングの有効性を検討すること。
• 可変なノイズスケールを備えた、分散を制御可能な新しい拡散SDEの族を設計することで、サンプリング効率を向上させること。
• 既存のODEソルバーと同等またはそれを上回る性能を示す効率的な線形多段階SDEソルバーを構築すること。
• 追加のモデル訓練を必要とせず、NFEが少ない設定でも最先端の生成品質を達成すること。

提案手法

• 元のSDEと同一の周辺分布を持つが、調整可能なパラメータτを介してノイズスケール（分散）を制御可能なSDEの族を提案する。
• [29]で提示された確率的アダムス法を拡張し、拡散過程に特有の指数的重み付き積分を組み込むことで、得られたSDEを解く。
• 提案された確率的アダムスソルバーの理論的収束次数解析を導入し、数値的安定性と正確性を保証する。
• 効率性を向上させ、サンプリングに必要な関数評価回数（NFE）を削減するために、線形多段階フレームワークを採用する。
• ノイズスケールτをキャリブレーションし、生成における多様性と品質のバランスを制御可能にする。
• 標準的なDPMと同一の逆過程を採用するが、提案されたSA-Solverによる強化された数値積分を適用する。

実験結果

リサーチクエスチョン

• RQ1 制御可能なノイズ分散を有する確率的サンプリングは、決定的ソルバーと比較して、拡散モデルにおける生成品質を向上させ得るか？
• RQ2 既存のODEソルバーと同等またはそれを上回る性能を示す、効率的な線形多段階SDEソルバーをどのように設計できるか？
• RQ3 異なるNFEレベルにおいて、ノイズスケール（τ）とサンプリング品質（FIDで測定）の最適なトレードオフは何か？
• RQ4 拡散モデルが完全に訓練されていない状態でも、提案手法は優れた性能を維持するか。これはモデル容量へのロバストネスを示唆するか？
• RQ5 トレーニングフリーで確率的ソルバーは、最小限の関数評価回数で最先端のFIDスコアを達成できるか？

主な発見

• SA-Solverは、CIFAR-10、LSUN-Bedroom、ImageNetのすべてのベンチマークデータセットで、適切な関数評価回数（NFE）の下で最先端のFIDスコアを達成し、決定的および確率的ベースラインを上回った。
• 極めて少ないNFE（例：5–10）の設定でも、DDIM、DPM-Solver、UniPCなどの既存の最先端手法と同等またはそれ以上のFIDスコアを達成した。
• LSUN-Bedroom 256x256データセットでは、NFE=60で最小FID 3.06を達成し、EDM(SDE)やUniPCを含むすべてのベースラインを上回った。
• 分類器フリー付きのImageNet 512x512 DiTモデルにおいて、NFE=60の条件下で、DDIM や UniPC と比較してより多様で詳細なサンプルを生成した。
• モデルが完全に訓練されていない状態でも、SA-Solverは決定的サンプラーと比較して優れた性能を維持しており、モデル容量へのロバストネスを示した。
• アブレーションスタディの結果、ノイズスケールτを増加（τ=1.0まで）させることで、NFEが少ない状況でのサンプル品質が向上し、多数のベンチマークでτ=0.8–1.0の範囲で最適な性能が得られた。