DisCo-Diff : Amélioration des modèles de diffusion continue par l'intégration de latents discrets

DisCo-Diff: Enhancing Continuous Diffusion Models with Discrete Latents

July 3, 2024

Auteurs: Yilun Xu, Gabriele Corso, Tommi Jaakkola, Arash Vahdat, Karsten Kreis

cs.AI

Résumé

Les modèles de diffusion (DMs) ont révolutionné l'apprentissage génératif. Ils utilisent un processus de diffusion pour encoder les données dans une distribution gaussienne simple. Cependant, encoder une distribution de données complexe, potentiellement multimodale, dans une seule distribution gaussienne continue représente sans doute un problème d'apprentissage inutilement difficile. Nous proposons les modèles de diffusion à variables latentes discrètes-continues (DisCo-Diff) pour simplifier cette tâche en introduisant des variables latentes discrètes complémentaires. Nous enrichissons les DMs avec des variables latentes discrètes apprenables, inférées par un encodeur, et entraînons le DM et l'encodeur de bout en bout. DisCo-Diff ne repose pas sur des réseaux pré-entraînés, ce qui rend le cadre universellement applicable. Les variables latentes discrètes simplifient considérablement l'apprentissage de la transformation complexe bruit-vers-données du DM en réduisant la courbure de l'équation différentielle ordinaire (ODE) générative du DM. Un transformeur autorégressif supplémentaire modélise la distribution des variables latentes discrètes, une étape simple car DisCo-Diff ne nécessite que peu de variables discrètes avec des petits codebooks. Nous validons DisCo-Diff sur des données jouets, plusieurs tâches de synthèse d'images ainsi que sur le docking moléculaire, et constatons que l'introduction de variables latentes discrètes améliore systématiquement les performances du modèle. Par exemple, DisCo-Diff atteint des scores FID de pointe sur les ensembles de données ImageNet-64/128 conditionnés par classe avec un échantillonneur ODE.

English

Diffusion models (DMs) have revolutionized generative learning. They utilize a diffusion process to encode data into a simple Gaussian distribution. However, encoding a complex, potentially multimodal data distribution into a single continuous Gaussian distribution arguably represents an unnecessarily challenging learning problem. We propose Discrete-Continuous Latent Variable Diffusion Models (DisCo-Diff) to simplify this task by introducing complementary discrete latent variables. We augment DMs with learnable discrete latents, inferred with an encoder, and train DM and encoder end-to-end. DisCo-Diff does not rely on pre-trained networks, making the framework universally applicable. The discrete latents significantly simplify learning the DM's complex noise-to-data mapping by reducing the curvature of the DM's generative ODE. An additional autoregressive transformer models the distribution of the discrete latents, a simple step because DisCo-Diff requires only few discrete variables with small codebooks. We validate DisCo-Diff on toy data, several image synthesis tasks as well as molecular docking, and find that introducing discrete latents consistently improves model performance. For example, DisCo-Diff achieves state-of-the-art FID scores on class-conditioned ImageNet-64/128 datasets with ODE sampler.