Diffusion Model 笔记

Generative Model

四种 DL-based 生成模型：

1. Variational Autoencoders (VAE)
2. Flow-based models
3. Generative Adversarial Networks.
4. Diffusion

这四种模型都有其缺点，如下图：

	VAE	FLOW	GAN	Diffusion
Pros	Fast Sampling rate. Diverse sample generation	Fast Sampling rate. Diverse sample generation	Fast Sampling rate. High sample generation quality.	High sample generation quality. Diverse sample generation
Cons	Low sample generation quality	Need specialized architecture, low sample generation quality	Unstable training, low sample generation diversity (Mode Collapse)	Low sampling rate

Diffusion Model

主要思想

…systematically and slowly destroy the structure in a data distribution through an iterative forward diffusion process. We then learn a reverse diffusion process that restores structure in data, yielding a highly flexible and tractable generative model of the data. This approach allows us to rapidly learn, sample from, and evaluate probabilities in deep generative models …

• Forward 过程（data to noise）：
  • 原始图像 $x_0$
  • 重复 $T$ 次加噪：$x_{t-1}+{ϵ}_{t-1}\to x_t$，最终得到 $x_T$
  • 无模型参与
• Backward 过程（noise to data）：
  • 有模型 $\theta$ 参与：$\theta :(x_t,t)\to {\widehat{ϵ}}_{t-1}$

DDPM

NOTE

DDPM 是 Diffusion Model 的一个子类，其特点在于：

• 预测噪声而非预测图像： 广义的扩散模型可以尝试在每一步预测上一张稍微清晰一点的图片。但 DDPM 发现，让神经网络（通常是 U-Net 架构）直接去预测当前图片上被添加了多少噪声（$ϵ$），效果出奇的好，数学优化上也更加稳定。
• 固定的加噪过程： DDPM 使用了基于马尔可夫链（Markov Chain）的设定，规定了每次加噪的幅度（Variance Schedule）是预先设定好的、固定的，不需要模型去学习，这大大降低了训练难度。
• 简化的损失函数： DDPM 的作者通过巧妙的数学化简，去除了复杂的变分下界（VLB）中一些难计算的项，得到了一个极其简单、优雅的均方误差损失函数（MSE）。模型只需要比较“真实的噪声”和“预测的噪声”的差异即可。

Forward 过程

$$\begin{array}{rl}q(x_1,x_2,\dots ,x_T\mid x_0)& =\prod _{t=1}^{T}q(x_t\mid x_{t-1})\\ q(x_t\mid x_{t-1})& =\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)\end{array}$$

$\sqrt{1-{\beta }_t}$ 这个系数是如何确定的

答案是为了保持方差。重参数化 $x_t$ 为 $x_t=\sqrt{1-{\beta }_t}{x}_{t-1}+\sqrt{{\beta }_t}z,z\sim \mathcal{N}(0,I)$，则 $x_t$ 的方差为：

$$Var(x_t)=Var(\sqrt{1-{\beta }_t}{x}_{t-1}+\sqrt{{\beta }_t}z)$$

假设 $x_{t-1}$ 已经是一个近似方差为 $I$ 的分布，且 $x_{t-1}$ 与新加入的噪声 $z$ 是相互独立的，根据方差的性质：

$$\begin{array}{rl}Var(x_t)& =(\sqrt{1-{\beta }_t}{)}^{2}Var(x_{t-1})+(\sqrt{{\beta }_t}{)}^{2}Var(z)\\ & =(1-{\beta }_t)\cdot I+{\beta }_t\cdot I\\ & =I\end{array}$$

这是最终 $x_t$ 能够顺利的收敛到 $\mathcal{N}(0,I)$ 的必要条件

在实际使用时，DDPM 的作者让 ${\beta }_t$ 取值为 $[0.001,0.02]$，总步数 $T=1000$，使用线性取值，如下图：

为了解决“获取 $x_t$ 就需要前向 $t$ 次的问题”，可以通过数学推导来求解，即求解 $q(x_t\mid x_0)$：

$$\begin{array}{rl}{x}_1& =\sqrt{1-{\beta }_0}{x}_0+\sqrt{{\beta }_0}{ϵ}_0\\ x_2& =\sqrt{1-{\beta }_1}{x}_1+\sqrt{{\beta }_1}{ϵ}_1\\ & =\sqrt{(1-{\beta }_0)(1-{\beta }_1)}{x}_0+\left(\sqrt{{\beta }_0(1-{\beta }_1)}{ϵ}_0+\sqrt{{\beta }_1}{ϵ}_1\right)\end{array}$$

这里 ${ϵ}_0,{ϵ}_1\sim \mathcal{N}(0,I)$ 且相互独立，则有

$${ϵ}_1^{\prime }=\sqrt{{\beta }_0(1-{\beta }_1)}{ϵ}_0+\sqrt{{\beta }_1}{ϵ}_1\sim \mathcal{N}(0,({\beta }_0(1-{\beta }_1)+{\beta }_1)I)$$

记 ${\alpha }_t=1-{\beta }_t$，则

$$\begin{array}{rl}{\beta }_0(1-{\beta }_1)+{\beta }_1& =(1-{\alpha }_0){\alpha }_1+(1-{\alpha }_1)=1-{\alpha }_0{\alpha }_1\\ (1-{\beta }_0)(1-{\beta }_1)& ={\alpha }_0{\alpha }_1\end{array}$$

再记 ${\bar{\alpha }}_t={\prod }_{i=1}^t{\alpha }_i$ 进而

$$x_2=\sqrt{{\bar{\alpha }}_1}{x}_0+\sqrt{1-{\bar{\alpha }}_1}{ϵ}_1^{\prime }$$

同理可以得到

$$\begin{array}{rl}{x}_t& =\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ\\ & ϵ\sim \mathcal{N}(0,I)\\ ⟹& q(x_t\mid x_0)=\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I)\end{array}$$

如此便可以一步得到 $x_t$，而无需反复进行前向过程

Backward 过程

(图片中 $q(x_t\mid x_0)$ 有误)

In 1949, W. Feller showed that, for gaussian (and binomial) distributions, the diffusion process’s reversal has the same functional form as the forward process. 这意味着可以假设后向过程 $(x_t,t)\to x_{t-1}$ 也遵循高斯分布

我们需要求得一个模型 $\theta$，即 $p_{\theta }(x_{t-1}\mid x_t)$，由上面这段话，可以假设

$$p_{\theta }(x_{t-1}\mid x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\Sigma }_{\theta }(x_t,t))$$

Loss Function

设计 loss 的 insight 为“最大化反向过程中生成的 $x_0$ 在原始数据分布中的 log-likelihood”，即

$$\mathcal{L}={\mathbb{E}}_{q(x_0)}[-\log p_{\theta }(x_0)]$$

这一点和 VAE 是相同的，可以仿照 VAE 引入变分推断的方法，即引入

$$D_{\text{KL}}\left(q(x_{1:T}\mid x_0)\mid p_{\theta }(x_{1:T}\mid x_0)\right)$$

NOTE

回忆 VAE 中推导 loss 时，引入了 $D_{\text{KL}}(q(z|x)\mid p_{\theta }(z|x))$。在 diffusion model 中，$x_{1;T}$ 可以视为 VAE 中的隐变量 $z$

$$\begin{array}{rl}0& \le D_{\text{KL}}\left(q(x_{1:T}\mid x_0)\mid p_{\theta }(x_{1:T}\mid x_0)\right)\\ & ={\mathbb{E}}_{q(x_{1:T}|x_0)}\left[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{1:T}|x_0)}\right]\\ & ={\mathbb{E}}_{q(x_{1:T}|x_0)}\left[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}+\log p_{\theta }(x_0)\right]\\ & ={\mathbb{E}}_{q(x_{1:T}|x_0)}\left[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right]+\log p_{\theta }(x_0)\\ ⟹-\log p_{\theta }(x_0)& \le {\mathbb{E}}_{q(x_{1:T}|x_0)}\left[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right]\\ ⟹\mathcal{L}& \le {\mathbb{E}}_{q(x_0)}\left[{\mathbb{E}}_{q(x_{1:T}|x_0)}\left[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right]\right]\\ & ={\mathbb{E}}_{q(x_{0:T})}\left[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right]\end{array}$$

而

可以将 loss 总结为：

看上去较为复杂，原作者进行了如下化简：

• 抛弃 $L_0$，在实验中获得了更好的结果
• 抛弃 $L_T$，因为该项不涉及神经网络参数

于是便只需要考虑 $L_{t-1}$

其中

$$\begin{array}{rl}q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)& =\mathcal{N}({\mathbf{x}}_{t-1};{\tilde{\mathbf{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0),{\tilde{\beta }}_t\mathbf{I}),\\ \mathrm{where}{\tilde{\mathbf{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)& \begin{array}{r}:=\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{\mathbf{x}}_0+\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{\mathbf{x}}_t\mathrm{and}{\tilde{\beta }}_t:=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t\end{array}\end{array}$$

在实验中，作者发现固定 $p_{\theta }$ 中的方差，对实验结果没有很大影响，即

$$p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\sigma }^{2}I)$$

进而可以推导得到

$$L_{t-1}={\mathbb{E}}_q\left[\frac{1}{2{\sigma }_t^2}\Vert {\tilde{\mu }}_t(x_t,x_0)-{\mu }_{\theta }(x_t,t){\Vert }^2\right]+C$$

通过将 $x_0$ 用 $x_t$ 和 $ϵ$ 表达，可以将上面的式子转化为每步对噪声 $ϵ$ 进行优化（这里去掉了范数前的系数）：

$${\mathbb{E}}_{t,x_0\sim q,ϵ\sim \mathcal{N}}\left[\Vert ϵ-{ϵ}_{\theta }(\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t){\Vert }^2\right]$$

也即

$${\mathbb{E}}_{t,x_0,ϵ}\left[\Vert ϵ-{ϵ}_{\theta }(x_t,t){\Vert }^2\right]$$

参考资料

• Introduction to Diffusion Models for Image Generation - A Comprehensive Guide
• InDepth Guide to Denoising Diffusion Probabilistic Models DDPM
• 由浅入深了解 Diffusion Model
• Diffusion Models: DDPM | Generative AI Animated - YouTube
• Diffusion models from scratch in PyTorch - YouTube