2. 扩散概率模型(diffusion probabilistic models)

随着 OpenAI 发布 DALL-E 模型,图像生成一下子火爆起来,其中的扩散模型算法也步入了众多工程师的视野。 扩散模型于2015年引入,其动机来自非平衡热力学(non-equilibrium thermodynamics), 本章我们介绍扩散模型的原理。

2.1. 扩散概率模型(diffusion probabilistic model)

早在2015年,就由 Jascha Sohl-DicksteinEric 等人提出了概率扩散模型(diffusion probabilistic model,DPM) [1],可以简称为扩散模型(diffusion model,DM)。 在原论文 [1] 中作者表示是受到了非平衡统计物理学(non-equilibrium statistical physics) 的启发进而提出了这个模型。当然这个劳什子 non-equilibrium statistical physics 大部分人包括我完全不了解, 所以我们不管它,我们从概率学的角度来阐述这个模型。

2.1.1. 马尔科夫分层自编码器(Markovian Hierarchical Variational Autoencoder,MHVAE)

事实上,从某个角度上看,扩散模型可以看做是VAE(变分自编码器)的一个扩展,因此我们从 VAE 的角度入手。 首先回顾一下VAE模型(如 图 2.1.1),其中有两个变量 \(X\)\(Z\)。 自编码的过程是把输入 \(x\) 编码成 \(z\),然后在从 \(z\) 解码回 \(x\)\(x\) 表示编码器的输入和解码器的输出,\(z\) 表示编码器的输出和解码器的输入。 就像图中所示那样,\(p(z|x)\) 对应编码过程,\(p(x|z)\) 对应解码过程, \(x\) 是可观测的,\(z\) 是不可观测的。不熟悉的读者可以复习一下上一章节 1变分自编码器

../_images/diffusion_vae.png

图 2.1.1 变分自编码器(VAE)的图表示 (图片来自 [2]

在 VAE 中编码过程只进行了一次,是否可以进行多次呢 (当然扩散模型的原论文 [1] 并不是从这个角度出发进而提出扩散模型的)? 如图 2.1.2 所示,把 VAE 的编码过程扩展成多次,从图上看,形如一个马尔科夫链的链式结构。 整个链是双向的,从左到右是一个逐步编码的过程,相当于把 VAE 的编码器循环执行了 \(T\) 次; 从右到左是一个逐步解码的过程,相当于把 VAE 的解码器循环执行了 \(T\) 次。 我们把编码过程称为前向过程(Forward Trajectory),解码过程称逆向过程(Reverse Trajectory), 无论前向还是逆向,每一个步骤(时刻) \(t\) 仅与它的上一个步骤(时刻)相关, 这是一个典型的马尔科夫过程(Markov chain)。

../_images/diffusion_hvae.png

图 2.1.2 马尔科夫分层自编码器(MHVAE)的图表示(图片来自 [2]

\(q(z_{t}|z_{t-1})\) 表示一个步骤的编码过程,\(p(z_{t}|z_{t-1})\) 表示一个步骤的解码过程。 模型所有的时刻的 \(z_{1:T}\) 构成了模型的隐变量集合, 整个模型的联合概率分布可以是表示为

(2.1.39)\[p(x,z_{1:T}) = p(z_T) p_{\theta}(x|z_1) \prod_{t=2}^T p_{\theta}(z_{t-1}|z_t)\]

隐变量 \(Z_{1:T}\) 的后验概率 \(q_{\phi}(z_{1:T}|x)\) 可以分解为

(2.1.40)\[q_{\phi}(z_{1:T}|x) = q_{\phi}(z_{1}|x) \prod_{t=2}^T q_{\phi}(z_{t}|z_{t-1})\]

在这个模型中,仍然是只有变量 \(X\) 是可观测的,观测样本的对数似然为 \(\ln p(x)\) ,和 VAE 中的推导是相同的,我们需要找到它的下界函数(ELBO)。

(2.1.41)\[ \begin{align}\begin{aligned} \ln p(x) &= \ln \int p(x,z_{1:T}) d z_{1:T}\\ &= \ln \int \frac{p(x,z_{1:T}) q_{\phi}(z_{1:T}|x) }{ q_{\phi}(z_{1:T}|x) } d z_{1:T} &(\text{乘上} 1= \frac{ q_{\phi}(z_{1:T}|x) }{ q_{\phi}(z_{1:T}|x) } \text{相当于没变化} )\\ &= \ln \mathbb{E}_{ q_{\phi}(z_{1:T}|x)} \left [ \frac{ p(x,z_{1:T}) }{ q_{\phi}(z_{1:T}|x) } \right ] &(\text{积分变期望})\\& \ge \mathbb{E}_{ q_{\phi}(z_{1:T}|x)} \left [ \ln \frac{ p(x,z_{1:T}) }{ q_{\phi}(z_{1:T}|x) } \right ] &(\text{应用 Jensen 不等式})\end{aligned}\end{align} \]

公式(2.1.41) 和 VAE 相比并没有太大变化,只是把 VAE 中的 \(z\) 换成了 \(z_{1:T}\) ,有关 EBLO 更详细的推导和介绍可以参考 VAE 的章节:1变分自编码器

2.1.2. 扩散模型

把 VAE 扩展成马尔科夫链式结构就得到了 MHVAE, MHVAE 是扩散模型的雏形,只需要在 MHVAE 的基础上再稍微调整一下就得到了扩散模型。

../_images/diffusion_vdm_base.png

图 2.1.3 扩散模型的图表示(图片来自 [2]

  • 首先,不再区分 \(x\)\(z\),反正都是对输入数据 \(x\) 进行变换, 每个节点都用符号 \(x\) 来表示。不用额外单独处理第一步了,这样整个过程可以是一个循环执行编码或者解码。 此时由于每一步的输出是下一步的输入,所以输入和输出数据的尺寸必须是一样的了,即所有 \(x_t\) 的尺寸必须是相同的, 这一点和原来的 VAE 不一样了,在 VAE 中 \(z\) 的尺寸一般是小于 \(x\) 的。

  • 其次,前向编码过程每一个步骤的编码器 \(q(x_t|x_{t-1})\) 不再通过神经网络去学习,而是直接固定为一个高斯线性变换。 这点很重要,在扩散模型中,编码器不再使用参数化的模型去学习拟合了,而是直接假设其是一个线性高斯变换,即后验条件概率 \(q(x_t|x_{t-1})\) 是一个已知的不含未知参数的高斯条件分布,具体方程稍后再介绍。

  • 最后,由于上一步编码器为线性高斯的假设,结合马尔科夫链的特征,理论上当 \(T \to \infty\) 时, \(x_T\) 是一个正态分布, 即随着 \(T\) 的增大,\(x_T\) 趋近于正态分布。通过为这个线性高斯设定一个小于 \(1\) 的渐系数, 可以使得 \(x_T\) 收敛到一个标准正态分布,即 \(\mathcal{N}(0,\textit{I})\) (均值为 \(0\) ,方差为单位方差 \(\textit{I}\)) 。

2.1.3. 前向-后向

图 2.1.3 所示,它是一个马尔科夫链式结构的贝叶斯网络(有向图网络), \(x_0\) 是模型的起点变量,它可以代表真实图片的分布(或者说随机变量)。 从 \(x_0\) 开始,按照这个链式结构逐步的演变成一个标准的高斯分布 \(x_T\)。 整个网络的联合概率分布可以表示为 \(p(x_{0:T})\), 如果按照从左到右的顺序对这个联合概率进行链式拆解

(2.1.42)\[p(x_{0:T}) = q(x_0) \prod_{t=1}^T q(x_t|x_{t-1})\]

我们把从左到右的演变过程称为前向过程,因为它是把一个真实图片 \(x_0\) 逐步添加噪声,并最终变成一个纯粹的随机高斯变量 \(x_T\) 的过程,所以也可以称其扩散过程。

如果按照从右到左的顺序,正好是前向扩散过程的逆过程,所以称之为逆过程, 它可以从一个纯粹的高斯噪声(标准高斯分布)随机变量 \(x_T\) 逐步演变,最终得到一个真实图片 \(x_0\) ,这个过程相当于生成了一张新图片,所以也称为(图像)生成过程 ,而这个过程在统计概率学上,本质是从一个联合概率分布进行采样的过程,所以也可以称为采样(sample)过程。 同样,可以按照从右到左的顺序对联合概率 \(p(x_{0:T})\) 进行链式分解

(2.1.43)\[p(x_{0:T}) = p(x_T) \prod_{t=T}^1 p(x_{t-1}|x_{t})\]

前向过程

然而,代表真实图像 \(x_0\) 的真实概率分布 \(p(x_0)\) 的概率密度函数是不知道的, 但我们能得到一批真实的图像样本,也就是我们有 \(x_0\) 的观测样本, 此时 \(x_0\) 是已知观测值,\(x_{1:T}\) 是未知的隐变量, 这时整个马尔科夫网络的联合概率变成了一个条件概率 \(q(x_{1:T}|x_0)\), 根据链式法则,可以展开为

(2.1.44)\[q(x_{1:T}|x_0) = q(x_1|x_0) \prod_{t=2}^T q(x_t|x_{t-1}) = \prod_{t=1}^T q(x_t|x_{t-1})\]

根据第二项假设,每一个 \(x_t\) 都是一个高斯变量, 前向过程每一个步骤的编码器 \(q(x_t|x_{t-1})\) 固定为一个线性高斯变换, 所谓线性高斯变换是指 \(x_t\) 的均值和 \(x_{t-1}\) 的值是线性关系。

在本论文中,定义 \(q(x_t|x_{t-1})\) 的方差与 \(x_{t-1}\) 是独立的, 并且为 \(\beta_t \textit{I}\), 其中 \(0 \lt \beta_1 \lt \beta_2 \lt \dots \beta_T \lt 1\) 。这么做的意义:前期方差较小,添加的噪声少,扩散速度慢;随着方差逐步加大,添加的噪声越来越多,扩散的速度加快。 \(\beta_t\) 作为一个超参数存在,人工指定它的值。

定义 \(q(x_t|x_{t-1})\) 的均值 \(\mu_{x_t}\)\(x_{t-1}\) 是线性关系, 这里设定另外一个系数 \(\alpha_t\),并且令 \(\alpha_t = 1-\beta_t\)\(\mu_{x_t}\)\(x_{t-1}\) 的关系定义为:

(2.1.45)\[\mu_{x_t} = \sqrt{\alpha_t} \ x_{t-1}\]

\(x_t\) 的方差定义成与 \(x_{t-1}\) 无关,而是经过缩放的单位方差, 定义为

(2.1.46)\[\Sigma_{x_t} = \beta_t \textit{I} = (1- \alpha_t ) \textit{I}\]

\(q(x_t|x_{t-1})\) 就是一个以 \(\sqrt{\alpha_t} \ x_{t-1}\) 为均值, 以 \((1- \alpha_t ) \textit{I}\) 为方差的高斯分布。 按照线性高斯的特征,如 公式(2.1.47) 所示, 它可以看做是在 \(\sqrt{\alpha_t} \ x_{t-1}\) 基础上加上一个 \(\mathcal{N} (0, (1- \alpha_t ) \textit{I} )\) 的随机高斯噪声。 这就相当于每一个步骤都在前一个步骤的基础上加上一个随机高斯噪声数据,随着\(t\) 的增加,\(x_t\) 逐步变成一个高斯噪声数据。 如图 2.1.3 所示,一张正常的图片逐步变成一张高斯噪声图。 这个过程好比把一滴墨水滴入一杯水中,随着时间的推移,墨水逐步扩散分不到各处,这杯水变成浑浊的一杯水, 原来的墨滴再也看不到了,这就是这个模型被称为扩散模型(diffusion model)的原因。

(2.1.47)\[q(x_t|x_{t-1}) = \mathcal{N} (\sqrt{\alpha_t} \ x_{t-1}, (1- \alpha_t ) \textit{I} )\]
(2.1.48)\[ \begin{align}\begin{aligned}x_{t} &=\sqrt{\alpha_t} \ x_{t-1} + \mathcal{N} (0, (1- \alpha_t ) \textit{I} )\\&=\sqrt{\alpha_t} \ x_{t-1} + \sqrt{1- \alpha_t } \ \epsilon \ \ \ ,\epsilon \sim \mathcal{N} (0, \textit{I} )\end{aligned}\end{align} \]

至于系数 \(\beta_t\) 的值可以作为一个超参数有人工指定[3] ,也可以作为一个可学习的参数由模型学习[4], 这里我们从直觉的角度来理解系数 \(\beta_t\) 的作用。

  1. 为什么要有系数 \(\beta_t\)

    前向扩散过程是希望 \(x_t\) 渐进趋近于标准正态分布,即均值趋近于 \(0\), 方差趋近于单位方差 \(\textit{I}\),并且 \(x_t\) 只和它的上一步 \(x_{t-1}\) 有关,同时 \(x_t\)\(x_{t-1}\) 是线性高斯变化,也就是说 \(\mu_{x_t}\)\(x_{t-1}\) 的关系必然是某个系数乘上 \(x_{t-1}\) 的线性关系。

  2. \(\alpha_t\) 为什么小于 \(1\)

    因为需要 \(\mu_{x_t} \to 0\),所以这个系数必然是小于 \(1\) 的。

  3. 什么要有根号?

    在满足 \(\mu_{x_t} \to 0\) 的同时,也要满足 \(\Sigma_{x_t} \to \textit{I}\)。 用一个系数同时控制均值和方差的变化,可以扩散过程更稳健。而 \(\mu_{x_t}\)\(\Sigma_{x_t}\) 并不是同一个量纲,方差是平方的,原论文是直接把 \(\alpha_t\) 定义在了方差上,所以在均值上就要开根号。

  4. 这个线性关系(线性高斯),为什么没有截距项?

    观察下 公式(2.1.47),其中的高斯噪声 \(\sqrt{1- \alpha_t } \ \epsilon\) 就相当于截距。

\(\alpha_t\) 并不是一个固定值,是可以随着 \(t\) 的增长逐渐变小的,这可以使得扩散过程中加的高斯噪声比重逐渐加重。 分析一下这样做的好处,前期如果加的噪声太多,会使得数据扩展的太快(比如突变),使得逆向还原变得困难; 同样因为后期数据本身已经接近随机噪声数据了,后期如果加的噪声不够多,相当于变化幅度小,扩散的太慢,这会使得链路变长需要的事件变多。 我们希望扩散的前期慢一点,后期快一点,因此 \(\alpha_t\) 最好是逐渐变小的。

总的来说,前向过程就是一个逐步添加高斯噪声,并且最终变成一个纯的高斯噪声数据。 另外前向过程中的编码分布 \(q(x_t|x_{t-1})\) 并没有进行参数化表示,而是假定是一个确定的线性高斯变换。

在前向过程我们需要计算每一个时刻的 \(x_t\),最直接的计算方式就是写一个循环,按照公式 公式(2.1.47) 逐步从 \(x_0\) 算到 \(x_T\),公式中的 \(\epsilon_{t}\) 就是从一个标准正态分布随机采样, 这和 VAE 中是一样的。这样用循环计算当然可以,但是当 \(T\) 较大时,明显是效率低下的。 事实上,得益于 \(q(x_t|x_{t-1})\) 是线性高斯的假设,这里有一个计算上的小技巧,可以直接从 \(x_0\) 一步计算任意的 \(t\),这样就可以并行计算全部的 \(x_t\)。 如下是具体的推导过程:

(2.1.49)\[ \begin{align}\begin{aligned}x_t &= \sqrt{\alpha_t} \ x_{t-1} + \sqrt{1-\alpha_t} \ \epsilon_{t}\\&= \sqrt{\alpha_t} \left( \sqrt{\alpha_{t-1}} \ x_{t-2} + \sqrt{1-\alpha_{t-1}} \ \epsilon_{t-1} \right ) + \sqrt{1-\alpha_t} \ \epsilon_{t}\\&= \sqrt{\alpha_t \alpha_{t-1} } \ x_{t-2} + \underbrace{ \sqrt{\alpha_t - \alpha_t \alpha_{t-1} }\ \epsilon_{t-1} + \sqrt{1- \alpha_t} \ \epsilon_{t} }_{\text{两个相互独立的0均值的高斯分布相加}}\\ &= \sqrt{\alpha_t \alpha_{t-1} } \ x_{t-2} + \underbrace{ \sqrt{ \sqrt{\alpha_t - \alpha_t \alpha_{t-1} }^2 + \sqrt{1- \alpha_t}^2 } \ \epsilon }_{\text{两个方差相加,用一个新的高斯分布代替}}\\&= \sqrt{\alpha_t \alpha_{t-1} } \ x_{t-2} + \sqrt{1- \alpha_t \alpha_{t-1}} \ \epsilon\\&= ...\\&= \sqrt{\prod_{i=1}^t \alpha_i} \ x_0 + \sqrt{1- \prod_{i=1}^t \alpha_i } \ \epsilon\\&= \sqrt{\bar{ \alpha}_t } \ x_0 + \sqrt{1- \bar{ \alpha}_t } \ \epsilon \ \ \ , \bar{\alpha} = \prod_{i=1}^t \alpha_i ,\ \ \epsilon \sim \mathcal{N}(0,\textit{I})\\&\sim \mathcal{N}(\sqrt{\bar{\alpha}_t } \ x_0, (1- \bar{ \alpha}_t) \textit{I})\end{aligned}\end{align} \]

前项过程的每一个步骤的转换条件概率 \(q(x_t|x_{t-1})\) 代表着从 \(x_{t-1})\)\(x_t\) 的转换方法(过程), 而且我们已经得到了条件概率分布 \(q(x_t|x_{t-1})\) 的具体形式( 公式(2.1.47) ) ,具体的,当有了 \(x_{t-1}\) 的值后,可以根据 公式(2.1.48) 计算 \(x_{t}\) 的值,就这样循环计算,直到计算出 \(x_{T}\) 的值。 得益于条件高斯分布的特性,我们推导出一个计算小技巧,公式(2.1.49), 可以直接从 \(x_0\) 计算出任意时刻的 \(x_t\) ,不用再逐步循环计算,这极大的提高了我们的计算速度。

最后的重点,我们发现只要设置了超参数 \(\alpha_{0:T}\) 的值,这个前向计算过程是可以直接解析(使用公式)计算的, 没有未知参数,不需要用一个模型学习这个过程。 然而逆向过程就不能直接计算的,是需要用模型学习的。

逆向过程

图 2.1.3,逆向过程是从右到左的解码过程,从 \(x_T\) (随机高斯噪声)开始, 逐步的解码成一个有意义的数据,比如生成一张图片。 解码过程的每一个时刻 \(t\) 要把 \(x_{t+1}\) 还原成 \(x_t\), 每一步骤对应的条件概率表示成 \(p(x_t|x_{t+1})\)。 按照逆向过程,对联合概率 \(p(x_{0:T})\) 进行分解为:

(2.1.50)\[p(x_{0:T}) = p(x_T) \prod_{t={T-1}}^0 p(x_{t}|x_{t+1})\]

其中 \(p(x_T)\) 的概率密度是知道的,它是一个标准高斯分布,即 \(p(x_T) \sim \mathcal{N}(0,\textit{I})\) 。然而 \(p_{\theta}(x_{t}|x_{t+1})\) 是难以计算的, 即使根据贝叶斯定理得出 \(p_{\theta}(x_{t}|x_{t+1})\) 的表达式, 它的分母部分(归一化项)是含有积分的,这个积分是没有解析解的,是非常难以计算的。 因此,我们想到用一个模型(神经网络)取拟合学习条件概率分布 \(p_{\theta}(x_{t}|x_{t+1})\) ,然后利用学习好的模型,充当 \(p_{\theta}(x_{t}|x_{t+1})\) ,这样就有了逆向过程的完整链式结构( 公式(2.1.50)) ,也就可以通过逆向过程,从一个随机高斯噪声 \(x_T\) 逐步生成一张真实的图片 \(x_0\)

为了符号区分, 定义符号 \(\theta\) 表示模型的学习参数, 不带 \(\theta\)\(p(x_t|x_{t+1})\) 表示解码器的真实分布, 带 \(\theta\)\(p_{\theta}(x_t|x_{t+1})\) 表示参数化模型学习的近似分布。 下一步就是如何学习 \(p_{\theta}(x_t|x_{t+1})\), 类似于 VAE 模型,通过最大化 ELBO 来进行参数学习。

2.1.4. 目标函数(ELBO)

我们知道学习一个概率分布的未知参数的常用算法是极大似然估计, 极大似然估计是通过极大化观测数据的对数概率(似然)实现的。 在这里,整个网络的联合概率分布是 \(p(x_{0:T})\) ,其中,我们只有 \(x_0\) 的观测数据(样本) ,没有 \(x_{1:T}\) 的观测样本, 因此我们极大化的是边缘分布 \(p(x_0)\),而不是联合分布 \(p(x_{0:T})\) 。边缘分布 \(p(x_0)\) 可以通过边际化得到

(2.1.51)\[p(x_0) = \int p(x_{0:T}) dx_{1:T}\]

显然扩散模型的学习过程和原始的VAE以及上一节介绍的 MHVAE 相似的,原始的观测数据对数似然 \(\ln p(x_0)\) 包含隐变量的积分操作,是无法直接极大化的,需要用它的下界函数(ELBO)进行替代。 接下来我们推导一下扩散模型的 ELBO 函数,我们在 公式(2.1.41) 的基础上继续推演, 推导过程如下:

(2.1.52)\[\begin{split}\begin{align} {\ln p(x_0)} &= {\ln \int p(x_{0:T}) dx_{1:T}}\\ &= {\ln \int \frac{p(x_{0:T})q(x_{1:T}|x_0)}{q(x_{1:T}|x_0)} dx_{1:T}}\\ &= {\ln \mathbb{E}_{q(x_{1:T}|x_0)}\left[\frac{p(x_{0:T})}{q(x_{1:T}|x_0)}\right]}\\ &\geq {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_{0:T})}{q(x_{1:T}|x_0)}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)\prod_{t=1}^{T}p_{x{\theta}}(x_{t-1}|x_t)}{\prod_{t = 1}^{T}q(x_{t}|x_{t-1})}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)p_{x{\theta}}(x_0|x_1)\prod_{t=2}^{T}p_{x{\theta}}(x_{t-1}|x_t)}{q(x_T|x_{T-1})\prod_{t = 1}^{T-1}q(x_{t}|x_{t-1})}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)p_{x{\theta}}(x_0|x_1)\prod_{t=1}^{T-1}p_{x{\theta}}(x_{t}|x_{t+1})}{q(x_T|x_{T-1})\prod_{t = 1}^{T-1}q(x_{t}|x_{t-1})}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)p_{x{\theta}}(x_0|x_1)}{q(x_T|x_{T-1})}\right] + \mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \prod_{t = 1}^{T-1}\frac{p_{x{\theta}}(x_{t}|x_{t+1})}{q(x_{t}|x_{t-1})}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln p_{x{\theta}}(x_0|x_1)\right] + \mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)}{q(x_T|x_{T-1})}\right] + \mathbb{E}_{q(x_{1:T}|x_0)}\left[ \sum_{t=1}^{T-1} \ln \frac{p_{x{\theta}}(x_{t}|x_{t+1})}{q(x_{t}|x_{t-1})}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln p_{x{\theta}}(x_0|x_1)\right] + \mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)}{q(x_T|x_{T-1})}\right] + \sum_{t=1}^{T-1}\mathbb{E}_{q(x_{1:T}|x_0)}\left[ \ln \frac{p_{x{\theta}}(x_{t}|x_{t+1})}{q(x_{t}|x_{t-1})}\right]}\\ &= \begin{aligned}[t] \mathbb{E}_{q(x_{1}|x_0)}\left[\ln p_{x{\theta}}(x_0|x_1)\right] &+ \mathbb{E}_{q(x_{T-1}, x_T|x_0)}\left[\ln \frac{p(x_T)}{q(x_T|x_{T-1})}\right] \\ &+ \sum_{t=1}^{T-1}\mathbb{E}_{q(x_{t-1}, x_t, x_{t+1}|x_0)}\left[\ln \frac{p_{x{\theta}}(x_{t}|x_{t+1})}{q(x_{t}|x_{t-1})}\right]\\ \end{aligned}\\ &= \begin{aligned}[t] {\underbrace{\mathbb{E}_{q(x_{1}|x_0)}\left[\ln p_{\theta}(x_0|x_1)\right]}_\text{reconstruction term}} &- {\underbrace{\mathbb{E}_{q(x_{T-1}|x_0)}\left[ \KL{q(x_T|x_{T-1})}{p(x_T)}\right]}_\text{prior matching term}} \\ &- {\sum_{t=1}^{T-1}\underbrace{\mathbb{E}_{q(x_{t-1}, x_{t+1}|x_0)}\left[ \KL {q(x_{t}|x_{t-1})}{p_{\theta}(x_{t}|x_{t+1})}\right]}_\text{consistency term}} \end{aligned} \end{align}\end{split}\]

这个推导过程看上去很复杂,但最后的结果却很简洁,由3个子项构成,现在分别介绍一下各项的含义。

  • \(\mathbb{E}_{q(x_{1}|x_0)}\left[\ln p_{\theta}(x_0|x_1)\right]\) 可以被看做是重建项, 回顾一下原始 VAE 模型的 ELBO 函数,它和原始 VAE 的第一项是一样的,从第一步的隐变量 \(x_1\) 重建成原来的数据 \(x_0\)

  • \(\mathbb{E}_{q(x_{T-1}|x_0)}\left[\KL{q(x_T|x_{T-1})}{p(x_T)}\right]\) 这一项是先验匹配项, 这一项和原始 VAE 也能对应上。稍微有点区别的是,这里这一项是没有学习参数的,\(q(x_T|x_{T-1})\) 是前向过程, 我们已经假设它是已知的线性高斯变化了,不需要再用模型学习了。当 \(T\) 足够大时,可以认为这一项就是 \(0\)

  • \(\mathbb{E}_{q(x_{t-1}, x_{t+1}|x_0)}\left[\KL{q(x_{t}|x_{t-1})}{p_{{\theta}}(x_{t}|x_{t+1})}\right]\) 是个 KL 散度度量,这一项称为一致项(consistency term),显然它使得在每一个时刻 \(t\),解码器 \(p_{{\theta}}(x_{t}|x_{t+1})\) 得到的 \(x_t\) 尽量和编码器 \(q(x_{t}|x_{t-1})\) 生成的内容保值一致。 \(q(x_{t}|x_{t-1})\) 是已知的线性高斯变换,不需要学习,而 \(p_{{\theta}}(x_{t}|x_{t+1})\) 是需要一个参数化的模型去学习的,也就是我们的模型学习的是 \(p_{{\theta}}(x_{t}|x_{t+1})\),其中 \(\theta\) 表示模型需要学习的参数集合。

../_images/diffusion_first_derivation.png

图 2.1.4 根据 ELBO( 公式(2.1.52))显示,对于每一个时刻 \(t\),最小化绿色和粉色箭头所代表的分布之间的差异。(图片来自 [2]

在极大化 公式(2.1.52) 时,第二项不包含任意可学习参数,因此可以忽略掉。 另外两项都是一个期望项,解决方法和 VAE 一样,可以通过采样法(MCMC)近似求解期望。 然而,第三项存在一个小问题,它的期望是关于两个变量的, \((x_{t-1},x_{t+1})\), 用采样法(MCMC)同时对两个随机变量进行采样,会导致更大的方差,这会使得优化过程不稳定,不容易收敛 [2]。 因此直接优化 公式(2.1.52) 并不是最佳的选择,我们继续看能不能改进一下。

改进方法的推导中依赖几个等式替换,在推导前先给出这几个依赖等式。 首先根据条件独立性原则,有下式成立

(2.1.53)\[q(x_t | x_{t-1}) = q(x_t | x_{t-1}, x_0)\]

备注

条件独立性

假设有三个随机变量,分别是 \(A,B,C\),三者之间的依赖关系为(有向图表示) \(A \to B \to C\) , 当变量 \(B\) 的值已知(确定)时,变量 \(A\) 和变量 \(B\) 是相互独立的,记作 \(A \ind C \mid B\), 读作:在 \(B\) 的条件下变量 \(A\) 和变量 \(C\) 相互独立, 此时有 \(P(C|B)=P(C|B,A)\) 成立。

\(q(x_t | x_{t-1}, x_0)\) 可以进一步根据贝叶斯定理求得,如下:

(2.1.54)\[q(x_t | x_{t-1}, x_0) = \frac{q(x_{t-1} \mid x_t, x_0)q(x_t \mid x_0)}{q(x_{t-1} \mid x_0)}\]

有了 公式(2.1.53)公式(2.1.54), 我再重新推导 ELBO 函数,推导过程如下:

(2.1.55)\[\begin{split}\begin{align} {\ln p(x_0)} &\geq {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_{0:T})}{q(x_{1:T}|x_0)}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)\prod_{t=1}^{T}p_{{\theta}}(x_{t-1}|x_t)}{\prod_{t = 1}^{T}q(x_{t}|x_{t-1})}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)p_{{\theta}}(x_0|x_1)\prod_{t=2}^{T}p_{{\theta}}(x_{t-1}|x_t)}{q(x_1|x_0)\prod_{t = 2}^{T}q(x_{t}|x_{t-1})}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)p_{{\theta}}(x_0|x_1)\prod_{t=2}^{T}p_{{\theta}}(x_{t-1}|x_t)}{q(x_1|x_0)\prod_{t = 2}^{T}q(x_{t}|x_{t-1}, x_0)}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p_{{\theta}}(x_T)p_{{\theta}}(x_0|x_1)}{q(x_1|x_0)} + \ln \prod_{t=2}^{T}\frac{p_{{\theta}}(x_{t-1}|x_t)}{q(x_{t}|x_{t-1}, x_0)}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)p_{{\theta}}(x_0|x_1)}{q(x_1|x_0)} + \ln \prod_{t=2}^{T}\frac{p_{{\theta}}(x_{t-1}|x_t)}{\frac{p(x_{t-1}|x_{t}, x_0)q(x_t|x_0)}{q(x_{t-1}|x_0)}}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)p_{{\theta}}(x_0|x_1)}{q(x_1|x_0)} + \ln \prod_{t=2}^{T}\frac{p_{{\theta}}(x_{t-1}|x_t)}{\frac{p(x_{t-1}|x_{t}, x_0)\cancel{q(x_t|x_0)}}{\cancel{q(x_{t-1}|x_0)}}}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)p_{{\theta}}(x_0|x_1)}{\cancel{q(x_1|x_0)}} + \ln \frac{\cancel{q(x_1|x_0)}}{q(x_T|x_0)} + \ln \prod_{t=2}^{T}\frac{p_{{\theta}}(x_{t-1}|x_t)}{q(x_{t-1}|x_{t}, x_0)}\right]}\\ &= {\mathbb{E}_{q(x_{1:T}|x_0)}\left[\ln \frac{p(x_T)p_{{\theta}}(x_0|x_1)}{q(x_T|x_0)} + \sum_{t=2}^{T}\ln\frac{p_{{\theta}}(x_{t-1}|x_t)}{q(x_{t-1}|x_{t}, x_0)}\right]}\\ &= \begin{aligned}[t] { \mathbb{E}_{ q(x_{1:T}|x_0) } \left[ \ln p_{{\theta}}(x_0|x_1) \right]\\ \\+ \mathbb{E}_{ q(x_{1:T}|x_0)} \left[ \ln \frac{p(x_T)}{q(x_T|x_0)} \right ]\\ \\+ \sum_{t=2}^{T} \mathbb{E}_{q(x_{1:T}|x_0)} \left[\ln\frac{p_{{\theta}}(x_{t-1}|x_t)}{q(x_{t-1}|x_{t}, x_0)}\right] }\\ \end{aligned}\\ &= \begin{aligned}[t]{ \mathbb{E}_{q(x_{1}|x_0)}\left[\ln p_{{\theta}}(x_0|x_1)\right] + \mathbb{E}_{q(x_{T}|x_0)}\left[\ln \frac{p(x_T)}{q(x_T|x_0)}\right] \\ + \sum_{t=2}^{T}\mathbb{E}_{q(x_{t}, x_{t-1}|x_0)}\left[\ln\frac{p_{{\theta}}(x_{t-1}|x_t)}{q(x_{t-1}|x_{t}, x_0)}\right] \\ }\end{aligned} \\ &= \begin{aligned}[t] \underbrace{\mathbb{E}_{q(x_{1}|x_0)}\left[\ln p_{{\theta}}(x_0|x_1)\right]}_\text{reconstruction term} &- \underbrace{\KL{q(x_T|x_0)}{p(x_T)}}_\text{prior matching term}\\ &- \sum_{t=2}^{T} \underbrace{\mathbb{E}_{q(x_{t}|x_0)}\left[\KL{q(x_{t-1}|x_t, x_0)}{p_{{\theta}}(x_{t-1}|x_t)}\right]}_\text{denoising matching term}\\ \end{aligned} \end{align}\end{split}\]

对比一下 公式(2.1.55)公式(2.1.52) 发生了一些变化,

  • \(\mathbb{E}_{q(x_{1}|x_0)}\left[\ln p_{{\theta}}(x_0|x_1)\right]\),这一项没有变化。

  • \(\KL{q(x_T|x_0)}{p(x_T)}\),这一项发生变化,其中由 \(q(x_T|x_{T-1})\) 变成了 \(q(x_T|x_0)\), 这其实没啥区别,都是约束最终的 \(x_{T}\) 要尽量接近先验 \(p(x_T)\) 的标准高斯分布。

  • 关键是第3项的变化,需要积分(采样)的变量只剩一个了,这样就简单了很多。 其中 \(q(x_{t-1}|x_t, x_0)\) 表示解码器的真实分布,带有 \(\theta\)\(p_{{\theta}}(x_{t-1}|x_t)\) 表示参数化的模型,二者的 KL 散度正好表达了参数化模型 \(p_{{\theta}}(x_{t-1}|x_t)\) 要尽量和真实分布 \(q(x_{t-1}|x_t, x_0)\) 接近,如图 图 2.1.5 所示。

../_images/diffusion_second_derivation.png

图 2.1.5 公式(2.1.55) 对应的可视化图形表示。粉色箭头表示根据贝叶斯公式得到的逆过程中解码器的真实分布 \(q(x_{t-1}|x_t, x_0)\), 绿色箭头表示需要学习的近似分布 \(p_{{\theta}}(x_{t-1}|x_t)\), 二者的 KL 散度表示必须领二者进尽可能的接近。(图片来自 [2]

要想极大化 公式(2.1.55) 形式的 ELBO 函数,还需要其中每一项展开成具体的参数化公式, 第二项不包含参数 \(\theta\),所以可以直接忽略。

reconstruction term

第一项 \(p_{{\theta}}(x_0|x_1)\) 和 VAE 中是解码器是类似的, 它是一个依赖 \(x_1\) 的条件高斯分布。 假设它的均值 \(\mu_{\theta}\) 是一个关于 \(x_1\)\(t=1\) 的参数化函数, 记作 \(\mu_{\theta}(x_1,t=1)\),根据以前的套路,我们可以用一个参数化的神经网络模型去拟合它。 和 VAE 不同的是,这里我们不再用模型去预测 \(p_{{\theta}}(x_0|x_1)\) 的方差, 而是假设它的方差是一个常量,暂时记作 \(\Sigma\) 。即模型输入 \(x_1\)\(t=1\),输出 \(\mu_{\theta}(x_1,t=1)\)。 根据高斯分布的概率密度函数,对 \(\ln p(x_0|x_1)\) 进行展开

(2.1.56)\[ \begin{align}\begin{aligned}\ln p(x_0|x_1) &=\ln \frac{1}{(2\pi)^{n/2} |\Sigma |^{1/2}}exp\{-\frac{1}{2}(x_0 - \mu_{\theta}(x_1,t=1))^{T}\Sigma^{-1}(x - \mu_{\theta}(x_1,t=1))\}\\& \propto -\frac{1}{2}(x_0 - \mu_{\theta}(x_1,t=1))^{T}(x_0 - \mu_{\theta}(x_1,t=1))\}\\& = -\frac{1}{2} \left\lVert x_0 - \mu_{\theta}(x_1,t=1) \right\rVert_2^2\end{aligned}\end{align} \]

从中可以看出极大化 \(\ln p(x_0|x_1)\) 就等价于极小化模型输出 \(\mu_{\theta}(x_1,t=1)\)\(x_0\) 的平方误差,模型输出其实就相当于在预测 \(x_0\) 的值,因此我们可以把模型输出 \(\mu_{\theta}(x_1,t=1)\) 用符号 \(\hat x_{0}(x_1,t=1)\) 来表示。

(2.1.57)\[ \begin{align}\begin{aligned}& {\operatorname{\arg\max}}_{\theta} \mathbb{E}_{q(x_{1}|x_0)}\left[\ln p_{{\theta}}(x_0|x_1)\right]\\&\Leftrightarrow {\operatorname{\arg\max}}_{\theta} \ \mathbb{E}_{q(x_{1}|x_0)} \left [ -\frac{1}{2} \left\lVert x_0 - \hat x_{0}(x_1,t=1) \right\rVert_2^2 \right ]\\&\Leftrightarrow {\operatorname{\arg\min}}_{\theta} \ \mathbb{E}_{q(x_{1}|x_0)} \left [ \left\lVert x_0 - \hat x_{0}(x_1,t=1) \right\rVert_2^2 \right ]\end{aligned}\end{align} \]

可以看到,这一项其实就相当于,我们用一个模型去预测样本 \(x_0\),损失函数等价于 均方误差。 至于其中的期望项 \(\mathbb{E}_{q(x_{1}|x_0)}\) ,是由于模型的输入 \(x_1\) 是一个随机变量, 并不是一个数值变量,所以需要对随机变量 \(x_1\) 求期望,而求期望时随机变量 \(x_1\) 的概率分布用的是条件概率 分布 \(q(x_{1}|x_0)\) ,这个概率分布用 公式(2.1.47) 即可求得。 当然按照 VAE 相似的套路,我们可以通过采样法近似求解期望。 这里可以暂时先不关注这一项的具体实现,接着看后面的内容,在最后的表达式中,这一项其实可以和第三项合并在一起。

prior matching term

第二项是先验匹配项(prior matching term),它是一个 KL 散度,\(q(x_T|x_0)\) 是前向扩散过程生成的 \(x_T\) 的后验分布(有了观察样本 \(x_0\) 之后,所以是后验); \(p(x_T)\) 是逆向生成过程中 \(x_T\) 的先验分布。这一项就是约束 \(x_T\) 的后验分布要和先验分布尽量接近, VAE 中也有类似的一项,起作用相当于正则项。不过在这里 \(q(x_T|x_0)\)\(p(x_T)\) 都没有包含未知参数,所以在极大化 ELBO 时,这一项相当于常数项,可以忽略。

denoising matching term

这一项也是一个 KL 散度,KL 散度的两项都是高斯分布, 只要能找到这两个高斯分布的均值和方差,就可以解析计算它们的 KL 散度了。 我们先来推导下真实分布 \(q(x_{t-1}|x_t, x_0)\), 在 公式(2.1.54) 的基础上继续推导展开,注意 \(q(x_{t-1}|x_t, x_0)\) 是条件概率分布, 其中 \(x_{t-1}\) 是未知量(分布的目标变量)\(x_{t},x_0\)条件变量

(2.1.58)\[\begin{split}\begin{align} {q(x_{t-1}|x_t, x_0)} &= {\frac{q(x_t | x_{t-1}, x_0)q(x_{t-1}|x_0)}{q(x_{t}|x_0)}}\\ &= {\frac{\mathcal{N}(x_{t} ; \sqrt{\alpha_t} x_{t-1}, (1 - \alpha_t)\textit{I})\mathcal{N}(x_{t-1} ; \sqrt{\bar\alpha_{t-1}}x_0, (1 - \bar\alpha_{t-1}) \textit{I})}{\mathcal{N}(x_{t} ; \sqrt{\bar\alpha_{t}}x_0, (1 - \bar\alpha_{t})\textit{I})}}\\ &\propto {\text{exp}\left\{-\left[\frac{(x_{t} - \sqrt{\alpha_t} x_{t-1})^2}{2(1 - \alpha_t)} + \frac{(x_{t-1} - \sqrt{\bar\alpha_{t-1}} x_0)^2}{2(1 - \bar\alpha_{t-1})} - \frac{(x_{t} - \sqrt{\bar\alpha_t} x_{0})^2}{2(1 - \bar\alpha_t)} \right]\right\}}\\ &= {\text{exp}\left\{-\frac{1}{2}\left[\frac{(x_{t} - \sqrt{\alpha_t} x_{t-1})^2}{1 - \alpha_t} + \frac{(x_{t-1} - \sqrt{\bar\alpha_{t-1}} x_0)^2}{1 - \bar\alpha_{t-1}} - \frac{(x_{t} - \sqrt{\bar\alpha_t} x_{0})^2}{1 - \bar\alpha_t} \right]\right\}}\\ &= {\text{exp}\left\{-\frac{1}{2}\left[\frac{(-2\sqrt{\alpha_t} x_{t}x_{t-1} + \alpha_t x_{t-1}^2)}{1 - \alpha_t} + \frac{(x_{t-1}^2 - 2\sqrt{\bar\alpha_{t-1}}x_{t-1} x_0)}{1 - \bar\alpha_{t-1}} + C(x_t, x_0)\right]\right\}} \\ &\propto {\text{exp}\left\{-\frac{1}{2}\left[- \frac{2\sqrt{\alpha_t} x_{t}x_{t-1}}{1 - \alpha_t} + \frac{\alpha_t x_{t-1}^2}{1 - \alpha_t} + \frac{x_{t-1}^2}{1 - \bar\alpha_{t-1}} - \frac{2\sqrt{\bar\alpha_{t-1}}x_{t-1} x_0}{1 - \bar\alpha_{t-1}}\right]\right\}}\\ &= {\text{exp}\left\{-\frac{1}{2}\left[(\frac{\alpha_t}{1 - \alpha_t} + \frac{1}{1 - \bar\alpha_{t-1}})x_{t-1}^2 - 2\left(\frac{\sqrt{\alpha_t}x_{t}}{1 - \alpha_t} + \frac{\sqrt{\bar\alpha_{t-1}}x_0}{1 - \bar\alpha_{t-1}}\right)x_{t-1}\right]\right\}}\\ &= {\text{exp}\left\{-\frac{1}{2}\left[\frac{\alpha_t(1-\bar\alpha_{t-1}) + 1 - \alpha_t}{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}x_{t-1}^2 - 2\left(\frac{\sqrt{\alpha_t}x_{t}}{1 - \alpha_t} + \frac{\sqrt{\bar\alpha_{t-1}}x_0}{1 - \bar\alpha_{t-1}}\right)x_{t-1}\right]\right\}}\\ &= {\text{exp}\left\{-\frac{1}{2}\left[\frac{\alpha_t-\bar\alpha_{t} + 1 - \alpha_t}{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}x_{t-1}^2 - 2\left(\frac{\sqrt{\alpha_t}x_{t}}{1 - \alpha_t} + \frac{\sqrt{\bar\alpha_{t-1}}x_0}{1 - \bar\alpha_{t-1}}\right)x_{t-1}\right]\right\}}\\ &= {\text{exp}\left\{-\frac{1}{2}\left[\frac{1 -\bar\alpha_{t}}{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}x_{t-1}^2 - 2\left(\frac{\sqrt{\alpha_t}x_{t}}{1 - \alpha_t} + \frac{\sqrt{\bar\alpha_{t-1}}x_0}{1 - \bar\alpha_{t-1}}\right)x_{t-1}\right]\right\}}\\ &= {\text{exp}\left\{-\frac{1}{2}\left(\frac{1 -\bar\alpha_{t}}{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}\right)\left[x_{t-1}^2 - 2\frac{\left(\frac{\sqrt{\alpha_t}x_{t}}{1 - \alpha_t} + \frac{\sqrt{\bar\alpha_{t-1}}x_0}{1 - \bar\alpha_{t-1}}\right)}{\frac{1 -\bar\alpha_{t}}{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}}x_{t-1}\right]\right\}}\\ &= {\text{exp}\left\{-\frac{1}{2}\left(\frac{1 -\bar\alpha_{t}}{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}\right)\left[x_{t-1}^2 - 2\frac{\left(\frac{\sqrt{\alpha_t}x_{t}}{1 - \alpha_t} + \frac{\sqrt{\bar\alpha_{t-1}}x_0}{1 - \bar\alpha_{t-1}}\right)(1 - \alpha_t)(1 - \bar\alpha_{t-1})}{1 -\bar\alpha_{t}}x_{t-1}\right]\right\}}\\ &= {\text{exp}\left\{-\frac{1}{2}\left(\frac{1}{\frac{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}{1 -\bar\alpha_{t}}}\right)\left[x_{t-1}^2 - 2\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)x_0}{1 -\bar\alpha_{t}}x_{t-1}\right]\right\}}\\ &\propto {\mathcal{N}(x_{t-1} ;} \underbrace{{\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)x_0}{1 -\bar\alpha_{t}}}}_{\mu_q(x_t, x_0)}, \underbrace{{\frac{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}{1 -\bar\alpha_{t}}\textit{I}}}_{{\Sigma}_q(t)}) \end{align}\end{split}\]

上面推导过程中的 \(C(x_t, x_0)\) 是与目标变量 \(x_{t-1}\) 无关的项,它仅仅与 \(x_0,x_{t},\alpha\) 有关, 所以可以把它看做是常量,可以不关注它。 总之最后的结果是 \(q(x_{t-1}|x_t, x_0)\) 同样是一个高斯分布,它均值记作 \(\mu_q(x_t, x_0)\),是一个关于 \(x_t, x_0\) 的函数。 方差记作 \(\Sigma_q(t)\),仅与 \(\alpha\) 有关,如果把 \(\alpha\) 当做一个超参数,那么显然方差是可以直接计算的。

(2.1.59)\[ \begin{align}\begin{aligned}\mu_q(x_t, x_0) &= { \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)x_0}{1 -\bar\alpha_{t}}}\\\Sigma_q(t) &= \frac{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}{ 1 -\bar\alpha_{t}} \textit{I} = \sigma_q^2(t) \textit{I}\end{aligned}\end{align} \]

我们的参数化模型分布 \(p_\theta(x_{t-1}|x_t)\),需要尽可能的接近真实高斯分布 \(q(x_{t-1}|x_t, x_0)\) ,显然 \(p_\theta(x_{t-1}|x_t)\) 也需要是一个高斯分布才行,我们用 \(\mu_{\theta}\) 表示它的均值, \(\Sigma_\theta\) 表示的方差。 真实分布 \(q(x_{t-1}|x_t, x_0)\) 的方差 \(\Sigma_q(t)\) 是一个仅与 \(\alpha\) 有关的量, 在本论文中把它看做一个不需要模型学习的常量,直接假设 \(\Sigma_\theta=\Sigma_q(t)\) 即可。 在本论文中,作者认为这个方差学习与否,最终效果差异不大,所以就没有用模型去学习。 然而后来,有其他学者发布新的论文,提出了学习方差的改进方法,认为学习方差效果会更好。 后面的章节会详细介绍这个改进方法,这里先假设方差是不需要模型学习的常量即可。

(2.1.60)\[p_{\theta} (x_{t-1}|x_t) \sim \mathcal{N}(x_{t-1} ;\mu_{\theta},\Sigma_q(t))\]

现在,真实分布 \(q(x_{t-1}|x_t, x_0)\) 和参数化模型学习分布 \(p_\theta(x_{t-1}|x_t)\) 各自的均值和方差表达式都确定了, 又已知两个高斯分布的 KL 散度计算公式如下:

(2.1.61)\[\begin{align} \KL{\mathcal{N}(x; {\mu}_x,{\Sigma}_x)}{\mathcal{N}(y; {\mu}_y,{\Sigma}_y)} &=\frac{1}{2}\left[\log\frac{|{\Sigma}_y|}{|{\Sigma}_x|} - d + \text{tr}({\Sigma}_y^{-1}{\Sigma}_x) + ({\mu}_y-{\mu}_x)^T {\Sigma}_y^{-1} ({\mu}_y-{\mu}_x)\right] \end{align}\]

根据这个公式,可以得到真实分布 \(q(x_{t-1}|x_t, x_0)\) 和参数化学习分布 \(p_\theta(x_{t-1}|x_t)\) 的 KL 散度为

(2.1.62)\[\begin{split}\begin{align} & \quad \,\KL{q(x_{t-1}|x_t, x_0)}{p_{{\theta}}(x_{t-1}|x_t)} \\ &= \KL{\mathcal{N}(x_{t-1}; {\mu}_q,{\Sigma}_q(t))}{\mathcal{N}(x_{t-1}; {\mu}_{{\theta}},{\Sigma}_q(t))}\\ &=\frac{1}{2}\left[\log\frac{|{\Sigma}_q(t)|}{|{\Sigma}_q(t)|} - d + \text{tr}({\Sigma}_q(t)^{-1}{\Sigma}_q(t)) + ({\mu}_{{\theta}}-{\mu}_q)^T {\Sigma}_q(t)^{-1} ({\mu}_{{\theta}}-{\mu}_q)\right]\\ &=\frac{1}{2}\left[\log1 - d + d + ({\mu}_{{\theta}}-{\mu}_q)^T {\Sigma}_q(t)^{-1} ({\mu}_{{\theta}}-{\mu}_q)\right]\\ &=\frac{1}{2}\left[({\mu}_{{\theta}}-{\mu}_q)^T {\Sigma}_q(t)^{-1} ({\mu}_{{\theta}}-{\mu}_q)\right]\\ &=\frac{1}{2}\left[({\mu}_{{\theta}}-{\mu}_q)^T \left(\sigma_q^2(t)\textbf{I}\right)^{-1} ({\mu}_{{\theta}}-{\mu}_q)\right]\\ &=\frac{1}{2\sigma_q^2(t)}\left[\left\lVert{\mu}_{{\theta}}-{\mu}_q\right\rVert_2^2\right] \end{align}\end{split}\]

极大化 ELBO 函数,等同于极小化 \(\KL{q(x_{t-1}|x_t, x_0)}{p_{{\theta}}(x_{t-1}|x_t)}\) ,又相当于极小化 \(\left[\left\lVert{\mu}_{{\theta}}-{\mu}_q\right\rVert_2^2\right]\) ,又是一个均方误差。 \(\mu_{\theta}\) 是参数化模型学习的分布 \(p_\theta(x_{t-1}|x_t)\) 的均值, 也是我们的模型需要预测(输出)的值, 显然 \(\mu_{\theta}\) 是和 \(x_t\) 相关的,即它肯定是 \(x_t\) 的一个函数, 也就是说我们的参数化模型输入 \(x_t\) ,并且输出 \(\mu_{\theta}\)。 前面我们已经推到出 \({\mu}_q\)

(2.1.63)\[\mu_q(x_t, x_0) = { \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)x_0}{1 -\bar\alpha_{t}}}\]

从中可以看出 \({\mu}_q\) 本身也是一个关于 \(x_t\) 的函数, 我们的目标是使得 \(\mu_{\theta}\)\({\mu}_q\) 尽量接近, 那么我们完全可以把 \({\mu}_{\theta}\) 人为定义成类似 \({\mu}_q\)公式(2.1.63)) 的计算方式,即

(2.1.64)\[\mu_{\theta}={\mu}_{{\theta}}(x_t, t) = \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)\hat x_{{\theta}}(x_t, t)}{1 -\bar\alpha_{t}}\]

模型不再直接预测 \({\mu}_q\) ,而是预测输出 \(\hat x_{{\theta}}(x_t, t)\) , 然后利用公式 公式(2.1.64) 计算得到 \({\mu}_q\) 的预测值, 这样一来,\(\left\lVert{\mu}_{{\theta}}-{\mu}_q\right\rVert_2^2\) 可以简化成

(2.1.65)\[ \begin{align}\begin{aligned}& \quad \left\lVert{\mu}_{{\theta}}-{\mu}_q \right\rVert_2^2\\&= \left\lVert \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)\hat x_{{\theta}}(x_t, t)}{1 -\bar\alpha_{t}} - { \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)x_0}{1 -\bar\alpha_{t}}} \right\rVert_2^2\\&= \left\lVert \frac{ \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)\hat x_{{\theta}}(x_t, t)}{1 -\bar\alpha_{t}} - { \frac{\sqrt{\bar\alpha_{t-1}}(1-\alpha_t)x_0}{1 -\bar\alpha_{t}}} \right\rVert_2^2\\&= \left\lVert \frac{ \sqrt{\bar\alpha_{t-1}}(1-\alpha_t) } {1 -\bar\alpha_{t}} ( \hat x_{{\theta}}(x_t, t) - x_0 ) \right\rVert_2^2\\&= \frac{ \sqrt{\bar\alpha_{t-1}}(1-\alpha_t) } {1 -\bar\alpha_{t}} \left\lVert ( \hat x_{{\theta}}(x_t, t) - x_0 ) \right\rVert_2^2\end{aligned}\end{align} \]

对于 \(2\le t \le T\) 的每一时刻 \(t\), 我们的参数化模型输入为 \(x_t, t\) , 预测输出 \(\hat x_{{\theta}}(x_t, t)\), 我们要最小化这个输出值和 \(x_0\) 的平方误差。 极大化 ELBO 函数( 公式(2.1.55) )中的第3项的负的 KL 散度, 就等价于极小化模型输出和 \(x_0\) 的平方误差。

(2.1.66)\[ \begin{align}\begin{aligned}& {\operatorname{\arg\max}}_{\theta} \left[ -\sum_{t=2}^{T} \mathbb{E}_{q(x_{t}|x_0)} \left[\KL{q(x_{t-1}|x_t, x_0)}{p_{{\theta}}(x_{t-1}|x_t)}\right] \right ]\\& \Leftrightarrow {\operatorname{\arg\min}}_{\theta} \left [ \sum_{t=2}^{T} \frac{1}{2\sigma_q^2(t)} \frac{ \sqrt{\bar\alpha_{t-1}}(1-\alpha_t) } {1 -\bar\alpha_{t}} \mathbb{E}_{q(x_{t}|x_0)} \left [ \left\lVert ( \hat x_{{\theta}}(x_t, t) - x_0 ) \right\rVert_2^2 \right ] \right ]\\ & \Leftrightarrow {\operatorname{\arg\min}}_{\theta} \left [ \sum_{t=2}^{T} \mathbb{E}_{q(x_{t}|x_0)} \left [ \left\lVert ( \hat x_{{\theta}}(x_t, t) - x_0 ) \right\rVert_2^2 \right ] \right ]\end{aligned}\end{align} \]

最终的目标函数

最后我们整合一下 公式(2.1.57)公式(2.1.66) 就得到了最终的目标函数,可以看到最终的目标函数其实非常简洁, 对任意 \(t \in [1,T]\), 参数化模型输入 \(x_t\)\(t\),输出 \(\hat x_{{\theta}}(x_t, t)\) ,模型的目标是优化预测值 \(\hat x_{{\theta}}(x_t, t)\) 和真实值 \(x_0\) 之间的平方误差。 公式中期望 \(\mathbb{E}_{q(x_{t}|x_0)}\) 的解决方法和 VAE 中一样,随机采样 \(x_t\) 即可,采样次数记为 \(L\) 次,\(L\) 可以设为一个超参数,从经验上看 \(L=1\) 其实就可以了。至于如何采样 \(x_t\) 的值,显然在前向过程中通过 公式(2.1.49) 可以直接计算得到。

(2.1.67)\[ \begin{align}\begin{aligned} & {\operatorname{\arg\max}}_{\theta} \quad \text{ELBO}\\ & \Leftrightarrow {\operatorname{\arg\max}}_{\theta} \left [ \mathbb{E}_{q(x_{1}|x_0)}\left[\ln p_{{\theta}}(x_0|x_1)\right] - \sum_{t=2}^{T} \mathbb{E}_{q(x_{t}|x_0)}\left[\KL{q(x_{t-1}|x_t, x_0)}{p_{{\theta}}(x_{t-1}|x_t)}\right] \right]\\ &\Leftrightarrow {\operatorname{\arg\min}}_{\theta} \left[ \mathbb{E}_{q(x_{1}|x_0)} \left [ \left\lVert x_0 - \hat x_{0}(x_1,t=1) \right\rVert_2^2 \right ] \right ] + \left [ \sum_{t=2}^{T} \mathbb{E}_{q(x_{t}|x_0)} \left [ \left\lVert( \hat x_{{\theta}}(x_t, t) - x_0 ) \right\rVert_2^2 \right ] \right ]\\ & \Leftrightarrow {\operatorname{\arg\min}}_{\theta} \sum_{t=1}^{T} \mathbb{E}_{q(x_{t}|x_0)} \left [ \left\lVert( \hat x_{{\theta}}(x_t, t) - x_0 ) \right\rVert_2^2 \right ]\end{aligned}\end{align} \]

在这个变分扩散模型中,前向过程不需要学习,模型只学习解码分布 \(p_{\theta}(x_{t-1}|x_t)\)。 在模训练阶段前向过程和逆向过程都要计算,而在模型推理(应用)阶段,不需要前向过程,仅有逆向即可。 在推理时,输入一个随机高斯噪声数据,然后迭代(循环)执行 \(p_{\theta}(x_{t-1}|x_t)\) 最终还原成(生成)一个新的有效图片。

2.1.5. 图片生成(采样)过程

训练完成的神经网络模型为 \(\hat{x}_{\theta}(x_t, t)\) ,它的输入 \(x_t\)\(t\) ,输出是对 \(x_0\) 的预测。有个它之后,就可以得到逆过程中的条件概率分布 \(p(x_{t-1}|x_{t})\) 的一个近似表示。

(2.1.68)\[ \begin{align}\begin{aligned}p(x_{t-1}|x_{t}) &\approx p_{\theta} (x_{t-1}|x_{t})\\&\approx q(x_{t-1}|x_t, x_0)\\&\sim \mathcal{N} (x_{t-1} ; \mu_q(x_t, x_0) , \Sigma_q(t) )\end{aligned}\end{align} \]

其中,均值 \(\mu_q(x_t, x_0)\) 和 方差 \(\Sigma_q(t)\) 分别为

(2.1.69)\[ \begin{align}\begin{aligned}\mu_q(x_t, x_0) &= { \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t) \hat{x}_{\theta}(x_t, t) }{1 -\bar\alpha_{t}}}\\\Sigma_q(t) &= \frac{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}{ 1 -\bar\alpha_{t}} \textit{I} = \sigma_q^2(t) \textit{I}\end{aligned}\end{align} \]

逆向生成(采样)过程可以简述为

  1. 设定 \(T=1000\)

  2. 从一个标准高斯分布中采样得到 \(x_T\)

  3. 利用神经网络模型计算 \(\hat{x}_{\theta}(x_t, t)\)

  4. 计算出 \(p(x_{t-1}|x_{t})\) 的均值 \(\mu_q(x_t, x_0)\) 和方差 \(\Sigma_q(t)\)

  5. \(p(x_{t-1}|x_{t})\) 随机采样的到 \(x_{t-1}\)

  6. 重复步骤3~5,直到 \(t=1\)

在每一个时刻 \(t\) 模型预测都是原始的 \(x_0\) ,从直觉上讲,感觉就不是很合理。当 \(t\) 比较大时,\(x_t\) 距离 \(x_0\) 应该是比较远的, 虽然表示时刻的 \(t\) 也作为模型的输入,但同样的模型参数要同时解决不同的 \(t\),其实是比较困难的。 事实也证明了 DPM 效果确实不是很靠谱,生成的图像质量通常并不高, 在论文[3] 中做了相应的对比实现,证明了这个这个结论。 所以虽然 DPM 在 2015年就提出了,但一直没有引发热度, 直到2020年,Ho et al 等人发表了论文 DPM 的改进模型 “Denoising Diffusion Probabilistic Models (DDPMs)” ,进一步 OPEN AI 发布了基于 DDPM 的 DALL-E ,DALL-E 的效果是非常显著的,这才是得扩散模型得到足够的关注。 下一节讨论 DDPM 做了哪些关键的改进。

2.2. 降噪扩散概率模型(Denoising diffusion probabilistic model,DDPM)

在 2020年 Ho et al 等人发表了论文 “Denoising Diffusion Probabilistic Models (DDPMs)” ,DDPM 对 DPM 做了关键的改进和优化,解决了 DPM 的一些不足,使得扩散模型生成的图像质量得到了大幅提升, 这才使得扩散模型在图像生成领域大放异彩。 DDPM 做的关键改进就是参数化模型预测的内容做了调整:

  • 不再是预测原始的 \(x_0\),而是预测每一个时刻添加的噪声,降低了模型的学习难度。

回顾下 公式(2.1.49),在前向过程中,我们得到 \(q(x_t|x_0)\) 的分布为:

(2.2.56)\[q(x_t|x_0) \sim \mathcal{N}(x_t; \sqrt{\bar{\alpha}_t } \ x_0, (1- \bar{ \alpha}_t) \textit{I})\]

\(x_t\) 的采样过程为

(2.2.57)\[x_t = \sqrt{\bar{\alpha}_t } \ x_0 + \sqrt{1- \bar{ \alpha}_t } \ \epsilon_t \ \ \ , \bar{\alpha}_t = \prod_{i=1}^t \alpha_i ,\ \ \epsilon_t \sim \mathcal{N}(0,\textit{I})\]

公式(2.2.57) 表达了从 \(x_0\) 计算 \(x_t\) 的关系, 现在可以把它反过来,用 \(x_t\) 得到 \(x_0\)

(2.2.58)\[x_0 = \frac{x_t -\sqrt{1- \bar{ \alpha}_t } \ \epsilon_t }{ \sqrt{\bar{\alpha}_t } } ,\ \ \epsilon_t \sim \mathcal{N}(0,\textit{I})\]

我们在 公式(2.1.62) 的基础上,重新推导下 ELBO 中第三项 KL 散度, 为方便阅读,把 公式(2.1.62) 复制到这里。

(2.2.59)\[\KL{q(x_{t-1}|x_t, x_0)}{p_{{\theta}}(x_{t-1}|x_t)} =\frac{1}{2\sigma_q^2(t)}\left[\left\lVert{\mu}_{{\theta}}-{\mu}_q\right\rVert_2^2\right]\]

其中 \({\mu}_q\) 代表逆过程真实分布 \(q(x_{t-1}|x_t, x_0)\) 的均值, \(\mu_{\theta}\) 代表需要学习的近似逆过程分布 \(p_{{\theta}}(x_{t-1}|x_t)\) 的均值,同时也是我们的参数化模型的预测输出值。 假定 \(q(x_{t-1}|x_t, x_0)\)\(p_{{\theta}}(x_{t-1}|x_t)\) 的方差是相同的,并且都是 \(\sigma_q^2(t)\)。 这次利用 公式(2.2.58)\({\mu}_q\) 重新推导

(2.2.60)\[\begin{split}\begin{align} {\mu}_q &= {\mu}_q(x_t, x_0) = \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)x_0}{1 -\bar\alpha_{t}}\\ &= \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)\frac{x_t - \sqrt{1 - \bar\alpha_t}\epsilon }{\sqrt{\bar\alpha_t}}}{1 -\bar\alpha_{t}}\\ &= \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + (1-\alpha_t)\frac{x_t - \sqrt{1 - \bar\alpha_t}\epsilon }{\sqrt{\alpha_t}}}{1 -\bar\alpha_{t}}\\ &= \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t}}{1 - \bar\alpha_t} + \frac{(1-\alpha_t)x_t}{(1-\bar\alpha_t)\sqrt{\alpha_t}} - \frac{(1 - \alpha_t)\sqrt{1 - \bar\alpha_t}\epsilon }{(1-\bar\alpha_t)\sqrt{\alpha_t}}\\ &= \left(\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1 - \bar\alpha_t} + \frac{1-\alpha_t}{(1-\bar\alpha_t)\sqrt{\alpha_t}}\right)x_t - \frac{(1 - \alpha_t)\sqrt{1 - \bar\alpha_t}}{(1-\bar\alpha_t)\sqrt{\alpha_t}}\epsilon \\ &= \left(\frac{\alpha_t(1-\bar\alpha_{t-1})}{(1 - \bar\alpha_t)\sqrt{\alpha_t}} + \frac{1-\alpha_t}{(1-\bar\alpha_t)\sqrt{\alpha_t}}\right)x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}\epsilon \\ &= \frac{\alpha_t-\bar\alpha_{t} + 1-\alpha_t}{(1 - \bar\alpha_t)\sqrt{\alpha_t}}x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}\epsilon \\ &= \frac{1-\bar\alpha_t}{(1 - \bar\alpha_t)\sqrt{\alpha_t}}x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}\epsilon \\ &= \frac{1}{\sqrt{\alpha_t}}x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}\epsilon \end{align}\end{split}\]

接下,把模型的预测输出 \(\mu_{\theta}\) ,也按照这个形式进行参数化

(2.2.61)\[\mu_{\theta}={\mu}_{{\theta}}(x_t, t) = \frac{1}{\sqrt{\alpha_t}}x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}} {\hat\epsilon}_{ {\theta}}(x_t, t)\]

\({\hat\epsilon}_{ {\theta}}(x_t, t)\) 表示模型的预测输出,然后通过 公式(2.2.61) 计算得到 \(\mu_{\theta}\)。 把 公式(2.2.60)公式(2.2.61) 一起带入 公式(2.2.59)

(2.2.62)\[\begin{split}\begin{align} & \KL{q(x_{t-1}|x_t, x_0)}{p_{{\theta}}(x_{t-1}|x_t)} \\ &=\frac{1}{2\sigma_q^2(t)}\left[\left\lVert{\mu}_{{\theta}}-{\mu}_q\right\rVert_2^2\right]\\ &= \frac{1}{2\sigma_q^2(t)}\left[\left\lVert\frac{1}{\sqrt{\alpha_t}}x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}{\hat\epsilon}_{{\theta}}(x_t, t) - \frac{1}{\sqrt{\alpha_t}}x_t + \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}\epsilon \right\rVert_2^2\right]\\ &= \frac{1}{2\sigma_q^2(t)}\left[\left\lVert \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}\epsilon - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}{\hat\epsilon}_{{\theta}}(x_t, t)\right\rVert_2^2\right]\\ &= \frac{1}{2\sigma_q^2(t)}\left[\left\lVert \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}(\epsilon - {\hat\epsilon}_{{\theta}}(x_t, t))\right\rVert_2^2\right]\\ &= \frac{1}{2\sigma_q^2(t)}\frac{(1 - \alpha_t)^2}{(1 - \bar\alpha_t)\alpha_t}\left[\left\lVert\epsilon - {\hat\epsilon}_{{\theta}}(x_t, t)\right\rVert_2^2\right] \end{align}\end{split}\]

其中 \(\epsilon\) 表示从时刻 \(t-1\) 到时刻 \(t\) 的前向过程中, 按照公式 公式(2.1.48) 添加到 \(x_{t-1}\) 的随机高斯噪声。

(2.2.63)\[ \begin{align}\begin{aligned}x_t &=\sqrt{\alpha_t} \ x_{t-1} + \sqrt{1- \alpha_t } \ \epsilon\\&= \sqrt{\bar{\alpha}_t } \ x_0 + \sqrt{1- \bar{ \alpha}_t } \ \epsilon \ \ \ , \bar{\alpha}_t = \prod_{i=1}^t \alpha_i ,\ \ \epsilon \sim \mathcal{N}(0,\textit{I})\\&\sim \mathcal{N}(\sqrt{\bar{\alpha_t} } \ x_0, (1- \bar{ \alpha_t}) \textit{I})\end{aligned}\end{align} \]

对于每一个时刻 \(t\):

  1. 从标准高斯分布 \(\epsilon\) 中随机采样一个噪声, 注意每一个时刻都是独立重新采样的随机高斯噪声,所以不同时刻 \(t\)\(\epsilon\) 是不一样的值。

  2. 代入上式中计算得到 \(x_t\)

  3. \(x_t\)\(t\) 输入参数神经网络模型,模型输出值 \(\hat \epsilon_{{\theta}}(x_t, t)\) 作为预测噪声 。

  4. 最小化 \(\hat \epsilon_{{\theta}}(x_t, t)\)\(\epsilon\) 之间的平方误差。

参数化模型直接输出值是 \({\hat\epsilon}_{ {\theta}}(x_t, t)\) ,这是就相当于 模型预测是每一步从 \(x_{t-1}\)\(t\) 添加的噪声。 由于预测噪声之后,再通过 公式(2.2.61) 就得到了逆向过程 \(p_{{\theta}}(x_{t-1}|x_t)\) 的期望。

从直觉上讲,模型预测每一步的噪声比直接预测 \(x_0\) 要容易很多, 从 DDPM 的论文中也证明了一点。 相比原来的 DPM,DDPM 能生成更高质量的图像数据。

论文中给出的训练过程伪代码:

../_images/diffusion_training.webp

图 2.2.1 DDPM训练过程伪代码(来自 [3]

论文中给出的采样(图像生成)过程伪代码如下:

../_images/diffusion_sampling.webp

图 2.2.2 DDPM采样过程伪代码(来自 [3]

其中 \(\beta_t = 1 - \alpha_t\)\(\sigma_{t}\) 是逆过程分布 \(q(x_{t-1}|x_t, x_0)\) 的标准差, 参照 公式(2.1.58)

在采样(图像生成)过程有一个细节,最后生成的图像(输出值)\(\tilde{x}_0\) 并不是通过 \(x_{t-1}\) 计算得到的,二者计算方式不一样。 逆过程中,\(x_{t-1}\) 的分布是 \(q(x_{t-1}|x_t, x_0)\), 它是一个高斯分布的随机变量,它需要通过随机采样得到, 所以上图伪代码中,是先计算出其均值 \(\tilde{\mu}\), 然后再加上一项随机高斯噪声 \(\sigma_t z\)

而最终生成的图像 \(x_0\) 是根据 公式(2.2.63) 的逆运算得到的。 其实当 \(t=1\) 时, \(\tilde{x_0}\)\(\tilde{\mu}\) 是相同的。 当 \(t=1\) 时,有 \(\bar{\alpha}_1 = \alpha_1\) ,此时

(2.2.64)\[ \begin{align}\begin{aligned}\tilde{\mu} &= \frac{1} {\sqrt{\alpha_t}} (x_t - \frac{\beta_t}{\sqrt{1-\bar{\alpha_t}} } \epsilon_{\theta}(x_t,t) )\\&= \frac{1} {\sqrt{\alpha_1}} (x_1 - \frac{1-\alpha_1}{\sqrt{1-\bar{\alpha_1}} } \bar{\epsilon} )\\&= \frac{1} {\sqrt{\alpha_1}} (x_1 - \frac{1-\alpha_1}{\sqrt{1-\bar{\alpha_1}} } \bar{\epsilon} )\\&= \frac{1} {\sqrt{\alpha_1}} (x_1 - \frac{1-\alpha_1}{\sqrt{1-\alpha_1} } \bar{\epsilon} )\\&= \frac{1} {\sqrt{\alpha_1}} (x_1 - \sqrt{1-\alpha_1} \bar{\epsilon} )\\&= \tilde{x}_0\end{aligned}\end{align} \]

也就是说我们最终输出的 \(\tilde{x}_0\)\(q(x_{t-1}|x_t, x_0,t=1)\) 的期望, 而不是 \(q(x_{t-1}|x_t, x_0,t=1)\) 的采样值,这和我们的 ELBO 函数的第一项是一致的。 事实上,所有的回归模型最终输出的都是期望值。

2.3. 基于分数的解释(Score-based DDPM)

在2019年,宋旸 [] 提出了一种基于分数的扩散模型 [5] , 实际上这种基于分数的扩散模型和本章基于马尔科夫的扩散模型是等价的, 论文 [] [2] 给出了这种等价说明。 后续的章节我们再单独详细的介绍宋旸论文中基于分数扩散模型的原理, 这里我们暂且按照论文 [] 的内容给出和 DDPM 等价的分数模型。

首先需要一个关于 Tweedie 公式(Tweedie’s Formula)的背景知识。 在统计学中,经常需要根据观测样本估计一个概率分布的未知参数,估计算法有很多,比如最大似然估计、贝叶斯估计等等。 Tweedie 公式是一种估计指数族分布均值参数的方法,如果你不知道什么是指数族,可以查看本博客中广义线性模型的内容。 高斯分布就是指数族分布中的一种,所以可以 Tweedie 公式估计高斯分布的均值参数。 这里省略 Tweedie 公式的证明过程,只关注怎么使用 Tweedie 公式。

假设有一个(多元)高斯随机变量 \(z \sim (\mathcal{N}{z;\mu_z,\Sigma_z})\)Tweedie 公式可以写为

(2.3.17)\[\mathbb{E}[\mu_z | z] = z + \Sigma_z \underbrace{ \nabla_z \log p(z)}_{\text{score}}\]

其中等式左边 \(\mathbb{E}[\mu_z | z]\) 表示:有 \(z\) 的样本条件下, \(\mu_z\) 的期望,直观地理解,就是用 \(z\) 的样本估计 \(\mu_z\) 的值。 等式右边就是具体计算方式,在这里只需要 \(z\) 的一条样本就可以。 \(\log p(z)\) 就是观测样本的对数似然, \(\nabla_z \log p(z)\) 就是观测样本的对数似然的一阶导数,又名梯度, 也叫 Stein 分数(Stein score),简称分数(score)。 注意,还有另外一个分数, fisher score ,二者不太一样,不要搞混了。 fisher score 是对数似然关于 参数 的梯度 \(\nabla_{\theta} \log p(z;\theta)\) ; Stein score 是对数似然关于 观测变量(样本) 的梯度 \(\nabla_z \log p(z)\)

回到 公式(2.1.49),在前向过程中,我们得到 \(q(x_t|x_0)\)

(2.3.18)\[q(x_t|x_0) = \mathcal{N}(x_{t} ; \sqrt{\bar\alpha_t} x_0, \left(1 - \bar\alpha_t\right)\textit{I})\]

\(q(x_t|x_0)\) 是一个高斯分布,虽然已知它的均值是 \(\sqrt{\bar\alpha_t} x_0\), 但我们仍可以利用 Tweedie 公式估计它的均值参数

(2.3.19)\[\sqrt{\bar\alpha_t} x_0 = \mathbb{E}\left[{\mu}_{x_t}|x_t\right] = x_t + (1 - \bar\alpha_t)\nabla_{x_t}\log p(x_t)\]

对上式进行移项可以得到一个新的关系式

(2.3.20)\[x_0 = \frac{x_t + (1 - \bar\alpha_t)\nabla\log p(x_t)}{\sqrt{\bar\alpha_t}}\]

这个关系式是一个分数表达 \(x_0\) 的等式。 现在回到逆过程中的真实分布 \(q(x_{t-1}|x_t,x_0)\) ,它的均值 \({\mu}_q(x_t, x_0)\)公式(2.1.63) 中已经给出,这里把 公式(2.3.20) 带入进去

(2.3.21)\[\begin{split}\begin{align} {\mu}_q(x_t, x_0) &= \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)x_0}{1 -\bar\alpha_{t}}\\ &= \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)\frac{x_t + (1 - \bar\alpha_t)\nabla\log p(x_t)}{\sqrt{\bar\alpha_t}}}{1 -\bar\alpha_{t}}\\ &= \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + (1-\alpha_t)\frac{x_t + (1 - \bar\alpha_t)\nabla\log p(x_t)}{\sqrt{\alpha_t}}}{1 -\bar\alpha_{t}}\\ &= \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t}}{1 - \bar\alpha_t} + \frac{(1-\alpha_t)x_t}{(1-\bar\alpha_t)\sqrt{\alpha_t}} + \frac{(1 - \alpha_t)(1 - \bar\alpha_t)\nabla\log p(x_t)}{(1-\bar\alpha_t)\sqrt{\alpha_t}}\\ &= \left(\frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1 - \bar\alpha_t} + \frac{1-\alpha_t}{(1-\bar\alpha_t)\sqrt{\alpha_t}}\right)x_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}}\nabla\log p(x_t)\\ &= \left(\frac{\alpha_t(1-\bar\alpha_{t-1})}{(1 - \bar\alpha_t)\sqrt{\alpha_t}} + \frac{1-\alpha_t}{(1-\bar\alpha_t)\sqrt{\alpha_t}}\right)x_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}}\nabla\log p(x_t)\\ &= \frac{\alpha_t-\bar\alpha_{t} + 1-\alpha_t}{(1 - \bar\alpha_t)\sqrt{\alpha_t}}x_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}}\nabla\log p(x_t)\\ &= \frac{1-\bar\alpha_t}{(1 - \bar\alpha_t)\sqrt{\alpha_t}}x_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}}\nabla\log p(x_t)\\ &= \frac{1}{\sqrt{\alpha_t}}x_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}} \nabla\log p(x_t) \end{align}\end{split}\]

重新按照个形式设定参数化分布 \(p_{\theta}(x_{t-1}|x_t)\) 均值 \(\mu_{\theta}\) 的表达式

(2.3.22)\[\mu_{\theta}(x_t, t) = \frac{1}{\sqrt{\alpha_t}}x_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}}s_{{\theta}}(x_t, t)\]

现在 \(s_{{\theta}}(x_t, t)\) 是模型学习预测的内容, 相当于模型直接预测的是分数 \(\nabla_{x_t} \log p(x_t)\), 最后我们推导下目标函数 ELBO 中相关的 KL 散度。

(2.3.23)\[\begin{split}\begin{align} & \quad \,\argmin_{{\theta}} \KL{q(x_{t-1}|x_t, x_0)}{p_{{\theta}}(x_{t-1}|x_t)} \nonumber \\ &= \argmin_{{\theta}}\KL{\mathcal{N}\left(x_{t-1}; {\mu}_q,{\Sigma}_q\left(t\right)\right)}{\mathcal{N}\left(x_{t-1}; {\mu}_{{\theta}},{\Sigma}_q\left(t\right)\right)}\\ &=\argmin_{{\theta}}\frac{1}{2\sigma_q^2(t)}\left[\left\lVert\frac{1}{\sqrt{\alpha_t}}x_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}}s_{{\theta}}(x_t, t) - \frac{1}{\sqrt{\alpha_t}}x_t - \frac{1 - \alpha_t}{\sqrt{\alpha_t}}\nabla\log p(x_t)\right\rVert_2^2\right]\\ &=\argmin_{{\theta}}\frac{1}{2\sigma_q^2(t)}\left[\left\lVert \frac{1 - \alpha_t}{\sqrt{\alpha_t}}s_{{\theta}}(x_t, t) - \frac{1 - \alpha_t}{\sqrt{\alpha_t}}\nabla\log p(x_t)\right\rVert_2^2\right]\\ &=\argmin_{{\theta}}\frac{1}{2\sigma_q^2(t)}\left[\left\lVert \frac{1 - \alpha_t}{\sqrt{\alpha_t}}(s_{{\theta}}(x_t, t) - \nabla\log p(x_t))\right\rVert_2^2\right]\\ &=\argmin_{{\theta}}\frac{1}{2\sigma_q^2(t)}\frac{(1 - \alpha_t)^2}{\alpha_t}\left[\left\lVert s_{{\theta}}(x_t, t) - \nabla\log p(x_t)\right\rVert_2^2\right] \label{eq:123} \end{align}\end{split}\]

最后的目标函数仍然是均方误差的形式,和原始扩散模型以及基于噪声的扩散模型是高度一致的。 公式中的 \(\nabla\log p(x_t)\) 是梯度的真实值,这个可以直接计算得到,pytorch 中可以计算指定变量的梯度。 \(s_{{\theta}}(x_t, t)\) 代表一个神经网络模型,它去学习和预测 \(\nabla\log p(x_t)\)

接下来看下梯度和噪声的关系,结合 公式(2.3.20)公式(2.2.58) 可得

(2.3.24)\[\begin{split}\begin{align} x_0 = \frac{x_t + (1 - \bar\alpha_t)\nabla\log p(x_t)}{\sqrt{\bar\alpha_t}} &= \frac{x_t - \sqrt{1 - \bar\alpha_t}{\epsilon}_t}{\sqrt{\bar\alpha_t}}\\ \therefore (1 - \bar\alpha_t)\nabla\log p(x_t) &= -\sqrt{1 - \bar\alpha_t}{\epsilon}_t\\ \nabla\log p(x_t) &= -\frac{1}{\sqrt{1 - \bar\alpha_t}}{\epsilon}_t \end{align}\end{split}\]

我们知道梯度的方向,是函数极值的方向。 梯度 \(\nabla\log p(x_t)\) 指向的是对数概率函数 \(\log p(x_t)\) 概率密度最大的方向, 它和添加的噪声相差一个负号,正好是相反的,梯度的方向是噪声的反方向,这十分符合直觉。

2.4. 扩散模型的三种等价表示

到这里我们一共介绍了扩撒模型的三种不同形式,1)直接预测初始样本;2)预测噪声;3)预测分数。 这三种形式理论是等价的,然后对于模型的学习难度以及灵活性的却是不同的。 这里我们把这三种形式进行总结对比一下,方便理解和复习。

逆向过程中真实的(后验)分布为 \(q(x_{t-1}|x_t,x_0)\) ,它是一个条件高斯分布,均值记为 \(\mu_q\), 方差记为 \(\Sigma_{q(t)}\)

(2.4.23)\[q(x_{t-1}|x_t,x_0) \sim \mathcal{N}(x_{t-1},\mu_q,\Sigma_{q(t)})\]

方差 \(\Sigma_{q(t)}\) 的表达式为

(2.4.24)\[\Sigma_q(t) = \frac{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}{ 1 -\bar\alpha_{t}} \textit{I} = \sigma_q^2(t) \textit{I}\]

方差看做一个已知的常量,默认模型不去学习它。 我们的参数化模型学习的分布记为 \(p_{\theta}(x_{t-1}|x_t)\), 我们要令 \(p_{\theta}(x_{t-1}|x_t)\) 尽量近似真实分布 \(q(x_{t-1}|x_t,x_0)\)。 实际实现时,模型学习拟合的是真实分布 \(q(x_{t-1}|x_t,x_0)\) 的均值 \(\mu_q\), 然而它可以推导出三种不同等价的形式,因此也就有三种不同的学习方式。

第一种形式:学习初始样本

均值 \(\mu_q\) 的第一种表达式为

(2.4.25)\[\mu_q(x_t, x_0) = { \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)x_0}{1 -\bar\alpha_{t}}}\]

按照个形式设定参数化分布 \(p_{\theta}(x_{t-1}|x_t)\) 的均值 \(\mu_{\theta}\)

(2.4.26)\[\mu_{\theta}={\mu}_{{\theta}}(x_t, t) = \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})x_{t} + \sqrt{\bar\alpha_{t-1}}(1-\alpha_t)\hat x_{{\theta}}(x_t, t)}{1 -\bar\alpha_{t}}\]

此时模型实际学习的是初始样本 \(x_0\), 显然这种形式对模型来说难度比较大,所以效果并不突出。

第二种形式:学习噪声

均值 \(\mu_q\) 的第二种表达式为

(2.4.27)\[\mu_q(x_t, x_0) = \frac{1}{\sqrt{\alpha_t}}x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}\ \epsilon\]

同理,按照个形式设定参数化分布 \(p_{\theta}(x_{t-1}|x_t)\) 的均值 \(\mu_{\theta}\)

(2.4.28)\[\mu_{\theta}={\mu}_{{\theta}}(x_t, t) = \frac{1}{\sqrt{\alpha_t}}x_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}} {\hat\epsilon}_{ {\theta}}(x_t, t)\]

此时模型学习预测的噪声 \(\epsilon\) ,相比直接学习初始样本 \(x_0\),每一步学习噪声简单了许多,模型的效果提升很大。

第三种形式:学习分数

依据 Tweedie 公式,均值 \(\mu_q\) 的表达式可以写成基于分数的形式

(2.4.29)\[{\mu}_q(x_t, x_0) = \frac{1}{\sqrt{\alpha_t}}x_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}}\nabla\log p(x_t)\]

同样,按照个形式设定参数化分布 \(p_{\theta}(x_{t-1}|x_t)\) 的均值 \(\mu_{\theta}\)

(2.4.30)\[{\mu}_q(x_t, x_0) = \frac{1}{\sqrt{\alpha_t}}x_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}} s_{\theta}(x_t,t)\]

此时模型学习预测的是分数(梯度) \(\nabla\log p(x_t)\)。 相比学习噪声,学习分数有个好处,他可以令我在逆过程采样图片时采用基于分数的采样算法, 而基于分数的采样算法种类很多,这极大的增加了算法的灵活性。

2.5. 改进降噪扩散概率模型(Improved Denoising Diffusion Probabilistic Models,IDDPM)

  1. 学习方差

  2. 线性的schedule 改成余弦schedule

待补充

源码 https://github.com/openai/improved-diffusion

2.6. 参考文献