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

초록

확산 모델(Diffusion Models, DMs)은 생성 학습 분야에 혁명을 일으켰습니다. 이 모델은 확산 과정을 활용하여 데이터를 단순한 가우시안 분포로 인코딩합니다. 그러나 복잡하고 잠재적으로 다중 모드를 가진 데이터 분포를 단일 연속 가우시안 분포로 인코딩하는 것은 불필요하게 어려운 학습 문제를 제기합니다. 우리는 이 과제를 단순화하기 위해 이산 잠재 변수를 도입한 이산-연속 잠재 변수 확산 모델(Discrete-Continuous Latent Variable Diffusion Models, DisCo-Diff)을 제안합니다. 우리는 학습 가능한 이산 잠재 변수를 확산 모델에 추가하고, 이를 인코더로 추론하며, 확산 모델과 인코더를 종단 간(end-to-end)으로 학습합니다. DisCo-Diff는 사전 학습된 네트워크에 의존하지 않아 프레임워크를 보편적으로 적용할 수 있습니다. 이산 잠재 변수는 확산 모델의 복잡한 노이즈-투-데이터 매핑 학습을 크게 단순화하며, 이는 확산 모델의 생성적 ODE의 곡률을 줄임으로써 가능합니다. 추가적으로, 자동회귀 트랜스포머는 이산 잠재 변수의 분포를 모델링하는데, 이는 DisCo-Diff가 작은 코드북을 가진 소수의 이산 변수만을 필요로 하기 때문에 간단한 과정입니다. 우리는 DisCo-Diff를 토이 데이터, 여러 이미지 합성 작업, 그리고 분자 도킹에 대해 검증하였고, 이산 잠재 변수를 도입함으로써 모델 성능이 일관적으로 향상됨을 확인했습니다. 예를 들어, DisCo-Diff는 ODE 샘플러를 사용하여 클래스 조건부 ImageNet-64/128 데이터셋에서 최첨단 FID 점수를 달성했습니다.

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.