DisCo-Diff: Mejorando Modelos de Difusión Continua con Latentes Discretos

DisCo-Diff: Enhancing Continuous Diffusion Models with Discrete Latents

July 3, 2024

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

cs.AI

Resumen

Los modelos de difusión (DMs) han revolucionado el aprendizaje generativo. Utilizan un proceso de difusión para codificar datos en una distribución Gaussiana simple. Sin embargo, codificar una distribución de datos compleja y potencialmente multimodal en una única distribución Gaussiana continua representa posiblemente un problema de aprendizaje innecesariamente desafiante. Proponemos Modelos de Difusión de Variables Latentes Discretas-Continuas (DisCo-Diff) para simplificar esta tarea mediante la introducción de variables latentes discretas complementarias. Ampliamos los DMs con variables latentes discretas aprendibles, inferidas con un codificador, y entrenamos el DM y el codificador de manera integral. DisCo-Diff no depende de redes pre-entrenadas, lo que hace que el marco sea universalmente aplicable. Las variables latentes discretas simplifican significativamente el aprendizaje del mapeo de ruido a datos del DM al reducir la curvatura de la Ecuación Diferencial Ordinaria (ODE) generativa del DM. Un transformador autoregresivo adicional modela la distribución de las variables latentes discretas, un paso sencillo porque DisCo-Diff solo requiere unas pocas variables discretas con diccionarios de códigos pequeños. Validamos DisCo-Diff en datos de juguete, varias tareas de síntesis de imágenes y acoplamiento molecular, y encontramos que la introducción de variables latentes discretas mejora consistentemente el rendimiento del modelo. Por ejemplo, DisCo-Diff logra puntuaciones FID de vanguardia en conjuntos de datos de ImageNet-64/128 condicionados por clase con muestreador 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.