扩散模型爆火,这是首篇综述与Github论文分类汇总

机器之心专栏

作者:杨灵

本综述(Diffusion Models: A Comprehensive Survey of Methods and Applications)来自加州大学 & Google Research 的 Ming-Hsuan Yang、北京大学崔斌实验室以及 CMU、UCLA、蒙特利尔 Mila 研究院等众研究团队,首次对现有的扩散生成模型(diffusion model)进行了全面的总结分析,从 diffusion model 算法细化分类、和其他五大生成模型的关联以及在七大领域中的应用等方面展开,最后提出了 diffusion model 的现有 limitation 和未来的发展方向。

文章链接:https://arxiv.org/abs/2209.00796

论文分类汇总 github 链接:https://github.com/YangLing0818/Diffusion-Models-Papers-Survey-Taxonomy

介绍

扩散模型(diffusion models)是深度生成模型中新的 SOTA。扩散模型在图片生成任务中超越了原 SOTA:GAN,并且在诸多应用领域都有出色的表现,如计算机视觉,NLP、波形信号处理、多模态建模、分子图建模、时间序列建模、对抗性净化等。此外,扩散模型与其他研究领域有着密切的联系,如稳健学习、表示学习、强化学习。然而,原始的扩散模型也有缺点,它的采样速度慢,通常需要数千个评估步骤才能抽取一个样本;它的最大似然估计无法和基于似然的模型相比;它泛化到各种数据类型的能力较差。如今很多研究已经从实际应用的角度解决上述限制做出了许多努力,或从理论角度对模型能力进行了分析。

然而,现在缺乏对扩散模型从算法到应用的最新进展的系统回顾。为了反映这一快速发展领域的进展,我们对扩散模型进行了首个全面综述。我们设想我们的工作将阐明扩散模型的设计考虑和先进方法,展示其在不同领域的应用,并指出未来的研究方向。此综述的概要如下图所示:

尽管 diffusion model 在各类任务中都有着优秀的表现,它仍还有自己的缺点,并有诸多研究对 diffusion model 进行了改善。为了系统地阐明 diffusion model 的研究进展,我们总结了原始扩散模型的三个主要缺点,采样速度慢,最大化似然差、数据泛化能力弱,并提出将的 diffusion models 改进研究分为对应的三类:采样速度提升、最大似然增强和数据泛化增强。我们首先说明改善的动机,再根据方法的特性将每个改进方向的研究进一步细化分类,从而清楚的展现方法之间的联系与区别。在此我们仅选取部分重要方法为例, 我们的工作中对每类方法都做了详细的介绍,内容如图所示:

在分析完三类扩散模型后,我们将介绍其他的五种生成模型 GAN,VAE,Autoregressive model, Normalizing flow, Energy-based model。考虑到扩散模型的优良性质,研究者们已经根据其特性将 diffusion model 与其他生成模型结合,所以为了进一步展现 diffusion model 的特点和改进工作,我们详细地介绍了 diffusion model 和其他生成模型的结合的工作并阐明了在原始生成模型上的改进之处。Diffusion model 在诸多领域都有着优异的表现,并且考虑到不同领域的应用中 diffusion model 产生了不同的变形,我们系统地介绍了 diffusion model 的应用研究,其中包含如下领域:计算机视觉,NLP、波形信号处理、多模态建模、分子图建模、时间序列建模、对抗性净化。对于每个任务,我们定义了该任务并介绍利用扩散模型处理任务的工作,我们将本项工作的主要贡献总结如下:

扩散模型基础

生成式建模的一个核心问题是模型的灵活性和可计算性之间的权衡。扩散模型的基本思想是正向扩散过程来系统地扰动数据中的分布,然后通过学习反向扩散过程恢复数据的分布,这样就了产生一个高度灵活且易于计算的生成模型。

A.Denoising Diffusion Probabilistic Models(DDPM)

一个 DDPM 由两个参数化马尔可夫链组成,并使用变分推断以在有限时间后生成与原始数据分布一致的样本。前向链的作用是扰动数据,它根据预先设计的噪声进度向数据逐渐加入高斯噪声,直到数据的分布趋于先验分布,即标准高斯分布。反向链从给定的先验开始并使用参数化的高斯转换核,学习逐步恢复原数据分布。用

表示原始数据及其分布, 则前向链的分布是可由下式表达:

这说明前向链是马尔可夫过程,x_t 是加入 t 步噪音后的样本,β_t 是事先给定的控制噪声进度的参数。当

趋于 1 时,x_T 可以近似认为服从标准高斯分布。当β_t 很小时,逆向过程的转移核可以近似认为也是高斯的:

我们可以将变分下界作为损失函数进行学习:

B.Score-Based Generative Models(SGM)

上述 DDPM 可以视作 SGM 的离散形式。SGM 构造一个随机微分方程(SDE)来平滑的扰乱数据分布,将原始数据分布转化到已知的先验分布:

和一个相应的逆向 SDE,来将先验分布变换回原始数据分布:

因此,要逆转扩散过程并生成数据,我们需要的唯一信息就是在每个时间点的分数函数。利用 score-matching 的技巧我们可以通过如下损失函数来学习分数函数:

对两种方法的进一步介绍和两者关系的介绍请参见我们的文章。

原始扩散模型的三个主要缺点,采样速度慢,最大化似然差、数据泛化能力弱。最近许多研究都在解决这些缺点,因此我们将改进的扩散模型分为三类:采样速度提升、最大似然增强和数据泛化增强。在接下来的三、四、五节我们将对这三类模型进行详细的介绍。

采样加速方法

在应用时,为了让新样本的质量达到最佳,扩散模型往往需要进行成千上万步计算来获取一个新样本。这限制了 diffusion model 的实际应用价值,因为在实际应用时,我们往往需要产生大量的新样本,来为下一步处理提供材料。研究者们在提高 diffusion model 采样速度上进行了大量的研究。我们对这些研究进行了详细的阐述。我们将其细化分类为三种方法:Discretization Optimization,Non-Markovian Process,Partial Sampling。

A.Discretization Optimization 方法优化求解 diffusion SDE 的方法。因为现实中求解复杂 SDE 只能使用离散解来逼近真正的解,所以该类方法试图优化 SDE 的离散化方法,在保证样本质量的同时减少离散步数。SGM 提出了一个通用的方法来求解逆向过程,即对前向和后向过程采取相同的离散方法。如果给定了前向 SDE 的离散方式:

那么我们就可以以相同的方式离散化逆向 SDE:

这种方法比朴素 DDPM 效果略好一点。进一步,SGM 向 SDE 求解器中加入了一个矫正器,从而让每一步生成的样本都有正确的分布。在求解的每一步,求解器给出一个样本后,矫正器都使用马尔可夫链蒙特卡罗方法来矫正刚生成的样本的分布。实验表明向求解器中加入矫正器比直接增加求解器的步数效率更高。

B.Non-Markovian Process方法突破了原有 Markovian Process 的限制,其逆过程的每一步可以依赖更多以往的样本来进行预测新样本,所以在步长较大时也能做出较好的预测,从而加速采样过程。其中主要的工作 DDIM,不再假设前向过程是马尔可夫过程,而是服从如下分布:

DDIM 的采样过程可以视为离散化的神经常微分方程,其采样过程更高效,并且支持样本的内插。进一步的研究发现 DDIM 可以视作流形上扩散模型 PNDM 的特例。

C.Partial Sampling方法通过在 generation process 中忽略一部分的时间节点,而只使用剩下的时间节点来生成样本,直接减少了采样时间。例如,Progressive Distillation 从训练好的扩散模型中蒸馏出效率更高的扩散模型。对于训练好的一个扩散模型,Progressive Distillation 会从新训练一个扩散模型,使新的扩散模型的一步对应于训练好的扩散模型的两步,这样新模型就可以省去老模型一半的采样过程。具体算法如下:

不断循环这个蒸馏过程就能让采样步骤指数级下降。

最大似然估计加强

扩散模型在最大似然估计的表现差于基于似然函数的生成模型,但最大化似然估计在诸多应用场景都有重要意义,比如图片压缩, 半监督学习, 对抗性净化。由于对数似然难以直接计算,研究主要集中在优化和分析变分下界(VLB)。我们对提高扩散模型最大似然估计的模型进行了详细的阐述。我们将其细化分类为三类方法:Objectives Designing,Noise Schedule Optimization,Learnable Reverse Variance。

A.Objectives Designing方法利用扩散 SDE 推倒出生成数据的对数似然与分数函数匹配的损失函数的关系。这样通过适当设计损失函数,就可以最大化 VLB 和对数似然。Song et al. 证明了可以设计损失函数的权重函数,使得 plug-in reverse SDE 生成样本的似然函数值小于等于损失函数值,即损失函数是似然函数的上界。分数函数拟合的损失函数如下:

我们只需将权重函数λ(t)设为扩散系数 g(t)即可让损失函数成为似然函数的 VLB,即:

B.Noise Schedule Optimization通过设计或学习前向过程的噪声进度来增大 VLB。VDM 证明了当离散步数接近无穷时,损失函数完全由信噪比函数 SNR(t)的端点决定:

那么在离散步数接近无穷时,可以通过学习信噪比函数 SNR(t)的端点最优化 VLB,而通过学习信噪比函数中间部分的函数值来实现模型其他方面的改进。

C.Learnable Reverse Variance方法学习反向过程的方差,从而较少拟合误差,可以有效地最大化 VLB。Analytic-DPM 证明,在 DDPM 和 DDIM 中存在反向过程中的最优期望和方差:

使用上述公式和训练好的分数函数,在给定前向过程的条件下,最优的 VLB 可以近似达到。

数据泛化增强

扩散模型假设数据存在于欧几里得空间,即具有平面几何形状的流形,并添加高斯噪声将不可避免地将数据转换为连续状态空间,所以扩散模型最初只能处理图片等连续性数据,直接应用离散数据或其他数据类型的效果较差。这限制了扩散模型的应用场景。数个研究工作将扩散模型推广到适用于其他数据类型的模型,我们对这些方法进行了详细地阐释。我们将其细化分类为两类方法:Feature Space Unification,Data-Dependent Transition Kernels。

A.Feature Space Unification方法将数据转化到统一形式的 latent space,然后再 latent space 上进行扩散。LSGM 提出将数据通过 VAE 框架先转换到连续的 latent space 上后再在其上进行扩散。这个方法的难点在于如何同时训练 VAE 和扩散模型。LSGM 表明由于潜在先验是 intractable 的,分数匹配损失不再适用。LSGM 直接使用 VAE 中传统的损失函数 ELBO 作为损失函数,并导出了 ELBO 和分数匹配的关系:

该式在忽略常数的意义下成立。通过参数化扩散过程中样本的分数函数,LSGM 可以高效的学习和优化 ELBO。

B.Data-Dependent Transition Kernels方法根据数据类型的特点设计 diffusion process 中的 transition kernels,使扩散模型可以直接应用于特定的数据类型。D3PM 为离散型数据设计了 transition kernel,可以设为 lazy random-walk,absorbing state 等。GEODIFF 为 3D 分子图数据设计了平移 - 旋转不变的图神经网络,并且证明了具有不变性的初分布和 transition kernel 可以导出具有不变性的边缘分布。假设 T 是一个平移 - 旋转变换,如:

那么生成的样本分布也有平移 - 旋转不变性:

和其他生成模型的联系

在下面的每个小节中,我们首先介绍其他五类重要的生成模型,并分析它们的优势和局限性。然后我们介绍了扩散模型是如何与它们联系起来的,并说明通过结合扩散模型来改进这些生成模型。VAE,GAN,Autoregressive model, Normalizing flow, Energy-based model 和扩散模型的联系如下图所示:

扩散模型的应用

在本节中,我们分别介绍了扩散模型在计算机视觉、自然语言处理、波形信号处理、多模态学习、分子图生成、时间序列以及对抗学习等七大应用方向中的应用,并对每类应用中的方法进行了细分并解析。例如在计算机视觉中可以用 diffusion model 进行图像补全修复(RePaint):

在多模态任务中可以用 diffusion model 进行文本到图像的生成(GLIDE):

还可以在分子图生成中用 diffusion model 进行药物分子和蛋白质分子的生成(GeoDiff):

应用分类汇总见表:

未来研究方向