SiT: Explorando Modelos Gerativos Baseados em Fluxo e Difusão com Transformadores de Interpolante Escaláveis

SiT: Exploring Flow and Diffusion-based Generative Models with Scalable Interpolant Transformers

January 16, 2024

Autores: Nanye Ma, Mark Goldstein, Michael S. Albergo, Nicholas M. Boffi, Eric Vanden-Eijnden, Saining Xie

cs.AI

Resumo

Apresentamos os Transformadores Interpolantes Escaláveis (SiT), uma família de modelos generativos construídos sobre a estrutura dos Transformadores de Difusão (DiT). O framework interpolante, que permite conectar duas distribuições de maneira mais flexível do que os modelos de difusão padrão, possibilita um estudo modular de várias escolhas de design que impactam modelos generativos baseados em transporte dinâmico: utilizar aprendizado em tempo discreto versus contínuo, decidir o objetivo que o modelo deve aprender, escolher o interpolante que conecta as distribuições e implementar um amostrador determinístico ou estocástico. Ao introduzir cuidadosamente os elementos mencionados, o SiT supera o DiT de forma uniforme em todos os tamanhos de modelos no benchmark condicional ImageNet 256x256, utilizando exatamente a mesma estrutura, número de parâmetros e GFLOPs. Ao explorar diversos coeficientes de difusão, que podem ser ajustados separadamente do aprendizado, o SiT alcança um score FID-50K de 2,06.

English

We present Scalable Interpolant Transformers (SiT), a family of generative models built on the backbone of Diffusion Transformers (DiT). The interpolant framework, which allows for connecting two distributions in a more flexible way than standard diffusion models, makes possible a modular study of various design choices impacting generative models built on dynamical transport: using discrete vs. continuous time learning, deciding the objective for the model to learn, choosing the interpolant connecting the distributions, and deploying a deterministic or stochastic sampler. By carefully introducing the above ingredients, SiT surpasses DiT uniformly across model sizes on the conditional ImageNet 256x256 benchmark using the exact same backbone, number of parameters, and GFLOPs. By exploring various diffusion coefficients, which can be tuned separately from learning, SiT achieves an FID-50K score of 2.06.