Allinea i Tuoi Passi: Ottimizzazione delle Pianificazioni di Campionamento nei Modelli di Diffusione

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

April 22, 2024

Autori: Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis

cs.AI

Abstract

I modelli di diffusione (DMs) si sono affermati come l'approccio più avanzato per la modellazione generativa nel dominio visivo e oltre. Un aspetto critico dei DMs è la loro lentezza nel campionamento, che si basa su molte valutazioni sequenziali di funzioni attraverso grandi reti neurali. Il campionamento dai DMs può essere visto come la risoluzione di un'equazione differenziale attraverso un insieme discretizzato di livelli di rumore noto come pianificazione del campionamento. Mentre i lavori precedenti si sono concentrati principalmente sulla derivazione di risolutori efficienti, poca attenzione è stata dedicata alla ricerca di pianificazioni di campionamento ottimali, e l'intera letteratura si affida a euristiche artigianali. In questo lavoro, per la prima volta, proponiamo un approccio generale e principiato per ottimizzare le pianificazioni di campionamento dei DMs al fine di ottenere output di alta qualità, chiamato Align Your Steps. Sfruttiamo metodi del calcolo stocastico e troviamo pianificazioni ottimali specifiche per diversi risolutori, DMs addestrati e dataset. Valutiamo il nostro nuovo approccio su diversi benchmark di sintesi di immagini, video e dati giocattolo 2D, utilizzando una varietà di campionatori diversi, e osserviamo che le nostre pianificazioni ottimizzate superano le pianificazioni artigianali precedenti in quasi tutti gli esperimenti. Il nostro metodo dimostra il potenziale inesplorato dell'ottimizzazione della pianificazione del campionamento, specialmente nel regime di sintesi a pochi passi.

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.