Alinea tus pasos: Optimización de los cronogramas de muestreo en modelos de difusión

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

April 22, 2024

Autores: Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis

cs.AI

Resumen

Los modelos de difusión (DMs) se han consolidado como el enfoque de modelado generativo más avanzado en el dominio visual y más allá. Un inconveniente crucial de los DMs es su lenta velocidad de muestreo, que depende de muchas evaluaciones secuenciales de funciones a través de grandes redes neuronales. El muestreo de los DMs puede verse como la resolución de una ecuación diferencial mediante un conjunto discretizado de niveles de ruido conocidos como el programa de muestreo. Si bien trabajos anteriores se centraron principalmente en derivar solucionadores eficientes, se ha prestado poca atención a encontrar programas de muestreo óptimos, y toda la literatura se basa en heurísticas diseñadas manualmente. En este trabajo, por primera vez, proponemos un enfoque general y fundamentado para optimizar los programas de muestreo de los DMs con el fin de obtener resultados de alta calidad, denominado Align Your Steps. Aprovechamos métodos del cálculo estocástico y encontramos programas óptimos específicos para diferentes solucionadores, DMs entrenados y conjuntos de datos. Evaluamos nuestro nuevo enfoque en varios puntos de referencia de síntesis de imágenes, videos y datos 2D de juguete, utilizando una variedad de muestreadores diferentes, y observamos que nuestros programas optimizados superan a los programas diseñados manualmente en casi todos los experimentos. Nuestro método demuestra el potencial no explotado de la optimización del programa de muestreo, especialmente en el régimen de síntesis con pocos pasos.

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.