DisCo-Diff: 이산 잠재변수를 통한 연속 확산 모델의 성능 향상

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