[論文レビュー] OMS-DPM: Optimizing the Model Schedule for Diffusion Probabilistic Models

Enshu Liu, Xuefei Ning+3 more|arXiv (Cornell University)|Jun 15, 2023

Advanced Neural Network Applications被引用数 6

ひとこと要約

OMS-DPMは、ノイズ除去ステップごとに異なる事前学習済みモデルを選択することで、拡散確率的モデルにおけるモデルスケジュール最適化のための新規手法を提案する。生成品質と速度の両方を顕著に向上させる。予測子ベースの進化的探索を用いることで、Stable Diffusionと同等の品質を維持しながら、2倍の高速化を達成。CIFAR-10、CelebA、ImageNet、LSUN-Churchの複数のベンチマークで最先端のFIDスコアを達成。

ABSTRACT

Diffusion probabilistic models (DPMs) are a new class of generative models that have achieved state-of-the-art generation quality in various domains. Despite the promise, one major drawback of DPMs is the slow generation speed due to the large number of neural network evaluations required in the generation process. In this paper, we reveal an overlooked dimension -- model schedule -- for optimizing the trade-off between generation quality and speed. More specifically, we observe that small models, though having worse generation quality when used alone, could outperform large models in certain generation steps. Therefore, unlike the traditional way of using a single model, using different models in different generation steps in a carefully designed \emph{model schedule} could potentially improve generation quality and speed \emph{simultaneously}. We design OMS-DPM, a predictor-based search algorithm, to optimize the model schedule given an arbitrary generation time budget and a set of pre-trained models. We demonstrate that OMS-DPM can find model schedules that improve generation quality and speed than prior state-of-the-art methods across CIFAR-10, CelebA, ImageNet, and LSUN datasets. When applied to the public checkpoints of the Stable Diffusion model, we are able to accelerate the sampling by 2$ imes$ while maintaining the generation quality.

研究の動機と目的

• 生成毎に何百〜何千回もニューラルネットワーク評価を必要とする、拡散確率的モデル（DPMs）の遅いサンプリング速度を解消すること。
• 生成品質と速度の両面で向上を図るための、これまで無視されてきた最適化次元「モデルスケジュール」を同定・活用すること。ここでは、ノイズ除去ステップごとに異なるモデルを用いることで、品質と速度を向上させる。
• 再訓練を一切行わず、事前学習済みモデルのみを用いて、任意の時間予算下での最適なモデルスケジュールを同定する手法を開発すること。
• Stable Diffusionなどの公開チェックポイントを含む、多様なデータセットおよび事前学習済みモデルに対し、効率的かつ汎用的なスケジューリングを可能にすること。
• ノイズ除去ステップにおけるモデルの非均質性を活用することで、従来手法よりも優れた品質-速度トレードオフを達成すること。

提案手法

• モデルスケジュールから生成品質を予測できるように学習する予測子ベースの探索アルゴリズム「OMS-DPM」を提案する。少量の訓練データを用いて学習する。
• 少数のモデルスケジュールのサブセット上で性能予測子を訓練し、未観測のスケジュールにも一般化可能にすることで、完全評価にかかるコストを削減する。
• 指数的に巨大なモデルスケジュール探索空間を、所定の時間予算下で効率的に探索するため、進化的アルゴリズムを用いる。
• DDIMのステップスキップやDPM-Solverの順序選択といった既存の加速技術を、探索のハイパーパrameterとして柔軟に統合可能である。
• 予測された性能に基づき、異なるサイズ・速度の事前学習済みモデルを、ノイズ除去ステップごとに異なるものに割り当てる。
• 1つのデータセットあたり数点のデータで予測子を学習することで、再訓練を最小限に抑えつつ、新しいデータセットへも再利用可能な汎用性を実現する。

実験結果

リサーチクエスチョン

• RQ1 異なるノイズ除去ステップに異なる事前学習済みモデルを動的に割り当てることで、DPMにおける生成品質とサンプリング速度の両方を向上させることができるか？
• RQ2 ノイズスケジュール、離散化、ソルバ設計に注目する従来の加速手法と比較して、モデルスケジュール最適化はどのように異なるか？
• RQ3 時間予算が制限された状況下で、モデルサイズと配置（初期ステップ vs. 後期ステップ）が最終的な生成品質に与える影響は何か？
• RQ4 少量のデータポイント（例：100点程度）で学習した軽量予測子は、多様なモデルスケジュールとデータセットに一般化可能か？これにより、完全評価なしで効率的な探索が可能か？
• RQ5 最適なモデルスケジュールの構造は、異なるデータセットやサンプリング予算によってどのように変化するか？

主な発見

• OMS-DPMは、Stable Diffusionモデルのサンプリングを2倍高速化しながら、FIDスコアを維持または向上させ、ベースライン手法を上回る。
• CIFAR-10では、1400msの遅延予算下でFIDをベースラインの3.56から3.19に低下させ、品質-速度トレードオフの顕著な改善を示した。
• 時間予算が厳しい状況では、最小で高速なモデル（例：2〜3モデル）のみが使用される。これは、ステップ数が少ないと、大きなモデルが誤差を引き起こすためである。
• LSUN-Churchでは、最適戦略として、ノイズ除去の初期段階で大きなモデルが使用されるが、ImageNet-64やCIFAR-10では、後期に大きなモデルがより効果的であることが示され、データセット依存のスケジューリングパターンが確認された。
• 100点程度のデータポイントで学習した予測子は、良好な一般化性能を示し、最小限の評価コストで効率的な探索が可能である。
• DPM-Solverでは、時間予算が厳しい状況では1次および2次ソルバが好まれるが、十分な時間が確保できる場合には3次ソルバが使用される。最適スケジュールにおける動的ソルバ選択の傾向が明らかになった。