Richten Sie Ihre Schritte aus: Optimierung von Probenahmeplänen in Diffusionsmodellen

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

April 22, 2024

papers.authors: Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis

cs.AI

papers.abstract

Diffusionsmodelle (DMs) haben sich als der modernste generative Modellierungsansatz im visuellen Bereich und darüber hinaus etabliert. Ein entscheidender Nachteil von DMs ist ihre langsame Abtastgeschwindigkeit, die auf vielen aufeinanderfolgenden Funktionsauswertungen durch große neuronale Netzwerke beruht. Das Abtasten von DMs kann als Lösung einer Differentialgleichung durch eine diskretisierte Reihe von Rauschniveaus, bekannt als das Abtastschema, betrachtet werden. Während vergangene Arbeiten hauptsächlich darauf abzielten, effiziente Solver abzuleiten, wurde wenig Aufmerksamkeit darauf verwendet, optimale Abtastpläne zu finden, und die gesamte Literatur stützt sich auf handgefertigte Heuristiken. In dieser Arbeit schlagen wir erstmals einen allgemeinen und prinzipiengeleiteten Ansatz zur Optimierung der Abtastpläne von DMs für hochwertige Ausgaben vor, genannt Align Your Steps. Wir nutzen Methoden aus der stochastischen Analysis und finden optimale Pläne, die spezifisch für verschiedene Solver, trainierte DMs und Datensätze sind. Wir evaluieren unseren neuartigen Ansatz an mehreren Bild-, Video- sowie 2D-Spielzeugdatensynthese-Benchmarks unter Verwendung verschiedener Sampler und stellen fest, dass unsere optimierten Pläne in fast allen Experimenten die bisherigen handgefertigten Pläne übertreffen. Unsere Methode zeigt das ungenutzte Potenzial der Optimierung von Abtastplänen auf, insbesondere im Regime der Synthese mit wenigen Schritten.

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.