17 天前

DisCo-Diff:通过离散潜变量增强连续扩散模型

Yilun Xu, Gabriele Corso, Tommi Jaakkola, Arash Vahdat, Karsten Kreis
DisCo-Diff:通过离散潜变量增强连续扩散模型
摘要

扩散模型(Diffusion Models, DMs)已彻底革新了生成式学习领域。它们通过扩散过程将数据编码为简单的高斯分布。然而,将复杂且可能具有多模态特性的数据分布编码为单一连续的高斯分布,无疑构成了一个过于艰巨的学习挑战。为此,我们提出离散-连续潜变量扩散模型(Discrete-Continuous Latent Variable Diffusion Models, DisCo-Diff),通过引入互补的离散潜变量来简化该任务。DisCo-Diff在传统扩散模型的基础上引入了可学习的离散潜变量,这些变量由编码器推断得到,并实现扩散模型与编码器的端到端联合训练。该框架不依赖预训练网络,因而具有广泛的适用性。离散潜变量显著简化了扩散模型中从噪声到数据映射的学习过程,其核心机制在于降低了生成过程所对应常微分方程(ODE)的曲率。此外,我们采用一个自回归Transformer模型来建模离散潜变量的分布,这一设计极为简便,因为DisCo-Diff仅需少量离散变量及小规模的码本(codebook)即可实现高效建模。我们在合成数据、多种图像生成任务以及分子对接任务上对DisCo-Diff进行了验证,结果表明,引入离散潜变量能持续提升模型性能。例如,在使用ODE采样器的类别条件ImageNet-64/128数据集上,DisCo-Diff取得了当前最优的FID得分,展现出卓越的生成质量与稳定性。