[論文レビュー] Fast Sampling of Diffusion Models via Operator Learning

Hongkai Zheng, Weili Nie+3 more|arXiv (Cornell University)|Nov 24, 2022

Advanced Neuroimaging Techniques and Applications被引用数 17

ひとこと要約

本論文では、1回の順方向プロパゲーションで高速かつ並列に拡散モデルをサンプリングするためのニューラルオペレータベースの新手法DSNOを提案する。拡散モデルのバックボーンにフーリエ領域における時系列畳み込み層を統合することで、DSNOは連続時間における逆過程の軌道を学習し、1回のモデル評価でCIFAR-10で3.78、ImageNet-64で7.83という最先端のFIDスコアを達成した。

ABSTRACT

Diffusion models have found widespread adoption in various areas. However, their sampling process is slow because it requires hundreds to thousands of network evaluations to emulate a continuous process defined by differential equations. In this work, we use neural operators, an efficient method to solve the probability flow differential equations, to accelerate the sampling process of diffusion models. Compared to other fast sampling methods that have a sequential nature, we are the first to propose a parallel decoding method that generates images with only one model forward pass. We propose diffusion model sampling with neural operator (DSNO) that maps the initial condition, i.e., Gaussian distribution, to the continuous-time solution trajectory of the reverse diffusion process. To model the temporal correlations along the trajectory, we introduce temporal convolution layers that are parameterized in the Fourier space into the given diffusion model backbone. We show our method achieves state-of-the-art FID of 3.78 for CIFAR-10 and 7.83 for ImageNet-64 in the one-model-evaluation setting.

研究の動機と目的

• 時間に敏感なアプリケーション向けに、リアルタイム推論を可能にするため、拡散モデルのサンプリングを高速化すること。
• 複数回の順方向プロパゲーションを必要とする従来の高速サンプリング手法に見られる逐次的ボトルネックを克服すること。
• ノイズからデータへの微分可能で連続時間の解軌道マッピングをニューラルオペレータを用いて開発すること。
• 推論時間を1回の順方向プロパゲーションに大幅に短縮しつつも、高いサンプル品質を維持すること。
• パラメータ効率が良く、フーリエ空間における時系列畳み込みを用いて時系列モデリングを拡散モデルに統合すること。

提案手法

• 初期のガウスノイズから連続時間における逆拡散過程の解軌道へマップするニューラルオペレータであるDSNOを提案する。
• フーリエ領域でパrameter化された時系列畳み込み層を導入し、軌道に沿った時系列相関をモデル化する。
• フーリエベースの演算を用いて全タイムステップの出力を同時に計算することで、逐次計算を回避し、並列デコードを可能にする。
• 既存のU-Netバックボーンに時系列畳み込みブロックを統合し、標準的な拡散モデルアーキテクチャと互換性を保つ。
• サンプルの忠実度を向上させるために、L1損失と知覚的（LPIPS）損失の組み合わせでモデルを訓練する。
• 2次関数的時間離散化とマルチスケールの監視を用いて、一般化性能とFIDスコアの向上を図る。

実験結果

リサーチクエスチョン

• RQ1 ニューラルオペレータを用いて、拡散モデルの連続時間における逆過程を効果的にモデル化し、高速なサンプリングを実現できるか？
• RQ2 フーリエ空間における時系列畳み込み層は、拡散軌道における長距離の時系列依存性を効率的に捉えることができるか？
• RQ3 フーリエベースの演算による並列デコードは、サンプル品質を損なわずに1ステップ生成を可能にするか？
• RQ4 提案手法は、1回の順方向プロパゲーションのみで最先端のFIDスコアを達成できるか？
• RQ5 損失関数の選択（例：L1対LPIPS）が、最終的なサンプル品質およびFIDスコアにどのように影響するか？

主な発見

• DSNOは、1回のモデル評価設定下でCIFAR-10で3.78という最先端のFIDスコアを達成した。
• DSNOは、1回のモデル評価でImageNet-64でも7.83という最先端のFIDスコアを達成した。
• VGGベースのLPIPS損失を用いることで、標準的なL1損失と比較してFIDスコアに顕著な改善が見られた。
• 時間分解能を4から8に引き上げることでFIDはわずかに改善されたが、限界的利得から4が効率性にとって十分であると示唆された。
• 提案された時系列畳み込みブロックは、モデルサイズにわずか10%の増加しかもたらさず、並列デコードと高い性能を実現した。
• 本手法は真の並列デコードを可能にし、従来の逐次的ソルバとは異なり、1回の順方向プロパゲーションで全タイムステップの出力を生成した。