生成式模型的基本目的
以图片生成和处理任务为例,假设我们的研究对象是一个16 * 16 的RGB图片,抛开对图片内容的限制,假如我们完全随机地生成这样一个尺寸的图片,每个像素的可能取值有 256 * 256 * 256 个,整个图片的可能取值情况有 N = ( 256 \cdot 256 \cdot 256 ) ^ {(16 \cdot 16 )} 个。
我们把可能随机生成的这样的一个图片设为(或者说,记为,表示为)一个随机变量 X,X的所有可能取值个数为 N,则我们可以将其样本空间表示为
$$
S_{X} = \{X_1, X_2, ..., X_N\}
$$
X 这个随机变量必然是服从某个分布的,我们将其记为
$$
X \sim p(X)
$$
(即我们随机生成的这样一个图片X,其取值为a的概率为 p(a) )
(即 P(X = a) = p(a) )
模型生成图片的过程,实际上就是对服从这个服从分布 p(X) 的随机变量X进行采样的过程,采样得到的 X' 就是一个属于 S_X 的图片。
但我们在许多图片生成任务中肯定不希望得到一个在 S_X 上完全均匀、没有特点、没有倾向的随机分布 p(X) ,不然生成出的图像大概率是一堆难以理解的马赛克。比如说,如果我们希望这个模型能生成一些与猫相关的图片,那这个模型对应的分布 p(X) 在与猫相关的图片对应的 X 值上的概率应该更高,而与猫无关的图片的概率应该更小,即:
$$
p(X_{与猫相关}) 较大 \\ p(X_{与猫无关}) 较小
$$
也就是说,在许多任务中,我们对这个分布 p(X) 是有特殊的要求和期望的,我们的目的就是求出这个满足我们要求的分布 p(X) 。只要我们能求出合适的分布 p(X),我们就可以通过对这个分布进行采样得到我们期望的图片了。(实际上不能说是对分布采样,应该说是对一个服从该分布的随机变量进行采样)
当然,不只是图片任务,在其他任务里,这里的 X 也可以是视频音频或文本。
我们怎么求我们希望的分布 p(X) 呢?如果 N 比较小,我们大可以人为地对 p(X) 按我们的偏好进行调整,但实际任务中 N 往往极大,难以直接获取 p(X) ,我们只能用各种方式来逼近这个分布 p(X) 了。
下面是几种估计 p(X) 的方法。
估计分布的方法之一:极大似然估计
极大似然法的原理很直观易懂。在有一定先验知识的情况下,我们可以对 p(X) 的形式进行一定的假设,而其中的待求参数为 \theta,记分布为 p_θ(X) 我们的任务就转化成了如何求目标分布对应的参数 θ。
(举个例子,我们假设已知 p(X) 是一个高斯分布,那我们只需要求该高斯分布对应的均值和方差即可,记 θ = (μ, σ^2) ,这里的 θ 就是我们的待求参数,只要求出 θ 就可以确定分布了)
在图片生成任务的例子中,假设我们有 k 个我们希望这个模型生成的图片样例,那么我们显然是希望这个模型生成这些图片的概率能足够的大,即我们希望最大化如下概率:(我们记样例组成的集合为 T_x)
$$
L_{\theta}(T_X) = \prod_{i=1}^{k}p_{\theta}(X_i)~~~(X_i \in T_X)
$$
那我们希望的参数 θ 应该为
$$
\hat{\theta} = \mathrm{argmax}_{\theta}L_{\theta}(S_X)
$$
这样我们就把分布 p_θ 的估计问题转化成了一个最优化问题,这就是极大似然估计(即希望我们要求的分布能使当下我们期望发生的事发生的概率尽可能地大)。
但实际上,我们首先很难对 p(X) 的形式有一个合理的预先认识和假设(大多数时候我们对其一无所知)。并且,在很多情况下,参数 \theta 对应的参数空间维度极大(比如16 * 16 的图片生成任务中,我们假设每个像素位置的取值为一个高斯分布,那么整个图片服从一个 16 * 16 维的多元高斯分布,一共有 16 * 16 * 2 个待求参数,即 \theta 是一个长度为512的向量,其参数空间的维度为512维),对这样一个高维的参数进行优化,需要相当相当多的数据才能取得有效的估计,数据成本极高。
估计分布的方法之二:隐变量估计
直接求解或估计 p(X) 还是有些困难,而直接用神经网络去拟合 p(X) 效果也往往欠佳。故我们引入一个隐变量 z ,先生成一个随机变量 z ,再由 z 经过某种变换生成最终的随机变量 X 。
(关于为什么要引入隐变量,更加理论化的解释请见另一篇关于贝叶斯方法和变分推断的笔记,这里给出一些更形象的描述https://zhuanlan.zhihu.com/p/355019238)
(因为实际上,观测数据只是我们可以直接观测到的样本指标或特征,而在观测数据之下蕴含的还有一些我们不能直接观测到的、潜在的、未知的特征,并且这些特征可能对观测数据产生影响,我们称其为隐藏特征。观测特征所在的空间为观测空间,隐藏数据所在的空间为隐藏空间。举例说明,考虑一个电商平台的用户行为分析。观测空间是用户在平台上的实际行为数据,如购买记录、浏览历史、搜索关键词等。而隐藏空间可以是用户的潜在兴趣或偏好。例如,假设一个用户购买了几本关于健康和健身的书籍,但观测数据中并没有直接记录用户的兴趣。用户的兴趣偏好会在一定程度上影响其在平台上的行为活动,但不会直接记录在观测数据中。用户在平台上的行为数据在形式上往往是比较复杂的,难以直接找到其一般的规律,那我们可以人为地添加潜在的主题变量来表示“用户的兴趣”,然后想办法构建一个由“用户的兴趣”到“用户的行为数据”的映射,因为“用户的兴趣”这一隐藏特征在表示和含义上远比观测特征简洁直白,可以很方便地分析理解数据,提供更精确和深入的推断和预测。)
(再举个图像的例子,一个16 * 16 的手写数字图像(假设其写的是9),作为直接的观测数据,数据量是较大的,但是其隐藏特征实际上就是一个数字 9 ,我们可以把隐藏特征纳入考虑对观测数据进行分析)
引入隐变量 z 后,我们记其服从于分布 s(z) , 我们可以从较为简单的分布 s(z) 出发,用积分公式间接地求出相对复杂的分布 p(x) :(对的,一般我们引入的这个隐变量就是相对简单的,不然不如不引)(这个公式的依据是 product rule)
$$
p(x) = \int p(x|z)s(z)\mathrm{d}z
$$
这里需要注意:
-
s(z) 分布的表达式形式一般比较简单,比如普通的多元高斯分布,并且往往是预先假设的一个没有未知量的确定分布。
-
z 的自由度一般比较低,或者说其维度比较低,应该是一个相对简单的变量。
-
在实践中,因为“从隐变量 z 经过某种变换生成最终的随机变量 X ” 往往是个确定的变换过程,即 x = f_{\theta}(z) (很多时候我们会用神经网络来实现 f_{\theta} ),所以有如下表达式:(没错,给定 z 的情况下其实际上是个关于 x 的冲激函数)
$$
p_{\theta}(x|z) = \left\{ \begin{array}{ll} 0 & \text{if } x \neq f_{\theta}(z) \\ \infty & \text{if } x = f_\theta (z) \\ \end{array} \right.
$$
-
在实践中从简单分布 s(z) 生成任意分布 p(x) 过程如下:假如我们想求 p(X') , 则先解出 x = f(z) 的解集 \{z_1, z_2, ..., z_m\} ,经过一定计算化简可以发现 p(X') = \sum_{i=1}^m s(z_i)
-
我们或许无法一开始就给出足够合理、正确的变换 f , 但是我们一般是用神经网络来实现 f,神经网络的反应能力是哈基米的七倍,它不舒服的时候会自己训练拟合到合适的函数(就是说我们可以通过训练来拟合到)(而且实际上,这里唯一的待求参数就是 f 这个函数的参数 \theta,所以隐变量模型实际上把一个概率分布的估计问题转化成了一个函数逼近的问题)。
隐变量模型背后的关键思想是:任何一个简单的概率分布经过一个足够复杂的函数后可以映射到任意概率分布。
但是,我们在训练这个模型(或者说,逼近求解 f )的时候,使用的还是极大似然的方法。
与前文类似,我们最大化的目标为:
$$
L_{\theta}(T_X) = \prod_{i=1}^{k}p_{\theta}(X_i)~~~(X_i \in T_X)
$$
其负对数为:
$$
l_{\theta}(T_X) = -\sum_{i=1}^n\ln{p_{\theta}(x_i)} = -\sum_{i=1}^n\ln{\int p_{\theta}(x_i|z)s(z)\mathrm{d}z}
$$
其关于 \theta 的梯度为
$$
\nabla_{\theta}~l_{\theta}(T_X) = -\sum_{i=1}^n \nabla_{\theta} \ln{\int p_{\theta}(x_i|z)s(z)\mathrm{d}z} = -\sum_{i=1}^n \frac {\int \nabla_{\theta} \ln{\int p_{\theta}(x_i|z)s(z)\mathrm{d}z}} {\int p_{\theta}(x_i|z)s(z)\mathrm{d}z}
$$
我们可以用这个梯度对模型参数 \theta 进行更新。
但是这个方法可能存在如下问题:
-
4 中的解集一般很难求,因为 f 在某些情况下并不可逆且较为复杂,不好直接求解。
-
在实践中,为了能让 z 表示更多的信息,我们往往假设其为一个连续的随机变量(比如高斯随机变量),这导致我们必须要计算连续空间上的积分(如果其为离散的我们就可以把积分转化为求和了),往往难以计算准确结果。
可见,这种方法在计算梯度时存在积分难以计算的问题,故我们通常采用蒙特卡洛采样的方法来计算其近似解。
原本的积分可以表示成期望的形式 \int p_\theta(x|z)s(z)\mathrm{d}z = \mathbb{E}_{z \sim s(z)} [p_\theta(x|z)] ,我们可以通过采样估计期望来估计积分。步骤如下:
-
对 z \sim s(z) 进行 m 次采样得到 z_1, z_2,...,z_m
-
可以表示出 p_\theta (x_j|z_j)
并且不难证明 \int p_\theta(x|z)s(z)\mathrm{d}z = \mathbb{E}_{z \sim s(z)} [p_\theta(x|z)] \approx \frac{1} {m} \sum_{j=1}^m p_\theta(x_j|z_j)
蒙特卡洛方法虽然简单,但是还是有两个可改进方向:
-
提高采样效率。这里的 m 要非常大才能得到相对满意的结果 ,并且 z 的取值范围中有一些值很可能产生坏样本,故我们尝试优化,缩小 z 的取值空间,缩小 s(z) 的方差 \sigma^2 ,这样 z 的采样范围也会缩小,采样次数也不需要很大,同时也可以排除一部分坏的z值。(VAE方向)。
-
这里我们将得到的 X' 和真实图片 X 做MSE,但很多真实情况表明MSE并不是很有效的度量,所以另外一种思路是用神经网络来学习一个更好的度量方式 D(x, x')。(GAN方向)
隐变量估计的改良之一:VAE
VAE 理论基础:变分推断
怎样缩小 z 的采样分布的方差呢,之前的采样方案是从 z 的先验概率分布 s(z) 中采样,现在可以考虑从 z 的后验概率分布 p(z|X) 中进行采样。
关于这点,可以这样理解:先验概率 p(z) 是没有限制的任意分布,而后验概率 p(z|X) 是给定 X 后的条件概率,所以更可能采样出能生成 X 的 z 。
但是我们一般无法直接求出后验分布,故我们在变分推断中选择使用另一个分布 q_\theta (z|X) 来近似地逼近、估计 p(z|X) ,拿从 q_\theta (z|X) 采样的结果近似地作为 p(z|X) 中得到的采样。
为什么这种方法叫变分推断呢,在这个问题中我们的优化对象是 q_\theta (z|X) (或者说是其参数 \theta ),而优化的目标是缩小 q_\theta (z|X) 与 p(z|X) 之间的“距离“,我们设其为 dist 。我们一般会表示出 dist = f(\theta) ,而我们的优化手段往往是求 f 关于 \theta 的导数然后使用梯度下降法,而在数学中,变分是一种函数的微小变化或泛函的微小变化的概念,与导数关系密切,故我们称其为变分推断。(说人话就是通过求导优化一个分布去逼近另一个分布)
在变分推断中,我们采用 KL散度 衡量两个概率分布的相似程度,作为两个分布之间的”距离“。
$$
\begin{flalign} &KL(q_\theta(z|x)||p(z|x))\\ =& \int q_\theta(z|x)\ln{\frac{q_\theta(z|x)}{p(z|x)}}\mathrm{d}z\\ =& \mathbb{E}_{z \sim q_\theta(z|x)}[\ln{\frac{q_\theta(z|x)}{p(z|x)}}]\\ =& \mathbb{E}_{z \sim q_\theta(z|x)}[\ln{q_\theta(z|x) - \ln{{p(z|x)}}}]\\ =& \mathbb{E}_{z \sim q_\theta(z|x)}[\ln{q_\theta(z|x) - \ln{\frac{p(x|z)p(z)}{p(x)}}}]\\ =& \mathbb{E}_{z \sim q_\theta(z|x)}[\ln{q_\theta(z|x)} - \ln{p(z)} - ln{p(x|z)}] + \ln{p(x)}\\ =& KL(q_\theta(z|x)||p(z)) - \mathbb{E}_{z \sim q_\theta(z|x)}[\ln{p(x|z)}] + \ln{p(x)} \end{flalign}
$$
整理后得到
$$
\ln{p(x)} - KL(q_\theta(z|x)||p(z|x)) = \mathbb{E}_{z \sim q_\theta(z|x)}[\ln{p(x|z)}] - KL(q_\theta(z|x)||p(z))
$$
这个式子是我们由纯粹推导得来的,但其中各部分又具有相当有实际意义的含义:
-
我们的整体目的其实是要最大化 \ln{p(x)} - KL(q_\theta(z|x)||p(z|x)), 在训练过程中,我喂给模型的图像 x 都是希望其生成的模型,其中 \ln {p(x)} 自然是模型生成图像 x 的对数似然值,我们肯定希望其尽可能大。另一方面, KL(q_\theta(z|x)||p(z|x)) 是我们人为设定的分布q_\theta(z|x)和希望其逼近的目标分布(即z 的后验分布)的KL散度,所以我们自然希望其尽可能的小,即其相反数尽可能的大。
-
等号右面的第一部分,\mathbb{E}_{z \sim q_\theta(z|x)}[\ln{p(x|z)}], z \sim q_\theta(z|x) 表示我们先给定一个 x ,然后我们从给定这个 x 后的 z 的后验分布 q_\theta(z|x) 中采样得到一个隐变量 z,这个过程其实就是给定 x 输出 z ,相当于把 x 编码成了 z ,我们称这个过程encode, 而 p(x|z) 表示给定 z 又生成回 x 的概率,反映的是把 z 解码为 x 的过程,我们称其为decode。用这种方式可以把这个整体理解为对于给定 x ,将其按分布 q_\theta(z|x) 编码为某个 z 后能再成功解码(重构)回 x 的对数似然函数值的期望。如果这个期望足够大,说明得到的 z 是 x 的一个足够好足够有效的表示,是一个能有效反映 x 的隐变量。我们称其为重构误差,希望其能足够的大(这个名字好奇怪,明明越大越好却起名为误差)。
-
对于等号右面的第二项 KL(q_\theta(z|x)||p(z)) ,其表示我们用来拟合后验分布的函数与实际的先验分布的函数之间的KL散度,即我们不希望我们得到的后验函数和原先验分布之间的差距过大。这可以视为一种正则化技巧(至于具体是怎么正则的,下面讲VAE结构时会详细说明)。
又因为 KL(q_\theta(z|x)||p(z|x)) \geq 0,故
$$
\ln{p(x)} \geq \mathbb{E}_{z \sim q_\theta(z|x)}[\ln{p(x|z)}] - KL(q_\theta(z|x)||p(z)) = ELBO
$$
ELBO(Evidence Lower Bound)称为对数似然的下界,此时可以把最大化对数似然转化为最大化ELBO(因为 p(x) 毕竟是我们所不知道的嘛,没法直接最小化KL(q_\theta(z|x)||p(z|x)))(而且只要ELBO这个下界足够大,说明 \ln{(p(x))} 基本不会小)
VAE 的结构与原理
我们上面通过理论推导得出,我们可以把最小化KL(q_\theta(z|x)||p(z|x)) 的目标转化为最大化ELBO。接下来我们只要从工程角度实现 encode 和 decode 的过程即可。(在模型中其实就是构造 encoder 和 decoder)
-
encoder的实现:
如同前面所说,对于每一个输入给定的 x ,我们都会有(而且需要有)一个分布 q_\theta(z|x) 来(通过对分布采样)得出一个与之对应的隐变量(或者说隐编码) z。而这个分布如何设计呢?我们假设其为正态分布,那么我们只需要决定其均值和方差即可。而这个均值和方差即可由神经网络输入 x 后输出得到,这样这个神经网络就帮助实现了 encode 的过程,我们称这一部分为 encoder。
-
KL(q_\theta(z|x)||p(z)) 的实现:
如果不对 q_\theta(z|x) 做约束,其方差很可能趋近于0,导致给定 x 后可能生成的 z 局限到某个常数附近,使 VAE 退化到 AE,这会大大削减模型生成新样本的能力和隐变量表达的丰富性。所以我们希望其方差适当地小但不等于0 。 我们可以通过使其与 N(0,I) 看齐来实现。我们可以直接假设 p(z) 为 N(0,I) ,计算KL(q_\theta(z|x)||p(z))。
(p(z)为N(0,I)其实也可以从另一个角度理解,即我们理想中的 p(z|x) 是 N(0,I) ,那我们在此假设下可以计算出:p(z) = \int p(z|x)p(x)\text{d}x = \int \mathcal{N}(0, I)p(x) \text{d}x = \mathcal{N}(0, I)\int p(x) \text{d}x = \mathcal{N}(0, I),所以说,我们理想中的 p(z) 也应该是 N(0,I)。)
在此基础上,我们可以进行如下计算(以一维情况为例):
$$
\begin{equation*} \begin{split} &KL(q_{\theta}(z|x)||\mathcal{N}(0, I)) \\ =& \int \frac{1}{\sqrt{2 \pi {\sigma}^2}} e^{-\frac{(x-{\mu})^2}{2{\sigma}^2}} \ln \frac{\frac{1}{\sqrt{2 \pi {\sigma}^2}} e^{-\frac{(x-{\mu})^2}{2{\sigma}^2}}}{\frac{1}{\sqrt{2 \pi}} e^{-\frac{x^2}{2}}} \text{d}x \\ =& \int \frac{1}{\sqrt{2 \pi {\sigma}^2}} e^{-\frac{(x-{\mu})^2}{2{\sigma}^2}} \ln \frac{1}{\sqrt{\sigma^2}}\times e^{\frac{x^2}{2} - \frac{(x-\mu)^2}{2\sigma^2}}\text{d}x \\ =& \int \frac{1}{\sqrt{2 \pi {\sigma}^2}} e^{-\frac{(x-{\mu})^2}{2{\sigma}^2}} [-\frac{1}{2} \ln \sigma^2 + \frac{1}{2}x^2 - \frac{1}{2}\frac{(x-\mu)^2}{\sigma^2}]\text{d}x \\ =& \frac{1}{2} \int \frac{1}{\sqrt{2 \pi {\sigma}^2}} e^{-\frac{(x-{\mu})^2}{2{\sigma}^2}} [-\ln \sigma^2 + x^2 - \frac{(x-\mu)^2}{\sigma^2}]\text{d}x \\ =& \frac{1}{2}(-\ln \sigma^2 + \mathbb{E}[x^2] - \frac{1}{\sigma^2}\mathbb{E}[(x-u)^2]) \\ =& \frac{1}{2}(-\ln \sigma^2 + \sigma^2 + \mu^2 - 1) \end{split} \end{equation*}
$$其中 \mathbb{E}[x^2] 的计算过程如下:
$$
\begin{align} \mathbb{E}X^2&=\frac{1}{\sqrt{2\pi}\sigma}\underbrace{\int_{-\infty}^{+\infty}x^2\mathrm{e}^{-\frac{(x-\mu)^2}{2\sigma^2}}\mathrm{d}x}_{x\to \sqrt{2}\sigma x+\mu}\\ &=\frac{1}{\sqrt{\pi}}\int_{-\infty}^{+\infty}(\sqrt{2}\sigma x+\mu)^2\mathrm{e}^{-x^2}\mathrm{d}x\\ &=\frac{2\sigma^2}{\sqrt{\pi}}\int_{-\infty}^{+\infty}x^2\mathrm{e}^{-x^2}\mathrm{d}x\\ &\quad+\frac{2\sqrt{2}\mu\sigma}{\sqrt{\pi}}\int_{-\infty}^{+\infty}x\mathrm{e}^{-x^2}\mathrm{d}x\\ &\quad+\frac{\mu^2}{\sqrt{\pi}}\int_{-\infty}^{+\infty}\mathrm{e}^{-x^2}\mathrm{d}x\\ &=\frac{2\sigma^2}{\sqrt{\pi}}\cdot\frac{\sqrt{\pi}}{2}+\frac{2\sqrt{2}\mu\sigma}{\sqrt{\pi}}\cdot 0+\frac{\mu^2}{\sqrt{\pi}}\cdot\sqrt{\pi}\\ &=\sigma^2+\mu^2. \end{align}
$$ -
重参数化技巧:
由于采样过程本身不可导,这里我们使用重参数化技巧(reparametrize),将从N(μ, σ^2)采样得到一个Z的过程视为为在N(0,1)上采样得到一个 \epsilon然后计算 Z = μ + ε × σ 的过程,这样就可以计算Z到μ、σ的导数了。(我有一点点疑惑这个概率是否会发生改变,待进一步证明)
-
多样本多维度情况下的目标函数:
上面介绍的重构误差是针对单个样本的,当样本数为 n 时,可以表示为 \frac{1}{n}\sum_{i=1}^{n} \ln p(x_i|z_i) ,并且可以转化为要求 x_i 和 \hat{x}_i 的MSE尽可能地小。
(为什么可以互相转化呢?我是抱有一定疑问的)
,即最小化 \sum_{i=1}^n ||x_i - \hat{x}_i||_2^2上面介绍的KL散度表示也是针对1维的隐变量 z 。由于正态分布是各向同性的,当维度为 d 时,可以表示为 \frac{1}{2}\sum_{i=1}^{d} (-\ln \sigma^2 + \sigma^2 + \mu^2 - 1) 。
两个加起来为: Loss = \frac{1}{n}\sum_{i=1}^{n} [||x_i - \hat{x}_i||_2^2 + \frac{1}{2}\sum_{i=1}^{d} (-\ln \sigma^2 + \sigma^2 + \mu^2 - 1)]
-
decoder
decoder本身是个确定的过程,只需要用一个神经网络来实现 x=f(z) 就行
隐变量模型的改良之二:GAN
GAN的工程实现思路看似简单,但其背后的数学内涵似乎是颇为精妙的,在此贴一篇博客,且待之后进一步研究。https://zhuanlan.zhihu.com/p/25071913
隐变量模型的改良之三:FLOW
在前面我们提到,传统的隐变量模型,如果不使用蒙特卡洛采样,对 x = f(z)这个方程进行求解是十分困难的(因为f往往十分复杂且不可逆)。
但在FLOW中,我们希望设计一个方便进行求解的函数 f,这样一个函数需要满足如下性质:
-
f 可逆。
-
f 为一一映射(其实可以由1得到),并且定义域和值域维度相同。
-
因为 f 可逆,容易得到 z = f^{-1}(x) \sim N(0,I)
且
$$
p(x)=\frac{1}{(2\pi)^{D/2}}\exp(-\frac{1}{2}||f^{-1}(x)||^2)\bigg |\det[\frac{\partial {f^{-1}}}{\partial x}]\bigg|
$$故我们应该要求 f^{-1}(x)的行列式也要易于计算。
这就是FLOW的理念,在这个理念上的代表作是 RevNet (Reversible Network),其详细原理的学习待进一步完善:细水长flow之NICE:流模型的基本概念与实现 - 科学空间|Scientific Spaces
Diffusion Model(DDPM)的基本思路
对于隐变量模型给出的基本思路:
p(x)=\int p(x|z) p(z) \mathrm{d}z
VAE给出的优化是提高 p(z)的采样效率:
-
encoder 用参数化的后验分布模型 q_\theta(z|x),直接预测均值\mu和方差\Sigma,使得 q_\theta(z|x) = N(\mu,\Sigma)
-
decoder 将从q_\theta(z|x)中采样得到的z经过参数化的 p_\theta(x|z) 生成最后的 \hat{x}
diffusion model 的大体思路其实与之类似,但是在diffusion model 中,q_\theta(z|x) 所对应的过程不是一步得到的,而是由P步连续变换得来的,即在encoder中:
$$
q_\theta(z,x_{1:T}|x)=q_\theta(z|x_{T-1})\cdot q_\theta(x_{T-1}|x_{T-2})\cdots q_\theta(x_{t}|x_{t-1}) \cdots q_\theta(x_2|x_1) \cdot q_\theta(x_1|x)
$$
同理,在decoder中有:
$$
p_\theta(x,x_{1:T}|z)=p_\theta(x|x_1)\cdot p_\theta(x_1|x_2)\cdots p_\theta(x_{t-1}|x_t) \cdots p_\theta(x_{T-2}|x_{T-1})\cdot p_\theta(x_{T-1}|z)
$$
(这种马尔科夫链式的分布加噪过程在某种程度是受非平衡热力学提出的)
其图解如下(符号略有不同,x 变成了 x_0 ,z 变成了 x_T):
-
从右向左是encoder过程,是一个无参数的 q(x_t|x_{t-1}) ,即这是一个纯粹人为的过程(比如从原始的清晰图片每次按照高斯分布进行一个映射)。
-
从左往右是decoder过程,是一个带参数的 p_\theta(x_{t-1}|x_t)。其并非像VAE那样直接预测 \hat{x} ,而是预测高斯噪声,在decode过程中逐渐减去高斯噪声,还原出清晰的图像。
-
在diffusion model 中隐变量 z (在图中是 x_T)和原始图片的维度是一样大的(但是理论上还是简单的(因为很像高斯噪声))。
Diffusion Model 的 encoder过程
首先定义递增的常量序列{\beta_t},满足0<\beta_1 \lt \beta_2 \lt \cdots \lt \beta_T<1 定义观测原图为随机变量x_0,然后定义从随机变量x_{t-1}到随机变量x_t的分布关系为
$$
\begin{align}q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-\beta_t}x_{t-1}, \beta_t )\end{align}
$$
即x_t是均值\sqrt{1-\beta_t}x_{t-1}且方差为\beta_t 的高斯分布,用类似重参数分解可以得到
\begin{align}x_t = \sqrt{1-\beta_t}x_{t-1} + \sqrt{\beta_t} \epsilon_{t-1}, \ \epsilon_i \sim \mathcal{N}(0, 1)\end{align}
对于公式(2)有以下一些解释
-
x_{t-1}和\epsilon_{t-1}前面的两个系数平方和等于1
-
\beta_t单调递增且\beta_t \in (0, 1),则可以保证t=0时候方差几乎为0,t=T时方差几乎为1
如果定义
-
\alpha_t=1-\beta_t。为了书写方便
-
\bar{\alpha}_t=\prod_{i=1}^t \alpha_i。为了书写方便
-
\bar{\epsilon}_k \sim \mathcal{N}(0, 1)。代表k个高斯分布合并之后的新高斯分布
那么递推展开可以得到
$$
\begin{align}x_t&=\sqrt{\alpha_t} x_{t-1} + \sqrt{1-\alpha_t} \epsilon_{t-1} \\ &=\sqrt{\alpha_t \alpha_{t-1}} x_{t-2} + \sqrt{1-\alpha_t \alpha_{t-1}} \bar{\epsilon}_2 \\&= \cdots \\&=\sqrt{\bar{\alpha_t}}x_0 +\sqrt{1-\bar{\alpha_t}}\bar{\epsilon}_t\end{align}
$$
注意
-
这里用了方差的性质,即两个高斯分布的和还是高斯分布,并且新方差等于这两个高斯分布的方差
-
\bar{\epsilon}_k是k个高斯分布合并之后的新高斯分布
-
x_0和\bar{\epsilon}_t前面两个系数的平方和仍然是1
(这几个性质真的好鸡儿巧妙,到底是怎么构造出来的,有什么说法吗)
观察公式(6)里的x_t可以发现
-
随着t \rightarrow T,\sqrt{\bar{\alpha_t}} \rightarrow 0且\sqrt{1-\bar{\alpha_t}} \rightarrow 1,因此x_t \rightarrow \mathcal{N}(0, 1),逐渐变成标准高斯分布,极端情况下x_T =\mathcal{N}(0, 1)
-
不仅q(x_t|x_{t-1})=\mathcal{N}(\sqrt{1-\beta_t}x_{t-1}, \beta_t )可以直接计算,并且q(x_t|x_0)=\mathcal{N}(\sqrt{\bar{\alpha_t}}x_0 , 1-\bar{\alpha_t})也可以直接计算
-
整个encoder过程的是完全透明的,可以高效地计算中间任意分布q(x_t|x_0)的方式
Diffusion Model 的 decoder 过程
扩散过程是将数据噪音化,那么反向过程就是一个去噪的过程,如果我们可以求出反向过程的每一步的真实分布q(\mathbf{x}_{t-1} \vert \mathbf{x}_t),那么从一个随机噪音\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})开始,逐渐去噪就能生成一个真实的样本,所以反向过程也就是生成数据的过程。
(打个比方,Diffusion Model 的原理类似于我想让模型学会盖楼,于是我们在模型面前把一栋完整的楼一锤一锤地砸烂成shit,但是每一锤都很小很简单,让模型学会把每一锤恢复,他就能把这一堆shit恢复成一栋楼。这样随便给模型一堆shit,模型就能以此盖一栋楼)
原论文认为,对于encoder过程的每一步 \begin{align}q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-\beta_t}x_{t-1}, \beta_t )\end{align} ,当 \beta_t 足够小的时候,其逆过程的分布 q(\mathbf{x}_{t-1} \vert \mathbf{x}_t) 也可以近似地看做一个正态分布。(具体的理论证明似乎比较地复杂,暂且按下)
那我们可以做出如下假设,并用神经网络来计算其中的参数:
$$
p_\theta(\mathbf{x}_{0:T}) = p(\mathbf{x}_T) \prod^T_{t=1} p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) \quad p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t))\\
$$
这里 p(\mathbf{x}_T)= \mathcal{N}(\mathbf{0}, \mathbf{I}),而 p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) 为参数化的高斯分布,它们的均值和方差由训练的网络 \boldsymbol{\mu}_\theta(\mathbf{x}_t, t) 和 \boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t) 给出。我们训练扩散的模型实际上就是要训练得到这些模型。我们要用神经网络计算出的这些分布来估计 q(x_{t-1}|x_t) 。
但是实际的分布 q(x_{t-1}|x_t) 是难以处理的,但是加上条件 x_0 的后验分布 q(x_{t-1}|x_t,x_0) 是方便处理的,我们不妨假设:
$$
q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = \mathcal{N}(\color{blue}{\tilde{\boldsymbol{\mu}}}(\mathbf{x}_t, \mathbf{x}_0), \color{red}{\tilde{\beta}_t} \mathbf{I}) \\
$$
下面我们对其中的参数进行推导计算如下:
首先,由贝叶斯分布,我们有:
$$
q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) = q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) }\\
$$
又因为encoder过程符合马尔科夫链特性,我们有 q(x_t|x_{t-1}, x_0) = q(x_t|x_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t\mathbf{I}) , 并且由前面 encoder过程中的推导,我们知道:
$$
q(\mathbf{x}_{t-1} \vert \mathbf{x}_0)=\mathcal{N}(\mathbf{x}_{t-1}; \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0, (1 - \bar{\alpha}_{t-1})\mathbf{I})\\q(\mathbf{x}_t \vert \mathbf{x}_0) = \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\mathbf{I}) \\
$$
(其实从这里就可以看出 q(x_{t-1}|x_t,x_0) 是货真价实的高斯分布了)带进原式子可得:
$$
\begin{aligned} q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) &= q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) } \\ &\propto \exp \Big(-\frac{1}{2} \big(\frac{(\mathbf{x}_t - \sqrt{\alpha_t} \mathbf{x}_{t-1})^2}{\beta_t} + \frac{(\mathbf{x}_{t-1} - \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0)^2}{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big) \\ &= \exp \Big(-\frac{1}{2} \big(\frac{\mathbf{x}_t^2 - 2\sqrt{\alpha_t} \mathbf{x}_t \color{blue}{\mathbf{x}_{t-1}} \color{black}{+ \alpha_t} \color{red}{\mathbf{x}_{t-1}^2} }{\beta_t} + \frac{ \color{red}{\mathbf{x}_{t-1}^2} \color{black}{- 2 \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_0} \color{blue}{\mathbf{x}_{t-1}} \color{black}{+ \bar{\alpha}_{t-1} \mathbf{x}_0^2} }{1-\bar{\alpha}_{t-1}} - \frac{(\mathbf{x}_t - \sqrt{\bar{\alpha}_t} \mathbf{x}_0)^2}{1-\bar{\alpha}_t} \big) \Big) \\ &= \exp\Big( -\frac{1}{2} \big( \color{red}{(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}})} \mathbf{x}_{t-1}^2 - \color{blue}{(\frac{2\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{2\sqrt{\bar{\alpha}_{t-1}}}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)} \mathbf{x}_{t-1} \color{black}{ + C(\mathbf{x}_t, \mathbf{x}_0) \big) \Big)} \end{aligned}\\
$$
这里的 C(\mathbf{x}_t, \mathbf{x}_0)是一个和\mathbf{x}_{t-1}无关的部分,所以省略。根据高斯分布的概率密度函数定义和上述结果,我们可以通过待定系数求出后验分布q(\mathbf{x}_{t-1} \vert \mathbf{x}_{t}, \mathbf{x}_0)的均值和方差:
$$
\begin{aligned} \tilde{\beta}_t &= 1/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}) \\&= 1/(\frac{\alpha_t - \bar{\alpha}_t + \beta_t}{\beta_t(1 - \bar{\alpha}_{t-1})}) \\&= \color{green}{\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t} \\ \tilde{\boldsymbol{\mu}}_t (\mathbf{x}_t, \mathbf{x}_0) &= (\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1} }}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0)/(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \bar{\alpha}_{t-1}}) \\ &= (\frac{\sqrt{\alpha_t}}{\beta_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1} }}{1 - \bar{\alpha}_{t-1}} \mathbf{x}_0) \color{green}{\frac{1 - \bar{\alpha}_{t-1}}{1 - \bar{\alpha}_t} \cdot \beta_t} \\ &= \frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t + \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t} \mathbf{x}_0\\ \end{aligned}\\
$$
可以看到方差是一个定量(扩散过程参数固定),而均值是一个依赖 \mathbf{x}_0 和 \mathbf{x}_t 的函数。这个分布将会被用于推导扩散模型的优化目标。
Diffusion Model 的优化目标
上面介绍了扩散模型的扩散过程和反向过程,现在我们来从另外一个角度来看扩散模型:如果我们把中间产生的变量看成隐变量的话,那么扩散模型其实是包含T个隐变量的隐变量模型(latent variable model),它可以看成是一个特殊的Hierarchical VAEs:
相比VAE来说,扩散模型的隐变量是和原始数据同维度的,而且encoder(即扩散过程)是固定的。既然扩散模型是隐变量模型,那么我们可以就可以基于变分推断来得到 variational lower bound(VLB,又称ELBO)作为最大化优化目标,这里有:
$$
\begin{aligned} \log p_\theta(\mathbf{x}_0) &=\log\int p_\theta(\mathbf{x}_{0:T}) d\mathbf{x}_{1:T}\\ &=\log\int \frac{p_\theta(\mathbf{x}_{0:T}) q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} d\mathbf{x}_{1:T}\\ &\geq \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}[\log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}]\\ \end{aligned}\\
$$
(这里最后一步使用了琴生不等式)
对于网络训练来说,我们的训练目标是VLB取负:
$$
L=-L_{\text{VLB}}=\mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}[-\log \frac{p_\theta(\mathbf{x}_{0:T})}{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}]=\mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}[\log \frac{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})}{p_\theta(\mathbf{x}_{0:T})}] \\
$$
我们对其进行进一步分解可以得到:
$$
\begin{aligned} L &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ \log\frac{q(\mathbf{x}_{1:T}\vert\mathbf{x}_0)}{p_\theta(\mathbf{x}_{0:T})} \Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ \log\frac{\prod_{t=1}^T q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{ p_\theta(\mathbf{x}_T) \prod_{t=1}^T p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t) } \Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=1}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} \Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1}, \mathbf{x}_{0})}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] & \text{ ;use } q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0)=q(\mathbf{x}_t \vert \mathbf{x}_{t-1})\\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \Big( \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)}\cdot \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1}\vert\mathbf{x}_0)} \Big) + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] & \text{ ;use Bayes' Rule }\\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_t \vert \mathbf{x}_0)}{q(\mathbf{x}_{t-1} \vert \mathbf{x}_0)} + \log\frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ -\log p_\theta(\mathbf{x}_T) + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} + \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{q(\mathbf{x}_1 \vert \mathbf{x}_0)} + \log \frac{q(\mathbf{x}_1 \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)} \Big]\\ &= \mathbb{E}_{q(\mathbf{x}_{1:T}\vert \mathbf{x}_{0})} \Big[ \log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_T)} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)} - \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1) \Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{T}\vert \mathbf{x}_{0})}\Big[\log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_T)}\Big]+\sum_{t=2}^T \mathbb{E}_{q(\mathbf{x}_{t}, \mathbf{x}_{t-1}\vert \mathbf{x}_{0})}\Big[\log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)}\Big] - \mathbb{E}_{q(\mathbf{x}_{1}\vert \mathbf{x}_{0})}\Big[\log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)\Big] \\ &= \mathbb{E}_{q(\mathbf{x}_{T}\vert \mathbf{x}_{0})}\Big[\log\frac{q(\mathbf{x}_T \vert \mathbf{x}_0)}{p_\theta(\mathbf{x}_T)}\Big]+\sum_{t=2}^T \mathbb{E}_{q(\mathbf{x}_{t}\vert \mathbf{x}_{0})}\Big[q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)\log \frac{q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)}{p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)}\Big] - \mathbb{E}_{q(\mathbf{x}_{1}\vert \mathbf{x}_{0})}\Big[\log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)\Big] \\ &= \underbrace{D_\text{KL}(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T))}_{L_T} + \sum_{t=2}^T \underbrace{\mathbb{E}_{q(\mathbf{x}_{t}\vert \mathbf{x}_{0})}\Big[D_\text{KL}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t))\Big]}_{L_{t-1}} -\underbrace{\mathbb{E}_{q(\mathbf{x}_{1}\vert \mathbf{x}_{0})}\log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1)}_{L_0} \end{aligned}
$$
可以看到最终的优化目标共包含 T+1 项,其中 L_0 可以看成是原始数据重建,优化的是负对数似然,L_0 可以用估计的\mathcal{N}(\mathbf{x}_0; \boldsymbol{\mu}_\theta(\mathbf{x}_1, 1), \boldsymbol{\Sigma}_\theta(\mathbf{x}_1, 1))来专门构建一个离散化的decoder来计算:
$$
p_{\theta}(\mathbf{x}_0\vert\mathbf{x}_1)=\prod^D_{i=1}\int ^{\delta_+(x_0^i)}_{\delta_-(x_0^i)}\mathcal{N}(x; \mu^i_\theta(x_1, 1), \Sigma^i_\theta(x_1, 1))\mathrm{d}x\\ \delta_+(x)= \begin{cases} \infty& \text{ if } x=1 \\ x+\frac{1}{255}& \text{ if } x <1 \end{cases} \\ \delta_+(x)= \begin{cases} -\infty& \text{ if } x=-1 \\ x-\frac{1}{255}& \text{ if } x >-1 \end{cases}
$$
(其中 i 是对像素点的编号),相当于计算由 \mathbf{x_1} 大致生成 \mathbf{x_0} 的所有像素的概率。
而L_T计算的是最后得到的噪音的分布和先验分布的KL散度,这个KL散度没有训练参数,近似为0,因为先验p(\mathbf {x}_T)=\mathcal{N}(\mathbf{0}, \mathbf{I})而扩散过程最后得到的随机噪音q(\mathbf{x}_{T}\vert \mathbf{x}_{0})也近似为\mathcal{N}(\mathbf{0}, \mathbf{I});而L_{t-1}则是计算的是估计分布p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t)和真实后验分布q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)的KL散度,这里希望我们估计的去噪过程和依赖真实数据的去噪过程近似一致。
DDPM对 p_\theta(\mathbf{x}_{t-1} \vert\mathbf{x}_t) 做了近一步简化,采用固定的方差:\boldsymbol{\Sigma}_\theta(\mathbf{x}_t, t)= \sigma_t^2\mathbf{I},这里的\sigma_t^2可以设定为\beta _t或者\tilde{\beta}_t(这其实是两个极端,分别是上限和下限,也可以采用可训练的方差,见论文Improved Denoising Diffusion Probabilistic Models和Analytic-DPM: an Analytic Estimate of the Optimal Reverse Variance in Diffusion Probabilistic Models)。这里假定\sigma_t^2=\tilde{\beta}_t,那么:
$$
q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)=\mathcal{N}(\mathbf{x}_{t-1}; {\tilde{\boldsymbol{\mu}}} (\mathbf{x}_t, \mathbf{x}_0), {\sigma_t^2} \mathbf{I}) p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), {\sigma_t^2} \mathbf{I})
$$
对于两个高斯分布的KL散度,其计算公式为(具体推导见生成模型之VAE):
$$
\text{KL}(p_1||p_2) = \frac{1}{2}(\text{tr}(\boldsymbol{\Sigma}_2^{-1}\boldsymbol{\Sigma}_1)+(\boldsymbol{\mu_2}-\boldsymbol{\mu_1})^{\top}\boldsymbol{\Sigma}_2^{-1}(\boldsymbol{\mu_2}-\boldsymbol{\mu_1})-n+\log\frac{\det(\boldsymbol{\Sigma_2})}{\det(\boldsymbol{\Sigma_1})}) \\
$$
那么就有:
$$
\begin{aligned} D_\text{KL}(q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0)\parallel p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t)) &=D_\text{KL}(\mathcal{N}(\mathbf{x}_{t-1}; {\tilde{\boldsymbol{\mu}}} (\mathbf{x}_t, \mathbf{x}_0), {\sigma_t^2} \mathbf{I}) \parallel \mathcal{N}(\mathbf{x}_{t-1}; \boldsymbol{\mu}_\theta(\mathbf{x}_t, t), {\sigma_t^2} \mathbf{I})) \\ &=\frac{1}{2}(n+\frac{1}{{\sigma_t^2}}\|\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) - {\boldsymbol{\mu}_\theta(\mathbf{x}_t, t)} \|^2 -n+\log1) \\ &=\frac{1}{2{\sigma_t^2}}\|\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) - {\boldsymbol{\mu}_\theta(\mathbf{x}_t, t)} \|^2 \end{aligned}\\
$$
那么优化目标L_{t-1}即为:
$$
L_{t-1}=\mathbb{E}_{q(\mathbf{x}_{t}\vert \mathbf{x}_{0})}\Big[ \frac{1}{2{\sigma_t^2}}\|\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) - {\boldsymbol{\mu}_\theta(\mathbf{x}_t, t)} \|^2\Big] \\
$$
从上述公式来看,我们其实是在希望网络学习到的均值 \boldsymbol{\mu}_\theta(\mathbf{x}_t, t)和后验分布的均值{\tilde{\boldsymbol{\mu}}} (\mathbf{x}_t, \mathbf{x}_0)尽可能地一致。不过DDPM发现,预测均值并不是最好的选择。根据前面得到的扩散过程的特性,我们有:
$$
\mathbf{x_t}(\mathbf{x_0},\mathbf{\epsilon})=\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{\epsilon} \quad \text{ where } \mathbf{\epsilon}\sim \mathcal{N}(\mathbf{0}, \mathbf{I}) \\
$$
将这个公式带入上述优化目标(注意这里的损失我们加上了对 \mathbf{x}_0 的数学期望),可以得到:
$$
\begin{aligned} L_{t-1}&=\mathbb{E}_{\mathbf{x}_{0}}\Big(\mathbb{E}_{q(\mathbf{x}_{t}\vert \mathbf{x}_{0})}\Big[ \frac{1}{2{\sigma_t^2}}\|\tilde{\boldsymbol{\mu}}_t(\mathbf{x}_t, \mathbf{x}_0) - {\boldsymbol{\mu}_\theta(\mathbf{x}_t, t)} \|^2\Big]\Big) \\ &=\mathbb{E}_{\mathbf{x}_{0},\mathbf{\epsilon}\sim \mathcal{N}(\mathbf{0}, \mathbf{I})}\Big[ \frac{1}{2{\sigma_t^2}}\|\tilde{\boldsymbol{\mu}}_t\Big(\mathbf{x_t}(\mathbf{x_0},\mathbf{\epsilon}), \frac{1}{\sqrt{\bar \alpha_t}} \big(\mathbf{x_t}(\mathbf{x_0},\mathbf{\epsilon}) - \sqrt{1 - \bar{\alpha}_t} \mathbf{\epsilon} \big)\Big ) - {\boldsymbol{\mu}_\theta(\mathbf{x_t}(\mathbf{x_0},\mathbf{\epsilon}), t)} \|^2\Big] \\ &=\mathbb{E}_{\mathbf{x}_{0},\mathbf{\epsilon}\sim \mathcal{N}(\mathbf{0}, \mathbf{I})}\Big[ \frac{1}{2{\sigma_t^2}}\|\Big (\frac{\sqrt{\alpha_t}(1 - \bar{\alpha}_{t-1})}{1 - \bar{\alpha}_t} \mathbf{x}_t(\mathbf{x_0},\mathbf{\epsilon}) + \frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1 - \bar{\alpha}_t} \frac{1}{\sqrt{\bar \alpha_t}} \big(\mathbf{x_t}(\mathbf{x_0},\mathbf{\epsilon}) - \sqrt{1 - \bar{\alpha}_t} \mathbf{\epsilon} \big) \Big) - {\boldsymbol{\mu}_\theta(\mathbf{x_t}(\mathbf{x_0},\mathbf{\epsilon}), t)} \|^2\Big] \\ &=\mathbb{E}_{\mathbf{x}_{0},\mathbf{\epsilon}\sim \mathcal{N}(\mathbf{0}, \mathbf{I})}\Big[ \frac{1}{2{\sigma_t^2}}\|\frac{1}{\sqrt{\alpha_t}}\Big( \mathbf{x}_t(\mathbf{x_0},\mathbf{\epsilon}) - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\mathbf{\epsilon}\Big) - {\boldsymbol{\mu}_\theta(\mathbf{x_t}(\mathbf{x_0},\mathbf{\epsilon}), t)} \|^2\Big] \end{aligned}\\
$$
近一步地,我们对\boldsymbol{\mu}_\theta(\mathbf{x_t}(\mathbf{x_0},\mathbf{\epsilon}), t)也进行重参数化,变成:
$$
\boldsymbol{\mu}_\theta(\mathbf{x_t}(\mathbf{x_0},\mathbf{\epsilon}), t)=\frac{1}{\sqrt{\alpha_t}}\Big( \mathbf{x}_t(\mathbf{x_0},\mathbf{\epsilon}) - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\mathbf{\epsilon}_\theta\big(\mathbf{x}_t(\mathbf{x_0},\mathbf{\epsilon}), t\big)\Big) \\
$$
这里的\mathbf{\epsilon}_\theta是一个基于神经网络的拟合函数,这意味着我们由原来的预测均值而换成预测噪音\mathbf{\epsilon}。我们将上述等式带入优化目标,可以得到:
$$
\begin{aligned} L_{t-1}&=\mathbb{E}_{\mathbf{x}_{0},\mathbf{\epsilon}\sim \mathcal{N}(\mathbf{0}, \mathbf{I})}\Big[ \frac{1}{2{\sigma_t^2}}\|\frac{1}{\sqrt{\alpha_t}}\Big( \mathbf{x}_t(\mathbf{x_0},\mathbf{\epsilon}) - \frac{\beta_t}{\sqrt{1 - \bar{\alpha}_t}}\mathbf{\epsilon}\Big) - {\boldsymbol{\mu}_\theta(\mathbf{x_t}(\mathbf{x_0},\mathbf{\epsilon}), t)} \|^2\Big] \\ &= \mathbb{E}_{\mathbf{x}_{0},\mathbf{\epsilon}\sim \mathcal{N}(\mathbf{0}, \mathbf{I})}\Big[ \frac{\beta_t^2}{2{\sigma_t^2}\alpha_t(1-\bar{\alpha}_t)}\| \mathbf{\epsilon}- \mathbf{\epsilon}_\theta\big(\mathbf{x}_t(\mathbf{x_0},\mathbf{\epsilon}), t\big)\|^2\Big]\\ &=\mathbb{E}_{\mathbf{x}_{0},\mathbf{\epsilon}\sim \mathcal{N}(\mathbf{0}, \mathbf{I})}\Big[ \frac{\beta_t^2}{2{\sigma_t^2}\alpha_t(1-\bar{\alpha}_t)}\| \mathbf{\epsilon}- \mathbf{\epsilon}_\theta\big(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{\epsilon}, t\big)\|^2\Big] \end{aligned}\\
$$
DDPM近一步对上述目标进行了简化,即去掉了权重系数,变成了:
$$
L_{t-1}^{\text{simple}}=\mathbb{E}_{\mathbf{x}_{0},\mathbf{\epsilon}\sim \mathcal{N}(\mathbf{0}, \mathbf{I})}\Big[ \| \mathbf{\epsilon}- \mathbf{\epsilon}_\theta\big(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\mathbf{\epsilon}, t\big)\|^2\Big]
$$
这里的t在[1, T]范围内取值(如前所述,其中取1时对应L_0)。由于去掉了不同t的权重系数,所以这个简化的目标其实是VLB优化目标进行了reweight。从DDPM的对比实验结果来看,预测噪音比预测均值效果要好,采用简化版本的优化目标比VLB目标效果要好。
What are Diffusion Models? | Lil'Log
[论文理解] Diffusion Models Beat GANs on Image Synthesis – sunlin-ai