SiT : Exploration des modèles génératifs basés sur les flux et la diffusion avec des Transformers interpolants scalables

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

January 16, 2024

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

cs.AI

Résumé

Nous présentons les Scalable Interpolant Transformers (SiT), une famille de modèles génératifs construits sur l'architecture des Diffusion Transformers (DiT). Le cadre d'interpolation, qui permet de relier deux distributions de manière plus flexible que les modèles de diffusion standard, rend possible une étude modulaire de divers choix de conception influençant les modèles génératifs basés sur le transport dynamique : l'utilisation d'un apprentissage en temps discret ou continu, la détermination de l'objectif à apprendre par le modèle, le choix de l'interpolant reliant les distributions, et le déploiement d'un échantillonneur déterministe ou stochastique. En introduisant soigneusement ces éléments, SiT surpasse DiT de manière uniforme pour toutes les tailles de modèles sur le benchmark conditionnel ImageNet 256x256, en utilisant exactement la même architecture, le même nombre de paramètres et les mêmes GFLOPs. En explorant divers coefficients de diffusion, qui peuvent être ajustés indépendamment de l'apprentissage, SiT atteint un 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.