本文介绍 score-based generative models,又称 diffusion models。 先简单介绍生成模型 (generative models) 的概念和几种典型的生成模型, 然后重点介绍 diffusion models。 本文的主要内容参考了 (Luo, Calvin, 2022), 以及几篇相关博客文章[2] [3] [4] [14]。 在 (Luo, Calvin, 2022) 的基础上,本文主要有以下两点改变:
- 对部分数学符号和统计学概念增加了更详细具体的解释;
- 对于部分数学推导增加了详细的中间步骤,以保证只有初等概率论知识背景的同学可以理解。
因为接下来的内容会从统计学上描述生成模型,为了方便阅读和理解,这里首先从图像生成的角度解释一些概念:
概念 | 符号 | 解释 |
---|---|---|
观测,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$ 后,推测其是由哪个模型采样出的“,
变量是概率模型或者是模型参数。
🎯概率和统计
概率是指已知模型推测结果:例如已知硬币两面均匀,求抛 100 次都是正面朝上的概率。
而统计是指已知实验结果反推模型:例如已知抛 100 次硬币都是正面朝上,求硬币两面均匀的概率。
变分贝叶斯 (Variational Bayes)
后面讲的很多内容涉及到变分贝叶斯的一些概念,为了理清这些概念方便大家理解,这里先回顾一下 variational Bayesian methods 的基本概念, 这部分的主要参考资料有: Berkely CS 285: Lecture 18, Variational Inference (Part 1 & Part 2), 维基百科,
什么是变分贝叶斯方法
很多地方中都会解释变分贝叶斯是什么,有什么用处, 由于这些解释各自的侧重点不同,因此各不相同,如同盲人摸象,无法提纲挈领。 总的来说就一句话,变分贝叶斯是用来逼近一些无法计算的 (intractable) 概率分布的,通常这些分布都涉及到对一些变量的积分。 或者更简单地,变分贝叶斯是用来逼近概率分布的。 在变分贝叶斯方法所涉及到的概率模型中,一般涉及到三种不同的变量:
- 数据 data $x$;
- 参数 parameter $\theta$;
- 隐变量 latent variable $z$。
变分贝叶斯是用来干什么的
变分贝叶斯有哪些应用呢?总的来说有两个:
- 逼近一些 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.)
- 获得观测数据的 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$。
假设我们观测到下图中的数据,我们如何计算它们的分布 $p(x)$? 显然 $p(x)$ 并不是一个常规的分布,是一个 intractable distribution。 但是我们观察到数据聚成了 3 个 clusters,每个 cluster 的数据似乎是从高斯分布中采样的。 因此我们定义隐变量 $z\in\{0, 1, 2\}$ 表示 $x$ 来自哪一个高斯, 注意 $z$ 是 离散的随机变量。 然后定义以下几个分布:
- $p(x|z)$:是高斯分布;
- $p(z)$: 隐变量的“先验分布” prior;
- $p(z|x)$: 隐变量的 “后验概率” posterior。
当然,实际条件下的隐变量 $z$ 远比例子中复杂,有的甚至是连续的。 但是高斯混合模型的例子已经可以将问题说明白:
- 我们观测到了一堆数据 $x$,其分布未可知;
- 为了得到 $p(z)$,我们设计了一个隐变量 $z$,使得 $p(x|z), p(z)$ 是 tractable 的。
- 通过 $p(x|z), p(z)$,我们可以求解 $p(x)$ 以及 $p(z|x)$。
这里我们用生成模型为例介绍变分贝叶斯中的三种随机变量。 如下图所示,一般而言生成模型(不论是 GAN 还是 VAE)的生成(或者采样)过程都是输入一个隐变量 $z$ 到模型中,模型参数为 $\theta$,然后得到生成的样本 $x$。
ELBO: Evidence Lower BOund
逼近后验概率 + 获得 evidence考虑隐变量 $z$ 关于数据 $x$ 的后验概率 (posterior),记为 $p(z|x)$。
下面我们推导如何用变分贝叶斯推断来逼近 (approximate) 隐变量 (latent) 的后验概率:$p(z|x)$。 我们假设有另外一个分布 $q_{\phi}(z|x)$,然后用 $q_{\phi}(z|x)$ 去逼近 $p(z|x)$: $q_{\phi}(z|x) \approx p(z|x)$。 通常会假设 $q_{\phi}(z|x)$ 属于某个常见分布 $\mathcal{Q}$(比如高斯分布),然后最小化 $q_{\phi}(z|x)$ 与 $p(z|x)$ 之间的 KL-divergence。 两个概率分布 $p$ 和 $q$ 之间的 KL Divergence 定义为: \begin{equation} D_{\text{KL}}(p \ \rvert\rvert \ q) = \int_x p(x)\log\frac{p(x)}{q(x)}dx. \label{eq:kld} \end{equation} 找到最佳的逼近 $q^*$ 可以表示成为最小化 KL Divergence: \begin{equation} q^* = \underset{q\in\mathcal{Q}}{\operatorname{arg min}} \ \ D_{\text{KL}}\big(q_{\phi}(z|x) \ || \ p(z|x)\big) \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 min}} \ \ D_{\text{KL}}(q_{\phi}(z|x) \ || \ p(z|x)) \end{equation}
那么我们如何最小化 $D_{\text{KL}}(q_{\phi}(z|x) \ || \ p(z|x))$ 呢?$p(z|x)$ 本身是 intractable 的,因此无法直接计算上式中的 KL Divergence。 通过接下来的推导,我们发现可以通过另外一种方式去最小化 KL Divergence。
- ELBO 是 $\log p(x)$ 的下界;
- 可以通过 maximize ELBO 来 maximize $p(x)$
接下来我们将介绍生成模型,并以 variational aotu-encoder (VAE) 为例,介绍如何通过优化 ELBO 来最大化 $p(x)$,进而训练模型。
生成模型
给定一些从某个概率分布中采样出的样本 $x\sim p(x)$,生成模型 (generative models) 的目的是建模这个概率分布 $p(x)$, 从而使我们能从这个概率分布中采样出新的样本 (也就是所谓的“生成”), 同时也能判断任意一个样本的似然 (likelihood, 也就是判断任意一个样本“真假”的程度)。
这里所谓的“从某个概率分布中采样出的样本”,是假定所有的样本,比如自然图像,都服从某个很复杂概率分布, 所有的观测到的样本(observation)都是从这个概率分布中采样出来的。 似然 (likelihood) 指的是某个样本是从某个特定的概率分布中采样出来的概率。 假设 $p(x)$ 是所有自然图像的概率分布,那么样本 $x_i$ 的似然则表明该样本是从“自然图像”中采样出来的概率。
现有的几种生成模型
这里简单介绍一些现有的生成模型。因为我们的重点在 diffusion models,因此对于其他的生成模型范式我们仅从 high-level 介绍其基本概念和区别。 现有的生成模型可以粗略地划分为以下几种:
- Generative Adversarial Network (GAN).
- likelihood-based.
- energy-based.
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: Variational Auto-Encoder
Variational Auto-Encoder
这节我们介绍 Variational Auto-encoder。不像很多其他文章首先介绍 VAE 的原理,再介绍 VAE 的结构。 我们先从 VAE 的具体实现和模型结构开始讲起,因为 VAE 的结构比较简单,而背后的原理相对比较复杂。
VAE[5] 采用 encoder-decoder 结构。如 图1 所示,整个 VAE 模型呈现一个瓶颈结构,包含一个 encoder 和一个 decoder。 Encoder 输入样本 $x$,输出一个高斯分布的均值 $\mu$ 和方差 $\sigma$ 从而得到高斯分布 $q(z|x)$; decoder 输入从该高斯分布采样出的一个 latent 向量 $z$, 输出生成的样本 $\hat{x}$。 训练时的损失函数为
从优化和 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 的定义出发,推导这两者之间的联系。
- 先验项 $D_{KL}\big(q_{\phi}(z|x) \ \rvert\rvert \ p(z)\big)$ 对应了 公式$\ref{eq:loss-vae}$ 中的 $D_{KL}\big(q_{\phi}(z|x) \ \rvert\rvert \ \mathcal{N}(0, \mathbf{I})\big)$。
- 重建项 $\mathbb{E}_{z\sim q_{\phi}(z|x)}\log p_{\theta}(x|z)$ 对应了 公式$\ref{eq:loss-vae}$ 中的 $\rvert\rvert x - \hat{x} \rvert\rvert_2$;
🎯先验项 (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-elbo-1}$ 导出。
假设 $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 的训练过程可以分为一下几步:
- 输入样本 $x$ 到 encoder,得到$\mu(x), \sigma(x)$;
- 从标准高斯分布中采样若干个 $\epsilon$: $\{\epsilon^1, ..., \epsilon^L \}$;
- 根据 公式$\ref{eq:epsilon-z}$ 得到隐变量 $z^1, ..., z^L$;
- 输入 $z$ 到 Decoder,得到输出 $\hat{x}^1, ..., \hat{x}^L$;
- 根据 公式 $\ref{eq:loss-vae}$ 计算损失;
- 反向传导,更新 $\phi, \theta$。
Diffusion Models
Diffusion models 是一种特殊的隐变量 (latent variable) 模型,它的特殊性体现在两个方面:
- Diffusion models 有多个 latent variables,而 VAE/GAN 这些只有一个 latent variable;
- Diffusion models 的隐变量与样本维度相同,而 VAE/GAN 的隐变量维度都远小于数据;
Diffusion model 包含两个过程: forward process 和 reverse process。
- 假设 $X_0$ 是从真实数据分布中采样的一个样本 $X_0\sim q(X)$, forward process 通过不断迭代地在 $X_0$ 上加很小的高斯噪声,得到一系列的 latent variables $X_1, X_2, ..., X_T$。 当 $T$ 足够大 ($\ge 1,000$) 时,最后得到的 $X_T$ 可近似认为服从 isotropic Gaussian distribution。
- 假设 $X_T$ 是从 isotropic Gaussian distribution 中采样出的一个样本点,reverse process 通过不断迭代地“去噪声”,可以逆向得到 $X_{T-1}, X_{T-2}, ..., X_0$。
Forward process 和 reverse process 都是通过在上一步的基础上不断迭代得到最终结果,因此可以用马尔科夫链来建模。
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”, 是马尔可夫模型的参数。
Diffusion models
前面讲过,Diffusion models 的 forward process (从 $X_0$ 到 $X_T$) 和 reverse process (从 $X_T$ 到 $X_0$) 都可以用马尔科夫链建模。 我们假设在 forward/reverse processes 中每一步都加的是高斯噪声,那么马尔科夫链的 Markev kernel 都是高斯分布。 下面我们分别形式化地介绍 forward process 和 reverse process 的马尔科夫链表示。
Forward Process
Forward process 从样本 $X_0$ 开始,每步加上一个很小的高斯噪声,最终得到 isotropic 高斯噪声 $X_T$。
→ Forward / diffusion process (从样本到噪声) →
从真实样本 $X_0$ 开始,逐渐在累加一个很小的高斯噪声,最后可以得到 isotropic 高斯分布 $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 为
需要推导一下和 公式$\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}$。
注意 公式$\ref{eq:q_xt_xt1}$ 中的 $\beta_t$ 通常是一个提前设置好的的超参数, 因此我们可以 analytically 从 $X_0$ 推导出任意时刻 $X_t$ 的分布 $q(X_t|X_0)$: \begin{equation} \begin{split} q(X_t|X_0) = \mathcal{N}\big(X_t; \sqrt{\overline{a}_t}X_0, (1-\overline{a}_t)\mathbf{I} \big) \\ \alpha_t = 1 - \beta_t, \ \ \overline{a}_t = \prod_{\tau=1}^t a_{\tau}. \end{split}\label{eq:q_xt_x0} \end{equation}
下面是 mmgeneration 对 公式$\ref{eq:q_xt_x0}$ 的实现:
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
Reverse Process
Reverse process 是 forward process 的反过程:从高斯噪声 $X_T$ 开始,每步加上一个很小的高斯噪声,最终得到 $X_0$。
← 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} Reverse process 的 Markov kernel 是:
注意 reverse process 的 Markov kernel 不再是常数,其均值 $\mu_{\theta}(X_t, t)$ 与方差 $\Sigma_{\theta}(X_t, t)$ 不仅与输入 $X_t$ 有关,且与参数 $\theta$ 有关。 在实现中,reverse process 的 Markov kernel 由神经网络来逼近,该神经网络参数为 $\theta$,输入为 $X_t$,输出为对应的均值和方差 $\mu_{\theta}(X_t, t), \Sigma_{\theta}(X_t, t)$。
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)$ |
优化目标
由于 diffusion models 的 latent variable 为 $X_{1:T}$,因此 VAE 中的 $p(Z)$ 对应着 diffusion models 中的 $p(X_{1:T})$。 下表列出了一般 latent variable models 和 diffusion models 几种分布的对照。
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)$ |
与 VAE 类似,diffusion models 训练时的优化目标也是最大化 ELBO。 根据上面的变量对照表,difusion models 中的 ELBO 应当为:
这里推导一下红色部分。当 $t \gt 1$ 时:
公式 $\ref{eq:diffusion-loss}$ 中第二项计算 $q(X_{t-1}|X_t, X_0)$ 和 $p_{\theta}(X_{t-1}|X_t)$ 之间的 KL Divergence。 其中 $q(X_{t-1}|X_t, X_0)$ 可以 analytically 推导出来:
KL Divergence 中的另一项 $p_{\theta}(X_{t-1} | X_t) := \mathcal{N}\big(X_{t-1}; \mu_{\theta}(X_t, t), \Sigma_{\theta}(X_t, t)\big)$ 是模型需要学习的 reverse 过程。 为了简化,假设 $\Sigma_{\theta}(X_t, t) = \sigma_t^2\mathbf{I}$。(Ho et al. 2020) 的实验结果表明 $\sigma_t = \tilde{\beta}_t = \frac{1-\overline{a}_{t-1}}{1-\overline{a}_t}\beta_t$ 和 $\sigma_t=\beta_t$ 的结果类似。 最终,公式 $\ref{eq:diffusion-loss}$ 第二项计算了以下两个高斯分布的 KL Divergence:
Gaussian | mean | std |
---|---|---|
$q(X_{t-1}|X_t, X_0)$ | $\tilde{\mu}_t(X_t, X_0)$ $\ref{eq:q_x_t1_xt_x0}$ | $\tilde{\beta}_t\mathbf{I}$ |
$p_{\theta}(X_{t-1}|X_t)$ | $\mu_{\theta}(X_t, t)$ | $\sigma_t\mathbf{I}$ |
两个高斯分布之间的 KL Divergence 可以解析地表达出来。 假设有 $p, q$ 两个高斯分布,均值方差分别是 $\mu_1. \mu_2, \sigma_1, \sigma_2$, 将高斯分布的概率密度函数带入 KL Divergence 定义很容易得出: \begin{equation*}\begin{aligned} D_{KL}(p, q) &= \int p(x)\log \frac{p(x)}{q(x)}dx \\ &= \frac{1}{2}\log\frac{\sigma_2}{\sigma_1} + \frac{\sigma_1^2 + (\mu_1 - \mu_2)^2}{2\sigma_2^2} - \frac{1}{2} \end{aligned}\end{equation*} 于是 公式$\ref{eq:diffusion-loss}$ 中第二项的 KL Divergence 可以表示为:
重参数技巧
到这里,最直接的方法就是用神经网络来直接预测 $\tilde{\mu}_t$,然后计算 公式$\ref{eq:l_t1_1}$ 中的损失项。 但是 (Ho et al. 2020) 采用了另外一种更复杂方式:重参数 (reparameterize) 公式 $\ref{eq:q_xt_x0}$ 里的 $p(X_t | X_0)$。 在正式介绍之前,这里我们先整理一下涉及到的变量和分布,方便回顾和检索。
符号 | 定义 | 相关分布 | |
---|---|---|---|
$\beta_t$ | 超参数 | $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:q_xt_xt1}$) | |
$a_t$ | $1 - \beta_t$ | 只是为了下面导出 $\overline{a}_t$ | |
$\overline{a}_t$ | $\overline{a}_t = \prod_{\tau=1}^t a_{\tau}$ | $q(X_t|X_0) = \mathcal{N}\big(X_t; \sqrt{\overline{a}_t}X_0, (1-\overline{a}_t)\mathbf{I} \big)$ | 公式$\ref{eq:q_xt_x0}$ |
$\tilde{\beta}_t$ | $\tilde{\beta}_t = \frac{1-\overline{a}_{t-1}}{1-\overline{a}_t}\beta_t$ | $q(X_{t-1}|X_t, X_0) = \mathcal{N}\big(X_{t-1}; \tilde{\mu}_t(X_t, X_0), \tilde{\beta}_t\mathbf{I} \big)$ | 公式 $\ref{eq:q_x_t1_xt_x0}$ |
$\tilde{\mu}_t$ | 详见 公式 $\ref{eq:q_x_t1_xt_x0}$ |
对 $p(X_t | X_0)$ 的 reparameterization 和上面 VAE 中的重参数技巧 类似。 根据 公式$\ref{eq:q_xt_x0}$: $$ q(X_t|X_0) = \mathcal{N}\big(X_t; \sqrt{\overline{a}_t}X_0, (1-\overline{a}_t)\mathbf{I} \big) $$ 假设变量 $\epsilon\sim \mathcal{N}(0, \mathbf{I})$ 服从标准正态分布,我们可以重新表示 $X_t$: \begin{equation}\begin{aligned} X_t(X_0, \epsilon) &= \epsilon\sqrt{1-\overline{a}_t} + \sqrt{\overline{a}_t}X_0 \\ X_0 &= \frac{X_t(X_0, \epsilon) - \epsilon\sqrt{1-\overline{a}_t}}{\sqrt{\overline{a}_t}} \end{aligned}\end{equation} 上式说明,均值为 $\sqrt{\overline{a}_t}X_0$,方差为 $1-\overline{a}_t$ 的高斯分布可以由标准高斯分布先缩放 $\sqrt{1-\overline{a}_t}$ 倍然后平移 $\sqrt{\overline{a}_t}X_0$ 得到。 带入 公式$\ref{eq:l_t1_1}$:
References
- Kevin P. Murphy. Probabilistic Machine Learning: Advanced Topics,
- Diederik P Kingma et al.. Auto-Encoding Variational Bayes, arXiv1312.6114
- Luo, Calvin. Understanding Diffusion Models: A Unified Perspective, arXiv2208.11970
- Jonathan Ho et al.. Denoising Diffusion Probabilistic Models, NeurIPS 2020
- Minh-Ngoc Tran et al.. A practical tutorial on Variational Bayes. arXiv2103.01327
- What are Diffusion Models? (Blog post)
- How diffusion models work: the math from scratch (Blog post)
- From Autoencoder to Beta-VAE (Blog post)
- Variational Bayesian methods (Wikipedia)
- Notation in probability and statistics (Wikipedia)
- Variational Bayesian methods (Part 1) (Berkeley CS-285 course)
- Variational Bayesian methods (Part 2) (Berkeley CS-285 course)
- Generative Modeling by Estimating Gradients of the Data Distribution (Blog post)
- diffusion model最近在图像生成领域大红大紫,如何看待它的风头开始超过GAN? - 我想唱high C的回答 - 知乎
- Langevin-type sampling methods
感谢阅读🤗 本文内容谢绝任何形式的转载,如果您想和朋友分享本文内容,请分享本文链接 kaizhao.net/blog/diffusion。 如果您发现文中的错误,或者有任何疑问,欢迎在下方留言交流 (留言功能基于 disqus,在中国大陆的读者可能需要一些技术手段才能连接🥲)。