DisCo-Diff: Verbetering van Continue Diffusiemodellen met Discrete Latenten

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

Samenvatting

Diffusiemodellen (DMs) hebben generatief leren gerevolutioneerd. Ze maken gebruik van een diffusieproces om gegevens te coderen in een eenvoudige Gaussische verdeling. Het coderen van een complexe, potentieel multimodale gegevensverdeling in een enkele continue Gaussische verdeling vormt echter een mogelijk onnodig uitdagend leerprobleem. Wij stellen Discrete-Continue Latente Variabele Diffusiemodellen (DisCo-Diff) voor om deze taak te vereenvoudigen door aanvullende discrete latente variabelen te introduceren. We verrijken DMs met leerbare discrete latente variabelen, die worden afgeleid met een encoder, en trainen het DM en de encoder end-to-end. DisCo-Diff is niet afhankelijk van vooraf getrainde netwerken, waardoor het framework universeel toepasbaar is. De discrete latente variabelen vereenvoudigen het leren van de complexe ruis-naar-gegevens mapping van het DM aanzienlijk door de kromming van de generatieve ODE van het DM te verminderen. Een aanvullende autoregressieve transformer modelleert de verdeling van de discrete latente variabelen, een eenvoudige stap omdat DisCo-Diff slechts weinig discrete variabelen met kleine codeboeken vereist. We valideren DisCo-Diff op speelgoedgegevens, verschillende beeld-synthesetaken en moleculair docken, en ontdekken dat de introductie van discrete latente variabelen de modelprestaties consistent verbetert. DisCo-Diff behaalt bijvoorbeeld state-of-the-art FID-scores op klasse-geconditioneerde ImageNet-64/128 datasets met een ODE-sampler.

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.