C/C++教程

Part2: DDPM as Example of Variational Inference

本文主要是介绍Part2: DDPM as Example of Variational Inference,对大家解决编程问题具有一定的参考价值,需要的程序猿们随着小编来一起学习吧!

很多次翻看DDPM,始终不太能理解论文中提到的\(\text{Variational Inference}\)到底是如何在这个工作中起到作用。五一假期在家,无意间又刷到徐亦达老师早些年录制的理论视频,没想到其中也有介绍这部分的内容。老师的上课方式总是娓娓道来,把每一步都讲解得很仔细。本文记录一下个人对开头问题的思考。

Background

如果需要简略地介绍一下DDPM这个工作,可能会用以下几句话简单地描述:DDPMMarkov的形式对数据(图片)“扩散过程”建模,使用神经网络进行训练拟合,学习数据的概率分布。

所以对于生成任务来说,希望从给定数据中学习到的是数据的潜在信息。比如图片生成,在给定一些图片后,模型学习到的是“正常图片长什么样子”,如:

  1. 一张包含手机正面的图片会有【手机屏幕】;
  2. 一张包含猫咪的图片会有人们观察到的猫咪模样;
  3. ...

对于图片中每个像素点和附近的像素点,进行“合理”布局,才能生成“符合人们认知的图片”。

图片生成能像常见的机器学习任务如分类任务、回归任务,能基于maximize likelihood的形式来训练么?

结论是很难,先回顾如何做maximum likelihood。给定一批数据,首先需要假定数据服从的分布,接着写出似然函数,之后直接通过解析解的形式或是梯度下降的形式,求出分布。

问题就出在假定分布这一步,没有人知道图片客观上服从什么分布。那如果使用神经网络直接拟合可以么?这好像也不现实,拿一张512*512*3的图片来说,网络输出层共有约75w的数值。

对于图片生成还有另外一个问题,世界上的图片太多了,目之所及稍做处理,皆为图片。即便使用神经网络能拟合,最后生成的图片很难存在多样性。

那目前图片生成模型都是怎么做的,比如VAE或是本文即将要介绍的Diffusion Model,它们学习的都是数据分布\(p(x)\),但直接求\(p(x)\)这么麻烦,需要怎么做?这其实也是\(\text{Variational Inference}\)的核心思想,“曲线救国”,通过引入其它分布,将原本难以优化的问题转变为可优化问题。

ELOB

先把上述提到的所有背景先抛开,研究一下\(p(x)\),看看能得到什么有意思的结论。

a. 基于条件概率分布,引入新的随机变量\(z\)\(p(x) = \frac{p(x, z)}{p(z\mid x)}\)

b. 对于两边同时取\(\ln\),等式依然成立,因此有:\(\ln{p(x)} = \ln{\frac{p(x, z)}{p(z \mid x)}}\)

c. 右边分子分母同乘以\(q(z)\)\(\ln{p(x)} = \ln{\frac{p(x, z) * q(z)}{p(z \mid x) * q(z)}} = \ln{\left(\frac{p(x, z)}{q(z)} * \frac{q(z)}{p(z \mid x)}\right)} = \ln{\frac{p(x, z)}{q(z)}} + \ln{\frac{q(z)}{p(z \mid x)}}\)

d. 再次,对于上式左右两边求关于\(q(z)\)的期望,等式依然成立:

\[\begin{aligned} &\mathbb{E}_{z\sim q(z)}{[\ln{p(x)}]} = \mathbb{E}_{z\sim q(z)}{(\ln{\frac{p(x, z)}{q(z)}} + \ln{\frac{q(z)}{p(z \mid x)}})} \\ \iff & \int_z q(z)\ln{p(x)}dz = \int_z q(z)\ln{\frac{p(x, z)}{q(z)}}dz + \int_z q(z)\ln{\frac{q(z)}{p(z \mid x)}}dz \\ \iff & \ln{p(x)} = \int_z q(z)\ln{\frac{p(x, z)}{q(z)}}dz + \int_z q(z)\ln{\frac{q(z)}{p(z \mid x)}}dz \end{aligned} \tag{1} \]

一系列变换后,\((1)\)式是最后的推导结果,等式右边由两个项组成。第二个项\(\int_z q(z)\ln{\frac{q(z)}{p(z \mid x)}}dz\),叫做KL散度,它被用来衡量两个分布之间的“距离”,性质是值不小于0

这样一来,通过\((1)\)可以得到不等式\((2)\)

\[\begin{equation*} \ln{p(x)} \geq \int_z q(z)\ln{\frac{p(x, z)}{q(z)}}dz \end{equation*} \tag{2} \]

\((1)\)式右边的第一项,同时也是\((2)\)式的右边项,被学者们叫做\(\text{ELBO(Evidence Lower Bound)}\)

Objective Function

上述推导的\((2)\)式可以被视作“定理”一般的存在,即对于某个分布的对数形式,总可以找到它的下界。

\((2)\)式可以用来做什么?在Background中提到,图片生成任务中的\(p(x)\)想要对它做maximum likelihood根本无法做起。目标依然是最大化\(p(x)\),但有了\((2)\)式,求解的目标可以转移到最大化它的下界\(\text{ELBO}\)

这也是论文中提到的:

This paper presents progress in diffusion probabilistic models. A diffusion probabilistic model (which we will call a “diffusion model” for brevity) is a parameterized Markov chain trained using variational inference to produce samples matching the data after finite time.

接下来,回到论文中,看看是如何一步步推导出DDPM的优化目标。\((3)\)式直接摘录于论文:

\[\begin{equation*} \ln{p(x)} \geq \int_z q(z)\ln{\frac{p(x, z)}{q(z)}}dz = \mathbb{E}_{z \sim q(z)}\left[\ln{\frac{p(x,z)}{q(z)}}\right] \end{equation*} \tag{2} \]

\[\begin{equation*} \mathbb{E}\left[-\log p_\theta\left(\mathbf{x}_0\right)\right] \leq \mathbb{E}_q\left[-\log \frac{p_\theta\left(\mathbf{x}_{0: T}\right)}{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_0\right)}\right]=\mathbb{E}_q\left[-\log p\left(\mathbf{x}_T\right)-\sum_{t \geq 1} \log \frac{p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)}{q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right)}\right]=: L \end{equation*} \tag{3} \]

下面一项项地对\((3)\) 进行拆解,并且将它与\((2)\)比对,能帮助更好地理解:

  1. \((3)\)不等号左边的\(\mathbb{E}\left[-\log p_\theta\left(\mathbf{x}_0\right)\right]\)进一步化简就是\(-\log p_\theta\left(\mathbf{x}_0\right)\)。其中,\(p_\theta\left(\mathbf{x}_0\right)\)便是模型要学习的最终目标:图像的分布,\(\theta\)是模型的参数,\(\mathbf{x}_0\)是图片;

  2. \((2)\)式的左右两边同时加上符号,\(\geq\)变为\(\leq\)

  3. \((3)\)不等式右边部分,\(\mathbb{E}_q\left[-\log \frac{p_\theta\left(\mathbf{x}_{0: T}\right)}{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_0\right)}\right]\)

    1. 很明显,\(q(\mathbf{x}_{1:T} \mid \mathbf{x}_0)\)相当于\((2)\)中引入的额外分布\(q(z)\)。对于\(z\),在生成模型中会给它一个称呼:隐变量\((\text{latent})\)。实际上,在diffusion models里,对\(\mathbf{x}_0\)加噪后的\(\mathbf{x}_1,\mathbf{x}_2,\ldots, \mathbf{x}_T\)就可以看作隐变量,那不妨记作\(z := \{\mathbf{x}_1,\mathbf{x}_2,\ldots, \mathbf{x}_T\}\)

    2. \(p_\theta\left(\mathbf{x}_{0: T}\right) = p_\theta\left(\mathbf{x}_{0}, \mathbf{x}_{1}, \ldots, \mathbf{x}_{T}\right)\),是关于\(\mathbf{x}_0, z\)的联合概率分布,因为选用马尔代夫链建模,那么依据马尔可夫链的性质,论文定义:

\[\begin{equation*} \begin{aligned} q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_0\right)&:=\prod_{t=1}^T q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right) \\ p_\theta\left(\mathbf{x}_{0: T}\right)&:=p\left(\mathbf{x}_T\right) \prod_{t=1}^T p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right) \end{aligned} \end{equation*} \tag{4} \]

  1. \((4)\)带入\((3)\)不等式右边的第一项,得到\(L\)

\[\begin{equation*} \begin{aligned} &\mathbb{E}_q\left[-\log \frac{p_\theta\left(\mathbf{x}_{0: T}\right)}{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_0\right)}\right] \\ =&\mathbb{E}_q\left[-\log \frac{p\left(\mathbf{x}_T\right) \prod_{t=1}^T p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)}{\prod_{t=1}^T q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right)}\right] \\ =&\mathbb{E}_q\left[-\log p\left(\mathbf{x}_T\right)-\sum_{t \geq 1} \log \frac{p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)}{q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right)}\right] := L \end{aligned} \end{equation*} \]

到目前为止,经过了很多轮的变换以及数学公式,先捋一遍,再往下。\(L\)是一个替代的优化目标,

\[\mathop{\arg\min}{(L)} \iff \mathop{\arg\min}{(-\ln{p}_{\theta}(\mathbf{x}_0))} \iff \mathop{\arg\max}{(\ln{p}_{\theta}(\mathbf{x}_0))} \]

接下来,论文中对\(L\)进行了重写,以下步骤直接摘录自论文\(\text{Appendix A}\)

\[\begin{equation*} \begin{aligned} L & =\mathbb{E}_q\left[-\log \frac{p_\theta\left(\mathbf{x}_{0: T}\right)}{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_0\right)}\right] \\ & =\mathbb{E}_q\left[-\log p\left(\mathbf{x}_T\right)-\sum_{t \geq 1} \log \frac{p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)}{q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right)}\right] \\ & =\mathbb{E}_q\left[-\log p\left(\mathbf{x}_T\right)-\sum_{t>1} \log \frac{p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)}{q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right)}-\log \frac{p_\theta\left(\mathbf{x}_0 \mid \mathbf{x}_1\right)}{q\left(\mathbf{x}_1 \mid \mathbf{x}_0\right)}\right] \\ &=\mathbb{E}_q\left[-\log p\left(\mathbf{x}_T\right)-\sum_{t>1} \log \left[\frac{p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)}{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0\right)} \cdot \frac{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_0\right)}{q\left(\mathbf{x}_t \mid \mathbf{x}_0\right)}\right]-\log \frac{p_\theta\left(\mathbf{x}_0 \mid \mathbf{x}_1\right)}{q\left(\mathbf{x}_1 \mid \mathbf{x}_0\right)}\right] \end{aligned} \end{equation*} \tag{5} \]

倒数两步的变换发生在第二项,具体依据为:

\[\begin{aligned} q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right) =& \frac{q\left(\mathbf{x}_t, \mathbf{x}_{t-1}\right)}{q\left(\mathbf{x}_{t-1}\right)} \\ =& \frac{q\left(\mathbf{x}_t, \mathbf{x}_{t-1} \mid \mathbf{x}_{0}\right) *q(\mathbf{x}_{0})}{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{0}\right) * q(\mathbf{x}_{0})} \\ =& \frac{q\left(\mathbf{x}_t, \mathbf{x}_{t-1} \mid \mathbf{x}_{0}\right) }{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_0\right)} \end{aligned} \quad \Rightarrow \quad \begin{aligned} &\sum_{t>1} \log \frac{p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)}{q\left(\mathbf{x}_t \mid \mathbf{x}_{t-1}\right)} \\ =& \sum_{t>1} \log \frac{p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)}{q\left(\mathbf{x}_t, \mathbf{x}_{t-1} \mid \mathbf{x}_{0}\right) } \cdot {q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_0\right)} \\ =& \sum_{t>1} \log \frac{p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)}{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0\right)} \cdot \frac{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_0\right)}{q\left(\mathbf{x}_t \mid \mathbf{x}_0\right)} \end{aligned} \]

接着对\((5)\)进行改写得到最终形式\((6)\)

\[\begin{aligned} L &=\mathbb{E}_q\left[-\log \frac{p\left(\mathbf{x}_T\right)}{q\left(\mathbf{x}_T \mid \mathbf{x}_0\right)}-\sum_{t>1} \log \frac{p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)}{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0\right)}-\log p_\theta\left(\mathbf{x}_0 \mid \mathbf{x}_1\right)\right] \\ &=\mathbb{E}_q[\underbrace{D_{\mathrm{KL}}\left(q\left(\mathbf{x}_T \mid \mathbf{x}_0\right) \| p\left(\mathbf{x}_T\right)\right)}_{L_T}+\sum_{t>1} \underbrace{D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t, \mathbf{x}_0\right) \| p_\theta\left(\mathbf{x}_{t-1} \mid \mathbf{x}_t\right)\right)}_{L_{t-1}} \underbrace{-\log p_\theta\left(\mathbf{x}_0 \mid \mathbf{x}_1\right)}_{L_0}] \end{aligned} \tag{6} \]

Summary

太好了,对于\((6)\)来说,它最起码是个可以优化的目标函数了,因为论文中定义马尔可夫链相邻状态的转变是服从高斯分布的。当然在论文中,\((6)\)还会进一步被改写,得到更加精简的\(\text{loss function}\)形式。
DDPM是应用\(\text{variational inference}\)进行优化求解的典型例子,很值得借鉴学习。

Reference

  • Variational Inference Basic, Xu, R.Y.D
  • Denoising Diffusion Probabilistic Models
这篇关于Part2: DDPM as Example of Variational Inference的文章就介绍到这儿,希望我们推荐的文章对大家有所帮助,也希望大家多多支持为之网!