Alinhe Seus Passos: Otimizando Agendamentos de Amostragem em Modelos de Difusão

Align Your Steps: Optimizing Sampling Schedules in Diffusion Models

April 22, 2024

Autores: Amirmojtaba Sabour, Sanja Fidler, Karsten Kreis

cs.AI

Resumo

Modelos de difusão (DMs) estabeleceram-se como a abordagem de modelagem generativa mais avançada no domínio visual e além. Uma desvantagem crucial dos DMs é sua velocidade lenta de amostragem, que depende de muitas avaliações sequenciais de funções através de grandes redes neurais. Amostrar a partir de DMs pode ser visto como resolver uma equação diferencial por meio de um conjunto discretizado de níveis de ruído conhecido como cronograma de amostragem. Embora trabalhos anteriores tenham se concentrado principalmente em derivar solucionadores eficientes, pouca atenção foi dada à busca de cronogramas de amostragem ótimos, e toda a literatura depende de heurísticas manuais. Neste trabalho, pela primeira vez, propomos uma abordagem geral e fundamentada para otimizar os cronogramas de amostragem de DMs visando saídas de alta qualidade, chamada Align Your Steps. Aproveitamos métodos do cálculo estocástico e encontramos cronogramas ótimos específicos para diferentes solucionadores, DMs treinados e conjuntos de dados. Avaliamos nossa nova abordagem em vários benchmarks de síntese de imagens, vídeos e dados bidimensionais simples, utilizando uma variedade de amostradores diferentes, e observamos que nossos cronogramas otimizados superam os cronogramas manuais anteriores em quase todos os experimentos. Nosso método demonstra o potencial inexplorado da otimização de cronogramas de amostragem, especialmente no regime de síntese com poucos passos.

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.