DisCo-Diff: Улучшение непрерывных моделей диффузии с дискретными латентными переменными

DisCo-Diff: Enhancing Continuous Diffusion Models with Discrete Latents

July 3, 2024

Авторы: Yilun Xu, Gabriele Corso, Tommi Jaakkola, Arash Vahdat, Karsten Kreis

cs.AI

Аннотация

Модели диффузии (DM) революционизировали генеративное обучение. Они используют диффузионный процесс для кодирования данных в простое гауссовское распределение. Однако кодирование сложного, потенциально мультимодального распределения данных в одно непрерывное гауссовское распределение, пожалуй, представляет собой излишне сложную задачу обучения. Мы предлагаем модели диффузии с дискретно-непрерывными скрытыми переменными (DisCo-Diff), чтобы упростить эту задачу путем введения дополнительных дискретных скрытых переменных. Мы дополняем DM обучаемыми дискретными скрытыми переменными, которые выводятся с помощью кодировщика, и обучаем DM и кодировщик совместно. DisCo-Diff не зависит от предварительно обученных сетей, что делает эту структуру универсально применимой. Дискретные скрытые переменные значительно упрощают обучение сложного отображения шума в данные DM путем уменьшения кривизны генеративного ОДУ DM. Дополнительный авторегрессионный трансформер моделирует распределение дискретных скрытых переменных, что является простым шагом, поскольку DisCo-Diff требует лишь нескольких дискретных переменных с небольшими кодовыми книгами. Мы проверяем DisCo-Diff на игрушечных данных, нескольких задачах синтеза изображений, а также молекулярном докинге и обнаруживаем, что введение дискретных скрытых переменных последовательно улучшает производительность модели. Например, DisCo-Diff достигает передовых показателей FID на класс-условных наборах данных ImageNet-64/128 с выборкой ОДУ.

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.