DisCo-Diff: Verbesserung von kontinuierlichen Diffusionsmodellen mit diskreten Latenten

DisCo-Diff: Enhancing Continuous Diffusion Models with Discrete Latents

July 3, 2024

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

cs.AI

Zusammenfassung

Diffusionsmodelle (DMs) haben das generative Lernen revolutioniert. Sie nutzen einen Diffusionsprozess, um Daten in eine einfache Gaußverteilung zu kodieren. Allerdings stellt die Kodierung einer komplexen, potenziell multimodalen Datenverteilung in eine einzige kontinuierliche Gaußverteilung möglicherweise ein unnötig herausforderndes Lernproblem dar. Wir schlagen Discrete-Continuous Latent Variable Diffusion Models (DisCo-Diff) vor, um diese Aufgabe zu vereinfachen, indem wir ergänzende diskrete latente Variablen einführen. Wir erweitern DMs um erlernbare diskrete Latente, die mit einem Encoder inferiert werden, und trainieren DM und Encoder end-to-end. DisCo-Diff ist nicht auf vorab trainierte Netzwerke angewiesen, was das Framework universell anwendbar macht. Die diskreten Latenten vereinfachen das Lernen der komplexen Rausch-zu-Daten-Abbildung des DMs erheblich, indem sie die Krümmung der generativen ODE des DMs reduzieren. Ein zusätzlicher autoregressiver Transformer modelliert die Verteilung der diskreten Latenten, ein einfacher Schritt, da DisCo-Diff nur wenige diskrete Variablen mit kleinen Codebüchern erfordert. Wir validieren DisCo-Diff an Spielzeugdaten, mehreren Bildsynthesetasks sowie molekularer Docking, und stellen fest, dass die Einführung diskreter Latenten die Modellleistung konsistent verbessert. Zum Beispiel erzielt DisCo-Diff state-of-the-art FID-Werte auf klassenbedingten ImageNet-64/128 Datensätzen mit 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.