SiT: Explorando Modelos Generativos Basados en Flujo y Difusión con Transformadores de Interpolación Escalables

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

Resumen

Presentamos los Transformadores de Interpolación Escalables (SiT), una familia de modelos generativos construidos sobre la base de los Transformadores de Difusión (DiT). El marco de interpolación, que permite conectar dos distribuciones de manera más flexible que los modelos de difusión estándar, posibilita un estudio modular de diversas decisiones de diseño que impactan en los modelos generativos basados en transporte dinámico: el uso de aprendizaje en tiempo discreto frente a continuo, la elección del objetivo que el modelo debe aprender, la selección del interpolante que conecta las distribuciones y la implementación de un muestreador determinista o estocástico. Al introducir cuidadosamente los elementos mencionados, SiT supera a DiT de manera uniforme en todos los tamaños de modelos en el benchmark condicional de ImageNet 256x256 utilizando exactamente la misma arquitectura base, número de parámetros y GFLOPs. Al explorar varios coeficientes de difusión, que pueden ajustarse de forma independiente al aprendizaje, SiT alcanza un puntaje 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.