ステップを揃える：拡散モデルにおけるサンプリングスケジュールの最適化

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

April 22, 2024

著者: Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis

cs.AI

要旨

拡散モデル（DMs）は、視覚領域をはじめとする生成モデリング手法の最先端として確立されています。DMsの重要な欠点は、大規模なニューラルネットワークを通じた多くの逐次的な関数評価に依存するため、サンプリング速度が遅いことです。DMsからのサンプリングは、サンプリングスケジュールとして知られる離散化されたノイズレベルを通じて微分方程式を解くことと見なすことができます。これまでの研究は主に効率的なソルバーの導出に焦点を当ててきましたが、最適なサンプリングスケジュールを見つけることにはほとんど注意が払われておらず、既存の文献は手作りのヒューリスティックに依存しています。本研究では、初めて、高品質な出力のためのDMsのサンプリングスケジュールを最適化する一般的で原理に基づいたアプローチを提案します。これを「Align Your Steps」と呼びます。我々は確率解析の手法を活用し、異なるソルバー、訓練済みのDMs、およびデータセットに特化した最適なスケジュールを見つけます。我々の新しいアプローチを、いくつかの画像、ビデオ、および2Dのトイデータ合成ベンチマークで評価し、さまざまなサンプラーを使用して、最適化されたスケジュールが従来の手作りのスケジュールをほぼすべての実験で上回ることを観察しました。我々の方法は、特に少ステップ合成の領域において、サンプリングスケジュール最適化の未開拓の可能性を示しています。

English

Diffusion models (DMs) have established themselves as the state-of-the-art generative modeling approach in the visual domain and beyond. A crucial drawback of DMs is their slow sampling speed, relying on many sequential function evaluations through large neural networks. Sampling from DMs can be seen as solving a differential equation through a discretized set of noise levels known as the sampling schedule. While past works primarily focused on deriving efficient solvers, little attention has been given to finding optimal sampling schedules, and the entire literature relies on hand-crafted heuristics. In this work, for the first time, we propose a general and principled approach to optimizing the sampling schedules of DMs for high-quality outputs, called Align Your Steps. We leverage methods from stochastic calculus and find optimal schedules specific to different solvers, trained DMs and datasets. We evaluate our novel approach on several image, video as well as 2D toy data synthesis benchmarks, using a variety of different samplers, and observe that our optimized schedules outperform previous hand-crafted schedules in almost all experiments. Our method demonstrates the untapped potential of sampling schedule optimization, especially in the few-step synthesis regime.