DisCo-Diff: Miglioramento dei Modelli di Diffusione Continua con Latenti Discreti

DisCo-Diff: Enhancing Continuous Diffusion Models with Discrete Latents

July 3, 2024

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

cs.AI

Abstract

I modelli di diffusione (DMs) hanno rivoluzionato l'apprendimento generativo. Utilizzano un processo di diffusione per codificare i dati in una semplice distribuzione gaussiana. Tuttavia, codificare una distribuzione di dati complessa e potenzialmente multimodale in un'unica distribuzione gaussiana continua rappresenta, senza dubbio, un problema di apprendimento eccessivamente impegnativo. Proponiamo i Modelli di Diffusione con Variabili Latenti Discrete-Continue (DisCo-Diff) per semplificare questo compito introducendo variabili latenti discrete complementari. Aumentiamo i DMs con variabili latenti discrete apprendibili, inferite tramite un encoder, e addestriamo il DM e l'encoder end-to-end. DisCo-Diff non si basa su reti pre-addestrate, rendendo il framework universalmente applicabile. Le variabili latenti discrete semplificano significativamente l'apprendimento della complessa mappatura rumore-dati del DM riducendo la curvatura dell'ODE generativo del DM. Un ulteriore trasformatore autoregressivo modella la distribuzione delle variabili latenti discrete, un passo semplice poiché DisCo-Diff richiede solo poche variabili discrete con piccoli codebook. Validiamo DisCo-Diff su dati giocattolo, diverse attività di sintesi di immagini e docking molecolare, e scopriamo che l'introduzione di variabili latenti discrete migliora costantemente le prestazioni del modello. Ad esempio, DisCo-Diff raggiunge punteggi FID all'avanguardia sui dataset ImageNet-64/128 con condizionamento di classe utilizzando un campionatore 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.