查看原文
其他

深度生成网络新思路:扩散概率模型

壁仞科技研究院 壁仞科技研究院 2021-09-19



摘要

自从Ian Goodfellow在2014年推出生成对抗网络(Generative Adversarial Network,GAN),生成模型获得了极大关注并迅速席卷全球,各种变体XX-GAN层出不穷,对抗训练也被Yann LeCun誉为“有史以来最酷的事情”,风头完全盖过了几乎在同一时期提出的生成模型另一分支:变分自编码器(Variational Auto-Encoder,VAE)。随着研究的推进和各种模型的尝试,VAE近几年逐渐崛起,其中“扩散模型”(Diffusion Model)更是被OpenAI的研究团队发现在图像合成方面有着超越GAN的采样质量,在无条件图像生成等任务上达到了当前最佳水平,也在后续的生成模型研究中具有广泛的应用前景和重要的启发作用。本文梳理了近年来扩散模型的发展轨迹,并基于所介绍的扩散模型的数学逻辑,探讨今后的研究方向。




介绍

机器学习(ML)发展至今,对数据的依赖日益增长,基于数据的决策大多都是以概率形式表达,希望在众多候选项中取最优解。从数学角度来看,最常见的概率模型有以下三种:1)判定模型,根据属性变量Z预测标记X,也就是求概率P(X|Z);2)有向图生成模型(贝叶斯模型),从观测值X中反推属性(隐变量)Z,即求概率P(Z|X);3)无向图生成模型(马尔可夫模型),对于隐变量Z和观测值X,求联合分布P(X,Z)然后采样。GAN和VAE作为生成模型的两大分支,目标是基本一致的:希望构建一个从隐变量Z到目标数据X的映射关系,例如最理想的情况是X=f(Z),就可以用Z来描述X。本质上来说,就是进行概率分布之间的变换。我们既希望模型能够对客观环境具备灵活的适应能力(flexible),也希望在数学上易于表达或计算上易于处理(analytically or computationally tractable),如何兼顾鱼与熊掌一直都是ML的概率模型的一大挑战。非参数高斯混合模型(non-parametric Gaussian mixture model)可以视作从tractable到flexible的一种平滑过渡,因为它即可代表具备优美数学表达的单一高斯分布(tractable,易于模拟),也可代表无穷个高斯分布的叠加(flexible,可以模拟任意的数据结构)。

基于这样的思路,基于非平衡态热力学的深度无监督学习[1]首先开发了扩散模型,将这种分布之间的变换与演化,抽象模拟成了类似于随机行走(Random Walk)或布朗运动(Brownian Motion)的一个马尔可夫链(Markov Chain,MC):从初始状态也是观测值的开始,一步步添加高斯噪音,直到步数T足够多之后,噪音将原始输入信号完全淹没,得到不同维度间相互独立的多维正态分布。由于每一步添加噪音都是tractable的概率,所以整体tractable,这样就得到了一条从输入信号到纯高斯噪声的MC。将此思路反过来,从纯粹的高斯噪声出发,如果我们能够对应的找到每一步t的噪音并将其从中移除,我们就可以从一堆噪音中恢复出最初始的真实信号。从的过程称为前向过程或扩散过程,去噪的逆过程就称为反向过程或生成过程。扩散模型要做的,就是针对该MC的扩散过程进行训练,然后通过反向过程进行解码,达到生成符合目标分布数据的目的。而如果以作为中间层将前向过程和后向过程拼接在一起,就构成了一个VAE的编码-解码网络结构。



概念与发展

去噪扩散概率模型(Denoising Diffusion Probabilistic Models,DDPM)[2]

DDPM将噪声形成的过程看作是一个多步的扩散过程。长度为T的正向过程的整个扩散轨迹若记为,每一步t的扩散行为都是添加高斯噪音,也就是乘上一个的条件概率分布(这里的β是控制扩散速率的一个参数,也称为diffusion schedule,每个时刻t都有自己的β)。反向过程可以写成类似的概率连乘表示,其中每一步都是高斯分布,但它的均值方差需要训练模型学习拟合。

由于在训练过程中,我们只有目标数据分布和高斯噪声,无法得知相关信息,因此无法直接用于训练。为解决该问题,DDPM提出了的近似公式:

即每一步的都可以看成是初始值和标准高斯噪音的一个线性组合,如果可以学习到各自的系数,就可以将每一步t中的中抽取出来,起到去噪的效果,本质上是训练模型从时刻t学习有效信息,建立信息到噪声的一个映射,再在反向过程中根据噪声值生成真实信号。对更多细节感兴趣的读者可以参阅原文[2]。

DDPM为了简便,对模型进行了一些简化,例如将前向过程的β设为超参而不进行学习,仅在0和T之间线性增长;目标函数也丢弃了一些近似项,并简化了保留项的权重。

DDPM 从噪音中恢复/生成有意义图像的效果见图1。此外,DDPM可以对隐空间中源图像进行插值,即先通过编码器q对图像,映射到对应的,,再用反向过程的p将,的插值进行解码映射到图像空间,实际效果见图2。


图1 渐进式生成(progressive generation)效果,摘自[2]

图2 插值效果,摘自[2]

 

去噪扩散隐式模型(Denoising Diffusion Implicit Models,DDIM)[3]

由DDPM可见,步数T是个重要参数,只有当T足够大,才能足够接近标准正态分布,但大T值会带来计算负担和瓶颈,因为整个过程无法并行,必须顺序迭代。DDIM 最大的创新在于观察到DDPM的目标函数经过简化后仅仅依赖于边际分布而无联合分布,且很多inference分布的联合分布具备同样的边际分布,所以尝试将非马尔可夫链的情况推导至同样的目标函数,并在DDPM的概率密度q的基础上改进设计了。给定有噪音的观测值时,可以先预测,再通过以及条件分布求得的采样,这样每一步对噪音预测的时候就可以不依赖于的信息。

DDIM提出了几个有趣的发现和证明:

- 证明了参数σ对应不同的生成过程,但共用同一个噪音模型,所以不必对不同参数设置进行重复训练。

- 当σ取某特殊正值的时候,前向过程满足MC性质,就是DDPM的情况;若σ=0时,整个前向过程变成确定性过程,模型就成了隐式概率模型(这也是DDIM名称的由来)。

- 模型效果随着步数T的增多而变好的特点比较显见,但在相同步数的情况下,参数σ越小,模型效果往往越好。结合前两条,意味着DDIM可以在相同采样步数的情况下取得比DDPM更优的结果,或者更少的训练步骤取得同样效果,训练效率更高。

- 最后,由于前向过程有 T 步,所以后向过程同样是 T 步;但由于模型只要固定了之后,本质上不依赖于特定的前向过程,所以DDIM可以采取远低于 T 步的前向过程(例如选取一个长度为的子序列),于是反向的生成过程采样步数随之降低,起到提高采样效率的结果。

 

改进的DDPM(Improved Diffusion Models,IDM)[4][5]

对数似然是生成模型中广泛使用的评价指标,针对对数似然的优化通常可以促使生成模型学习捕捉数据分布的所有特征信息。此外很多工作显示,对数似然方面的一点小改进可能会显著提高采样质量和特征表征的学习效果。DDPM 没有达到其他基于似然逻辑的模型的对数似然水平,不禁令人怀疑这是否暗藏着某些缺陷——例如数据信息覆盖度不足,所以IDM从对数似然入手做了一些尝试与改进,在不牺牲采样质量(例如 FID 得分没有损失太多)的前提下提高了对数似然,最后跟 GAN 在精确(precision)和召回(recall)方面进行了比较,发现改进的模型可以学习到分布的(几乎)全部 modes,样本多样性很好。

IDM关于DDPM主要做了如下的改进与调整:

1. 首先,从公式(1)中可以发现,β取0和取1代表着模型的两个极端:

- β=0,条件概率密度变成Delta函数,也就是确定性函数——的值,且方差为0,对应着反向过程熵的下界。

- β=1,条件概率密度变成标准正态分布,对应着反向过程熵的上界。

但很奇怪的是对两种对立的极端情况DDPM却取得了几乎相同的采样结果,所以值得思考为什么这个选择没有影响到采样。IDM将方差参数化,去尝试学习

2.修改目标函数,IDM认为DDPM简化的目标函数可能会损失部分信息,所以在DDPM的优化目标基础上添加了之前未简化版本作为尾项(但配以0.001的权重),目的是在保留DDPM简化项作为主导的同时,适当保留原始项的信息。

3.DDPM 中扩散系数β的线性变化逻辑虽然适用于高分辨率对图像,但对小图像(64x64 和 32x32)却并非最优。所以IDM将其改为了cosine逻辑(准确来说是),这样可以在训练中期对公式(1)中的依旧保持接近线性的衰减而非出现线性逻辑的骤降情况。实验发现在达到相同的终止状态下,线性逻辑前期会快速“恶化”,而余弦逻辑则相对较平滑地从真实信号过渡至噪音(见图3)。

4.重要度采样(Importance Sampling)。DDPM中对时间t进行了均匀采样,但由于DDPM的目标函数原始表达式中对不同的t项拥有不同的量级,可能会引入计划外的噪音,所以IDM加入了重要度采样来进行调整。

5.与DDIM类似,IDM也对步长T进行了抽取长度为S的子序列的操作。

6.增加注意力机制(attention)。实验发现增加注意力头的数量或减少每个头的通道数都可以改善FID。

7.使用BigGAN的残差模块(residual module)对激活进行上采样和下采样。


图3 cosine逻辑(下)对比线性逻辑(上)的噪音效果,摘自[4]

 

改进后的IDM在所有任务中都得到了最优的FID评分,绝大多数任务上也都得到了最优的sFID评分;模型结构在LSUN和ImageNet 64×64上可以得到SOTA的图像生成结果。对于更高分辨率的ImageNet,模型要优于最佳的GAN。更重要的是,这可以通过抽取S仅有25步的子序列达到,将[1][2]中动辄一千步的迭代压力显著降低。

图4显示了生成模型的样本质量之间的比较,左边为是BigGAN-deep的结果,中间是IDM,右边是训练集中的图像。在图像质量差不多的情况下,IDM的图像采样多样性明显高于GAN。作者认为IDM更高的召回率意味着可以学到更完整的数据分布信息。


图4 GAN和IDM的采样效果比较,摘自[5]



总结

在GAN模型面临着难训练、难收敛且容易出现模式崩溃等问题的情况下,VAE和扩散模型凭借着出色的采样质量得到了越来越多的关注,经过改进的模型结构足以在无条件图像生成任务上实现SOTA,分类器引导技术使模型可以进一步提高相关任务的样本质量。最后,通过将引导与上采样相结合,可以使高分辨率条件下图像合成的结果达到SOTA。

但现阶段扩散模型在训练计算量上仍然高于GAN,采样速度也不如GAN,在单步模型上尚无法与GAN相竞争。IDM也没有提供有效的策略在未标记的数据集中做准确性和多样性之间的平衡。如何增加数据并行,降低内存和迭代次数消耗,提高整体运算效率,是扩散模型今后需要解决的问题,也是很有启发的研究方向。



参考文献

[1] Sohl-Dickstein, Jascha, et al. "Deep unsupervised learning using nonequilibrium thermodynamics." InternationalConference on Machine Learning. PMLR, 2015.

[2] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoisingdiffusion probabilistic models." arXiv preprint arXiv:2006.11239 (2020).

[3] Song, Jiaming, Chenlin Meng, and Stefano Ermon."Denoising diffusion implicit models." arXiv preprintarXiv:2010.02502 (2020).

[4] Nichol, Alex, and Prafulla Dhariwal. "Improveddenoising diffusion probabilistic models." arXiv preprintarXiv:2102.09672 (2021).

[5] Dhariwal, Prafulla, and Alex Nichol. "Diffusion models beat gans on image synthesis." arXiv preprint arXiv:2105.05233 (2021).



 往期推荐

1、「热力学计算」:从Landauer边界到终极计算机

2、神经网络的图结构

3、深度低照度图像增强-损失函数详解



关于壁仞科技研究院


壁仞科技研究院作为壁仞科技的前沿研究部门,旨在研究新型智能计算系统的关键技术,重点关注新型架构,先进编译技术和设计方法学,并将逐渐拓展研究方向,探索未来智能系统的各种可能。壁仞科技研究院秉持开放的原则,将积极投入各类产学研合作并参与开源社区的建设,为相关领域的技术进步做出自己的贡献。

扫码关注我们


: . Video Mini Program Like ,轻点两下取消赞 Wow ,轻点两下取消在看

您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存