"> "> 深度学习-扩散模型 | Yufei Luo's Blog

深度学习-扩散模型

简介

扩散模型是受非平衡热力学的启发提出的一种生成模型,最早出现于2015年的文章Deep Unsupervised Learning using Nonequilibrium Thermodynamics。它定义了一个扩散步骤的马尔可夫链,逐渐向数据添加随机噪声,然后学习逆扩散过程,从噪声中构建所需的数据样本。

目前所使用的扩散模型大都是基于2020年的去噪扩散概率模型(Denoising Diffusion Probabilistic Model,DDPM)这一工作,它对之前的扩散模型进行了简化,并通过变分推断来进行建模。本文主要内容为DDPM的相关知识,不涉及在此基础上的改进工作。

背景知识

ELBO

在贝叶斯体系中,通常假设观测到的数据\(\boldsymbol{x}\)是由隐变量\(\boldsymbol{z}\)来控制生成的,可以用一个联合概率分布\(p(\boldsymbol{x},\boldsymbol{z})\)来描述。如果想求出所有观测数据的概率\(p(\boldsymbol{x})\),可以有两种方式,一种办法是计算积分

\[ p(\boldsymbol{x})=\int p(\boldsymbol{x},\boldsymbol{z})d\boldsymbol{z} \]

另一种办法是使用条件概率公式

\[ p(\boldsymbol{x})=\frac{p(\boldsymbol{x},\boldsymbol{z})}{p(\boldsymbol{z}|\boldsymbol{x})} \]

如果要使用精确推断的方法计算出\(p(\boldsymbol{x})\),或是需要对所有的隐变量\(\boldsymbol{z}\)进行积分,又或者是能够知道\(p(\boldsymbol{z}|\boldsymbol{x})\)。对于复杂模型来说这两种办法都不可取,但是从这两个公式中都可以推导出证据下界(Evidence Lower Bound,ELBO),从而使用近似推断的方式求解。

在贝叶斯体系中,推断(Inference)指的是利用已知变量推测未知变量的分布,即我们在已经输入变量\(\boldsymbol{x}\)后,如何获得未知变量\(\boldsymbol{z}\)的分布\(p(\boldsymbol{z}|\boldsymbol{x})\)。精确推断方法指的是准确地计算\(p(\boldsymbol{z}|\boldsymbol{x})\),但往往需要很大的计算开销或者很难实现,现实应用中近似推断更为常用。

近似推断的方法往往分为两大类,第一类是采样,常见的是MCMC方法,第二类是使用另一个分布近似\(p(\boldsymbol{z}|\boldsymbol{x})\),典型代表就是变分推断。变分推断(Variational Inference)是近似推断方法的一种,是一大类通过简单分布来近似复杂分布,从而求解推断问题方法的总称。

从上面两个公式都可以推导出如下的不等式,其中不等号右侧为证据下界:

\[ \log p(\boldsymbol{x}) \geq \mathbb{E}_{q_{\phi}(\boldsymbol{z} | \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\right] \]

不等式的推导过程分别如下:

\[ \begin{aligned} \log p(\boldsymbol{x}) & =\log \int p(\boldsymbol{x}, \boldsymbol{z}) d \boldsymbol{z} \\ & =\log \int \frac{p(\boldsymbol{x}, \boldsymbol{z}) q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})} d \boldsymbol{z} \\ & =\log \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}\left[\frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\right] \\ & \geq \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} \mid \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} |\boldsymbol{x})}\right] \end{aligned} \]

\[ \begin{aligned} \log p(\boldsymbol{x}) & =\log p(\boldsymbol{x}) \int q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x}) dz ~~\left(\int q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x}) dz=1 \right)\\ & =\int q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})(\log p(\boldsymbol{x})) d z \\ & =\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}[\log p(\boldsymbol{x})] \\ & =\mathbb{E}_{q_{\phi}(\boldsymbol{z} | \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{x}, \boldsymbol{z})}{p(\boldsymbol{z} | \boldsymbol{x})}\right] \\ & =\mathbb{E}_{q_{\phi}(\boldsymbol{z} | \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{x}, \boldsymbol{z}) q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}{p(\boldsymbol{z} | \boldsymbol{x}) q_{\phi}(\boldsymbol{z} | \boldsymbol{x})}\right] \\ & =\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\right]+\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\left[\log \frac{q_{\phi}(\boldsymbol{z} | \boldsymbol{x})}{p(\boldsymbol{z} | \boldsymbol{x})}\right] \\ & =\mathbb{E}_{q_{\phi}(\boldsymbol{z} | \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\right]+D_{\mathrm{KL}}\left(q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x}) \| p(\boldsymbol{z} | \boldsymbol{x})\right) \\ & \geq \mathbb{E}_{q_{\phi}(\boldsymbol{z} | \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\right]~~\left( \text{KL Divergence always} \ge 0\right) \end{aligned} \]

从上述推导过程可以看出,证据即\(\log p(\boldsymbol{x})\)可以被分解为ELBO加上估计后验概率\(q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})\)和真实后验概率\(p(\boldsymbol{z}|\boldsymbol{x})\)之间的KL散度。由于KL散度总为非负数,加之\(\log p(\boldsymbol{x})\)关于\(\phi\)为一个常数,因此如果最大化ELBO,其实也就是等价于最小化估计后验概率和真实后验概率之间的KL散度,这样便可以间接使得估计后验概率尽可能地接近真实后验概率。

变分自编码器

变分自编码器(Variational Autoencoder,VAE)的训练目标就是直接最大化ELBO。这里的变分指的是从一系列以\(\phi\)作为参数的后验分布族中优化得到最优的后验分布\(q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})\);而自编码器则指的是采用了传统自编码器的框架,后验分布\(q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})\)对应于编码器,它将观测数据转化为隐变量的概率分布;同时,\(p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})\)对应于一个确定性的解码器,它将一个给定的隐变量解码为观测数据。

ELBO可以进一步被写为如下形式:

\[ \begin{aligned} \mathbb{E}_{q_{\phi}(\boldsymbol{z} | \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{x}, \boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\right] & =\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\left[\log \frac{p_{\boldsymbol{\theta}}(\boldsymbol{x} |\boldsymbol{z}) p(\boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\right] \\ & =\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x} |\boldsymbol{z})\right]+\mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\left[\log \frac{p(\boldsymbol{z})}{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\right] \\ & =\mathbb{E}_{q_{\phi}(\boldsymbol{z} | \boldsymbol{x})}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x} | \boldsymbol{z})\right]-D_{\mathrm{KL}}\left(q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x}) \| p(\boldsymbol{z})\right) \end{aligned} \]

其中第一项\(\mathbb{E}_{q_{\phi}(\boldsymbol{z}|\boldsymbol{x})}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x}|\boldsymbol{z})\right]\)代表重建项,用于衡量解码器从隐变量重建数据的概率似然,目的是保证编码器能够学习到有效的隐变量分布;第二项\(D_{\mathrm{KL}}\left(q_{\boldsymbol{\phi}}(\boldsymbol{z}|\boldsymbol{x})\| p(\boldsymbol{z})\right)\)代表先验匹配项,用于衡量学到的隐变量分布与先验假设之间的相似程度,最小化这一项希望模型能够真正地学到一个概率分布而不是坍缩为狄拉克分布。

通常来说,VAE的编码器试图去学习一个协方差矩阵为对角阵的多变量高斯分布,而先验分布通常被选择为一个标准的多变量高斯分布:

\[ \begin{aligned} q_{\phi}(\boldsymbol{z} | \boldsymbol{x}) & =\mathcal{N}\left(\boldsymbol{z} ; \boldsymbol{\mu}_{\phi}(\boldsymbol{x}), \boldsymbol{\sigma}_{\boldsymbol{\phi}}^{2}(\boldsymbol{x}) \mathbf{I}\right) \\ p(\boldsymbol{z}) & =\mathcal{N}\left(\boldsymbol{z} ; \mathbf{0}, \mathbf{I}\right) \end{aligned} \]

但是如果直接从隐变量分布中随机采样,则会导致整个计算过程不可微分。因此,通常使用如下的重参数化技巧:

\[ \boldsymbol{z}=\boldsymbol{\mu}_{\phi}(\boldsymbol{x})+\boldsymbol{\sigma}_{\phi}(\boldsymbol{x}) \odot \boldsymbol{\epsilon} ,~ \boldsymbol{\epsilon} \sim \mathcal{N}(\boldsymbol{\epsilon} ; \mathbf{0}, \mathbf{I}) \]

也就是说,VAE的编码器试图将数据\(\boldsymbol{x}\)编码为正态分布的均值和方差\(\boldsymbol{\mu}_{\phi}(\boldsymbol{x}),\boldsymbol{\sigma}_{\phi}(\boldsymbol{x})\),解码器则尝试从这个正态分布的采样中还原数据\(\boldsymbol{x}\)。此时,ELBO中的KL散度可以被直接计算出来,而重建损失则可以使用蒙特卡洛采样来代替计算期望,从而训练目标可以写为:

\[ \begin{aligned} &\underset{\boldsymbol{\phi}, \boldsymbol{\theta}}{\arg \max } \mathbb{E}_{q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x})}\left[\log p_{\boldsymbol{\theta}}(\boldsymbol{x} | \boldsymbol{z})\right]-D_{\mathrm{KL}}\left(q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x}) \| p(\boldsymbol{z})\right) \\ \approx & \underset{\boldsymbol{\phi}, \boldsymbol{\theta}}{\arg \max } \sum_{l=1}^{L} \log p_{\boldsymbol{\theta}}\left(\boldsymbol{x} | \boldsymbol{z}^{(l)}\right)-D_{\mathrm{KL}}\left(q_{\boldsymbol{\phi}}(\boldsymbol{z} | \boldsymbol{x}) \| p(\boldsymbol{z})\right) \end{aligned} \]

层级VAE

层级VAE(Hierarchical Variational Autoencoder,HVAE)是VAE的泛化形式,它将VAE堆叠为多层,从而构建了多层级的隐变量。也就是说,浅层的隐变量被看作是从深层更抽象的隐变量中采样而得。

在HVAE中,第\(t\)层的隐变量可以构造成\(1\sim t-1\)层所有隐变量的条件概率的形式,即\(p(\boldsymbol{z}_t|\boldsymbol{z}_1,\dots,\boldsymbol{z}_{t-1})\)。为简便起见,下面讨论马尔可夫HVAE这一简单情形,即HVAE的每一层具有马尔可夫性,它可以理解为将多个VAE简单地堆叠起来,如下图所示:

image-20221207163001095

此时,联合概率分布和隐变量的分布可以写成下面的形式:

\[ \begin{aligned} p\left(\boldsymbol{x}, \boldsymbol{z}_{1: T}\right) & =p\left(\boldsymbol{z}_{T}\right) p_{\boldsymbol{\theta}}\left(\boldsymbol{x} | \boldsymbol{z}_{1}\right) \prod_{t=2}^{T} p_{\boldsymbol{\theta}}\left(\boldsymbol{z}_{t-1} | \boldsymbol{z}_{t}\right) \\ q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1: T} | \boldsymbol{x}\right) & =q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1} | \boldsymbol{x}\right) \prod_{t=2}^{T} q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{t} | \boldsymbol{z}_{t-1}\right) \end{aligned} \]

从而ELBO也可以写为:

\[ \log p(\boldsymbol{x})\geq \mathbb{E}_{q_{\phi}\left(\boldsymbol{z}_{1: T} | \boldsymbol{x}\right)}\left[\log \frac{p\left(\boldsymbol{x}, \boldsymbol{z}_{1: T}\right)}{q_{\boldsymbol{\phi}}\left(\boldsymbol{z}_{1: T} | \boldsymbol{x}\right)}\right] \]

原理

前向过程与后向过程

在HVAE模型的基础之上,加上如下的三条限制,便可以得到扩散模型:

  1. 隐变量的维度与数据的维度完全相同
  2. 编码器是预先定义好的高斯分布,不再包含可学习参数
  3. 通过精心设计编码器中高斯分布的参数,使得最后一层变为标准的正态分布

由此,便可以给出扩散模型的数学表述:

扩散模型是具有形式\(p_{\theta}(\boldsymbol{x}_0):=\int p_\theta(\boldsymbol{x}_{0:T})d\boldsymbol{x}_{1:T}\)的隐变量模型,其中\(\boldsymbol{x}_1,\dots,\boldsymbol{x}_T\)是与数据\(\boldsymbol{x}_0\)具有相同维度的隐变量。联合概率分布\(p_\theta(\boldsymbol{x}_{0:T})\)被称为反向过程(Reverse Process),它被定义为一个具有高斯转移概率的马尔可夫链:

\[ \begin{aligned} & p(\boldsymbol{x}_T)=\mathcal{N}(\boldsymbol{x}_T;\boldsymbol{0},\boldsymbol{I}) \\ & p_\theta(\boldsymbol{x}_{0:T}):=p(\boldsymbol{x}_T)\prod_{t=1}^{T}p_\theta(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t) \\ & p_\theta(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t):=\mathcal{N}(\boldsymbol{x}_{t-1};\boldsymbol{\mu}_{\theta}(\boldsymbol{x}_t,t),\boldsymbol{\Sigma}_{\theta}(\boldsymbol{x}_t,t)) \end{aligned} \]

而相应的近似先验概率分布\(q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0)\)被称为前向过程(Forward Process)或者扩散过程,它是一个固定的马尔可夫链,按照一系列的方差系数\(\beta_1,\dots,\beta_T\)逐渐向数据中添加高斯噪声,直到最终的概率分布已经完全等同于高斯白噪声:

\[ \begin{aligned} & q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0):= \prod_{t=1}^{T} q(\boldsymbol{x}_t|\boldsymbol{x}_{t-1}) \\ & q(\boldsymbol{x}_{t}|\boldsymbol{x}_{t-1}):=\mathcal{N}(\boldsymbol{x}_{t};\sqrt{1-\beta_t}\boldsymbol{x}_{t-1},\beta_t\boldsymbol{I}) \end{aligned} \]

其中\(\beta_t\)通常被设置为常数(当然也可以设置为可学习参数),而且通常会适当选取使得隐变量的方差处于一个相近的尺度。这样的编码过程具有方差保持性(Variance-preserving)。

因此,扩散模型的训练目标就是学习反向过程中条件概率分布\(p_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_{t})\)的参数。在模型训练完成之后,就可以从高斯噪声中采样,然后不断地使用\(p_{\theta}(\boldsymbol{x}_{t-1}|\boldsymbol{x}_{t})\)进行\(T\)步的去噪,从而生成样本\(\boldsymbol{x}_{0}\)

仿照HVAE,扩散模型可以被简单地表示为下图所示的形式:

image-20221204153348386

损失函数

模型的训练过程为最小化如下的变分上界:

\[ \begin{aligned} &-\log p_{\theta}\left(\boldsymbol{x}_{0}\right) \\ =& -\log \int p_\theta(\boldsymbol{x}_{0:T})d\boldsymbol{x}_{1:T} \\ =& -\log \int \frac{p_\theta(\boldsymbol{x}_{0:T})q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0)}{q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0)}d\boldsymbol{x}_{1:T}\\ =& -\log \mathbb{E}_q \left[\frac{p_\theta(\boldsymbol{x}_{0:T})}{q(\boldsymbol{x}_{1:T}|\boldsymbol{x}_0)}\right] \\ \leq &\mathbb{E}_{q}\left[-\log \frac{p_{\theta}\left(\boldsymbol{x}_{0: T}\right)}{q\left(\boldsymbol{x}_{1: T} | \boldsymbol{x}_{0}\right)}\right] ~~(\text{Jenson's formula}) \\ =& \mathbb{E}_{q}\left[-\log \frac{p\left(\boldsymbol{x}_{T}\right)\prod_{t=1}^T p_{\theta}\left(\boldsymbol{x}_{t-1}|\boldsymbol{x}_{t}\right)}{\prod_{t=1}^T q\left(\boldsymbol{x}_{t} | \boldsymbol{x}_{t-1}\right)}\right]\\ =& \mathbb{E}_{q}\left[-\log p\left(\boldsymbol{x}_{T}\right)-\sum_{t \geq 1} \log \frac{p_{\theta}\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}\right)}{q\left(\boldsymbol{x}_{t} | \boldsymbol{x}_{t-1}\right)}\right] \end{aligned} \]

为简便起见,定义:

\[ L:=\mathbb{E}_{q}\left[-\log p\left(\boldsymbol{x}_{T}\right)-\sum_{t \geq 1} \log \frac{p_{\theta}\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}\right)}{q\left(\boldsymbol{x}_{t} | \boldsymbol{x}_{t-1}\right)}\right] \]

前向过程的一个重要性质是它允许采样任意时刻\(t\)的数据\(\boldsymbol{x}_t\)。记\(\alpha_t:=1-\beta_t\)\(\bar{\alpha}_t:=\prod_{s=1}^t \alpha_s\),有

\[ q(\boldsymbol{x}_t|\boldsymbol{x}_0)=\mathcal{N}(\boldsymbol{x}_t;\sqrt{\bar{\alpha}_0}\boldsymbol{x}_0,(1-\bar{\alpha}_t)\boldsymbol{I}) \]

根据\(q(\boldsymbol{x}_{t}|\boldsymbol{x}_{t-1}):=\mathcal{N}(\boldsymbol{x}_{t};\sqrt{1-\beta_t}\boldsymbol{x}_{t-1},\beta_t\boldsymbol{I})\)的定义,可得:

\[ \boldsymbol{x}_{t}=\sqrt{1-\beta_t}\boldsymbol{x}_{t-1}+\sqrt{\beta_t}\boldsymbol{\epsilon}, ~\boldsymbol{\epsilon}\sim \mathcal{N}(0,\boldsymbol{I}) \]

进一步将\(\boldsymbol{x}_{t-1}\)替换为\(\boldsymbol{x}_{t-2},\dots,\boldsymbol{x}_0\),有

\[ \begin{aligned} \boldsymbol{x}_{t}=&\sqrt{1-\beta_t}\boldsymbol{x}_{t-1}+\sqrt{\beta_t}\boldsymbol{\epsilon} \\ =& \sqrt{1-\beta_t}(\sqrt{1-\beta_{t-1}}\boldsymbol{x}_{t-2}+\sqrt{\beta_{t-1}}\boldsymbol{\epsilon})+\sqrt{\beta_t}\boldsymbol{\epsilon} \\ =& \sqrt{(1-\beta_t)(1-\beta_{t-1})}\boldsymbol{x}_{t-2}+\sqrt{\sqrt{(1-\beta_t)\beta_{t-1}}^2+\sqrt{\beta_t}^2}\boldsymbol{\epsilon} ~~(\text{The additivity of Gaussian dist.})\\ =& \sqrt{(1-\beta_t)(1-\beta_{t-1})}\boldsymbol{x}_{t-2}+\sqrt{(1-\beta_t)\beta_{t-1}+\beta_t}\boldsymbol{\epsilon}\\ =& \sqrt{(1-\beta_t)(1-\beta_{t-1})}\boldsymbol{x}_{t-2}+\sqrt{1-(1-\beta_t)(1-\beta_{t-1})}\boldsymbol{\epsilon} \\ =&\cdots \\ =&\sqrt{\prod_{i=1}^t (1-\beta_i)}\boldsymbol{x}_0+\sqrt{1-\prod_{i=1}^t(1-\beta_i)}\boldsymbol{\epsilon} \\ =& \sqrt{\bar{\alpha}_t}\boldsymbol{x}_0+\sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon} ~\sim \mathcal{N}(\boldsymbol{x}_{t};\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0,(1-\bar{\alpha}_t)\boldsymbol{I}) \end{aligned} \]

为了方便使用梯度下降的方法进行训练,可以使用上述结论将\(L\)进一步改写为:

$$ \[\begin{aligned} L:=&\mathbb{E}_{q}\left[-\log p\left(\boldsymbol{x}_{T}\right)-\sum_{t \geq 1} \log \frac{p_{\theta}\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}\right)}{q\left(\boldsymbol{x}_{t} | \boldsymbol{x}_{t-1}\right)}\right] \\ =&\mathbb{E}_{q}\left[-\log p\left(\boldsymbol{x}_{T}\right)-\sum_{t \geq 1} \log \frac{p_{\theta}\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}\right)}{q\left(\boldsymbol{x}_{t} | \boldsymbol{x}_{t-1},\boldsymbol{x}_{0}\right)}\right]~~ \left(q\left(\boldsymbol{x}_{t} | \boldsymbol{x}_{t-1}\right)=q\left(\boldsymbol{x}_{t} | \boldsymbol{x}_{t-1},\boldsymbol{x}_{0}\right)\right) \\ =&\mathbb{E}_{q}\left[-\log p\left(\boldsymbol{x}_{T}\right)-\sum_{t > 1} \log \frac{p_{\theta}\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}\right)}{q\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t},\boldsymbol{x}_{0}\right)}\cdot \frac{q\left(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0\right)}{q\left(\boldsymbol{x}_{t} |\boldsymbol{x}_0\right)} - \log \frac{p_{\theta}\left(\boldsymbol{x}_{0}|\boldsymbol{x}_{1}\right)}{q\left(\boldsymbol{x}_{1}|\boldsymbol{x}_{0}\right)}\right]~~\left(q\left(\boldsymbol{x}_{t}|\boldsymbol{x}_{t-1},\boldsymbol{x}_0\right)=\frac{q\left(\boldsymbol{x}_{t-1}|\boldsymbol{x}_{t},\boldsymbol{x}_0\right)q\left(\boldsymbol{x}_{t}|\boldsymbol{x}_0\right)}{q\left(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0\right)}\right) \\ =& \mathbb{E}_{q}\left[-\log p\left(\boldsymbol{x}_{T}\right)-\log \prod_{t>1} \frac{p_{\theta}\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}\right)}{q\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t},\boldsymbol{x}_{0}\right)}\cdot \frac{q\left(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0\right)}{q\left(\boldsymbol{x}_{t} |\boldsymbol{x}_0\right)} - \log {p_{\theta}\left(\boldsymbol{x}_{0}|\boldsymbol{x}_{1}\right)}\right] \\ =& \mathbb{E}_{q}\left[-\log \frac{p\left(\boldsymbol{x}_{T}\right)}{q(\boldsymbol{x}_{T}|\boldsymbol{x}_{0})}-\sum_{t > 1} \log \frac{p_{\theta}\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}\right)}{q\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t},\boldsymbol{x}_{0}\right)} - \log {p_{\theta}\left(\boldsymbol{x}_{0}|\boldsymbol{x}_{1}\right)}\right] \left(\prod_{t>1}\frac{q\left(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0\right)}{q\left(\boldsymbol{x}_{t} |\boldsymbol{x}_0\right)}=\frac{q\left(\boldsymbol{x}_{1}|\boldsymbol{x}_0\right)}{q\left(\boldsymbol{x}_{T} |\boldsymbol{x}_0\right)} \right) \\ =& \mathbb{E}_{q} \left[ D_{KL}(q(\boldsymbol{x}_T|\boldsymbol{x}_0)\|p(\boldsymbol{x}_T))+ \sum_{t>1} D_{KL}(q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)\|p_\theta(\boldsymbol{x}_{t-1} |\boldsymbol{x}_{t}))- \log {p_{\theta}\left(\boldsymbol{x}_{0}|\boldsymbol{x}_{1}\right)}\right] \end{aligned}\]

$$

其中,第一项为先验匹配项,即最终产生的加噪样本\(\boldsymbol{x}_T\)与标准正态分布有多接近;第二项代表去噪匹配项,也就是衡量两个相邻步\(t\)\(t-1\)去噪结果的匹配程度;第三项可以被解释为重建项,也就是从带噪样本\(\boldsymbol{x}_1\)重建得到去噪样本\(\boldsymbol{x}_0\)的效果。

在上面的推导过程中使用了贝叶斯公式\(q\left(\boldsymbol{x}_{t}|\boldsymbol{x}_{t-1},\boldsymbol{x}_0\right)=\frac{q\left(\boldsymbol{x}_{t-1}|\boldsymbol{x}_{t},\boldsymbol{x}_0\right)q\left(\boldsymbol{x}_{t}|\boldsymbol{x}_0\right)}{q\left(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0\right)}\),至于为什么要这样替换,首先在不使用贝叶斯公式的情况下进行推导:

\[ \begin{aligned} &\mathbb{E}_{q}\left[-\log \frac{p\left(\boldsymbol{x}_{T}\right)\prod_{t=1}^T p_{\theta}\left(\boldsymbol{x}_{t-1}|\boldsymbol{x}_{t}\right)}{\prod_{t=1}^T q\left(\boldsymbol{x}_{t} | \boldsymbol{x}_{t-1}\right)}\right] \\ =& \mathbb{E}_{q}\left[-\log \frac{p\left(\boldsymbol{x}_{T}\right)p_{\theta}\left(\boldsymbol{x}_{0}|\boldsymbol{x}_{1}\right)\prod_{t=1}^{T-1} p_{\theta}\left(\boldsymbol{x}_{t}|\boldsymbol{x}_{t+1}\right)}{q\left(\boldsymbol{x}_{T} | \boldsymbol{x}_{T-1}\right)\prod_{t=1}^{T-1} q\left(\boldsymbol{x}_{t} | \boldsymbol{x}_{t-1}\right)}\right]\\ =& \mathbb{E}_{q}\left[-\log p_{\theta}\left(\boldsymbol{x}_{0}|\boldsymbol{x}_{1}\right)- \log\frac{p_{\theta}\left(\boldsymbol{x}_{0}|\boldsymbol{x}_{1}\right)}{q\left(\boldsymbol{x}_{T} | \boldsymbol{x}_{T-1}\right)} -\log\frac{\prod_{t=1}^{T-1} p_{\theta}\left(\boldsymbol{x}_{t}|\boldsymbol{x}_{t+1}\right)}{\prod_{t=1}^{T-1} q\left(\boldsymbol{x}_{t} | \boldsymbol{x}_{t-1}\right)}\right] \\ =& -\mathbb{E}_{q}\left[\log p_{\theta}\left(\boldsymbol{x}_{0}|\boldsymbol{x}_{1}\right)\right]+\mathbb{E}_{q}\left[D_{\mathrm{KL}}\left(q\left(\boldsymbol{x}_{T} | \boldsymbol{x}_{T-1}\right) \| p\left(\boldsymbol{x}_{T}\right)\right)\right]+ \sum_{t=1}^{T-1} {\mathbb{E}_{q}\left[D_{\mathrm{KL}}\left(q\left(\boldsymbol{x}_{t} | \boldsymbol{x}_{t-1}\right) \| p_{\theta}\left(\boldsymbol{x}_{t} | \boldsymbol{x}_{t+1}\right)\right)\right]} \end{aligned} \]

其中第一项可以被看成重建项,与VAE中的处理方法类似;第二项可以被看成先验匹配项,衡量最终的加噪样本与标准正态分布有多接近;最后一项可以被看作一致性衡量项,它试图让\(\boldsymbol{x}_t\)在前向加噪和后向去噪过程中的概率分布尽可能地一致。

但是在最后一项中,同时包含有两个随机变量\(\boldsymbol{x}_{t-1}\)\(\boldsymbol{x}_{t+1}\),这样会导致用蒙特卡洛采样方法计算期望时,这一项的方差可能会偏大。尤其是在\(T\)的值较大时,方差会进一步地累加。而如果使用贝叶斯公式进行替换,则只需要使用一个随机变量,从而使得方差降低。

上式中出现了概率\(q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)\),它可以通过使用贝叶斯公式转换为如下公式计算而得:

$$ \[\begin{aligned} q\left(\boldsymbol{x}_{t-1}|\boldsymbol{x}_{t},\boldsymbol{x}_0\right)= & \frac{q\left(\boldsymbol{x}_{t}|\boldsymbol{x}_{t-1},\boldsymbol{x}_0\right)q\left(\boldsymbol{x}_{t-1}|\boldsymbol{x}_0\right)}{q\left(\boldsymbol{x}_{t}|\boldsymbol{x}_0\right)} \\ = & \frac{\mathcal{N}(\boldsymbol{x}_{t};\sqrt{\alpha_t}\boldsymbol{x}_{t-1},(1-\alpha_t)\boldsymbol{I})\cdot \mathcal{N}(\boldsymbol{x}_{t-1};\sqrt{\bar{\alpha}_{t-1}}\boldsymbol{x}_{0},(1-\bar{\alpha}_{t-1})\boldsymbol{I})}{\mathcal{N}(\boldsymbol{x}_{t};\sqrt{\bar{\alpha}_{t}}\boldsymbol{x}_{0},(1-\bar{\alpha}_{t})\boldsymbol{I})} \\ \propto & \exp\left\{-\left[ \frac{(\boldsymbol{x}_{t}-\sqrt{\alpha_t}\boldsymbol{x}_{t-1})^2}{2(1-\alpha_t)}+ \frac{(\boldsymbol{x}_{t-1}-\sqrt{\bar{\alpha}_{t-1}}\boldsymbol{x}_{0})^2}{2(1-\bar{\alpha}_{t-1})}-\frac{(\boldsymbol{x}_{t}-\sqrt{\bar{\alpha}_{t}}\boldsymbol{x}_{0})^2}{2(1-\bar{\alpha}_{t})}\right]\right\} \\ =& \cdots \\ =& \exp \left\{-\frac{1}{2}\left(\frac{1}{\frac{\left(1-\alpha_{t}\right)\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_{t}}}\right)\left[\boldsymbol{x}_{t-1}^{2}-2 \frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right) \boldsymbol{x}_{t}+\sqrt{\bar{\alpha}_{t-1}}\left(1-\alpha_{t}\right) \boldsymbol{x}_{0}}{1-\bar{\alpha}_{t}} \boldsymbol{x}_{t-1}\right]\right\} \\ \propto & \mathcal{N}\left(\boldsymbol{x}_{t-1} ; \frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right) \boldsymbol{x}_{t}+\sqrt{\bar{\alpha}_{t-1}}\left(1-\alpha_{t}\right) \boldsymbol{x}_{0}}{1-\bar{\alpha}_{t}}, \frac{\left(1-\alpha_{t}\right)\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_{t}} \boldsymbol{I}\right) \end{aligned}\]

$$

为方便起见,记\(\tilde{\mu}_t(\boldsymbol{x}_t,\boldsymbol{x}_0):=\frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right) \boldsymbol{x}_{t}+\sqrt{\bar{\alpha}_{t-1}}\left(1-\alpha_{t}\right) \boldsymbol{x}_{0}}{1-\bar{\alpha}_{t}}\)\(\tilde{\beta}_t:=\frac{\left(1-\alpha_{t}\right)\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_{t}}\)

\(L\)中的三项处理方式分布如下:

  1. 先验匹配项\(D_{KL}(q(\boldsymbol{x}_T|\boldsymbol{x}_0)\|p(\boldsymbol{x}_T))\)

    当假设\(\beta_t\)都为固定值时,前向过程没有可训练参数,因此这部分为常数,可以将其忽略掉。

  2. 去噪匹配项\(\sum_{t>1} D_{KL}(q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)\|p_\theta(\boldsymbol{x}_{t-1} |\boldsymbol{x}_{t}))\)

    由于

    \[ q\left(\boldsymbol{x}_{t-1}|\boldsymbol{x}_{t},\boldsymbol{x}_0\right)\propto \mathcal{N}\left(\boldsymbol{x}_{t-1} ; \frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right) \boldsymbol{x}_{t}+\sqrt{\bar{\alpha}_{t-1}}\left(1-\alpha_{t}\right) \boldsymbol{x}_{0}}{1-\bar{\alpha}_{t}}, \frac{\left(1-\alpha_{t}\right)\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_{t}} \boldsymbol{I}\right) \\ p_\theta(\boldsymbol{x}_{t-1}|\boldsymbol{x}_{t})=\mathcal{N}\left(\boldsymbol{x}_{t-1};\mathbf{\mu}_0(\boldsymbol{x}_{t},t),\sigma^2_t\mathbf{I}\right) \]

    因此

    \[ D_{KL}(q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)\|p_\theta(\boldsymbol{x}_{t-1} |\boldsymbol{x}_{t}))=\mathbb{E}_q\left[\frac{1}{2\sigma_t^2}\| \tilde{\boldsymbol{\mu}}_t(\boldsymbol{x}_t,\boldsymbol{x}_0)-\boldsymbol{\mu}_\theta(\boldsymbol{x}_t,t) \|^2 \right]+C \]

    进一步地,我们可以使用\(\boldsymbol{x}_t(\boldsymbol{x}_0,\boldsymbol{\epsilon})=\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0+\sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon},~\boldsymbol{\epsilon}\sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I})\),以及\(\tilde{\mu}_t(\boldsymbol{x}_t,\boldsymbol{x}_0)\)\(\tilde{\beta}_t\)的表达式,得到下面的结果:

    $$

    \[\begin{aligned} & \mathbb{E}_q\left[\frac{1}{2\sigma_t^2}\| \tilde{\boldsymbol{\mu}}_t(\boldsymbol{x}_t,\boldsymbol{x}_0)-\boldsymbol{\mu}_\theta(\boldsymbol{x}_t,t) \|^2 \right] \\ = & \mathbb{E}_{\boldsymbol{x}_0,\boldsymbol{\epsilon}}\left[ \frac{1}{2\sigma_t^2}\| \tilde{\boldsymbol{\mu}}_t(\boldsymbol{x}_t,\boldsymbol{x}_0)-\boldsymbol{\mu}_\theta(\boldsymbol{x}_t,t) \|^2 \right] \\ =& \mathbb{E}_{\boldsymbol{x}_{0}, \boldsymbol{\epsilon}}\left[\frac{1}{2 \sigma_{t}^{2}}\left\|\tilde{\boldsymbol{\mu}}_{t}\left(\boldsymbol{x}_{t}\left(\boldsymbol{x}_{0}, \boldsymbol{\epsilon}\right), \frac{1}{\sqrt{\bar{\alpha}_{t}}}\left(\boldsymbol{x}_{t}\left(\boldsymbol{x}_{0}, \boldsymbol{\epsilon}\right)-\sqrt{1-\bar{\alpha}_{t}} \boldsymbol{\epsilon}\right)\right)-\boldsymbol{\mu}_{\theta}\left(\boldsymbol{x}_{t}\left(\boldsymbol{x}_{0}, \boldsymbol{\epsilon}\right), t\right)\right\|^{2}\right] \\ =& \mathbb{E}_{\boldsymbol{x}_{0}, \boldsymbol{\epsilon}}\left[\frac{1}{2 \sigma_{t}^{2}}\left\|\frac{1}{\sqrt{\alpha_{t}}}\left(\boldsymbol{x}_{t}\left(\boldsymbol{x}_{0}, \boldsymbol{\epsilon}\right)-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \boldsymbol{\epsilon}\right)-\boldsymbol{\mu}_{\theta}\left(\boldsymbol{x}_{t}\left(\boldsymbol{x}_{0}, \boldsymbol{\epsilon}\right), t\right)\right\|^{2}\right] \end{aligned}\]

    $$

    也就是说,给定\(\boldsymbol{x}_t\)之后,\(\boldsymbol{\mu}_\theta\)需要能够预测\(\frac{1}{\sqrt{\alpha_{t}}}\left(\boldsymbol{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \boldsymbol{\epsilon}\right)\)的值。因此可以考虑将参数化的模型构造为如下形式:

    \[ \boldsymbol{\mu}_\theta=\frac{1}{\sqrt{\alpha_{t}}}\left(\boldsymbol{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \boldsymbol{\epsilon}_\theta(\boldsymbol{x}_{t},t)\right) \]

    也就是说,让\(\boldsymbol{\epsilon}_\theta\)\(\boldsymbol{x}_t\)中去预测加入的噪声\(\boldsymbol{\epsilon}\)。按照这一定义,损失函数可以进一步优化为:

    \[ \mathbb{E}_{\boldsymbol{x}_{0}, \boldsymbol{\epsilon}}\left[\frac{\beta_t^2}{2\sigma_t^2\alpha_t(1-\bar{\alpha}_t)}\left\|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_\theta\left(\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0+\sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon},t\right)\right\|^2\right] \]

    或者是使用一个更加简单的形式:

    \[ \mathbb{E}_{\boldsymbol{x}_{0}, \boldsymbol{\epsilon},t}\left[\left\|\boldsymbol{\epsilon}-\boldsymbol{\epsilon}_\theta\left(\sqrt{\bar{\alpha}_t}\boldsymbol{x}_0+\sqrt{1-\bar{\alpha}_t}\boldsymbol{\epsilon},t\right)\right\|^2\right] \]

    这一表达式的形式十分接近于去噪得分匹配模型。

  3. 重建项\(\log {p_{\theta}\left(\boldsymbol{x}_{0}|\boldsymbol{x}_{1}\right)}\)

    这一项与VAE中的重建误差类似,可以用相同的方法计算。简便起见有时也可以将其并入到去噪匹配项中(即\(t=1\)

等价形式

关于扩散模型的损失函数,也可以等价地构造为从任意噪声水平恢复原始图像,这是因为去噪匹配项可以等价地构造成如下的形式:

$$ \[\begin{aligned} &D_{KL}(q(\boldsymbol{x}_{t-1}|\boldsymbol{x}_t,\boldsymbol{x}_0)\|p_\theta(\boldsymbol{x}_{t-1} |\boldsymbol{x}_{t})) \\ =&\mathbb{E}_q\left[\frac{1}{2\sigma_t^2}\| \tilde{\boldsymbol{\mu}}_t(\boldsymbol{x}_t,\boldsymbol{x}_0)-\boldsymbol{\mu}_\theta(\boldsymbol{x}_t,t) \|^2 \right]+C \\ =&\mathbb{E}_q\left[\frac{1}{2\sigma_t^2}\left\| \frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right) \boldsymbol{x}_{t}+\sqrt{\bar{\alpha}_{t-1}}\left(1-\alpha_{t}\right) \boldsymbol{x}_{0}}{1-\bar{\alpha}_{t}}-\frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right) \boldsymbol{x}_{t}+\sqrt{\bar{\alpha}_{t-1}}\left(1-\alpha_{t}\right) \hat{\boldsymbol{x}}_{\theta}(\boldsymbol{x}_t,t)}{1-\bar{\alpha}_{t}} \right\|^2 \right]+C \\ =&\mathbb{E}_q\left[ \frac{1}{2\sigma_t^2}\cdot \frac{\bar{\alpha}_{t-1}\left(1-\alpha_{t}\right)^2}{(1-\bar{\alpha}_{t})^2}\left\| \boldsymbol{x}_{0}-\hat{\boldsymbol{x}}_{\theta}(\boldsymbol{x}_t,t) \right\|^2 \right]+C \\ \end{aligned}\]

$$

其中,\(\hat{\boldsymbol{x}}_{\theta}(\boldsymbol{x}_t,t)\)是一个参数化的神经网络,用于从噪声图像中预测原始图像。

扩散模型与得分匹配

得分匹配结合朗之万动力学(Score matching with Langevin dynamics,SMLD)的生成方法指的是用模型估计每个噪声尺度下的得分(对数概率密度在数据点处的梯度值),然后基于这一模型使用朗之万动力学不断减小噪声尺度来进行采样。去噪扩散概率模型(Denoising diffusion probabilistic modeling,DDPM)指的是训练一个概率模型序列,使之能够对不断给数据添加高斯噪声对数据进行破坏的正向过程做去噪的反向建模。这二者是从不同角度出发得到的生成模型,但是可以被放在一个统一的框架下去看待。

等价损失函数

Tweedie公式说明,给定一个指数族分布以及一组来自该分布的样本,概率分布的均值可以由样本均值加上一个修正项来估计。例如给定一个高斯分布的样本$ ( ; {z}, {z}) $,Tweedie公式可以表示为如下形式:

\[ \mathbb{E}\left[\boldsymbol{\mu}_{z} |\boldsymbol{z}\right]=\boldsymbol{z}+\boldsymbol{\Sigma}_{z} \nabla_{\boldsymbol{z}} \log p(\boldsymbol{z}) \]

其中,\(\nabla_{\boldsymbol{z}} \log p(\boldsymbol{z})\)通常被称为得分函数。

在扩散模型的推导过程中,有:

\[ q(\boldsymbol{x}_t|\boldsymbol{x}_0)=\mathcal{N}(\boldsymbol{x}_t;\sqrt{\bar{\alpha}_0}\boldsymbol{x}_0,(1-\bar{\alpha}_t)\boldsymbol{I}) \]

而对于\(\boldsymbol{x}_t\),使用Tweedie公式可以得到均值\(\boldsymbol{\mu}_{\boldsymbol{x}_t}=\sqrt{\bar{\alpha}_0}\boldsymbol{x}_0\)的期望:

\[ \mathbb{E}\left[\boldsymbol{\mu}_{\boldsymbol{x}_t} |\boldsymbol{x}_t\right]=\boldsymbol{x}_t+(1-\bar{\alpha}_t)\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t) \]

因此有:

\[ \boldsymbol{x}_t+(1-\bar{\alpha}_t)\nabla_{\boldsymbol{x}_t} \log p(\boldsymbol{x}_t)=\sqrt{\bar{\alpha}_0}\boldsymbol{x}_0 \]

将上式代入到扩散模型中\(\tilde{\mu}_t(\boldsymbol{x}_t,\boldsymbol{x}_0)\)的表达式中,有:

$$ \[\begin{aligned} \tilde{\mu}_t(\boldsymbol{x}_t,\boldsymbol{x}_0)=&\frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right) \boldsymbol{x}_{t}+\sqrt{\bar{\alpha}_{t-1}}\left(1-\alpha_{t}\right) \boldsymbol{x}_{0}}{1-\bar{\alpha}_{t}} \\ =&\frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right) \boldsymbol{x}_{t}+\sqrt{\bar{\alpha}_{t-1}}\left(1-\alpha_{t}\right) \frac{\boldsymbol{x}_{t}+\left(1-\bar{\alpha}_{t}\right) \nabla \log p\left(\boldsymbol{x}_{t}\right)}{\sqrt{\bar{\alpha}_{t}}}}{1-\bar{\alpha}_{t}} \\ =&\frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right) \boldsymbol{x}_{t}+\left(1-\alpha_{t}\right) \frac{\boldsymbol{x}_{t}+\left(1-\bar{\alpha}_{t}\right) \nabla \log p\left(\boldsymbol{x}_{t}\right)}{\sqrt{\alpha_{t}}}}{1-\bar{\alpha}_{t}} \\ =& \cdots \\ =& \frac{1}{\sqrt{\alpha_{t}}} \boldsymbol{x}_{t}+\frac{1-\alpha_{t}}{\sqrt{\alpha_{t}}} \nabla \log p\left(\boldsymbol{x}_{t}\right) \end{aligned}\]

$$

因此,同样可以定义\(\boldsymbol{\mu}_\theta(\boldsymbol{x}_t,t)\)为如下形式:

\[ \boldsymbol{\mu}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t}, t\right)=\frac{1}{\sqrt{\alpha_{t}}} \boldsymbol{x}_{t}+\frac{1-\alpha_{t}}{\sqrt{\alpha_{t}}} \boldsymbol{s}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t}, t\right) \]

从而扩散模型中的去噪匹配项可以写成:

\[ \begin{aligned} &\underset{\boldsymbol{\theta}}{\arg \min } D_{\mathrm{KL}}\left(q\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}, \boldsymbol{x}_{0}\right) \| p_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}\right)\right) \\ =& \underset{\boldsymbol{\theta}}{\arg \min } D_{\mathrm{KL}}\left(\mathcal{N}\left(\boldsymbol{x}_{t-1} ; \boldsymbol{\mu}_{q}, \boldsymbol{\Sigma}_{q}(t)\right) \| \mathcal{N}\left(\boldsymbol{x}_{t-1} ; \boldsymbol{\mu}_{\boldsymbol{\theta}}, \boldsymbol{\Sigma}_{q}(t)\right)\right) \\ =& \underset{\boldsymbol{\theta}}{\arg \min } \frac{1}{2 \sigma_{q}^{2}(t)}\left[\left\|\frac{1}{\sqrt{\alpha_{t}}} \boldsymbol{x}_{t}+\frac{1-\alpha_{t}}{\sqrt{\alpha_{t}}} \boldsymbol{s}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t}, t\right)-\frac{1}{\sqrt{\alpha_{t}}} \boldsymbol{x}_{t}-\frac{1-\alpha_{t}}{\sqrt{\alpha_{t}}} \nabla \log p\left(\boldsymbol{x}_{t}\right)\right\|_{2}^{2}\right] \\ = &\underset{\boldsymbol{\theta}}{\arg \min } \frac{1}{2 \sigma_{q}^{2}(t)} \frac{\left(1-\alpha_{t}\right)^{2}}{\alpha_{t}}\left[\left\|\boldsymbol{s}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t}, t\right)-\nabla \log p\left(\boldsymbol{x}_{t}\right)\right\|_{2}^{2}\right]\\ \end{aligned} \]

其中\(\boldsymbol{s}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t}, t\right)\)为一个参数化的网络,它被用于估计得分函数的值。

由于

\[ \boldsymbol{x}_{0}=\frac{\boldsymbol{x}_{t}+\left(1-\bar{\alpha}_{t}\right) \nabla \log p\left(\boldsymbol{x}_{t}\right)}{\sqrt{\bar{\alpha}_{t}}}=\frac{\boldsymbol{x}_{t}-\sqrt{1-\bar{\alpha}_{t}} \boldsymbol{\epsilon}_{0}}{\sqrt{\bar{\alpha}_{t}}} \]

因此有

\[ \nabla \log p\left(\boldsymbol{x}_{t}\right)=-\frac{1}{\sqrt{1-\bar{\alpha}_{t}}} \boldsymbol{\epsilon}_{0} \]

得分函数的含义是如何在数据空间中移动以取得最大概率,因此去噪的过程也可以理解为将数据往概率密度较高的方向移动。

在去噪得分匹配中,损失函数的表达式为:

\[ \underset{\boldsymbol{\theta}}{\arg \min } \sum_{t=1}^{T} \lambda(t) \mathbb{E}_{p_{\sigma_{t}}\left(\boldsymbol{x}_{t}\right)}\left[\left\|\boldsymbol{s}_{\boldsymbol{\theta}}(\boldsymbol{x}, t)-\nabla \log p_{\sigma_{t}}\left(\boldsymbol{x}_{t}\right)\right\|_{2}^{2}\right] \]

它与上面推导得到的扩散模型中的去噪匹配项具有相近的格式。此外,在SMLD的框架下,数据的生成使用的是带有退火的朗之万动力学,即采样的过程中逐步降低噪声等级\(t\),从而生成的样本会逐步接近真实样本,这也与扩散模型的采样过程有类似之处。

SDE下的统一框架

伊藤方程

在连续的状态空间中,DDPM的训练目标中暗含了在不同的噪声等级上面计算得分,因此它与SMLD可以被统称为基于得分的生成模型(Score-based generative models)。它们可以用伊藤随机微分方程(SDE)构造为一个统一的框架:

image-20221128101113630

给定一个扩散过程\(\{\boldsymbol{x}_{t}\}_{t=0}^T\)​,其中\(t\in [0,T]\)​为一个连续的时间变量,\(\boldsymbol{x}_0\sim p_0\)​代表初始概率分布,\(\boldsymbol{x}_T\sim p_T\)​代表一个先验分布。通常来说,\(p_T\)是一个非结构化的先验分布,并且不包含关于\(p_0\)的任何信息,例如固定均值和方差的高斯分布。扩散过程可以用伊藤方程来描述,即如下的随机微分方程:

\[ d\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{x},t)dt+g(t)d\boldsymbol{w} \]

其中\(\boldsymbol{w}\)是一个标准的维纳过程,即布朗运动;\(\boldsymbol{f}(\cdot,t):\mathbb{R}^d\rightarrow \mathbb{R}^d\)是一个向量函数,被称为\(\boldsymbol{x}_{t}\)的漂移系数;\(g(\cdot):\mathbb{R}\rightarrow \mathbb{R}\)是一个标量函数,被称为\(\boldsymbol{x}_{t}\)的扩散系数,此处为了简便起见,假设扩散系数是个标量,而且与\(\boldsymbol x\)无关。如果SDE的系数在状态\(\boldsymbol{x}\)和时间\(t\)上都满足全局Lipschitz性,则它有唯一的强解。

如果从\(\boldsymbol{x}_T\sim p_T\)开始,将上述的扩散过程反过来进行,可以得到样本\(\boldsymbol{x}_0\sim p_0\)。这样的反过程同样是个扩散过程,可以用反向时间的SDE来描述:

\[ d\boldsymbol{x}=[\boldsymbol{f}(\boldsymbol{x},t)-g(t)^2\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})]dt+g(t)d\bar{\boldsymbol{w}} \]

其中\(\bar{\boldsymbol{w}}\)是一个标准的反向维纳过程,对应于时间反向从\(T\)到0的过程;\(dt\)是一个无限小的负时间步长。此时,如果知道所有时间\(t\)的得分\(\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})\),就可以从反向SDE中对\(p_0\)进行采样。

SDE与DDPM

对于DDPM,它的条件概率分布\(\{p_{\alpha_i}(\boldsymbol{x}|\boldsymbol{x}_0)\}_{i=1}^N\)对应于如下的离散马尔可夫链:

\[ \boldsymbol{x}_i=\sqrt{1-\beta_i}\boldsymbol{x}_{i-1}+\sqrt{\beta_i}\boldsymbol{z}_{i-1},i=1,\dots,N \]

定义\(\bar{\beta}_i=N\beta_i\),上式可以改写为:

\[ \boldsymbol{x}_i=\sqrt{1-\frac{\bar{\beta}}{N}}\boldsymbol{x}_{i-1}+\sqrt{\frac{\bar{\beta}}{N}}\boldsymbol{z}_{i-1},i=1,\dots,N \]

\(N\rightarrow \infty\)时,\(\bar{\beta}_i\)就变为了一个自变量为\(t\in [0,1]\)的函数\(\beta(t)\),即\(\beta(\frac{i}{N})=\bar{\beta}_i\)。此外令\(\boldsymbol{x}(\frac{i}{N})=\boldsymbol{x}_i\)\(\boldsymbol{z}(\frac{i}{N})=\boldsymbol{z}_i\),上式可以进一步变为:

\[ \begin{aligned} \boldsymbol{x}(t+\Delta t) =& \sqrt{1-\beta(t+\Delta t)\Delta t}\boldsymbol{x}(t)+\sqrt{\beta(t+\Delta t)\Delta t}\boldsymbol{z}(t) \\ \approx & \boldsymbol{x}(t)-\frac{1}{2}\beta(t+\Delta t)\Delta t\boldsymbol{x}(t)+\sqrt{\beta(t+\Delta t)\Delta t}\boldsymbol{z}(t) \\ \approx & \boldsymbol{x}(t)-\frac{1}{2}\beta(t)\Delta t\boldsymbol{x}(t)+\sqrt{\beta(t)\Delta t}\boldsymbol{z}(t) \end{aligned} \]

\(N\rightarrow \infty\)\(\Delta t\rightarrow 0\)时,表达式变为下面的SDE:

\[ d\boldsymbol{x}=-\frac{1}{2}\beta(t)\boldsymbol{x}dt+\sqrt{\beta(t)}d\boldsymbol{w} \]

如果初始状态的方差为1,上述的SDE会得到一个方差不变的过程,因此也被称为保持方差(Variance Preserving)的SDE。

SDE与SMLD

使用去噪得分匹配的方法训练一个时间相关的得分模型\(\boldsymbol{s_\theta}(\boldsymbol{x},t)\)就可以对\(\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})\)进行估计,模型的训练目标为:

\[ \boldsymbol{\theta}^*=\arg\min_{\boldsymbol{\theta}} \mathbb{E}_t\left\{ \lambda(t)\mathbb{E}_{\boldsymbol{x}_0}\mathbb{E}_{\boldsymbol{x}_t|\boldsymbol{x}_0}\left[ \left\| \boldsymbol{s_\theta}(\boldsymbol{x}_t,t)-\nabla_{\boldsymbol{x}_t}\log_{p_{0t}}(\boldsymbol{x}_t|\boldsymbol{x}_0) \right\|_2^2 \right] \right\} \]

其中\(\lambda:[0,T]\rightarrow \mathbb{R}^+\)是一个函数值为正的权重函数,\(t\)\([0,T]\)上均匀采样,\(\boldsymbol{x}_{0}\sim p_0(\boldsymbol{x})\)\(\boldsymbol{x}_{t}\sim p_{0t}(\boldsymbol{x}_{t}|\boldsymbol{x}_{0})\)。通常来说,\(p_{0t}(\boldsymbol{x}_{t}|\boldsymbol{x}_{0})\)要求是一个已知的概率分布。如果\(\boldsymbol{f}(\cdot,t)\)为一个仿射函数,那么这个概率分布总是正态分布,它的均值和方差具有闭式解。

基于训练好的得分模型,可以用朗之万动力学的方法做采样来生成数据。使用\(N\)个不同的噪声等级\(\sigma_i,i=1,\dots,N\),每个不同的噪声等级所对应的干扰样本\(\boldsymbol{x}_i\)及其条件概率\(p_{\sigma_i}(\boldsymbol{x}|\boldsymbol{x}_0)\)满足如下的马尔可夫链:

\[ \boldsymbol{x}_i=\boldsymbol{x}_{i-1}+\sqrt{\sigma_i^2-\sigma_{i-1}^2}\boldsymbol{z}_{i-1},i=1,\dots,N \]

其中,\(\boldsymbol{z}_{i-1}\sim \mathcal{N}(0,I)\)。此外为了简便起见,假设\(\sigma_0=0\)。当\(N\rightarrow \infty\)时,使用连续变量\(t\in[0,1]\)代替下标\(i\),此时\(\{\sigma_i\}_{i=1}^N\)就变为函数\(\sigma(t)\)\(\boldsymbol{z}_i\)变为\(\boldsymbol{z}(t)\),上述的马尔可夫链也就变成了连续的随机过程\(\{\boldsymbol{x}_{t}\}_{t=0}^1\)

\(\boldsymbol{x}(\frac{i}{N})=\boldsymbol{x}_i\)\(\sigma(\frac{i}{N})=\sigma(t)\)\(\boldsymbol{z}(\frac{i}{N})=\boldsymbol{z}_i\),令\(\Delta t=\frac{1}{N}\),上述马尔可夫链可以改写为:

\[ \boldsymbol{x}(t+\Delta t)=\boldsymbol{x}(t)+\sqrt{\sigma^2(t+\Delta t)-\sigma^2(t)}\boldsymbol{z}(t) \approx \boldsymbol{x}(t)+\sqrt{\frac{d[\sigma^2(t)]}{dt}\Delta t}\boldsymbol{z}(t) \]

\(N\rightarrow \infty\)时,也就是\(\Delta t\rightarrow 0\)时,就变为下面的SDE:

\[ d\boldsymbol{x}=\sqrt{\frac{d[\sigma^2(t)]}{dt}}d\boldsymbol{w} \]

\(t\rightarrow \infty\)时,这一SDE对应于一个方差爆炸的过程,因此被称为方差爆炸(Variance Exploding)的SDE。

概率流ODE

考虑一个一般形式的SDE:

\[ d\boldsymbol{x}=\boldsymbol{f}(\boldsymbol{x},t)dt+\boldsymbol{G}(\boldsymbol{x},t)d\boldsymbol{w} \]

其中\(\boldsymbol{f}(\cdot,t):\mathbb{R}^d\rightarrow \mathbb{R}^d\)\(\boldsymbol{G}(\cdot,t):\mathbb{R}^d\rightarrow \mathbb{R}^{d\times d}\)。边际概率密度函数\(p_t(\boldsymbol{x}_{t})\)的演化服从Fokker-Planck方程(推导过程可以参考Fokker-Planck方程 - 知乎 (zhihu.com)):

\[ \begin{aligned} \frac{\partial p_t(\boldsymbol{x})}{\partial t}=&-\sum_{i=1}^d \frac{\partial}{\partial x_i}[f_i(\boldsymbol{x},t)p_t(\boldsymbol{x})]+\frac{1}{2}\sum_{i=1}^d\sum_{j=1}^d\frac{\partial^2}{\partial x_i \partial x_j}\left[\sum_{k=1}^d G_{ik}(\boldsymbol{x},t)G_{jk}(\boldsymbol{x},t)p_t(\boldsymbol{x})\right] \\ =& -\sum_{i=1}^d \frac{\partial}{\partial x_i}[f_i(\boldsymbol{x},t)p_t(\boldsymbol{x})]+\frac{1}{2}\sum_{i=1}^d\frac{\partial}{\partial x_i }\left[ \sum_{j=1}^d\frac{\partial}{\partial x_j}\left[\sum_{k=1}^d G_{ik}(\boldsymbol{x},t)G_{jk}(\boldsymbol{x},t)p_t(\boldsymbol{x})\right]\right] \end{aligned} \]

其中,

\[ \begin{aligned} & \sum_{j=1}^d\frac{\partial}{\partial x_j}\left[\sum_{k=1}^d G_{ik}(\boldsymbol{x},t)G_{jk}(\boldsymbol{x},t)p_t(\boldsymbol{x})\right] \\ =& \sum_{j=1}^d\frac{\partial}{\partial x_j}\left[\sum_{k=1}^d G_{ik}(\boldsymbol{x},t)G_{jk}(\boldsymbol{x},t)\right]p_t(\boldsymbol{x}) + \sum_{j=1}^d\sum_{k=1}^d G_{ik}(\boldsymbol{x},t)G_{jk}(\boldsymbol{x},t)p_t(\boldsymbol{x})\frac{\partial}{\partial x_j} \log p_t(\boldsymbol{x}) \\ =& p_t(\boldsymbol{x}) \nabla\cdot [\boldsymbol{G}(\boldsymbol{x},t)\boldsymbol{G}(\boldsymbol{x},t)^T]+p_t(\boldsymbol{x})\boldsymbol{G}(\boldsymbol{x},t)\boldsymbol{G}(\boldsymbol{x},t)^T\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x}) \end{aligned} \]

从而有

\[ \frac{\partial p_t(\boldsymbol{x})}{\partial t}=-\sum_{i=1}^{d}\frac{\partial}{\partial x_i}\left\{ f_i(\boldsymbol{x},t)p_t(\boldsymbol{x})-\frac{1}{2}p_t(\boldsymbol{x})\left[ \nabla\cdot [\boldsymbol{G}(\boldsymbol{x},t)\boldsymbol{G}(\boldsymbol{x},t)^T]+\boldsymbol{G}(\boldsymbol{x},t)\boldsymbol{G}(\boldsymbol{x},t)^T\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})\right]\right\} \]

定义

\[ \tilde{\boldsymbol{f}}(\boldsymbol{x},t):= \boldsymbol{f}(\boldsymbol{x},t)-\frac{1}{2}\nabla\cdot [\boldsymbol{G}(\boldsymbol{x},t)\boldsymbol{G}(\boldsymbol{x},t)^T]+\boldsymbol{G}(\boldsymbol{x},t)\boldsymbol{G}(\boldsymbol{x},t)^T\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x}) \]

从而可以得到一个形式更简单的ODE:

\[ d\boldsymbol{x}=\tilde{\boldsymbol{f}}(\boldsymbol{x},t)dt \]

这也就是说,对于任意的扩散过程,都对应于一个确定性的ODE,它的轨迹与SDE有着同样的概率密度\(\{p_t(\boldsymbol{x})\}_{t=0}^T\)。上述方程被称为概率流ODE(Probability flow ODE),通过它可以实现很多功能:

  • 可以计算出任意输入数据的概率似然。即:

    \[ \log p_0(\boldsymbol{x}(0))=\log p_T(\boldsymbol{x}(T))+\int_{0}^{T} \nabla\cdot \tilde{\boldsymbol{f}}(\boldsymbol{x}(t),t)dt \]

  • 通过积分可以实现隐变量到数据的解码,因此可以对隐变量进行修改从而实现图像插值等图像编辑操作。

  • 由于前向SDE中并没有可训练参数,因此如果得分估计完全准确,概率流ODE就可以产生完全相同的轨迹。这也就是说,当训练数据、模型表示能力、优化精度都充足的情况下,对于数据的编码也是唯一的。

  • 使用ODE求解器,可以实现高效采样。

可控生成

定义

可控生成指的是从一个条件概率分布\(p_0(\boldsymbol{x}_{0}|\boldsymbol{y})\)中进行采样。例如从黑白图像生成彩色图像,这里的条件就是黑白图像;生成特定类别的图像,这里的条件就是类别;从被损坏的图像中复原出完整图像,这里的条件就是被损坏的图像;使用文本生成图片任务中,这里的条件就是文本的特征向量;等等。

在扩散模型中,可控生成可以通过简单地在每个转移步中加上条件\(\boldsymbol{y}\)来实现:

\[ p\left(\boldsymbol{x}_{0: T} | y\right)=p\left(\boldsymbol{x}_{T}\right) \prod_{t=1}^{T} p_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t-1} | \boldsymbol{x}_{t}, y\right) \]

根据上面的推导,此时同样有如下三种等价的模型实现方式:

\[ \begin{aligned} &\hat{\boldsymbol{x}}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t}, t, y\right) \approx \boldsymbol{x}_{0}\\ &\hat{\boldsymbol{\epsilon}}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t}, t, y\right) \approx \boldsymbol{\epsilon}_{0}\\ &\boldsymbol{s}_{\boldsymbol{\theta}}\left(\boldsymbol{x}_{t}, t, y\right) \approx \nabla \log p\left(\boldsymbol{x}_{t} | y\right) \end{aligned} \]

以及对应的SDE:

\[ d\boldsymbol{x}=\left\{\boldsymbol{f}(\boldsymbol{x},t)-g(t)^2\nabla_\boldsymbol{x}\log p_t(\boldsymbol{x}|\boldsymbol{y})\right\}dt+g(t)d\boldsymbol{w} \]

但是这种直接实现的方法泛化性较差,因此通常会使用指导(Guidance)的方式来显式地控制模型中的条件。

加入指导

根据贝叶斯公式

\[ p_t(\boldsymbol{x}_t|\boldsymbol{y})=\frac{ p_t(\boldsymbol{x}_t)p(\boldsymbol{y}|\boldsymbol{x}_t)}{p(\boldsymbol{y})} \]

可以得到得分函数的一个等价形式:

\[ \nabla_{\boldsymbol{x}_t}\log p_t(\boldsymbol{x}_t|\boldsymbol{y})=\nabla_{\boldsymbol{x}_t}\log p_t(\boldsymbol{x}_t)+\nabla_{\boldsymbol{x}_t}\log p_t(\boldsymbol{y}|\boldsymbol{x}_t) \]

由此可得对应的SDE:

\[ d\boldsymbol{x}=\left\{\boldsymbol{f}(\boldsymbol{x},t)-g(t)^2[\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x})+\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{y}|\boldsymbol{x})]\right\}dt+g(t)d\bar{\boldsymbol{w}} \]

上式中的\(\nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{y}|\boldsymbol{x})\)可以使用另外一个独立的模型去估计,或者是借助一个得分函数来近似计算,即\(\nabla_\boldsymbol{x}\log p_t(\boldsymbol{x}_{t}|\boldsymbol{y})\approx \boldsymbol{s_{\theta^*}}(\boldsymbol{x}_{t},t)+\nabla_{\boldsymbol{x}}\log p_t(\hat{\boldsymbol{y}}(t)|\boldsymbol{x}_{t})\),其中\(\hat{\boldsymbol{y}}(t)\)是从\(p(\boldsymbol{y}(t)|\boldsymbol{y})\)中采集到的一个样本。

使用得分函数近似计算的推导过程:

假设有一个通过SDE对\(\boldsymbol{x}\)不断做干扰所得到的一个扩散过程\(\{\boldsymbol{x}_{t}\}_{t=0}^T\),以及一个训练好的得分模型\(\boldsymbol{s_{\boldsymbol{\theta}^*}}(\boldsymbol{x}_{t},t)\)

根据贝叶斯定理,有

\[ \nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x}(t)|\boldsymbol{y})=\nabla_{\boldsymbol{x}}\log \int p_t(\boldsymbol{x}(t)|\boldsymbol{y}(t),\boldsymbol{y})p(\boldsymbol{y}(t)|\boldsymbol{y})d\boldsymbol{y}(t) \]

其中\(\boldsymbol{y}(t)\)\(\boldsymbol{x}_{t}\)和前向过程\(p(\boldsymbol{y}(t)|\boldsymbol{x}_{t})\)共同定义。

假设如下两个条件成立:

  1. \(p(\boldsymbol{y}(t)|\boldsymbol{y})\)是可追踪的。这个概率分布通常可以从前向过程和SDE之间的交互来推导出来。
  2. \(p_t(\boldsymbol{x}_{t}|\boldsymbol{y}(t),\boldsymbol{y})\approx p_t(\boldsymbol{x}_{t}|\boldsymbol{y}(t))\)。对于较小的\(t\)\(\boldsymbol{y}(t)\)几乎与\(\boldsymbol{y}\)相同,因此近似成立;而对于较大的\(t\)\(\boldsymbol{y}\)将会与\(\boldsymbol{x}_{t}\)远离,从而对于\(\boldsymbol{x}_{t}\)的影响很小。此外,对于较大的\(t\),近似误差在最后的采样结果中也影响较小,因为它在采样过程的早期就已经被使用。

在这两个假设下,有

\[ \begin{aligned} \nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x}(t)|\boldsymbol{y})& \approx \nabla_{\boldsymbol{x}}\log \int p_t(\boldsymbol{x}(t)|\boldsymbol{y}(t))p(\boldsymbol{y}(t)|\boldsymbol{y})d\boldsymbol{y} \\ & \approx \nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x}(t)|\hat{\boldsymbol{y}}(t)) \\ & = \nabla_{\boldsymbol{x}}\log p_t(\boldsymbol{x}(t))+\nabla_{\boldsymbol{x}}\log p_t(\hat{\boldsymbol{y}}(t)|\boldsymbol{x}(t)) \\ & \approx \boldsymbol{s_{\boldsymbol{\theta}^*}}(\boldsymbol{x}(t),t)+\nabla_{\boldsymbol{x}}\log p_t(\hat{\boldsymbol{y}}(t)|\boldsymbol{x}(t)) \end{aligned} \]

参考文献

  1. Denoising Diffusion Probabilistic Models (arxiv.org)
  2. Understanding Diffusion Models: A Unified Perspective (arxiv.org)
  3. Score-Based Generative Modeling through Stochastic Differential Equations (arxiv.org)