Diffusion Models

2022-08-28

本文介绍 score-based generative models,又称 diffusion models。 先简单介绍生成模型 (generative models) 的概念和几种典型的生成模型, 然后重点介绍 diffusion models。 本文的主要内容参考了文献[1], 以及几篇相关博客文章[2] [3] [4] [x]。 因为 generative models 和 diffusion models 涉及到一些统计学概念以及必要的公式推导, 作为非理论研究人员和工程师,理解这些统计学概念的物理意义可能比熟悉公式推导更重要。 因此,在文献[1]的基础上,本文主要有以下两点改变:

  1. 对部分数学符号和统计学概念增加了更详细具体的解释;
  2. 对于部分数学推导,省略了推导过程,只给出最终结论。

因为接下来的内容会从统计学上描述生成模型,为了方便阅读和理解,这里首先从图像生成的角度解释一些概念:

概念 符号 解释
观测,observation $x$ 从真实数据中观测到的样本,也即真实的数据,比如一张自然图像
采样 从学习到的概率分布中采样出新的样本,或者更通俗的讲“生成新的样本“的过程
似然 likelihood $p_{\theta}(x)$ 给定样本之后,似然指的是该样本的可能性,似然是概率模型参数 $\theta$ 的函数;
隐变量 latent variable $z$ 描述样本属性的一个变量,比如图像 $x$ 拍摄时的光照/天气/时间,或者图像中描述的内容。隐变量根据任务不同而不同。
编码器 encoder $q_{\phi}(z|x)$ 从样本 $x$ 到隐变量 $z$ 的过程称为编码,对应的模型称为编码器,模型参数为 $\phi$
解码器 decoder $p_{\theta}(x|z)$ 从隐变量 $z$ 到样本 $x$ 的过程称为解码,对应的模型称为解码器,模型参数为 $\theta$

🎯概率和似然
这里需要解释一下 似然 概率 的区别,虽然他们都用符号 $p_{\theta}(x)$ 表示。 概率 (probability) 是指”已知模型后,采样出某个样本 $x$ 的概率“,变量是 $x$; 而似然 (likelihood) 是指”已知样本 $x$ 后,推测其是由哪个模型采样出的“。

🎯概率和统计
这里也区分一下概率 (probability) 和 统计 (statistics) 的区别。 概率是指已知模型推测结果,例如已知骰子各面朝上的可能性均等,求抛两次和大于 10 的概率。 而统计是指已经知晓实验结果反过来推测模型。例如已知抛 100 次筛子的实验结果,推测筛子在一次实验中各面朝上的概率。

概率 (probability) 和统计 (statistic) 的区别。前者已知模型推测结果,后者以已知结果推断模型(从而预测未来的结果)。

变分贝叶斯 (Variational Bayes)

后面讲的很多内容涉及到变分贝叶斯的一些概念,为了理清这些概念方便大家理解,这里先回顾一下 variational Bayesian methods 的基本概念, 这部分的主要参考资料为:

什么是变分贝叶斯方法

很多地方中都会解释变分贝叶斯是什么,有什么用处, 由于这些解释各自的侧重点不同,因此各不相同,如同盲人摸象,无法提纲挈领。 总的来说就一句话,变分贝叶斯是用来逼近一些无法计算的 (intractable) 概率分布的,通常这些分布都涉及到对一些变量的积分。 或者更简单地,变分贝叶斯是用来逼近概率分布的。 在变分贝叶斯方法所涉及到的概率模型中,一般涉及到三种不同的变量:

  1. 数据 data $x$;
  2. 参数 parameter $\theta$;
  3. 隐变量 latent variable $z$。
其中数据 $x$ 称为 "observed variable",参数 $\theta$ 和 隐变量$z$ 称为 unobserved variables。

变分贝叶斯是用来干什么的

变分贝叶斯有哪些应用呢?总的来说有两个:

  1. 逼近一些 unobserved variables 的后验概率 (posterior) $p(z|x)$。(To provide an analytical approximation to the posterior probability of the unobserved variables, in order to do statistical inference over these variables.)
  2. 获得观测数据的 marginal likelihood (又称为 evidence) $p(x) = \int_z p(x|z)\cdot p(z)dz$;(To derive a lower bound for the marginal likelihood (sometimes called the evidence) of the observed data (i.e. the marginal probability of the data given the model, with marginalization performed over unobserved variables).
上面第二点不容易理解,看不懂没关系,后面会结合实例解释。

这里先解释一个概念叫 intractable。我们说一个分布是 intractable,指的是该分布不可以解析地 (analytically) 计算,只能通过逼近或者采样的方式间接估计。 如果知道某个随机变量服从高斯分布,只是参数未知,那么这个分布仍然可以解析地估计出来,不属于 intractable。 一般数据的分布是 intractable 的,而一些人为假设服从某些特定分布的隐变量 (latent variable) 是 tractable 的。

高斯混合模型:简单的隐变量模型

我们用高斯混合模型来解释 observed data $x$ 和 unobserved latent variable $z$。

一个简单的高斯混合模型。数据 $x$ 的分布较为复杂且 intractable。我们可以使用一个隐变量 $z$,使得 $p(x|z), p(z)$ 都是 tractable 的。 在这个例子 $x$ 聚集成 3 个 Gaussian distributions,隐变量 $z\in \{0, 1, 2\}$ 是离散的。

假设我们观测到下图中的数据,我们如何计算它们的分布 $p(x)$?

显然 $p(x)$ 并不是一个常规的分布,是一个 intractable distribution。 但是我们观察到数据聚成了 3 个 clusters,每个 cluster 的数据似乎是从高斯分布中采样的。 因此我们定义隐变量 $z\in\{0, 1, 2\}$ 表示 $x$ 来自哪一个高斯, 注意 $z$ 是 离散的随机变量。 然后定义以下几个分布:

当然,实际条件下的隐变量 $z$ 远比例子中复杂,有的甚至是连续的。 但是高斯混合模型的例子已经可以将问题说明白:

这里我们用生成模型为例介绍变分贝叶斯中的三种随机变量。 如下图所示,一般而言生成模型(不论是 GAN 还是 VAE)的生成(或者采样)过程都是输入一个隐变量 $z$ 到模型中,模型参数为 $\theta$,然后得到生成的样本 $x$。

从变分贝叶斯角度解释一个简单的生成模型。输入 latent variable $z$,经过参数为 $\theta$ 的模型,得到生成的样本 $x$。

ELBO: Evidence Lower BOund

Variational Bayesian 通过 $q_{\phi}(z|x)$ 来逼近 $p(z|x)$。 假设有分布 $q_{\phi}(z|x)\in\mathcal{Q}$ 属于某个分布 (a family of distributions)。 以高斯分布为例,$\mathcal{Q}$ 指的是全体高斯分布的集合 (均值方差不同), $q^i(z)\in\mathcal{Q}$ 指的是高斯分布集体中的一个个体。 可以通过最小化 $D_{KL}(q_{\phi}(z|x) \ || \ p(z|x))$ 来找到最佳的 approximation $q^*(z)$。
逼近后验概率 + 获得 evidence

考虑隐变量 $z$ 关于数据 $x$ 的后验概率 (posterior),记为 $p(z|x)$。

\begin{equation} p(z|x) = \frac{p(x|z) \cdot p(z)}{p(x)} = \frac{p(x|z) \cdot p(z)}{\int_z p(x, z)} \label{eq:bayes-1} \end{equation}
公式$\ref{eq:bayes-1}$ 等号右边分母 $p(x) = \int_z p(x, z)$ 就是 "intractable" 的,因为数据 $x$ 的分布不可知,也无法对所有的隐变量 $z$ 求积分。

下面我们推导如何用变分贝叶斯推断来逼近 (approximate) 隐变量 (latent) 的后验概率:$p(z|x)$。 我们假设有另外一个分布 $q_{\phi}(z|x)$,然后用 $q_{\phi}(z|x)$ 去逼近 $p(z|x)$: $q_x(z) \approx p(z|x)$。 通常会假设 $q_{\phi}(z|x)$ 属于某个常见分布 $\mathcal{Q}$(比如高斯分布),然后最小化 $q_{\phi}(z|x)$ 与 $p(z|x)$ 之间的 KL-divergence 来得到最佳的近似 $q^*$: \begin{equation} q^* = \underset{q\in\mathcal{Q}}{\operatorname{arg max}} \ \ D_{KL}(p(z|x) \ || \ q_{\phi}(z|x)) \label{eq:opt-qz} \end{equation}

很多其它的资料中会用符号 $q(z)$ 表示我们假设的分布,这里为了和后面的 encoder-decoder 保持符号的一致,我们用 $q_{\phi}(z|x)$ 来表示。 实际上这个分布是 depend on 数据 $x$ 的,同时也会被 $\phi$ 参数化,因此用 $q_{\phi}(z|x)$ 更准确。 因此,更准确地说,公式$\ref{eq:opt-qz}$ 实际上是寻找最优的参数 $\phi$: \begin{equation} \phi^* = \underset{\phi}{\operatorname{arg max}} \ \ D_{KL}(p(z|x) \ || \ q_{\phi}(z|x)) \end{equation}

那么我们如何最小化 $D_{KL}(q_{\phi}(z|x) \ || \ p(z|x))$ 呢?$p(z|x)$ 本身是 intractable 的,因此无法直接计算 KL divergence。 通过接下来的推导,我们发现可以通过另外一种方式去最小化 KL divergence。

\begin{equation} \begin{aligned} D_{KL}\Big(q_{\phi}(z|x) \ || \ p(z|x)\Big) &= \int_z q_{\phi}(z|x) \Big[\log \frac{q_{\phi}(z|x)}{p(z|x)}\Big] dz \\ &= \int_z q_{\phi}(z|x) \Big[\log \frac{q_{\phi}(z|x)p(x)}{p(z,x)}\Big] dz \\ &= {\color{orange}{\int_z q_{\phi}(z|x) \Big[ \log p(x)\Big]dz}} + \int_z q_{\phi}(z|x) \Big[ \log \frac{q_{\phi}(z|x)}{p(z,x)}\Big]dz \\ &= \log p(x) + \mathbb{E}_{z\sim q_{\phi}(z|x)}\Big[ \log \frac{q_{\phi}(z|x)}{p(z,x)}\Big] \\ &= \log p(x) - \mathbb{E}_{z\sim q_{\phi}(z|x)}\Big[ \log \frac{p(z,x)}{q_{\phi}(z|x)}\Big] \\ &= \log p(x) - \color{red}{\text{ELBO}} \end{aligned}\label{eq:elbo-in-var-bayesian-0} \end{equation}
公式$\ref{eq:elbo-in-var-bayesian-0}$ 等式右边第一项 $\log p(x)$ 与 $z$ 无关,可以当作常数。 因此最小化 $D_{KL}\Big(q_{\phi}(z|x) \ || \ p(z|x)\Big)$ 等价于最大化 $\mathbb{E}_{z\sim q_{\phi}(z|x)}\Big[ \log \frac{p(z,x)}{q_{\phi}(z|x)}\Big]$。 公式$\ref{eq:elbo-in-var-bayesian-0}$ 通过简单变换可以得到:
\begin{equation} \begin{aligned} \log p(x) &= \mathbb{E}_{z\sim q_{\phi}(z|x)}\Big[ \log \frac{p(z,x)}{q_{\phi}(z|x)}\Big] + D_{KL}\Big(q_{\phi}(z|x) \ || \ p(z|x)\Big) \\ \log p(x) &\ge \mathbb{E}_{z\sim q_{\phi}(z|x)}\Big[ \log \frac{p(z,x)}{q_{\phi}(z|x)}\Big] = \color{red}{\text{ELBO}} \end{aligned}\label{eq:elbo-inequal} \end{equation}
公式$\ref{eq:elbo-inequal}$ 中的 ELBO 称为 Evidence Lower BOund, 很多其他资料中用 $\mathbf{L}(q)$ 来表示,因为它实际上是 $q$ 的函数(因为数据 $x$ 已经固定)。 在变分贝叶斯中,$x$ 作为唯一能被直接观测到的变量,被称为 "evidence",公式$\ref{eq:elbo-inequal}$ 告诉我们 ELBO 是 $\log p(x)$ 下界,因为被称为 evidence lower bound。 请记住 公式$\ref{eq:elbo-in-var-bayesian-0}$,$\ref{eq:elbo-inequal}$ 告诉我们的两个重要的结论:
  1. ELBO 是 $\log p(x)$ 的下界;
  2. 可以通过 maximize ELBO 来 maximize $p(x)$

下面我们将介绍生成模型,并以 variational aotu-encoder (VAE) 为例,介绍如何通过优化 ELBO 来最大化 $p(x)$,进而训练模型。

生成模型

给定一些从某个概率分布中采样出的样本 $x$ (比如图片),生成模型 (generative models) 的目的是建模这个概率分布 $p(x)$, 从而使我们能从这个概率分布中采样出新的样本 (也就是所谓的“生成”), 同时也能判断任意一个样本的似然(likelihood, 也就是判断任意一个样本“真假”的程度)。

这里所谓的“从某个概率分布中采样出的样本”,是假定所有的样本,比如自然图像,都服从某个很复杂概率分布, 所有的观测到的样本(observation)都是从这个概率分布中采样出来的。 似然 (likelihood) 指的是某个样本是从某个特定的概率分布中采样出来的概率。 假设 $p(x)$ 是所有自然图像的概率分布,那么样本 $x_i$ 的似然则表明该样本是从“自然图像”中采样出来的概率。

现有的几种生成模型

这里简单介绍一些现有的生成模型。因为我们的重点在 diffusion models,因此对于其他的生成模型范式我们仅从 high-level 介绍其基本概念和区别。 现有的生成模型可以粗略地划分为以下几种:

Generative Adversarial Network (GAN) 通过“对抗训练”的方式学习一个复杂分布的采样过程。 其中生成器(generator)负责采样(或者叫生成),判别器(discriminator)负责区分生成的样本和真实的样本。 Likelihood-based models 直接建模样本的 likelihood,并对真实观测的样本赋予很高的 likelihood 以学到真实样本的分布。 Energy-based models 用任意形式的能量函数来学习分布函数,并通过 normalizing 将其转换为一个概率分布函数。

这里解释一下”任意形式的能量函数“和”通过 normalizing 转化为概率分布函数“。 概率分布函数 $p(x)$ 需要满足: $p(x) \le 1$ 以及 $\int p(x) = 1$, 而能量函数没有任何约束。一个无约束的函数可以通过 normalizing flow 转化为概率分布函数,具体可以参考 https://en.wikipedia.org/wiki/Flow-based_generative_model

VAE, and Hierarchical VAE

Variational Auto-Encoder

这节我们介绍 Variational Auto-encoder。不像很多其他文章首先介绍 VAE 的原理,再介绍 VAE 的结构。 我们先从 VAE 的具体实现和模型结构开始讲起,因为 VAE 的结构比较简单,而背后的原理相对比较复杂。

VAE[5] 采用 encoder-decoder 结构。如 图1 所示,整个 VAE 模型呈现一个瓶颈结构,包含一个 encoder 和一个 decoder。

Fig.1: VAE 模型结构,包含一个 encoder,decoder。encoder 输入样本 $x$,输出隐变量 $z$ 的分布; decoder 输入 $z$,生成样本 $x'$。
Encoder 输入样本 $x$ 得到 $\mu, \sigma$ 从而得到 $q(z|x)$; decoder 输入 latent $z$ 输出生成的样本 $\hat{x}$。 训练时的损失函数为
\begin{equation} \mathcal{L}_{VAE} = \underbrace{\rvert\rvert x - \hat{x} \rvert\rvert_2}_{\text{reconstruction term}} + \underbrace{D_{KL}\big(q_{\phi}(z|x) \ \rvert\rvert \ \mathcal{N}(0, \mathbf{I})\big)}_{\text{prior term}} \label{eq:loss-vae} \end{equation}
公式$\ref{eq:loss-vae}$ 中 $D_{KL}(\cdot \rvert\rvert \cdot)$ 为概率分布间的 KL-divergence,$\mathcal{N}(0, \mathbf{I})$ 为标准高斯分布,由于 $q_{\phi}(z|x)$ 是均值为 $\mu$ 方差为 $\sigma$ 的高斯分布, 因此可以解析的表达出来。

从优化和 auto-encoder 的角度来说,VAE 就是一个带有正则约束的自编码器。 损失函数第一项为自编码器中的重建损失; 第二项为正则项,约束后验概率 $q_{\phi}(z|x)$ 为标准高斯分布。 接下来我们将用变分贝叶斯理论推导出 VAE 的目标函数实际上是在最大化 ELBO。

从 ELBO 角度理解 VAE 的损失函数

从 variational Bayes 的角度,VAE 是在给定数据集 $$\mathcal{D}=\big\{x_i \big\}_{i=1, ..., N}$$ 的情况下最大化 ELBO: \begin{equation} \underset{\phi, \theta}{\operatorname{arg max}} \ \ \mathbb{E}_{z\sim q_{\phi}(z|x)}\log \frac{p(x, z)}{q_{\phi}(z|x)} \label{eq:loss-vae-elbo-1} \end{equation}

那么 公式$\ref{eq:loss-vae-elbo-1}$ 和 公式$\ref{eq:loss-vae}$ 有什么联系呢? 我们还是从 ELBO 的定义出发,推导这两者之间的联系。

\begin{equation}\begin{aligned} \text{ELBO} &= \mathbb{E}_{z\sim q_{\phi}(z|x)}\Big[ \log \frac{p(z,x)}{q_{\phi}(z|x)}\Big] \\ &= \mathbb{E}_{z\sim q_{\phi}(z|x)}\Big[ \log \frac{p_{\theta}(x|z)\cdot p(z)}{q_{\phi}(z|x)}\Big] \\ &= \mathbb{E}_{z\sim q_{\phi}(z|x)}\log p_{\theta}(x|z) - \mathbb{E}_{z\sim q_{\phi}(z|x)}\Big[ \log \frac{q_{\phi}(z|x)}{p(z)}\Big] \\ &= \underbrace{\mathbb{E}_{z\sim q_{\phi}(z|x)}\log p_{\theta}(x|z)}_{\text{reconstruction term}} - \underbrace{D_{KL}\big(q_{\phi}(z|x) \ \rvert\rvert \ p(z)\big)}_{\text{prior term}} \end{aligned}\label{eq:loss-vae-elbo} \end{equation}
公式 $\ref{eq:loss-vae-elbo}$ 将 ELBO 拆分成了两项:

🎯先验项 (prior term)
先验项好理解, 我们只需要假设隐变量 $z$ 的先验分布 $p(z)$ 为标准高斯分布 $\mathcal{N}(0, \mathbf{I})$,这里就和 公式$\ref{eq:loss-vae}$ 的第二项完全相等。

🎯重建项 (reconstruction)
重建项的物理意义是说,给定一个隐变量 $z$ 和模型 $p_{\theta}$,我们可以计算任意样本的似然 $p_{\theta}(x|z)$, 但是我们想要真实样本 $x_i$ 的似然最大。 公式$\ref{eq:loss-vae-elbo}$ 等号右边第一项是样本 $x$ 的似然,因此这项最大化真实样本的似然。 利用最大似然估计 (maximum likelihood estimation, MLE),我们可以推导出 公式$\ref{eq:loss-vae}$ 可以由 公式$\ref{eq:loss-vae}$ 在假设 $p_{\theta}(x|z)$ 导出。

假设 $p_{\theta}(x|z)$ 服从高斯分布且对任意的 $z$ 方差都是 $\sigma$, Decoder 的输出 $\hat{x}$ 是高斯分布的均值 。 那么真实样本 $x_i$ 的似然为: $$ p(X=x_i|z) = \sqrt{2\pi\sigma}\cdot e^{\frac{(x_i-\hat{x})^2}{\sigma^2}}. $$ 带入 公式$\ref{eq:loss-vae-elbo}$ 得到: $$ \log p_{\theta}(x|z) = \log\big(\sqrt{2\pi\sigma}\cdot e^{\frac{(x-\hat{x})^2}{\sigma^2}}\big) \propto (x-\hat{x})^2, $$ 也就是 公式$\ref{eq:loss-vae}$ 中的形式。

如果 $p_{\theta}(x|z)$ 服从拉普拉斯分布 (Laplacian Distribution): $$p(x) = \frac{1}{2b}\exp(-\frac{|x-\mu|}{b}).$$ $x$ 的似然变为: $$ p(X=x_i|z) = \frac{1}{2b}\exp(-\frac{|x-\hat{x}|}{b}) $$ 因此 公式$\ref{eq:loss-vae-elbo}$ 中的重建项变成了: $$ \log\big(\frac{1}{2b}\exp(-\frac{|x-\hat{x}|}{b})\big) \propto \rvert x-\hat{x}\rvert, $$ 也就是将 公式$\ref{eq:loss-vae}$ 中的 $L_2$ 变成了 $L_1$。

这里也可以看到,不同的损失函数 (L1 还是 L2) 实际上对应了 最大似然估计 (MLE) 中对变量 $x$ 概率分布的假设 (高斯或者拉普拉斯),尽管都只预测均值。 这里也可以看出,从优化的角度 (公式 $\ref{eq:loss-vae}$) 理解 VAE 要远比从概率统计的角度理解 (公式 $\ref{eq:loss-vae-elbo}$) 简单。 我个人觉得部分原因可能是我们对初等代数/微积分的熟悉程度要比概率统计的深。

🎯求期望 $\mathbb{E}_{z\sim q_{\phi}(z|x)}$
这里还需要注意一个细节:公式$\ref{eq:loss-vae-elbo-1}$ 中重建项有一个期望符号。 这个期望是通过对同一个输入 $x$ 采样多个隐变量 $z$ 实现的。 因为输入样本 $x$ 后我们实际上是得到了 $z$ 的 posterior $p_{\phi}(z|x)$,我们可以从这个 posterior 中采样若干个 $z$ 来计算 公式$\ref{eq:loss-vae-elbo-1}$ 中 的期望。

重参数技巧

上面围绕 公式$\ref{eq:loss-vae-elbo}$ 解释了 VAE 的损失函数和 ELBO 的关系,但是还有一个问题没有解决, 那就是 encoder 只输出了隐变量的分布 $q_{\phi}(z|x)$,然后在分布中”采样“得到了一个 $z$, 进而作为 decoder 的输出并得到 $\hat{x}$。 问题是,这里的”采样“操作并不是可导 (differentiable) 的, 不可导的操作是无法在神经网络中优化的。 VAE 使用 重参数技巧 (re-parameterize trick) 来解决这个问题。 具体地,encoder 网络不直接输出隐变量 $z$,而是输出隐变量分布的参数。 例如,我们假设隐变量 $z$ 服从多元高斯分布,那么 encoder 的输出就是该分布的参数 (均值和协方差)。

假设有样本 $x$ 和 从标准高斯分布 中随机采样出的向量 $\epsilon$, 先将 $x$ 输入到 encoder 中 得到 $\mu(x), \sigma(x)$。 然后我们通过将 $\epsilon$ 平移和缩放 来得到隐变量: \begin{equation} z = \epsilon \odot \sigma(x) + \mu(x) \label{eq:epsilon-z} \end{equation} 上式中 $\odot$ 表示element-wise multiply,这时候 $z$ 就是从均值为 $\mu$ 协方差矩阵为 $\sigma$ 的高斯分布中采样出来的。同时隐变量 $z$ 对输入 $x$ 的导数 $\frac{\partial z}{\partial x}$ 是可求的: $$ \frac{\partial z}{\partial x} = \epsilon\frac{\partial \sigma(x)}{\partial x} + \frac{\partial \mu(x)}{\partial x}. $$ 这里的重参数技巧 (re-parameterization trick) 相当于把”采样“这个不可导的操作前置了,提前采样出 $\epsilon$,然后通过平移和缩放得到 $z$。 这样整个网络从损失 $\mathcal{L}$ 到输入 $x$ 中间的所有操作都是可导的。

整个 VAE 的训练过程可以分为一下几步:

  1. 输入样本 $x$ 到 encoder,得到$\mu(x), \sigma(x)$;
  2. 从标准高斯分布中采样若干个 $\epsilon$: $\{\epsilon^1, ..., \epsilon^L \}$;
  3. 根据 公式$\ref{eq:epsilon-z}$ 得到隐变量 $z^1, ..., z^L$;
  4. 输入 $z$ 到 Decoder,得到输出 $\hat{x}^1, ..., \hat{x}^L$;
  5. 根据 公式 $\ref{eq:loss-vae}$ 计算损失;
  6. 反向传导,更新 $\phi, \theta$。
在实际实现中,当 batch 训练时 batch 足够大时,可以每次只采样一个 $\epsilon$ (也即 $L=1$)。

Diffusion Models

Markov Chain 马尔可夫链

假设 $X_t$ 包含了 $t$ 时刻某系统的所有信息,那么可以认为 $X_t$ 已经包含了足够的信息来预测未来的状态: \begin{equation} p(X_{t+1}|X_t) = p(X_{t+1}|X_t, X_{t-1}), \label{eq:markov} \end{equation} 这里的“包含了足够的信息”指的是不依赖于过去的信息 (比如 $X_{t-1}$),上面的公式形式化地表达了同样的意思。 也可以说 $t+1$ 时刻的状态 $X_{t+1}$ 只由 $t$ 时刻的状态 $X_t$ 决定,与任何 $X_{\tau}, \tau\lt t$ 都无关。 上面的假设称为“马尔可夫假设 (Markov assumption)”,也称满足马尔可夫假设的随机变量“具有马尔科夫性”。 $X_1, X_2, ..., X_T$ 又被称为马尔可夫链,马尔可夫模型,在随机过程中也被称作马尔可夫过程。

条件概率 $p(X_t|X_{t-1})$ 称为“transition function” 或者 “transition kernel” 或者 “Markov kernel”, 是马尔可夫模型的参数。 在离散条件下,$p(X_t|X_{t-1})$ 可以用一个矩阵来表示,这个矩阵被称为“马尔可夫转移矩阵”。 有了 Markov kernel,给定任意时刻的状态 $X_t$ 都可以推断出将来某一时刻的状态 $p(X_{t+n}), n\gt 0$。 如果 $p(X_t|X_{t-1})$ 与时间 $t$ 无关,则这个马尔可夫模型是“time-invariant 时间无关” 的。 时间无关的假设极大地减少了马尔可夫模型的参数数目,因为不同时刻共享了同一个 transition kernel。

Diffusion models

Diffusion models 是一种特别的 latent variable model。与其它 latent variable models (比如 VAE) 不同的是,diffusion models 的 latent variable 是一系列和图像 $X_0$ 同样维度的变量。 假设图像 $X_0\in\mathbb{R}^{H\times W\times 3}$,那么隐变量 $X_{t=1,...,T}\in\mathbb{R}^{H\times W\times 3}$ 和 $X_0$ 维度相同。 在 diffusion models 中,可以将这些隐变量的 joint distribution $p(X_1, ..., X_T) = p(X_{1:T})$ 看作是前文中的 $p(Z)$

从 $X_0$ 到 $X_T$从 $X_T$ 到 $X_0$ 的变换过程都可以通过不断在输入上累加高斯噪声完成,这个”不断在输入上累加高斯噪声”的过程可以描述为一个马尔科夫链。 $X_0\rightarrow X_T$ 称为 diffusion process 或者是 forward process; $X_T\rightarrow X_0$ 称为 reverse process。

原图来自[xx]
上面的话用 variation Bayessian 的语言描述就是: diffusion models 的联合概率分布 $p(X, Z)=p(X_{0:T})$ (从 $X_T$ 到 $X_0$) 和后验概率 $p(Z | X_0)=p(X_{1:T} | X_0)$ (从 $X_0$ 到 $X_T$) 都是马尔科夫链。 下面我们分别形式化地介绍 forward process 和 reverse process 的马尔科夫链表示。

Forward Process

→ Forward / diffusion process (从样本到噪声) →

从 $X_0$ 开始,逐渐在输入上累加上一个很小的高斯分布,最后可以得到 $X_T\in\mathcal{N}(\mathbf{0}, \mathbf{I})$: \begin{equation} q(X_{1:T} | X_0) = \prod_{t=1}^T q(X_{t-1} | X_t) \label{eq:diffusion-forward-markov} \end{equation} 公式$\ref{eq:diffusion-reverse-markov}$ 的 Markov kernel 为 $$ q(X_t | X_{t-1}) := \mathcal{N}\big(X_t; \sqrt{1-\beta_t} \ X_{t-1}, \beta_t\mathbf{I} \big) $$

需要解释一下和 公式$\ref{eq:diffusion-forward-markov}$ 中的连乘为什么成立,下面 reverse process 中的 公式$\ref{eq:diffusion-reverse-markov}$ 也类似。 如果 $X_{T-2}, X_{T-1}, X_{T}$ 满足“马尔科夫性”,也即 公式$\ref{eq:markov}$ 成立, 那么 \begin{equation} \begin{aligned} & \ p(X_{T-2}|X_{T-1}) \cdot p(X_{T-1}|X_{T}) \\ =& p(X_{T-2}|X_{T-1}, X_{T}) \cdot p(X_{T-1}|X_{T}) \\ =& \ \frac{p(X_{T-2}, X_{T-1}, X_{T})}{p(X_{T}, X_{T-1})} \cdot \frac{p(X_{T}, X_{T-1})}{p(X_{T})}\\ =& \ \frac{p(X_{T-2}, X_{T-1}, X_T)}{p(X_{T})} \\ =& \ p(X_{T-1}, X_{T-2}|X_{T}) \\ \end{aligned} \end{equation} 把上式推广到 $X_0\sim X_T$ 就得到 公式$\ref{eq:diffusion-forward-markov}$。

Reverse Process

← Reverse process (从噪声到样本) ←

从高斯噪声 $X_T\in\mathcal{N}(\mathbf{0}, \mathbf{I})$ 开始,逐渐在输入上累加上一个很小的高斯分布,最后可以生成样本 $X_0$: \begin{equation} p_{\theta}(X_{0:T}) = p(X_T)\prod_{t=1}^T p(X_{t-1} | X_t), \label{eq:diffusion-reverse-markov} \end{equation} 其中 $X_T\in\mathcal{N}(\mathbf{0}, \mathbf{I})$ reverse process 马尔科夫链的 Markov kernel 是:

$$ p(X_{t-1} | X_t) := \mathcal{N}\big(X_{t-1}; \mu_{\theta}(X_t, t), \Sigma_{\theta}(X_t, t)\big) $$
其中 $\theta$ 是 diffusion model 的参数,得到了 $\theta$,就可以从 $X_T$ 开始利用 公式$\ref{eq:diffusion-reverse-markov}$ 生成样本。 reverse process 从一个噪声开始最终生成 样本 $X_0$。

reverse process 中每一步的均值 $\mu_{\theta}(X_t, t)$ 和方差 $\Sigma_{\theta}(X_t, t)$ 是用神经网络来逼近的。

其次,Forward process 有一个很有用的特性,就是多步的结果 $p(X_t|X_0)$ 可以由 $X_0$ 一步直接导出: \begin{equation} q(X_t|X_0) = \mathcal{N}\big(X_t; \sqrt{\overline{a}_t}X_0, (1-\overline{a}_t)\mathbf{I} \big) \label{eq:qxt-x0} \end{equation} 其中 \begin{equation} a_t = 1 - \beta_t, \ \ \overline{a}_t = \prod_{s=1}^t a_{s}. \end{equation} 关于这部分的具体推导详见 DDPM 论文的附录部分。

Processes Markov kernels Outcome distributions
Forward Process $q(X_t | X_{t-1}) := \mathcal{N}\big(X_t; \sqrt{1-\beta_t} \ X_{t-1}, \beta_t\mathbf{I} \big)$ $q(X_{1:T} | X_0) = \prod_{t=1}^T q(X_{t-1} | X_t)$
Reverse Process $p(X_{t-1} | X_t) := \mathcal{N}\big(X_{t-1}; \mu_{\theta}(X_t, t), \Sigma_{\theta}(X_t, t)\big)$ $p_{\theta}(X_{0:T}) = p(X_T)\prod_{t=1}^T p(X_{t-1} | X_t)$
Table: Diffusion models 的 forward 和 reverse 过程。

优化目标

diffusion models 的 latent variable 是 $X_1, ..., X_T$, 因此 $p(X_1, ..., X_T) = p(X_{1:T})$ 相当于 公式$\ref{eq:elbo-in-var-bayesian-0}$ 中的 $p(Z)$。 这里将 diffusion model 中的各种分布和一般 latent variable model 的分布作一个对照:

Latent models diffusion models
$p(z)$ $p(X_{1:T}) = p(X_1, ..., X_T)$
$p(z|x)$ $p(X_{1:T}|X_0) = p(X_1, ..., X_T | X_0)$
$p(x, z)$ $p(X_{0:T}) = p(X_0, ..., X_T)$
Table: 一般 latent models 和 diffusion models 的变量对照。

与 VAE 类似,diffusion models 训练时的优化目标也是最大化 ELBO。 根据上面的变量对照表,difusion models 中的 ELBO 应当为:

\begin{equation} \mathbb{E}_{z\sim q_{\phi}(z|x)}\Big[ \log \frac{p(z,x)}{q_{\phi}(z|x)}\Big] = \mathbb{E}_{z\sim q_{\phi}(z|x)}\Big[ \log \frac{p(X_{0:T})}{q_{\phi}(X_{1:T}|X_0)}\Big] \end{equation}
因此 diffusion models 在的损失函数为:
\begin{equation}\begin{aligned} \mathcal{L} &:= \mathbb{E}_q \Big[-\log \frac{p_{\theta(X_{0:T})}}{q(X_{1:T}|X_0)}\Big] \\ &= \mathbb{E}_q \Big[-\log \frac{p(X_T)\prod_{t=1}^{T} p_{\theta}(X_{t-1}|X_t)}{\prod_{t=1}^{T} q(X_t|X_{t-1})}\Big] \\ &= \mathbb{E}_q \Big[-\log p(X_T) - \log\frac{\prod_{t=1}^{T} p_{\theta}(X_{t-1}|X_t)}{\prod_{t=1}^{T} q(X_t|X_{t-1})}\Big] \\ &= \mathbb{E}_q \Big[-\log p(X_T) - \sum_{t=1}^T\log\frac{p_{\theta}(X_{t-1}|X_t)}{q(X_t|X_{t-1})}\Big] \\ \end{aligned}\end{equation}
参考 公式$\ref{eq:elbo-inequal}$ 中的不等式:
\begin{equation} \log p(x) \le \mathbb{E}_{z\sim q_{\phi}(z|x)}\Big[ \log \frac{p(z,x)}{q_{\phi}(z|x)}\Big] = \mathbb{E}_{z\sim q_{\phi}(z|x)}\Big[\log\frac{p(X_{0:T})}{q_{\phi}(X_{1:T}|X_0)}\Big] \end{equation}

\begin{equation} p_{\theta}(X_0, ..., X_T) = p_{\theta}(X_{0:T}) = p(X_T)\prod_{t=1}^T p_{\theta}(X_{t-1}|X_t) \end{equation}
\begin{equation} p_{\theta}(X_{t-1}X_t) := \mathcal{N}(X_{t-1}; \mu_{\theta}(X_t, t), \Sigma_{\theta}(X_t, t)) \end{equation}
\begin{equation} \begin{aligned} \mathbb{E}\Big[-\log p_{\theta}(X_0)\Big] \le \mathbb{E}_q \Big[-\log \frac{p_{\theta(X_{0:T})}}{q(X_{1:T}|X_0)}\Big] = \mathbb{E}_q \Big[-\log p(X_T) - \sum_{t>1} \log \frac{p_{\theta}(X_{t-1}|X_t)}{q(X_t|X_{t-1})}\Big] := \mathcal{L} \end{aligned} \end{equation}
$$ \begin{aligned} p(z) &= p(X_{1:T}) \\ \end{aligned} $$
\begin{equation} q(X_t|X_0) = \mathcal{N}\big(X_t; \sqrt{\overline{a}_t}X_0, (1-\overline{a}_t)\mathbf{I} \big) \end{equation} $$ \begin{aligned} a_t &= 1 - \beta_t \\ \overline{a}_t &= \prod_{s=1}^t a_{s} \end{aligned} $$
mmgen.models.BasicGaussianDiffusion.q_sample
def q_sample(self, x_0, t, noise=None):
    r"""Get diffusion result at timestep `t` by `q(x_t | x_0)`.
    Args:
        x_0 (torch.Tensor): Original image without diffusion.
        t (torch.Tensor): Target diffusion timestep.
        noise (torch.Tensor, optional): Noise used in reparameteration
            trick. Default to None.
    Returns:
        torch.tensor: Diffused image `x_t`.
    """
    device = get_module_device(self)
    num_batches = x_0.shape[0]
    tar_shape = x_0.shape
    noise = self.get_noise(noise, num_batches=num_batches)
    mean = var_to_tensor(self.sqrt_alphas_bar, t, tar_shape, device)
    std = var_to_tensor(self.sqrt_one_minus_alphas_bar, t, tar_shape,
                        device)

    return x_0 * mean + noise * std

References

  1. Kevin P. Murphy. Probabilistic Machine Learning: Advanced Topics,
  2. Diederik P Kingma et al.. Auto-Encoding Variational Bayes, arXiv1312.6114
  3. Luo, Calvin. Understanding Diffusion Models: A Unified Perspective, arXiv2208.11970
  4. Jonathan Ho et al.. Denoising Diffusion Probabilistic Models, NeurIPS 2020
  5. Minh-Ngoc Tran et al.. A practical tutorial on Variational Bayes. arXiv2103.01327
  6. What are Diffusion Models? (Blog post)
  7. How diffusion models work: the math from scratch (Blog post)
  8. From Autoencoder to Beta-VAE (Blog post)
  9. Variational Bayesian methods (Wikipedia)
  10. Notation in probability and statistics (Wikipedia)
  11. Variational Bayesian methods (Part 1) (Berkeley CS-285 course)
  12. Variational Bayesian methods (Part 2) (Berkeley CS-285 course)
  13. Generative Modeling by Estimating Gradients of the Data Distribution (Blog post)
  14. diffusion model最近在图像生成领域大红大紫,如何看待它的风头开始超过GAN? - 我想唱high C的回答 - 知乎
  15. Langevin-type sampling methods

感谢阅读🤗 本文内容谢绝任何形式的转载,如果您想和朋友分享本文内容,请分享本文链接 kaizhao.net/blog/diffusion。 如果您发现文中的错误,或者有任何疑问,欢迎在下方留言交流 (留言功能基于 disqus,在中国大陆的读者可能需要一些技术手段才能连接🥲)。