1 前言

扩散模型的出现让我眼前一亮。它的原理与GAN截然不同，生成质量的天花板似乎更高，而且背后的原理啃起来也很费劲。

这次分享下我阅读Understanding Diffusion Models: A Unified Perspective这篇文章的笔记。这篇文章尝试从一个统一的视角分析扩散模型，介绍了扩散模型、VAE、基于分数的生成模型之间的联系。文章的数学推理很详细，基本没有跳步，很适合数学基础一般的同学。

本文基本上是原论文的大致翻译和概括。笔记难免存在一些信息上的简略，感兴趣的读者可以阅读原文。

2 背景

3 证据下界

样本和隐变量构成联合分布\(p(\vec x, \vec z)\)。在基于似然度的方法中，我们希望用模型最大化所有观测样本\(\vec x\)的似然度\(p(\vec x)\)。

我们可以利用证据下界（Evidence Lower Bound，ELBO）： \[ \mathbb E_{q_\phi(\vec z|\vec x)}\left[\log \frac{p(\vec x, \vec z)}{q_\phi(\vec z|\vec x)}\right] \tag{1}\] 我们称\(\log p(\vec x)\)为证据（的大小），证据下界和证据的关系是： \[ \log p(\vec x) \geq \mathbb E _{q_\phi(\vec z|\vec x)} \left[\log \frac{p(\vec x, \vec z)}{q_\phi(\vec z| \vec x)}\right] \] 其中\(q_\phi(\vec z| \vec x)\)就是需要我们优化的将观测样本\(\vec x\)映射到隐变量\(\vec z\)的编码器，是对真实后验\(p(\vec z|\vec x)\)的近似，而\(\phi\)是其参数。

后面我们将会看到，我们通过优化参数\(\phi\)，最大化公式 1，就能获得取得真实样本的数据分布，并能从中采样。

现在让我们理清为什么要最大化ELBO。 \[ \begin{aligned} \log p(\vec x) &= \log \int p(\vec x, \vec z) d\vec z & \text{对联合分布的边缘化} \\ &= \log \int \frac{p(\vec x, \vec z)q_\phi(\vec z|\vec x)}{q_\phi(\vec z|\vec x)} dz & \text{分子分母乘以同一个数} \\ &= \log \mathbb E _{q_\phi(\vec z|\vec x)} \left[\frac{p(\vec x, \vec z)}{q_\phi(\vec z|\vec x)}\right] & \text{期望的定义} \\ &\geq \mathbb E _{q_\phi(\vec z|\vec x)}\left[ \log \frac{p(\vec x, \vec z)}{q_\phi(\vec z|\vec x)} \right] & \text{Jensen不等式} \end{aligned} \tag{2}\] 以上是一种利用Jensen不等式的推导方式。还有一种证明方式，稍显冗长，但是能提供更多为何使用ELBO的直觉。

\[ \begin{aligned} \log p(\vec x) &= \log p(\vec x) \int q_\phi(\vec z|\vec x) dz & \text{乘以}1 = \int q_\phi(\vec z| \vec x)d\vec z\\ &= \int q_\phi(\vec z|\vec x)(\log p(\vec x)) d\vec z & \log p(\vec x)\text{移到积分符号后} \\ &= \mathbb E _{q_\phi(\vec z|\vec x)} \left[\log p(\vec x)\right] & \text{期望的定义}\\ &= \mathbb E _{q_\phi(\vec z|\vec x)}\left[\log\frac{p(\vec x, \vec z)}{p(\vec z|\vec x)}\right] & \text{链式法则}\\ &= \mathbb E _{q_\phi(\vec z|\vec x)} \left[\log\frac{p(\vec x, \vec z)q_\phi(\vec z|\vec x)}{p(\vec z|\vec x)q_\phi(\vec z|\vec x)}\right] & \text{分子分母同乘以一个数}\\ &= \mathbb E _{q_\phi(\vec z|\vec x)} \left[\log\frac{p(\vec x,\vec z)}{q_\phi(\vec z|\vec x)}\right] + \mathbb E_{q_\phi(\vec z|\vec x)}\left[\log\frac{q_\phi(\vec z|\vec x)}{p(\vec z|\vec x)}\right] & \text{期望的拆分}\\ &= \mathbb E _{q_\phi(\vec z|\vec x)} \left[\log \frac{p(\vec x, \vec z)}{q_\phi(\vec z|\vec x)}\right] + D_\text{KL} (q_\phi(\vec z|\vec x)\Vert p(\vec z|\vec x)) & \text{KL散度的定义} \\ &\geq \mathbb E _{q_\phi(\vec z|\vec x)} \left[\log \frac{p(\vec x, \vec z)}{q_\phi(\vec z|\vec x)}\right] & \text{KL散度非负} \end{aligned} \tag{3}\]

虽然公式 2和公式 3都证明了ELBO是证据下界，但是对于公式 3的理解是更关键的，即：

1. 两者的差值恰好为\(q_\phi(\vec z|\vec x)\)和\(p(\vec z|\vec x)\)的KL散度
2. 公式 3的左边实际上是一个常数，这个常数等于ELBO与KL散度的和。
3. 进一步的，最大化ELBO的过程也就是最小化\(q_\phi(\vec z|\vec x)\)和\(p(\vec z|\vec x)\)的KL散度的过程。由于我们不知道\(p(\vec z|\vec x)\)的真值，所以我们无法直接最小化KL散度，但是ELBO允许我们间接实现这项优化。
4. 一旦训练完毕，ELBO可以用于估计生成样本的似然度。因为ELBO的优化目标是逼近\(\log p(\vec x)\)。

4 变分自动编码器

变分自动编码器（Variational Autoencoder，VAE）的名字中有“变分”两个字，是因为它的目标是从所有可能的后验分布中寻找最优的\(q_phi(\vec z|\vec x)\)。它被称为编码器，是因为它具备传统编码器模型的特质，即它尝试将输入数据压缩为低维向量然后试图再还原为原来的输入。

VAE的目标是直接最大化ELBO。为了更清楚地分析VAE，我们继续拆解分析ELBO： \[ \begin{aligned} \mathbb E_{q_\phi(\vec z|\vec x)} \left[\log\frac{p(\vec x, \vec z)}{q_\phi(\vec z|\vec x)}\right] &= \mathbb E_{q_\phi(\vec z|\vec x)} \left[\log\frac{p_\theta(\vec x|\vec z)p(\vec z)}{q_\phi(\vec z|\vec x)}\right] & \text{链式法则}\\ &= \mathbb E_{q_\phi(\vec z|\vec x)}\left[\log p_\theta(\vec x|\vec z)\right] + \mathbb E_{q_\phi(\vec z|\vec x)}\left[\log \frac{p(\vec z)}{q_\phi(\vec z|\vec x)}\right] & \text{期望的拆分}\\ &= \underbrace{\mathbb E_{q_\phi(\vec z|\vec x)}\left[\log p_\theta(\vec x|\vec z)\right]}_{重构项} - \underbrace{D_\text{KL}(q_\phi(\vec z|\vec x)\Vert p(\vec z))}_{先验匹配项} & \text{KL散度的定义} \end{aligned} \tag{4}\] 在这个公式中，\(q_\phi(\vec z|\vec x)\)被称为编码器（encoder），\(p_\theta(\vec x|\vec z)\)被称为解码器（decoder）。

公式 4被拆分成两项，其中重构项度量了学习到的样本分布是否建模了真正的分布，而先验匹配项学习到的隐变量是否服从预设的先验分布。

在VAE模型中，encoder通常选择将隐变量设置为服从具有对角协方差矩阵的多元高斯分布，将先验分布设置为标准正态分布： \[ \begin{aligned} q_\phi(\vec z|\vec x) &= \mathcal N(\vec z; \vec \mu_{\vec\phi}(\vec x), \vec \sigma^2_{\vec \phi}(\vec x) \vec I)\\ p(\vec z) &= \mathcal N(\vec z; \vec 0, \vec I) \end{aligned} \]

在这样的设置下，先验匹配项中的KL散度是可以解析计算的，重构项可以通过蒙特卡洛估计方法近似。优化目标可以重写为： \[ \argmax{\vec \phi, \vec \theta}{\mathbb E _{q_\phi(\vec z|\vec x)}}\left[\log p_{\vec \theta}(\vec x|\vec z)\right] - D_{KL}(q_{\vec \phi}(\vec z|\vec x)\Vert p(\vec z)) \approx \argmax{\vec \phi, \vec \theta} \sum_{l=1}^L \log p_{\vec \theta}(\vec x|\vec z^{(l)}) - D_{KL}(q_{\vec \phi}(\vec z|\vec x) \Vert p(\vec z)) \]

其中\(\left\{\vec z^{(l)}\right\}_{l=1}^L\)是对于每个数据集中的样本\(\vec x\)从\(q_\phi(\vec z|\vec x)\)中采样的\(L\)个隐变量。这其中随机采样的过程通常是不可微的，幸运的是我们可以采用以下的“重参数化”技巧： \[ \vec z = \vec \mu_\pi(\vec x) + \vec \sigma_\phi(\vec x) \odot \vec \epsilon, \text{with} \vec\epsilon \sim \mathcal N(\vec \epsilon; \vec 0, \vec I) \] 其中\(\odot\)表示逐元素乘积。得益于重参数化方法，我们可以实现损失函数对\(\vec \phi\)和\(\vec \theta\)的求导。

5 分层的变分自动编码器

分层变分自动编码器（Hierarchical Variational Autoencoder，HVAE）将VAE推广到具有多层隐变量的情形。每一层的隐变量都以前一层隐变量为条件生成。论文讨论了HVAE的特例，马尔科夫HVAE（Markovian HVAE，MHVAE）。在MHVAE中，\(\vec z_t\)的生成只依赖\(\vec z_{t+1}\)，而不用考虑\(\vec z_{t+2}\)。MHVAE的联合分布和后验分布是： \[ p(\vec x, \vec z_{1:T})=p(\vec z_T)p_{\vec \theta}(\vec x|\vec z_1)\prod_{t=2}^T p_{\vec \theta}(\vec z_{t-1}|\vec z_t) \tag{5}\] \[ q_{\vec \phi} (\vec z_{1:T}\vert \vec x)=q_\phi(\vec z_1|\vec x)\prod_{t=2}^T q_{\vec \phi}(\vec z_t\vert z_{t-1}) \tag{6}\] 这时ELBO公式可以写为 \[ \begin{aligned} \log p(\vec x) &= \log \int p(\vec x, \vec z_{1:T}) d \vec z_{1:T} & \text{对联合分布的边缘化}\\ &= \log \int \frac{p(\vec x, \vec z_{1:T})q_{\vec \phi}(\vec z_{1:T}\vert \vec x)}{q_{\vec \phi}(\vec z_{1:T}\vert \vec x)} d\vec z_{1:T} & \text{分子分母乘以同一个数} \\ &= \log \mathbb E_{q_{\vec \phi}(\vec z_{1:T}\vert \vec x)} \left[\frac{p(\vec x, \vec z_{1:T})}{q_{\vec \phi}(\vec z_{1:T}\vert \vec x)}\right] &\text{期望的定义} \\ &\geq \mathbb E _{q_{\vec \phi}(\vec z_{1:T}\vert \vec x)} \left[\log \frac{p(\vec x, \vec z_{1:T})}{q_{\vec \phi}(\vec z_{1:T} \vert \vec x)}\right] & \text{Jensen不等式} \end{aligned} \] 然后再将联合分布公式 5和后验分布公式 6代入得到 \[ \mathbb E _{q_{\vec \phi}(\vec z_{1:T}\vert \vec x)} \left[\log \frac{p(\vec x, \vec z_{1:T})}{q_{\vec \phi}(\vec z_{1:T} \vert \vec x)}\right] = \mathbb E_{q_{\vec \phi}(\vec z_{q:T}\vert \vec x)} \left[ \log \frac{ p(\vec z_T)p_{\vec \theta}(\vec x|\vec z_1)\prod_{t=2}^T p_{\vec \theta}(\vec z_{t-1}|\vec z_t) }{ q_\phi(\vec z_1|\vec x)\prod_{t=2}^T q_{\vec \phi}(\vec z_t\vert z_{t-1}) } \right] \]

后面我们将会看到，在讨论变分扩散模型时，这个目标函数将可以分解为更多可解释的部件。

6 变分扩散模型

变分扩散模型（Variational Diffusion Model，VDM）可以视为施加了以下三个条件的MHVAE：

1. 隐变量的维度大小总是和样本数据的维度大小一样
2. 每一步的encoder不是可学习的，而是预定义的线性高斯模型
3. 每一步的encoder的参数随著步骤\(t\)变化，使得最后一层的隐变量最终服从标准正态分布。

因为有第一条限制，所以我们可以引入符号\(\vec x_t\)同时表示隐变量和原始数据。当\(t=0\)时，它代表原始数据，当\(t\in[1, T]\)时，它代表第\(t\)层的隐变量。

于是VDM的后验分布可以重写为： \[ q(\vec x_{1:T}\vert \vec x_0) = \prod_{t=1}^Tq(\vec x_t \vert x_{t-1}) \] 与普通的MHVAE不同，VDM的encoder不是可学习的（条件2），而是被通常被定义为均值\(\mu_t(\vec x_t)=\sqrt{\alpha_t} \vec x_{t-1}\)，协方差矩阵\(\Sigma_t(\vec x_t) = (1 - \alpha_t)\mat I\). 这里\(\alpha_t\)是随层级\(t\)变化的参数，可以是预设的，或者是可学习的。

编码器的模型可以记为： \[ q(\vec x_t|\vec x_{t-1}) =\mathcal N(\vec x_t;\sqrt{a_t}\vec x_{t-1}, (1- \alpha_t) \mat I) \] 根据条件3，经过若干层这样的编码器，最终的\(p(\vec x_T)\)将会服从标准正态分布。于是MHVAE的联合分布公式 5变为 \[ p(\vec x_{0:T}) = p(\vec x_T)\prod_{t=1}^T p_{\vec \theta}(\vec x_{t-1}|\vec x_t) \] 其中 \[ p(\vec x_T) = \mathcal N(\vec x_T; \mat 0, \mat I) \] 总的来说，VDM假设的几个条件要求每一步“编码”都是持续增加噪声，直到数据成为完全的高斯噪声的过程。

在VDM中，\(q(\vec x_t|\vec x_{t-1})\)是完全固定的，因此我们关系的就只剩\(p_{\vec \theta}(\vec x_{t-1}|\vec x_t)\)的学习了。

如果你已经完成了一个VDM的训练，那么从模型中采样图片的过程就是从\(p(\vec x_T)\)中随机采样一个数据，然后迭代地运行\(p_{\vec \theta}(\vec x_{t-1}|\vec x_t)\)，直到生成\(\vec x_0\).

与其它HVAE相似，VDM同样可以用ELBO作为优化目标： \[ \begin{aligned} \log p(\vec x) &= \log \int p(\vec x_{0:T} d\vec x_{1:T}) &\text{对联合分布的边缘化}\\ &= \log \int \frac{p(\vec x_{0:T})q(\vec x_{1:T}|\vec x_0)}{q(\vec x_{1:T}|\vec x_0)} d\vec x_{1:T} & \text{分子分母乘以同一个数} \\ &= \log \E _{q(\vec x_{1:T}|\vec x_0)} \left[\frac{p(\vec x_{0:T})}{q(\vec x_{1:T}|\vec x_0)}\right] &\text{期望的定义}\\ &\geq \mathbb E_{q(\vec x_{1:T}|\vec x_0)}\left[\log \frac{p(\vec x_{0:T})}{q(\vec x_{1:T}|\vec x_0)}\right] & \text{Jensen不等式}\\ &= \mathbb E_{q(\vec x_{1:T}|\vec x_0)}\left[\log\frac{p(\vec x_T)\prod_{t=1}^Tp_{\vec \theta}(\vec x_{t-1}|\vec x_t)}{\prod_{t=1}^T q(\vec x_t|\vec x_{t-1})}\right] & \text{链式法则、马尔科夫性质}\\ &= \mathbb E_{q(\vec x_{1:T}|\vec x_0)} \left[\log\frac{p(\vec x_T)p_{\vec \theta}(\vec x_0|\vec x_1)\prod_{t=2}^Tp_{\vec\theta}(\vec x_{t-1}|\vec x_t)}{q(\vec x_T|\vec x_{T-1})\prod_{t=1}^{T-1}q(\vec x_t|\vec x_{t-1})}\right] \\ &= \mathbb E_{q(\vec x_{1:T}|\vec x_0)} \left[\log\frac{p(\vec x_T)p_{\vec \theta}(\vec x_0|\vec x_1)\prod_{t=1}^{T-1}p_{\vec\theta}(\vec x_t|\vec x_{t+1})}{q(\vec x_T|\vec x_{T-1})\prod_{t=1}^{T-1}q(\vec x_t|\vec x_{t-1})}\right] \\ &= \mathbb E_{q(\vec x_{1:T}|\vec x_0)} \left[\log\frac{p(\vec x_T)p_{\vec \theta}(\vec x_0|\vec x_1)}{q(\vec x_T|\vec x_{T-1})}\right]+ \mathbb E_{q(\vec x_{1:T}|\vec x_0)} \left[\log\prod_{t=1}^{T-1}\frac{p_{\vec \theta}(\vec x_t|\vec x_{t+1})}{q(\vec x_t|\vec x_{t-1})}\right] & \text{数学期望有线性性质} \\ &= \mathbb E_{q(\vec x_{1:T}|\vec x_0)} \left[\log p_{\vec\theta}(\vec x_0|\vec x_1)\right]+ \mathbb E_{q(\vec x_{1:T}|\vec x_0)} \left[\log\frac{p(\vec x_T)}{q(\vec x_T | \vec x_{T-1})}\right] + \mathbb E_{q(\vec x_{1:T}|\vec x_0)} \left[\sum_{t=1}^{T-1}\log \frac{p_{\vec\theta}(\vec x_t|\vec x_{t+1})}{q(\vec x_t|\vec x_{t-1})}\right] \\ &= \mathbb E_{q(\vec x_{1:T}|\vec x_0)} \left[\log p_{\vec\theta}(\vec x_0|\vec x_1)\right]+ \mathbb E_{q(\vec x_{1:T}|\vec x_0)} \left[\log\frac{p(\vec x_T)}{q(\vec x_T | \vec x_{T-1})}\right] + \sum_{t=1}^{T-1}\mathbb E_{q(\vec x_{1:T}|\vec x_0)}\left[\log\frac{p_\vtheta(\vec x_t|\vec x_{t+1})}{q(\vx_t|\vx_{t-1})}\right] \\ &= \E_{q(\vx_1|\vx_0)} \left[\log p_{\vec\theta}(\vec x_0|\vec x_1)\right] + \E_{q(\vx_{T-1},\vx_T|\vx_0)}\left[\log\frac{p(\vec x_T)}{q(\vec x_T | \vec x_{T-1})}\right] + \sum_{t=1}^{T-1}\E_{q(\vx_{t-1}, \vx_t,\vx_{t+1}|\vx_0)} \left[\log\frac{p_\vtheta(\vec x_t|\vec x_{t+1})}{q(\vx_t|\vx_{t-1})}\right] \\ & \text{按数学期望的定义展开} \\ &= \underbrace{\E_{q(\vx_1|\vx_0)} \left[\log p_{\vec\theta}(\vec x_0|\vec x_1)\right]}_{重构项} + \int \int \log \frac{p(\vx_T)}{q(\vx_T|\vx_{T-1})} q(\vx_{T-1}, \vx_T|\vx_0)d\vx_{T}d\vx_{T-1} + \sum_{t=1}^{T-1}\E_{q(\vx_{t-1}, \vx_t,\vx_{t+1}|\vx_0)} \left[\log\frac{p_\vtheta(\vec x_t|\vec x_{t+1})}{q(\vx_t|\vx_{t-1})}\right] \\ &= \underbrace{\E_{q(\vx_1|\vx_0)} \left[\log p_{\vec\theta}(\vec x_0|\vec x_1)\right]}_{重构项} + \int \left( \int \log \frac{p(\vx_T)}{q(\vx_T|\vx_{T-1})} q(\vx_{T}|\vx_{T-1})d\vx_{T}\right)q(\vx_{T-1}|\vx_0)d\vx_{T-1} + \sum_{t=1}^{T-1}\E_{q(\vx_{t-1}, \vx_t,\vx_{t+1}|\vx_0)} \left[\log\frac{p_\vtheta(\vec x_t|\vec x_{t+1})}{q(\vx_t|\vx_{t-1})}\right] \\ & \text{应用KL散度的定义} \\ &= \underbrace{\E_{q(\vx_1|\vx_0)} \left[\log p_{\vec\theta}(\vec x_0|\vec x_1)\right]}_{重构项} + \int \Big(- D_\text{KL}\left(q(\vx_T\vert \vx_{T-1})\Vert p(\vx_T)\right)\Big)q(\vx_{T-1}|\vx_0)d\vx_{T-1} + \sum_{t=1}^{T-1}\E_{q(\vx_{t-1}, \vx_t,\vx_{t+1}|\vx_0)} \left[\log\frac{p_\vtheta(\vec x_t|\vec x_{t+1})}{q(\vx_t|\vx_{t-1})}\right] \\ & \text{期望的定义} \\ &= \underbrace{\E_{q(\vx_1|\vx_0)} \left[\log p_{\vec\theta}(\vec x_0|\vec x_1)\right]}_{重构项} - \mathbb E_{q(\vx_{T-1}|\vx_0)} \left[ D_\text{KL}\left(q(\vx_T\vert \vx_{T-1})\Vert p(\vx_T)\right)\right] + \sum_{t=1}^{T-1}\E_{q(\vx_{t-1}, \vx_t,\vx_{t+1}|\vx_0)} \left[\log\frac{p_\vtheta(\vec x_t|\vec x_{t+1})}{q(\vx_t|\vx_{t-1})}\right] \\ & \text{对于最后一项也是同理}\\ &= \underbrace{\E_{q(\vx_1|\vx_0)} \left[\log p_{\vec\theta}(\vec x_0|\vec x_1)\right]}_{重构项} - \underbrace{\mathbb E_{q(\vx_{T-1}|\vx_0)} \left[ D_\text{KL}\left(q(\vx_T\vert \vx_{T-1})\Vert p(\vx_T)\right)\right]}_{先验匹配项} - \sum_{t=1}^{T-1}\underbrace{\E_{q(\vx_{t-1},\vx_{t+1}|\vx_0)} \left[D_\text{KL}(q(\vx_t|\vx_{t-1})\Vert p_\vtheta (\vx_t|\vx_{t+1}))\right]}_{一致性约束项} \\ \end{aligned} \]

可以看到ELBO可以拆解为三项：

1. 重构项要求用第一层隐变量恢复原始数据的似然度最大化；
2. 先验匹配项要求最终的隐变量接近高斯先验分布。这一项其实不会体现在损失函数里，因为这一步没有任何可训练参数。实践中我们会通过一些设计和参数选择，使得最后一层隐变量接近高斯先验；
3. 一致性约束项要求隐变量在前向加噪和逆向去噪的两个方向产生的隐变量分布一致。

那么VDM下，ELBO的所有项都可以用蒙特卡洛采样的方式优化。但是注意到最后一项涉及到两个随机变量的采样\(\{\vx_{t-1}, \vx_{t+1}\}\)，这可能导致估计结果的方差偏大。为了优化这个问题，我们可以考虑重新推导ELBO，写出一个一次只需采样一个随机变量的版本。

注意到由于马尔科夫链的性质，有 \[ q(\vx_t|\vx_{t-1}, \vx_0) = \frac{ q(\vx_{t-1}|\vx_t,\vx_0) q(\vx_t|\vx_0) } { q(\vx_{t-1}|\vx_0) }, \] 利用这个式子，我们可以重新推导ELBO： \[ \begin{aligned} \log p(\vx) &\geq \E _{q(\vx_{1:T}|\vx_0)}\left[\log\frac{p(\vx_{0:T})}{q(\vx_{1:T}|\vx_0)}\right] \\ &= \mathbb E_{q(\vec x_{1:T}|\vec x_0)}\left[\log\frac{p(\vec x_T)\prod_{t=1}^Tp_{\vec \theta}(\vec x_{t-1}|\vec x_t)}{\prod_{t=1}^T q(\vec x_t|\vec x_{t-1})}\right] & \text{链式法则、马尔科夫性质}\\ &= \mathbb E_{q(\vec x_{1:T}|\vec x_0)}\left[\log \frac{p(\vx_T)p_\vtheta (\vx_0|\vx_1)\prod_{t=2}^T p_\vtheta (\vx_{t-1}|\vx_t)}{q(\vx_1|\vx_0)\prod_{t=2}^Tq(\vx_t|\vx_{t-1})}\right] \\ &= \mathbb E_{q(\vec x_{1:T}|\vec x_0)}\left[\log \frac{p(\vx_T)p_\vtheta (\vx_0|\vx_1)\prod_{t=2}^T p_\vtheta (\vx_{t-1}|\vx_t)}{q(\vx_1|\vx_0)\prod_{t=2}^Tq(\vx_t|\vx_{t-1}, \vx_0)}\right] & \text{应用马尔科夫链的性质} \\ &= \mathbb E_{q(\vec x_{1:T}|\vec x_0)}\left[\log \frac{p(\vx_T)p_\vtheta (\vx_0|\vx_1)}{q(\vx_1|\vx_0)} + \log \prod_{t=2}^T\frac{ p_\vtheta (\vx_{t-1}|\vx_t)}{q(\vx_t|\vx_{t-1}, \vx_0)}\right] \\ &= \mathbb E_{q(\vec x_{1:T}|\vec x_0)}\left[\log \frac{p(\vx_T)p_\vtheta (\vx_0|\vx_1)}{q(\vx_1|\vx_0)} + \log \prod_{t=2}^T\frac{ p_\vtheta (\vx_{t-1}|\vx_t)}{\frac{q(\vx_{t-1}|\vx_t,\vx_0)\textcolor{red}{q(\vx_t|\vx_0)}}{\textcolor{red}{q(\vx_{t-1}|\vx_0)}}}\right] & \text{贝叶斯定理}\\ & \text{红色部分在连乘的相邻项间被抵消} \\ &=\E _{q(\vx_{1:T}|\vx_0)}\left[\log \frac{p(\vx_T)p_\vtheta(\vx_0|\vx_1)}{\cancel{q(\vx_1|\vx_0)}} + \log \frac{\cancel{q(\vx_1|\vx_0)}}{q(\vx_T|\vx_0)}+\log\prod_{t=2}^T \frac{p_\vtheta(\vx_{t-1}|\vx_t)}{q(\vx_{t-1}|\vx_t, \vx_0)}\right] \\ &=\E _{q(\vx_{1:T}|\vx_0)}\left[\log \frac{p(\vx_T)p_\vtheta(\vx_0|\vx_1)}{q(\vx_T|\vx_0)}+\sum_{t=2}^T \log \frac{p_\vtheta(\vx_{t-1}|\vx_t)}{q(\vx_{t-1}|\vx_t, \vx_0)}\right] \\ &=\E _{q(\vx_{1:T}|\vx_0)}[\log p_\vtheta(\vx_0|\vx_1)] + \E _{q(\vx_{1:T}|\vx_0)}\left[ \log\frac{p(\vx_T)}{q(\vx_T|\vx_0)} \right] + \sum_{t=2}^T\E _{q(\vx_{1:T}|\vx_0)} \left[\log\frac{p_\vtheta(\vx_{t-1}|\vx_t)}{q(\vx_{t-1}|\vx_t, \vx_0)}\right] \\ &=\E_{q(\vx_1|\vx_0)}[\log p_\vtheta(\vx_0|\vx_1)] + \E_{q(\vx_T|\vx_0)}\left[\log\frac{p(\vx_T)}{q(\vx_T|\vx_0)}\right] + \sum_{t=2}^T \E_{q(\vx_t, \vx_{t-1}|\vx_0)}\left[\log\frac{p_\vtheta(\vx_{t-1}|\vx_t)}{q(\vx_{t-1}|\vx_t, \vx_0)}\right] \\ &=\underbrace{\E_{q(\vx_1|\vx_0)}[\log p_\vtheta(\vx_0|\vx_1)]}_{重构项} - \underbrace{D_\text{KL}(q(\vx_T|\vx_0)\Vert p(\vx_T))}_{先验匹配项} - \sum_{t=2}^T \underbrace{\E_{q(\vx_t|\vx_0)}[D_\text{KL}(q(\vx_{t-1} | \vx_t, \vx_0)\Vert p_\vtheta (\vx_{t-1}|\vx_t))]}_{去噪匹配项} \end{aligned} \tag{7}\] 经过这样的推导，我们将ELBO拆解为3项：

1. 重构项：可以用蒙特卡洛方法采样估计
2. 先验匹配项：不包含可训练参数。由于VDM的假设，这一项的值为0
3. 去噪匹配项：在这一项中，\(q(\vx_{t-1}|\vx_t, \vx_0)\)为ground-truth信号，它定义了已知真实样本\(\vx_0\)时，从\(\vx_t\)到\(\vx_{t-1}\)的去噪过程是怎么样的。我们用\(p_\vtheta(\vx_{t-1}|\vx_t)\)去近似它。

注意到由于我们没有应用马尔科夫性质以外的假设，所以其实以上推导对于任意MHVAE都是适用的。也适用于\(T=1\)时的MHVAE，这时MHVAE退化成普通的VAE。

ELBO中，\(D_\text{KL}(q(\vx_{t-1} | \vx_t, \vx_0)\Vert p_\vtheta (\vx_{t-1}|\vx_t))\)是优化的重点。在一般的MHVAE中，由于encoder可能是可学习的任意的函数，这个优化很难实现。而在VDM中，我们将encoder设置为固定的线性高斯模型，应用重参数化技巧，encoder过程可以重写为： \[ \vx_t = \sqrt{\alpha_t} \vx_{t-1} + \sqrt{1 - \alpha_t} \vec\epsilon, \text{其中}\vec\epsilon \sim \mathcal N(\vec\epsilon; \mZero, \mI) \] 这是一个递归的过程。那么对于任意的\(t\)，\(\vx_t\sim q(\vx_t|\vx_0)\)可以写为（式子中所有的\(\vec \epsilon\)都独立同分布地采样于\(\mathcal N(\vec \epsilon;\mZero, \mI)\)）： \[ \begin{aligned} \vx_t &= \sqrt{\alpha_t} \vx_{t-1} + \sqrt{1 - \alpha_t} \vec\epsilon^*_{t-1} \\ &= \sqrt{\alpha_t}(\sqrt{\alpha_{t-1}}\vx_{t-2} + \sqrt{1 - \alpha_{t-1}}\vec\epsilon^*_{t-2}) + \sqrt{1 - \alpha_t} \vec\epsilon^*_{t-1}\\ &= \sqrt{\alpha_t\alpha_{t-1}} \vx_{t-2} + \sqrt{\alpha_t - \alpha_t \alpha_{t-1}} \vec\epsilon^*_{t-2} + \sqrt{1 - \alpha_t}\vec\epsilon^*_{t-1} \\ &\text{两个高斯变量相加} \\ &= \sqrt{\alpha_t\alpha_{t-1}} \vx_{t-2} + \sqrt{\sqrt{\alpha_t - \alpha_t \alpha_{t-1}}^2 + \sqrt{1 - \alpha_t}^2} \vec\epsilon_{t-2} \\ &= \sqrt{\alpha_t\alpha_{t-1}} \vx_{t-2} + \sqrt{\alpha_t - \alpha_t\alpha_{t-1} + 1 - \alpha_t}\vec\epsilon_{t-2} \\ &= \sqrt{\alpha_t \alpha_{t-1}} \vx_{t-2} + \sqrt{1 - \alpha_t\alpha_{t-1}} \vec\epsilon_{t-2} \\ &= \dots \\ &= \sqrt{\prod_{i=1}^t \alpha_i} \vx_0 + \sqrt{1 - \prod_{i=1}^t\alpha_i}\vec\epsilon_0 \\ &= \sqrt{\bar{\alpha}_t} \vx_0 + \sqrt{1 - \bar{\alpha}_t} \vec\epsilon_0 \\ &\sim \mathcal N(\vx_t; \sqrt{\bar\alpha_t} \vx_0, (1 - \bar\alpha_t)\mI) \end{aligned} \tag{8}\] 这里应用了“两个独立高斯随机变量之和仍服从高斯分布，其均值为原分布均值之和，方差为原分布方差之和”这个知识点。

这样，我们就发现\(q(\vx_{t}|\vx_0)\)也服从高斯分布，知道了它的表达式。我们也可由此得知\(q(\vx_{t-1}|\vx_0)\)的表达式。

接下来的工作就是得到\(q(\vx_{t-1}|\vx_t, \vx_0)\)： \[ \begin{aligned} q(\vx_{t-1}|\vx_t, \vx_0) &= \frac{q(\vx_t|\vx_{t-1}, \vx_0) q(\vx_{t-1}|\vx_0)}{q(\vx_t|\vx_0)} \\ &= \frac{\mathcal N(\vx_t; \sqrt{\alpha_t}\vx_{t-1}, (1 - \alpha_t)\mI)\mathcal N(\vx_{t-1}; \sqrt{\bar{\alpha}_{t-1}}\vx_0, (1 - \bar\alpha_{t-1})\mI)}{\mathcal N(\vx_t;\sqrt{\bar\alpha_t}\vx_0, (1 - \bar\alpha_t)\mI)} \\ &\propto \exp\left\{-\left[ \frac{(\vx_t - \sqrt{\alpha_t}\vx_{t-1})^2}{2(1 - \alpha_t)} + \frac{(\vx_{t-1} - \sqrt{\bar\alpha_{t-1}}\vx_0)^2}{2(1 - \bar\alpha_{t-1})} - \frac{(\vx_t - \sqrt{\bar\alpha_t}\vx_0)^2}{2(1 - \bar\alpha_t)} \right] \right\} \\ &= \exp\left\{-\frac{1}{2}\left[ \frac{(\vx_t - \sqrt{\alpha_t}\vx_{t-1})^2}{1 - \alpha_t} + \frac{(\vx_{t-1} - \sqrt{\bar\alpha_{t-1}}\vx_0)^2}{1 - \bar\alpha_{t-1}} - \frac{(\vx_t - \sqrt{\bar\alpha_t}\vx_0)^2}{1 - \bar\alpha_t} \right] \right\} \\ & \text{把与}\vx_t,\vx_0\text{有关的常数项摘出来} \\ &= \exp\left\{ -\frac{1}{2}\left[\frac{-2\sqrt{\alpha_t} \vx_t \vx_{t-1} + \alpha_t \vx^2_{t-1}}{1 - \alpha_t} + \frac{\vx^2_{t-1} - 2\sqrt{\bar\alpha_{t-1}}\vx_{t-1}\vx_0}{1 - \bar\alpha_{t-1}} + C(\vx_t, \vx_0)\right] \right\} \\ &\propto \exp\left\{-\frac{1}{2} \left[ -\frac{2\sqrt{\alpha_t}\vx_t\vx_{t-1}}{1 - \alpha_t} + \frac{\alpha_t\vx^2_{t-1}}{1 - \alpha_t} + \frac{\vx_{t-1}^2}{1 - \bar\alpha_{t-1}} - \frac{2\sqrt{\bar\alpha_{t-1}}\vx_{t-1}\vx_0}{1 - \bar\alpha_{t-1}} \right]\right\} \\ &= \exp\left\{-\frac{1}{2} \left[(\frac{\alpha_t}{1 - \alpha_t} + \frac{1}{1 - \bar\alpha_{t-1}})\vx^2_{t-1} - 2 \left(\frac{\sqrt{\alpha_t}\vx_t}{1 - \alpha_t} + \frac{\sqrt{\bar\alpha_{t-1}}\vx_0}{1 - \bar\alpha_{t-1}}\right)\vx_{t-1}\right] \right\}\\ &= \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})}\vx^2_{t-1} - 2 \left(\frac{\sqrt{\alpha_t}\vx_t}{1 - \alpha_t} + \frac{\sqrt{\bar\alpha_{t-1}}\vx_0}{1 - \bar\alpha_{t-1}}\right)\vx_{t-1}\right] \right\}\\ &\text{注意}\alpha_t \bar\alpha_{t-1} = \bar\alpha_t \\ &= \exp\left\{-\frac{1}{2} \left[\frac{1 - \bar\alpha_t}{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}\vx^2_{t-1} - 2 \left(\frac{\sqrt{\alpha_t}\vx_t}{1 - \alpha_t} + \frac{\sqrt{\bar\alpha_{t-1}}\vx_0}{1 - \bar\alpha_{t-1}}\right)\vx_{t-1}\right] \right\}\\ &= \exp\left\{ -\frac{1}{2} \left(\frac{1 - \bar\alpha_t}{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}\right)\left[\vx_{t-1}^2 - 2 \frac{ \left(\frac{\sqrt{\alpha_t}\vx_t}{1 - \alpha_t} + \frac{\sqrt{\bar\alpha_{t-1}}\vx_0}{1 - \bar\alpha_{t-1}}\right)}{\frac{1 - \bar\alpha_t}{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}}\vx_{t-1}\right] \right\}\\ &= \exp\left\{ -\frac{1}{2} \left(\frac{1 - \bar\alpha_t}{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}\right)\left[\vx_{t-1}^2 - 2 \frac{ \left(\frac{\sqrt{\alpha_t}\vx_t}{1 - \alpha_t} + \frac{\sqrt{\bar\alpha_{t-1}}\vx_0}{1 - \bar\alpha_{t-1}}\right)(1 - \alpha_t)(1 - \bar\alpha_{t-1})}{1 - \bar\alpha_t}\vx_{t-1}\right] \right\}\\ &= \exp\left\{ -\frac{1}{2} \left(\frac{1}{\frac{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}{1 - \bar\alpha_t}}\right)\left[\vx_{t-1}^2 - 2 \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})\vx_t + \sqrt{\bar\alpha_{t-1}}(1 - \alpha_t)\vx_0}{1 - \bar\alpha_t}\vx_{t-1}\right] \right\}\\ &\propto \mathcal N(\vx_{t-1}; \underbrace{\frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})\vx_t + \sqrt{\bar\alpha_{t-1}}(1 - \alpha_t)\vx_0}{1 - \bar\alpha_t}}_{\vec\mu_q(\vx_t, \vx_0)},\underbrace{\frac{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}{1 - \bar\alpha_t}\mI}_{\mathbf\Sigma_q(t)}) \end{aligned} \tag{9}\]

于是我们可以看到\(\vx_{t-1}\sim q(\vx_{t-1}\vert\vx_t,\vx_0)\)服从正态分布。\(\vec\mu_q(\vx_t, \vx_0)\)是\(\vx_t,\vx_0\)的函数，而\(\mathbf \Sigma_q(t)\)是\(\alpha\)参数的函数。

设\(\mathbf \Sigma_q(t)=\sigma^2_q(t)\mI\)，其中 \[ \sigma^2_q(t)= \frac{(1 - \alpha_t)(1 - \bar\alpha_{t-1})}{1 - \bar\alpha_t} \tag{10}\]

我们要让\(p_\vtheta(\vx_{t-1}|\vx_t)\)尽可能接近\(q(\vx_{t-1}|\vx_t,\vx_0)\)。既然\(q(\vx_{t-1}|\vx_t,\vx_0)\)是高斯分布，我们也可以用高斯分布建模\(p_\vtheta(\vx_{t-1}|\vx_t)\)。不妨让它的协方差矩阵也为\(\mathbf \Sigma_q(t)=\sigma^2_q(t)\mI\)，均值则设为\(\vec\mu_\vtheta(\vx_t, t)\)，是\(\vx_t\)的函数，但不依赖\(\vx_0\)，毕竟decoder无法获得\(\vx_0\)的真值。

要让两个高斯分布相近，需要考虑它们的KL散度： \[ D_{\text{KL}} (\mathcal N(\vx;\vec\mu_\vx, \mathbf\Sigma_x)\Vert \mathcal N(\vec y;\vec\mu_y,\mathbf \Sigma_y)) = \frac{1}{2}\left[ \log \frac{|\mathbf \Sigma_y|}{|\mathbf \Sigma_x|} - d + \text{tr}(\mathbf \Sigma_y^{-1}\mathbf \Sigma_x) + (\vec\mu_y-\vec\mu_x)^T\mathbf \Sigma_y^{-1}(\vec\mu_y-\vec\mu_x) \right] \] 这里\(d\)是数据的维度。

因为我们将协方差矩阵设置为相同的，因此上式将只与均值的差有关。 \[ \begin{aligned} &\argmin{\vtheta} D_\text{KL}(q(\vx_{t-1}|\vx_t,\vx_0)\Vert p_\vtheta(\vx_{t-1}|\vx_t)) \\ =~& \argmin{\vtheta} D_\text{KL}(\mathcal N(\vx_{t-1};\vec\mu_q,\mathbf\Sigma_q(t))\Vert\mathcal N(\vx_{t-1}; \vec\mu_\vtheta, \mathbf\Sigma_q(t))) \\ =~&\argmin{\vtheta} \frac{1}{2} \left[\log\frac{|\mSigma_q(t)|}{|\mSigma_q(t)|} - d + \text{tr}(\mSigma_q(t)^{-1}\mSigma_q(t)) + (\vec\mu_\vtheta -\vec\mu_q)^T\mSigma_q(t)^{-1}(\vec\mu_\vtheta -\vec\mu_q)\right]\\ =~& \argmin{\vtheta} \frac{1}{2}[\log 1 - d + d+ (\vec\mu_\vtheta - \vec\mu_q)^T\mSigma_q(t)^{-1}(\vec\mu_\vtheta -\vec\mu_q)]\\ =~&\argmin{\vtheta} \frac{1}{2}\left[ (\vec\mu_\vtheta - \vec\mu_q)^T\mSigma_q(t)^{-1}(\vec\mu_\vtheta -\vec\mu_q) \right]\\ =~&\argmin{\vtheta} \frac{1}{2} \left[(\vec\mu_\vtheta - \vec\mu_q)^T(\sigma_q^2(t)\mI)^{-1}(\vec\mu_\vtheta -\vec\mu_q) \right] \\ =~&\argmin{\vtheta} \frac{1}{2\sigma^2(t)}\left[\left\Vert \vec\mu_\vtheta - \vec\mu_q \right\Vert_2^2\right] \end{aligned} \] 上面的式子中，\(\vec\mu_\vtheta\)是\(\vec\mu_q(\vec x_t, \vx_0)\)的缩写， \[ \vec\mu_q(\vx_t,\vx_0)=\frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})\vx_t + \sqrt{\bar\alpha_{t-1}}(1 - \alpha_t)\vx_0}{1 - \bar\alpha_t} \]

而\(\vec\mu_\vtheta\)是\(\vec\mu_\vtheta (\vx_t, t)\)的缩写。为了让\(\vec\mu_\vtheta\)趋近\(\vec\mu_q\)，可以这样设计： \[ \vec\mu_\vtheta(\vx_t, t) = \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})\vx_t + \sqrt{\bar\alpha_{t-1}}(1 - \alpha_t)\hat \vx_\vtheta(\vx_t, t)}{1 - \bar\alpha_t} \] 其中\(\vx_\vtheta(\vx_t, t)\)是基于\(\vx_t\)和\(t\)对\(\vx_0\)做出的预测。综合以上推理，优化问题变为： \[ \begin{aligned} &\argmin \vtheta D_\text{KL} (q(\vx_{t-1}|\vx_t, \vx_0) \Vert p_\vtheta(\vx_{t-1}|\vx_t)) \\ =~&\argmin \vtheta D_\text{KL} (\mathcal N(\vx_{t-1}; \vec\mu_q, \mSigma_q(t)) \Vert \mathcal N(\vx_{t-1}; \vec\mu_\vtheta, \mSigma_q(t)))\\ =~&\argmin \vtheta \frac{1}{2\sigma_q^2(t)}\left[\left\Vert {\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})\vx_t + \sqrt{\bar\alpha_{t-1}}(1 - \alpha_t)\hat\vx_\vtheta(\vx_t, t) \over 1 - \bar\alpha_t} - {\sqrt{\alpha_t} (1 - \bar\alpha_{t-1})\vx_t + \sqrt{\bar\alpha_{t-1}}(1 - \alpha_t)\vx_0\over 1 - \bar\alpha_t}\right\Vert_2^2\right] \\ =~&\argmin \vtheta \frac{1}{2\sigma_q^2(t)} \left[\left\Vert {\sqrt{\bar\alpha_{t-1}}(1 - \alpha_t)\hat\vx_\vtheta(\vx_t, t) \over 1 - \bar\alpha_t} - {\sqrt{\bar\alpha_{t-1}(1 - \alpha_t)\vx_0 \over 1 - \bar\alpha_t}} \right\Vert_2^2 \right]\\ =~&\argmin\vtheta \frac{1}{2\sigma_q^2(t)} \left[\left\Vert {\sqrt{\bar\alpha_{t-1}}(1 - \alpha_t)\over 1 - \bar\alpha_t}(\hat\vx_\vtheta(\vx_t, t) - \vx_0) \right\Vert_2^2\right]\\ =~&\argmin\vtheta \frac{1}{2\sigma_q^2(t)} {{\bar\alpha_{t-1}}(1 - \alpha_t)^2\over (1 - \bar\alpha_t)^2} \left[\left\Vert \hat\vx_\vtheta(\vx_t, t) - \vx_0 \right\Vert_2^2\right] \end{aligned} \tag{11}\] 于是，VDM的优化问题可以归结为用一个神经网络从带噪声的图像中恢复原始图像。

对公式 7中的求和项的优化可以近似为对在所有时间步\(t\)上如下期望值的优化： \[ \argmin\vtheta \E _{t\sim U\left\{2, T\right\}}\left[ \E_{q(\vx_t|\vx_0)} \left[ D_\text{KL}(q(\vx_{t-1}|\vx_t, \vx_0)\Vert p_\vtheta(\vx_{t-1}|\vx_t)) \right] \right] \]

7 学习噪声的参数

本节讨论影响VDM噪声的参数\(\alpha_t\)要如何学习得到。比较容易想到的办法是使用以\(\vec\eta\)为参数的模型\(\hat \alpha_{\vec\eta}(t)\)作预测。这样做是低效的，因为推理时，你需要在每一步\(t\)都预测对应的\(\bar\alpha_t\)。当然你可以提前把计算结果存下来。但是下文将介绍另一种方法。

将公式 10带入公式 11，我们得到： \[ \begin{aligned} \frac{1}{2\sigma^2_q(t)}\frac{\bar\alpha_{t-1}(1 - \alpha_t)^2}{(1 - \bar\alpha_t)^2}\left[\left\Vert \hat\vx_\vtheta(\vx_t, t) - \vx_0\right\Vert_2^2\right] &= \frac{1}{2\frac{(1 - \alpha_t)(1- \bar\alpha_{t-1})}{1 - \bar\alpha_t}} \frac{\bar\alpha_{t-1}(1 - \alpha_t)^2}{(1 - \bar\alpha_t)^2} \left[\left\Vert \hat\vx_\vtheta(\vx_t, t) - \vx_0\right\Vert_2^2\right] \\ &= \frac{1}{2} \frac{1 - \bar\alpha_t}{(1 - \alpha_t)(1 - \bar\alpha_{t-1})} \frac{\bar\alpha_{t-1}(1 - \alpha_t)^2}{(1 - \bar\alpha_t)^2} \left[\Vert \hat\vx_\vtheta(\vx_t, t) - \vx_0\Vert_2^2\right] \\ &= \frac{1}{2} \frac{\bar\alpha_{t-1}(1 - \alpha_t)}{(1 - \bar\alpha_{t-1})(1 - \bar\alpha_t)} \left[\Vert \hat\vx_\vtheta(\vx_t, t) - \vx_0\Vert_2^2\right]\\ &= \frac{1}{2} \frac{\bar\alpha_{t-1} - \bar\alpha_t}{(1 - \bar\alpha_{t-1})(1 - \bar\alpha_t)} \left[\Vert \hat\vx_\vtheta(\vx_t, t) - \vx_0\Vert_2^2\right] \\ &= \frac{1}{2} \frac{\bar\alpha_{t-1} -\bar\alpha_{t-1}\bar\alpha_t +\bar\alpha_{t-1}\bar\alpha_t - \bar\alpha_t}{(1 - \bar\alpha_{t-1})(1 - \bar\alpha_t)} \left[\Vert \hat\vx_\vtheta(\vx_t, t) - \vx_0\Vert_2^2\right] \\ &= \frac{1}{2} \frac{\bar\alpha_{t-1}(1 - \bar\alpha_t) - \bar\alpha_t(1 - \bar\alpha_{t-1})}{(1 - \bar\alpha_{t-1})(1 - \bar\alpha_t)} \left[\Vert \hat\vx_\vtheta(\vx_t, t) - \vx_0\Vert_2^2\right] \\ &= \frac{1}{2} \left(\frac{\bar\alpha_{t-1}}{1 - \bar\alpha_{t-1}} - \frac{\bar\alpha_t}{1 - \bar\alpha_t}\right) \left[\Vert \hat\vx_\vtheta(\vx_t, t) - \vx_0\Vert_2^2\right] \\ \end{aligned} \tag{12}\]

回想起\(q(\vx_t|\vx_0)\)是形为\(\mathcal N(\vx_t; \sqrt{\bar\alpha_t}\vx_0, (1 - \bar\alpha_t)\mI)\)，根据SNR（信噪比）的定义，\(\text{SNR} = \frac{\mu^2}{\sigma^2}\)，时间步\(t\)的SNR为： \[ \text{SNR}(t) = \frac{\bar\alpha_t}{1 - \bar\alpha_t} \] 那么公式 12可以进一步简化为： \[ \frac{1}{2\sigma^2_q(t)}\frac{\bar\alpha_{t-1}(1 - \alpha_t)^2}{(1 - \bar\alpha_t)^2}\left[\left\Vert \hat\vx_\vtheta(\vx_t, t) - \vx_0\right\Vert_2^2\right] = \frac{1}{2}(\text{SNR}(t-1)-\text{SNR}(t)) \left[\Vert \hat\vx_\vtheta(\vx_t, t) - \vx_0\Vert_2^2\right] \] 在VDM中，SNR应该随著时间步\(t\)增加而增加，因为\(\vx_t\)会随著\(t\)增加，从原图逐渐变成标准正态分布。

那么不妨将SNR函数设计为 \[ \text{SNR}(t) = \frac{\bar\alpha_t}{1 - \bar\alpha_t}= \exp(-\omega_{\vec\eta}(t)) \] 所以 \[ \bar\alpha_t = \text{sigmoid}(-\omega_\vec\eta(t)) \] 其中\(\vec\eta\)是可学习的模型参数。

8 VDM的三种等效形式

如前文所述，VDM可以设计为从\(\vx_t\)预测\(\vx_0\)的模型。但是，VDM还有两种其它等效形式。

重写公式 8的重参数化技巧，得到： \[ \vx_0 = \frac{\vx_t - \sqrt{1 - \bar\alpha_t}\vec\epsilon_0}{\sqrt{\bar\alpha_t}} \] 代入公式 9中得到的\(\vec\mu_q(\vx_t, \vx_0)\)，得到 \[ \begin{aligned} \vec\mu_q(\vx_t, \vx_0) &= \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})\vx_t + \sqrt{\bar\alpha_{t-1}}(1 - \alpha_t)\vx_0}{1 - \bar\alpha_t} \\ &= \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})\vx_t + \sqrt{\bar\alpha_{t-1}}(1 - \alpha_t) \frac{\vx_t - \sqrt{1 - \bar\alpha_t}\vec\epsilon_0}{\sqrt{\bar\alpha_t}}}{1 - \bar\alpha_t} \\ &= \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})\vx_t + (1 - \alpha_t) \frac{\vx_t - \sqrt{1 - \bar\alpha_t}\vec\epsilon_0}{\sqrt{\alpha_t}}}{1 - \bar\alpha_t} \\ &= \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})\vx_t}{1 - \bar\alpha_t} + \frac{(1 - \alpha_t)\vx_t}{(1 - \bar\alpha_t)\sqrt{\alpha_t}} - \frac{(1 - \alpha_t)\sqrt{1 - \bar\alpha_t}\vec\epsilon_0}{(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)\vx_t - \frac{(1 - \alpha_t)\sqrt{1 - \bar\alpha_t}}{(1 - \bar\alpha_t)\sqrt{\alpha_t}} \vec\epsilon_0\\ &= \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)\vx_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}} \vec\epsilon_0\\ &= \frac{1 - \bar\alpha_t}{(1 - \bar\alpha_t)\sqrt{\alpha_t}}\vx_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}} \vec\epsilon_0\\ &= \frac{1}{\sqrt{\alpha_t}}\vx_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}} \vec\epsilon_0\\ \end{aligned} \] 因此，另一种\(\vec\mu_\vtheta(\vx_t, t)\)的等效设计是： \[ \vec\mu_\vtheta(\vx_t, t) = \frac{1}{\sqrt{\alpha_t}} \vx_t - \frac{1-\alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}\hat{\vec\epsilon}_\vtheta(\vx_t, t) \] 对应的，优化问题就变为： \[ \begin{aligned} &\argmin \vtheta D_\text{KL} (q(\vx_{t-1}|\vx_t, \vx_0)\Vert p_\vtheta(\vx_{t-1}|\vx_t)) \\ =~&\argmin\vtheta D_\text{KL}(\mathcal N(\vx_{t-1};\vec\mu_q, \mat\Sigma_q(t))\Vert \mathcal N(\vx_{t-1};\vec \mu_\vtheta, \mat\Sigma_q(t)))\\ =~&\argmin\vtheta \frac{1}{2\sigma_q^2(t)}\left[\left\Vert \frac{1}{\sqrt{\alpha_t}}\vx_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}} \hat{\vec\epsilon}_\vtheta(\vx_t, t) - \frac{1}{\sqrt{\alpha_t}}\vx_t + \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}\vec\epsilon_0\right\Vert_2^2\right] \\ =~&\argmin\vtheta \frac{1}{2\sigma_q^2(t)}\left[\left\Vert \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt\alpha_t}\vec\epsilon_0 - \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt{\alpha_t}}\hat{\vec\epsilon}_\vtheta(\vx_t, t)\right\Vert_2^2\right] \\ =~&\argmin\vtheta \frac{1}{2\sigma_q^2(t)}\left[\left\Vert \frac{1 - \alpha_t}{\sqrt{1 - \bar\alpha_t}\sqrt\alpha_t}(\vec\epsilon_0 - \hat{\vec\epsilon}_\vtheta(\vx_t, t))\right\Vert_2^2\right] \\ =~&\argmin\vtheta \frac{1}{2\sigma_q^2(t)}\frac{(1 - \alpha_t)^2}{(1 - \bar\alpha_t)\alpha_t}\left[\left\Vert \vec\epsilon_0 - \hat{\vec\epsilon}_\vtheta(\vx_t, t)\right\Vert_2^2\right] \\ \end{aligned} \] 式子中\(\hat{\vec\epsilon}_\vtheta\)是用于预测噪声\(\vec\epsilon_0\sim\mathcal N(\vec\epsilon; \vec 0, \mI)\)，从而将\(\vec x_t\)恢复为\(\vec x_0\)的模型。

这样，我们看到理论上预测噪声和预测原始图像在理论上是等价的。但是许多工作表明实际上预测噪声效果更好。

第三种等价形式的推导要用到特威迪公式（Tweedies’s Formula）。对于高斯变量\(\vec z\sim \mathcal N(\vec z; \vec \mu_z, \mat \Sigma_z)\)，特威迪公式表明： \[ \E[\vec\mu_{\vec z}|\vec z] = \vec z + \Sigma_z\nabla_{\vec z}\log p(\vec z) \] 已知\(q(\vec x_t|\vec x_0)=\mathcal N(\vec x_t; \sqrt{\bar\alpha_t} \vec x_0, (1 - \bar\alpha_t)\mI)\)，那么，根据特威迪公式有 \[ \E[\vec\mu_{\vx_t}|\vx_t] = \vx_t + (1 - \bar\alpha_t)\nabla_{\vx_t}\log p(\vx_t) \] 后面为了方便，将\(\nabla_{\vx_t} \log p(\vx_t)\)简写为\(\nabla \log p(\vx_t)\). 已知\(\vec\mu_{\vx_t}=\sqrt{\bar\alpha_t} \vx_0\)，因此 \[ \begin{aligned} \sqrt{\bar\alpha_t}\vx_0 = \vx_t + (1 - \bar\alpha_t)\nabla\log p(\vx_t)\\ \therefore \vx_0 = \frac{\vx_t + (1 - \bar\alpha_t)\nabla\log p(\vx_t)}{\sqrt{\bar\alpha_t}} \end{aligned} \] 再次将\(\vx_0\)代入ground-truth的去噪过程公式 9，得到： \[ \begin{aligned} \vec\mu_q(\vx_t, \vx_0) &= \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})\vx_t + \sqrt{\bar\alpha_{t-1}}(1 - \alpha_t)\vx_0}{1 - \bar\alpha_t}\\ &= \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})\vx_t + \sqrt{\bar\alpha_{t-1}}(1 - \alpha_t)\frac{\vx_t + (1 - \bar\alpha_t)\nabla\log p(\vx_t)}{\sqrt{\bar\alpha_t}}}{1 - \bar\alpha_t}\\ &= \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})\vx_t +(1 - \alpha_t)\frac{\vx_t + (1 - \bar\alpha_t)\nabla\log p(\vx_t)}{\sqrt{\alpha_t}}}{1 - \bar\alpha_t}\\ &= \frac{\sqrt{\alpha_t}(1 - \bar\alpha_{t-1})\vx_t }{1 - \bar\alpha_t} + \frac{(1 - \alpha_t)\vx_t}{(1 - \bar\alpha_t)\sqrt{\alpha_t}} + \frac{(1 - \alpha_t)(1 - \bar\alpha_t)\nabla\log p(\vx_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)\vx_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}}\nabla\log p(\vx_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)\vx_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}}\nabla\log p(\vx_t)\\ &= \frac{1 - \bar\alpha_t}{(1 - \bar\alpha_t)\sqrt{\alpha_t}} \vx_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}}\nabla\log p(\vx_t)\\ &= \frac{1}{\sqrt{\alpha_t}} \vx_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}}\nabla\log p(\vx_t) \\ \end{aligned} \] 再一次，类似的将\(\vec\mu_\vtheta(\vx_t, t)\)设计为： \[ \vec\mu_\vtheta(\vx_t, t)=\frac{1}{\sqrt{\alpha_t}} \vx_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}} \vec s_\vtheta(\vx_t, t) \] 对应的优化过程变成： \[ \begin{aligned} &\argmin \vtheta D_\text{KL} (q(\vx_{t-1}|\vx_t, \vx_0)\Vert p_\vtheta(\vx_{t-1}\vert \vx_t))\\ =~& \argmin\vtheta D_\text{KL}(\mathcal N(\vx_{t-1}; \vec \mu_q, \mSigma_q(t))\Vert \mathcal N(\vx_{t-1};\vec\mu_\vtheta, \mSigma_a(t))) \\ =~& \argmin\vtheta \frac{1}{2\sigma_q^2(t)} \left[\left\Vert \frac{1}{\sqrt{\alpha_t}} \vx_t + \frac{1 - \alpha_t}{\sqrt{\alpha_t}} \vec s_\vtheta(\vx_t, t) - \frac{1}{\sqrt{\alpha_t}} \vx_t - \frac{1 - \alpha_t}{\sqrt{\alpha_t}}\nabla\log p(\vx_t) \right\Vert_2^2\right] \\ =~& \argmin\vtheta \frac{1}{2\sigma_q^2(t)} \left[\left\Vert \frac{1 - \alpha_t}{\sqrt{\alpha_t}} \left(\vec s_\vtheta(\vx_t, t) - \nabla\log p(\vx_t) \right) \right\Vert_2^2\right] \\ =~& \argmin\vtheta \frac{1}{2\sigma_q^2(t)} \frac{(1 - \alpha_t)^2}{\alpha_t} \left[\Vert \vec s_\vtheta(\vx_t, t) - \nabla\log p(\vx_t) \Vert_2^2\right] \end{aligned} \tag{13}\] 式子中\(\vec s_\vtheta(\vx_t, t)\)是一个预测分数函数（score function）\(\nabla \log p(\vx_t)\)的模型，本质上是在时间步\(t\)对于\(\vx_t\)的梯度。

分数函数\(\nabla \log p(\vx_t)\)与噪声\(\vec\epsilon_0\)显然很相似。不难看出 \[ \begin{aligned} & \vx_0 = \frac{\vx_t + (1 - \bar\alpha_t)\nabla\log p(\vx_t)}{\sqrt{\bar\alpha_t}} = \frac{\vx_t - \sqrt{1 - \bar\alpha_t}\vec\epsilon_0}{\sqrt{\bar\alpha_t}}\\ & \therefore \nabla \log p(\vx_t) = -\frac{1}{\sqrt{1 - \bar\alpha_t}}\vec\epsilon_0 \end{aligned} \] 结果显示分数函数和噪声的区别在于一个由\(t\)决定的系数。直觉上，噪声是施加于原图像使其变得 随机的过程，而我们证明了分数函数通过建模反方向的噪声来恢复图像。

总的来看，我们得到了三种等价的优化目标：直接预测\(\vx_0\), 预测噪声\(\vec \epsilon_0\)，预测分数\(\nabla \log p(\vx_t)\).

9 基于分数的生成式模型

我们展示了VDM可以通过优化\(\vec s_\vtheta(\vx_t, t)\)，预测\(\nabla \log p(\vx_t)\)来实现。然而这个基于特威迪公式的推导没有充分展示出设计背后的信息。分数函数（score function）到底是什么呢。为什么它值得我们去建模呢。幸运的是我们可以学习另一类生成式模型，基于分数的生成模型（Score-based Generative Models），了解分数函数的意义。我们会看到VDM可以解释为一种基于分数的生成模型。

在此之前我们先大概了解一下基于能量的模型（energy-based model）。首先，任意概率分布可以重写为这样的形式： \[ p_\vtheta(x) = \frac{1}{Z_\vtheta}e^{-f_\vtheta(\vx)}, \] 其中\(f_\vtheta(\vx)\)就是能量函数。\(Z_\vtheta\)是用于使得\(\int p_\vtheta(\vx)d\vx=1\)成立的系数。\(Z_\vtheta=\int e^{-f_\vtheta(\vx)}d\vx\)并不一定容易计算。如果\(f_\vtheta(\vx)\)很复杂，\(Z_\vtheta\)就容易变得不可解。

一种避免建模\(Z_\vtheta\)的方法是使用\(\vec s_\vtheta(\vx)\)来学习分数函数\(\nabla \log p(\vx)\)。注意到， \[ \begin{aligned} \nabla_\vx\log p_\vtheta(\vx) &= \nabla_\vx \log(\frac{1}{Z_\vtheta}e^{-f_\vtheta(\vx)}) \\ &= \nabla_\vx\log\frac{1}{Z_\vtheta} + \nabla_\vx \log e^{-f_\vtheta(\vx)}\\ &=-\nabla_\vx f_\vtheta(\vx) \\ &\approx \vec s_\vtheta(\vx) \end{aligned} \]

因此\(\nabla_\vx \log p_\vtheta(\vx)\)可以用一个神经网络近似。分数函数可以通过最小化费雪散度（Fisher Divergence）学习： \[ \mathbb E_{p(x)}\left[ \Vert \vec s_\vtheta(\vx) - \nabla \log p(\vx) \Vert_2^2 \right] \tag{14}\] 分数函数表示的是样本\(\vx\)在数据空间中往哪个方向移动能最大化其对数似然。

一旦学到了这样的分数函数，我们就可以用如下的过程来生成样本： \[ \vx_{i+1}\leftarrow \vx_i + c \nabla \log p(\vx_i) + \sqrt{2c}\vec\epsilon, ~ i=0,1,\dots,K, \] 其中\(\vx_0\)是空间中随机采样的一点，噪声\(\vec\epsilon\sim\mathcal N(\vec\epsilon;\mZero,\mI)\)的作用是防止采样总是收敛于同一模式。这个采样过程被称为朗之万动力学。

目标函数公式 14需要我们得到真实的分数函数，但在建模复杂分布（比如真实图像）的时候是做不到的。score matching技术允许我们在不知道真实分布的情况下最小化费雪散度。

这就是基于分数的生成式模型的原理。但是原始的基于分数的生成模型有几个问题：

1. 如果数据是高维空间中的低维流形时，分数函数不是良定义的。如果一个点落在低维流形外，这个点的概率为0，对数函数在此无定义。而自然图像就被认为是一种高维空间中的低维流形。
2. 学习到的分数函数在数据的低密度区域可能不准确。因为我们训练公式 14时，对于见得少或者没见过的数据，模型收不到很多监督信号。而采样却是从空间中的随机点开始的，不准确的分数函数将导致采样结果落在非最优点
3. 朗之万动力学采样不支持对混合分布的采样。例如对于 \[ p(\vx) = c_1 p_1(\vx) + c_2 p_2(\vx) \] 从特定位置初始化的采样点，可能会以均等的机会落到两个分布中，即使\(c_1 != c_2\).

而以上几个缺点可以同时用VDM的方法解决——往数据里加不同大小的噪声：

1. 高斯噪声的加入将使得空间中每一点的概率都不为0.
2. 高斯噪声的增大使得空间中每一点在训练中被采样到的几率变得更加均匀。流形的低密度区能得到更好的训练。
3. 逐步加强的高斯噪声能形成一种“中间态”的分布，允许我们的采样能够遵循分布的混合系数。具体的，我们可以定义不同时间步下不同的噪声等级\(\left\{\sigma_t\right\}_{t=1}^T\)，并定义每个时间步\(t\)的数据分布 \[ p_{\sigma_t}(\vx_t) = \int p(\vx) \mathcal N(\vx_t; \vx, \sigma_t^2\mI) d\vx \]

神经网络\(\vec s_\vtheta(\vx, t)\)将学习不同时间步下的分数函数： \[ \argmin \vtheta \sum_{t=1}^T \lambda(t)\mathbb E_{p_{\sigma_t}(\vx_t)} \left[ \Vert \vec s_\vtheta(\vx, t) - \nabla \log p_{\sigma_t}(\vx_t) \Vert_2^2 \right], \tag{15}\] 其中\(\lambda(t)\geq 0\)是权重系数。在噪声大的时候，上面的目标函数使模型能够学习不同分布模式的比例；在噪声小的时候，模式逐渐分离，分数函数更精准地学到每个模式的细节。在采样时，对于每个\(t=T, T-1, \dots, 2, 1\)，我们先从高噪音模式开始，然后逐渐降低噪音，直到样本收敛于某个具体模式。

注意到公式 15和公式 13的形式一致。至此，我们建立起了VDM和基于分数的生成模型之间的联系。

从基于分数的生成模型的视角出发，我们还可以发现当MHVAE的时间步数\(T\rightarrow \infty\)时，相当于将离散的随机过程变为连续的随机过程，这时可以用SDE（stochastic differential equation）来描述这个过程，而采样可以通过求逆向的SDE来完成。

10 引导信息

前文只讨论了\(p(x)\). 但我们有时会希望用引导信息控制生成的图像，需要条件概率\(p(\vx|y)\). 这是图像超分辨率、文生图模型的基石。

一种自然的方式是在每一时间步加上条件信息，将公式 \[ p(\vx_{0:T}) = p(\vx_T)\prod_{t=1}^T p_\vtheta(\vx_{t-1}|\vx_t) \] 转变为 \[ p(\vx_{0:T}|y) = p(\vx_T)\prod_{t=1}^T p_\vtheta(\vx_{t-1}|\vx_t, y) \] 然后我们可以预测\(\hat \vx_\vtheta(\vx_t, t, y)\approx \vx_0\)，或者\(\hat{\vec\epsilon}_\vtheta(\vx_t, t, y)\approx \vec\epsilon_0\)，或者\(\vec s_\vtheta(\vx_t, t, y)\approx \nabla \log p(\vx_t|y)\)，从而构造一个VDM。

这种方法可能的缺点是，模型可能忽略或者不充分重视条件信息。

为了解决这个问题，可以使用一些引导技巧。两种常见的引导技巧包括分类器引导和免分类器引导。

11 总结

1. VAE模型是MHVAE的一个特例
2. 介绍了VDM的三种等效优化目标：

• 预测原始图像
• 预测噪声
• 预测分数函数

3. VDM有以下缺点，值得进一步思考

• VDM没有遵循或者模拟人类通常生成数据的方式
• VDM不产生可解释的隐变量。而VAE则有可能产生一些有意义的隐变量。
• 隐变量和原始数据的尺寸被限定为相同的，因此无法学到压缩的、有意义的隐变量。
• 采样过程比较昂贵，需要采样多步。