DisCo-Diff: Melhorando Modelos de Difusão Contínua com 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

Resumo

Os modelos de difusão (DMs) revolucionaram a aprendizagem generativa. Eles utilizam um processo de difusão para codificar dados em uma distribuição Gaussiana simples. No entanto, codificar uma distribuição de dados complexa e potencialmente multimodal em uma única distribuição Gaussiana contínua representa, possivelmente, um problema de aprendizagem desnecessariamente desafiador. Propomos Modelos de Difusão de Variáveis Latentes Discretas-Contínuas (DisCo-Diff) para simplificar essa tarefa, introduzindo variáveis latentes discretas complementares. Nós aprimoramos os DMs com variáveis latentes discretas aprendíveis, inferidas com um codificador, e treinamos o DM e o codificador de ponta a ponta. O DisCo-Diff não depende de redes pré-treinadas, tornando o framework universalmente aplicável. As variáveis latentes discretas simplificam significativamente a aprendizagem do mapeamento de ruído para dados do DM, reduzindo a curvatura da EDO generativa do DM. Um transformador autoregressivo adicional modela a distribuição das variáveis latentes discretas, um passo simples, pois o DisCo-Diff requer apenas algumas variáveis discretas com pequenos codebooks. Validamos o DisCo-Diff em dados de brinquedo, várias tarefas de síntese de imagens, bem como acoplamento molecular, e descobrimos que a introdução de variáveis latentes discretas melhora consistentemente o desempenho do modelo. Por exemplo, o DisCo-Diff alcança pontuações FID de última geração em conjuntos de dados ImageNet-64/128 condicionados por classe com amostrador de EDO.

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.