MEGChai MEGChai
  • 文章
    • 随笔
    • 笔记
    • 教程
  • 关于

扩散模型(一):DDPM 基本原理与 MegEngine 实现

Chai
2022-06-12 17:23:00笔记阅读 2,089

2022-07-07 更新:MegEngine 实现代码已经发布在 https://github.com/MegEngine/MegDiffusion

前言

大部分与 DDPM 有关的教程喜欢直接从 OpenAI 的 Beats GAN 论文与 Pytorch 版本实现开始介绍,笔者猜测可能是因为 DDPM Ho20 的 Tensorflow 版本代码比较难阅读、而 15 年开山论文所用数学符号比较晦涩等等原因。笔者希望这个系列能够从最简单的 Ho20 版本开始介绍,并给出相应的 MegEngine 实现(代码风格实现偏 Educational 而非 Engineering),其中 MegDiffusion/megdiffusion/model/ddpm.py 中的 UNet 精准复现了 DDPM 官方 Tensorflow 版本的模型前向结果,并提供了转换后的预训练模型。本文高度服务于那些希望复现 DDPM 论文实验的读者,如果只是想要学习 DDPM 原理,可以直接阅读《What are Diffusion Models? 》,否则感受不到本文的意义。

这篇文章是接触 Diffusion Model 领域的第一篇笔记,也略带有些 Step-by-step 教程的性质,尽可能地做到 Self-contained, 里面记录的观点也是基于 20 年 DDPM 原始论文的认知给出的,未必解释得合理;而现今前沿的一些研究可能会以更加严谨、系统而全面的角度来对 Diffusion Model 进行解读,因此本文的将假想读者定位为和我一样的领域新手(尚处于读论文做复现的阶段),且不保证会有系列后续的文章。

什么是扩散模型?

相较于 GANs, VAEs 与 Normalizing Flows 等生成式模型,扩散模型(Diffusion Model)的思路非常简单:向数据中逐步地加入随机噪声,然后学习其逆向过程,希望最终能够从噪声中重建所需的数据样本 —— 有意思的地方在于,我们可以训练一个神经网络模型来学习逆向扩散过程。

扩散模型(一):DDPM 基本原理与 MegEngine 实现-MEGChai

在逆向扩散的每一步,神经网络模型的输入是当前的样本 $\mathbf{x}_{t}$ 和加噪程度 $t$ ,想一想预期输出可以是什么呢?有多种可能:比如采取最粗暴的一种做法,即一步到位,要求模型直接对 $\mathbf{x}_{0}$ 进行预测,发现实验效果不佳;想一下扩散的样子... 逐步地,渐渐地变化,因此扩散模型要求对去噪的下一个状态进行预测,即对 $\mathbf{x}_{t-1}$ 进行预测。而在实际实验的过程中,作者发现:先让模型对去噪过程中的噪声 $\mathbf{\epsilon}$ 进行预测,再通过相关计算得到 $\mathbf{x}_{t-1}$ 会有更好的效果。如果你对整个流程还不是还清楚,不用担心,下面我们将详细介绍一下扩散模型的前向和逆向过程,附带代码实现。

前向扩散过程

给定某个采样自真实数据分布的数据点(比如一张图片) $\mathbf{x}_{0} \sim q(\mathbf{x})$ , 前向扩散过程(Forward diffusion process) $q(\mathbf{x}_{1:T} \vert \mathbf{x}_0)$ 指在有限的 $T$ 步内(理想情况下应当满足 $T \rightarrow \infty$, 但这样就无法处理了,论文中使用 $T=1000$ 近似,因此这个值只要足够大即可),逐渐不断地向其中加入 Gaussian 随机噪声,将得到一个由带噪声样本组成的序列 $\mathbf{x}_{1}, \ldots, \mathbf{x}_{T}$. 其中每一步的加噪程度 $\beta$ 由事先生成的 $\{\beta_t \in (0, 1)\}_{t=1}^T$ 决定(满足 $0<\beta_{1}<\beta_{2}<\ldots<\beta_{T}<1$ ),写成条件概率公式如下:

$$
q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) =
\mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t\mathbf{I})
$$

根据马尔可夫链的无记忆性(下一个状态的概率分布只与当前状态有关),可以得到扩散过程的联合概率分布:

$$
q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) = \prod^T_{t=1} q(\mathbf{x}_t \vert \mathbf{x}_{t-1})
$$

在这个过程中,随着 $t$ 不断增大(直到 $t=T$ ),原始数据样本 $\mathbf{x}_{0}$ 将逐渐地丢失掉原有的特征,最终得到的 $\mathbf{x}_{T}$ 将会是一个各向同性的高斯分布(Isotropic Gaussian distribution, 指其协方差矩阵为单位阵乘以正的常数,即方差不因为方向而变化),即纯粹的噪声数据。

回忆一下 重参数/参数重整化技巧(Reparameterization trick) —— 假定我们希望从某个高斯分布中通过采样生成噪声 $\boldsymbol{\epsilon} \sim \mathcal{N}(\boldsymbol{\mu}, \boldsymbol{\sigma}^{2})$, 由于采样操作本身是不可导的,这样做会导致当 $\boldsymbol{\mu}$ 和 $\boldsymbol{\sigma}$ 是神经网络的参数时,梯度下降算法不可用。因此我们可以先从标准高斯分布中采样 $\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ 并看作是一个常数,再得到 $\boldsymbol{\epsilon} = \boldsymbol{\mu} + \boldsymbol{\sigma} \cdot \mathbf{z}$ 作为采样的等价过程。这样就将随机性转移到了 $\mathbf{z}$ 上,而 $\boldsymbol{\mu}$ 和 $\boldsymbol{\sigma}$ 在上述形式的仿射变换中是可导的。只要你在计算 Loss 的整个链路中出现了 “采样” 操作,为了能够进行梯度计算完成反向传播,都可以使用重参数技巧。因此使用参数重整化技巧,$q(\mathbf{x}_t \vert \mathbf{x}_{t-1})$ 的采样过程等价于 ${\color{blue} \mathbf{x}_t = \sqrt{1 - \beta_t} \mathbf{x}_{t-1} + \sqrt{\beta_t} \odot \mathbf{z}_{t-1}}$, 而独立的随机变量 $\mathbf{z}_{t-1} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ 保证了采样的随机性。

注意:上式中完全不含有可训练的参数,也即是说扩散模型的前向扩散过程不需要对网络进行训练。只要给定 $\mathbf{x}_0$ 和 $\beta$, 通过不断迭代,就可以算出 $t$ 为任意时刻所对应的 $q(\mathbf{x}_t)$ 分布。实际上在逆扩散的过程中网络也没有进行训练,我们做的仅仅是使用训练好的网络对噪声进行 “采样”。这也是扩散模型的特点之一:你可以将不同的扩散策略与不同的预测模型进行组合使用,只需提供一致的维度与超参数等信息(比如 $T$ 值)。那么我们的神经网络模型是如何进行参数训练的呢?在设计目标函数的时候将揭开谜底,在此之前让我们先搞清楚并实现最基本的扩散过程。

代码实现:逐步加噪

我们用单张图片+纯 NumPy 代码来实现一下前向扩散(加噪)过程,帮助你更好地理解逻辑:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline
image = plt.imread("chai.jpg")
plt.imshow(image)
plt.show()
 
def transform(image):
    """Uint8 [0, 255] -> Float [-1, 1]"""
    return (image / 255 * 2) - 1
 
def transform_rev(image):
    """Float [-1, 1] -> Uint8 [0, 255]"""
    return ((image + 1) / 2 * 255).astype("uint8")
 
image = transform(image)
# plt.imshow(image)  # Error
# plt.imshow(transform_rev(image))  # OK
扩散模型(一):DDPM 基本原理与 MegEngine 实现-MEGChai
我们用于演示扩散模型流程的样例图 $\mathbf{x}_0$

注意 Uint8 图片的数值范围是 $[0,255]$, 我们在加噪前需要将其缩放到噪声所属的标准正态分布 $[-1, 1]$ 的范围,确保逆向过程 $p\left(\mathbf{x}_{T}\right)$ 为标准正态先验。而在需要使用图片(比如可视化)时,需将范围变换回来。这也是扩散模型经典的前后处理逻辑。

1
2
3
4
5
6
7
8
9
10
11
12
13
timesteps = 1000
 
def linear_schedule(timesteps, start=1e-4, end=2e-2):
    return np.linspace(start, end, timesteps)
 
def sigmoid_schedule(timesteps, start=1e-4, end=2e-2):
    betas = linear_schedule(start, end, timesteps)
    sigmoid = 1 / (1 + np.exp(-betas))
    return sigmoid * (end - start) + start
 
# def xxx_schedule(...):
 
betas = linear_schedule(timesteps)

使用 Variance Schedule 生成固定的方差序列 betas, 当需要获取第 $t$ 步的方差时,只需要通过 betas[t] 取得。除了 linear_schedule 这种线性实现外,还可以使用其它的生成方式如 sigmoid_schedule, cosine_schedule, quadratic_schedule 等,只需要确保 $0<\beta_{1}<\ldots<\beta_{T}<1$ 即可。

这里的 timesteps 即 $T$, 可看作是一个超参数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def q_sample_single_step(x, t, noise=None):
    """q(x_t | x_{t-1})"""
    if noise is None:
        noise = np.random.normal(0, 1, x.shape)
    mean = np.sqrt(1-betas[t]) * x
    var = betas[t]
    return mean + np.sqrt(var) * noise
 
def q_sample_step_by_step(x, t, noise=None):
    """q(x_{1:T} | x_{0})"""
    for i in range(t):
        x = q_sample_single_step(x, i)
    return x
 
noisy_image = q_sample_step_by_step(image, timesteps)
plt.imshow(transform_rev(noisy_image))
plt.show()

实现和演示从 $\mathbf{x}_0$ 逐步加噪得到 $\mathbf{x}_{1}, \ldots, \mathbf{x}_{T}$ 的过程,单步过程即 $\mathbf{x}_t = \sqrt{1 - \beta_t} \mathbf{x}_{t-1} + \sqrt{\beta_t} \cdot \mathbf{z}_{t-1}$, 只要加噪的次数足够多,最终图片中的原始信息将会所剩无几,变成一个各向同性的高斯分布。

扩散模型(一):DDPM 基本原理与 MegEngine 实现-MEGChai
在 t=1000 时,已经接近纯噪声
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
import matplotlib.animation as animation
 
image = plt.imread("chai.jpg")
image = transform(image)
 
fig = plt.figure()
ims = []
 
for i in range(timesteps):
    image = q_sample_single_step(image, i)
    im = plt.imshow(transform_rev(image), animated=True)
    ims.append([im])
    
animate = animation.ArtistAnimation(
    fig, ims, interval=50, blit=True, repeat_delay=1000)
animate.save('chai-forward.gif')
plt.show()
扩散模型(一):DDPM 基本原理与 MegEngine 实现-MEGChai
可视化整个加噪过程

实际执行这段代码体验一下,你会发现逐步加噪过程的计算代价太大了,计算量随着 $T$ 的设定线性增长。

优化:直接计算加噪后图片

事实上,由于我们选择将加噪过程定义为 $q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) =
\mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t\mathbf{I})$ 并重参数化成 $\mathbf{x}_t = \sqrt{1 - \beta_t} \mathbf{x}_{t-1} + \sqrt{\beta_t} \cdot \mathbf{z}_{t-1}$ 这种形式,在数学上具有一定的好处 —— 我们可以直接根据 $\mathbf{x}_0$ 和 $\beta$ 计算出任意时刻的 $q(\mathbf{x}_t)$.

为了方便推导,我们令 $\alpha_{t}=1-\beta_{t} $ 以及 $\bar{\alpha}_{t}=\prod_{i=1}^{T} \alpha_{i}$, 则有:

$$
\begin{aligned}
\mathbf{x}_{t} &=\sqrt{\alpha_{t}} \mathbf{x}_{t-1}+\sqrt{1-\alpha_{t}} \mathbf{z}_{t-1} \\
&= \sqrt{\alpha_{t}} \left( \sqrt{\alpha_{t-1}} \mathbf{x}_{t-2}+\sqrt{1-\alpha_{t-1}} \mathbf{z}_{t-2} \right) +\sqrt{1-\alpha_{t}} \mathbf{z}_{t-1} \\
&=\sqrt{\alpha_{t} \alpha_{t-1}} \mathbf{x}_{t-2} + {\color{blue} \left( \sqrt{\alpha_{t} -\alpha_{t} \alpha_{t-1}} \mathbf{z}_{t-2} + \sqrt{1-\alpha_{t}} \mathbf{z}_{t-1} \right) } \\
&=\sqrt{\alpha_{t} \alpha_{t-1}} \mathbf{x}_{t-2}+ {\color{blue} \sqrt{1-\alpha_{t} \alpha_{t-1}} \overline{\mathbf{z}}_{t-2} } \\
&=\sqrt{\alpha_{t} \alpha_{t-1} \alpha_{t-2}} \mathbf{x}_{t-3}+\sqrt{1-\alpha_{t} \alpha_{t-1} \alpha_{t-2}} \overline{\mathbf{z}}_{t-3} \\
&=\ldots \\
&=\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}
\end{aligned}
$$

在 蓝色部分 推导中利用到了正态分布的叠加特性:对于两个正态分布 $X \sim \mathcal{N} \left(\boldsymbol{\mu}_1 , \boldsymbol{\sigma}_1^2 \right) $ 和 $Y \sim \mathcal{N} \left(\boldsymbol{\mu}_2 , \boldsymbol{\sigma}_2^2 \right) $ , 叠加后得到的 $aX+bY$ 分布满足 $aX+bY \sim \mathcal{N} \left( a \boldsymbol{\mu}_1 + b \boldsymbol{\mu}_2 , a \boldsymbol{\sigma}_1^2 + b \boldsymbol{\sigma}_2^2 \right)$. 而前文提到 $\mathbf{z} \sim \mathcal{N} (\mathbf{0}, \mathbf{I})$, 所以 $\left( \sqrt{\alpha_{t} -\alpha_{t} \alpha_{t-1}} \mathbf{z}_{t-2} + \sqrt{1-\alpha_{t}} \mathbf{z}_{t-1} \right)$ 合并后服从 $\mathcal{N} (\mathbf{0}, \sqrt{1-\alpha_{t} \alpha_{t-1}})$, 使用重整化技巧用 $\mathbf{0} + \overline{\mathbf{z}}_{t-2} \cdot \sqrt{1-\alpha_{t} \alpha_{t-1}} $ 等价,即 $\sqrt{1-\alpha_{t} \alpha_{t-1}} \overline{\mathbf{z}}_{t-2}$, 其中 $ \overline{\mathbf{z}}_{t-2} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$.

以此递推,最终有重整化形式的 $\mathbf{x}_{t} = \sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \mathbf{z} $ , 写成条件概率即:

$$
q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)=\mathcal{N}\left(\mathbf{x}_{t} ; \sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0},\left(1-\bar{\alpha}_{t}\right) \mathbf{I}\right)
$$

在设计 $\beta$ 时为了避免早期一次性加入太多噪声有 $0<\beta_{1}<\beta_{2}<\ldots<\beta_{T}<1$, 因此对应地 $1>\bar{\alpha}_{1}>\cdots>\bar{\alpha}_{T}>0$, 当设定的 $T$ 很大时,$\mathbf{x}_{t}$ 最终接近 $\mathbf{z}$, 而我们知道 $\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$. 佐证了加噪到最后将变成纯粹的 Gaussian 噪声。下面来实现一下:

1
2
3
4
5
6
7
8
9
10
11
# define alphas and alphas_cumprod
alphas = 1. - betas
alphas_cumprod = np.cumprod(alphas, axis=0)
 
def q_sample(x_start, t, noise=None):
    """q(x_t | x_0)"""
    if noise is None:
        noise = np.random.normal(0, 1, x_start.shape)
    mean = np.sqrt(alphas_cumprod[t]) * x_start
    var = (1. - np.sqrt(alphas_cumprod[t]))
    return mean + np.sqrt(var) * noise

按照公式写出 $q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)$ 的重参数化实现,注意到我们其实已经提前计算好了 betas, alphas, alphas_cumprod 与 $\beta$, $\alpha$, $\bar{\alpha}$ 对应,这些值会被多次用到,因此提前计算好会比较方便。实际上,在后续的计算中还有许多中间计算结果会被重复使用,如 $\sqrt{\bar{\alpha}_{t}}$, $\sqrt{1-\bar{\alpha}_{t}}$ 等, 为了避免重复计算,这些值也可以提前计算好,届时使用 $t$ 作为下标去直接索引。

下面我们提前计算好这些值,给出新的 q_sample 实现:

1
2
3
4
5
6
7
8
9
10
# calculations for diffusion q(x_t | x_{t-1}) and others
sqrt_alphas_cumprod = np.sqrt(alphas_cumprod)
sqrt_one_minus_alphas_cumprod = np.sqrt(1. - np.sqrt(alphas_cumprod))
 
def q_sample(x_start, t, noise=None):
    if noise is None:
        noise = np.random.normal(0, 1, x_start.shape)
    mean = sqrt_alphas_cumprod[t] * x_start
    std = sqrt_one_minus_alphas_cumprod[t]
    return mean + std * noise

这样我们就能够在需要使用到对应值时,直接从存储好的变量中去取,从而减少计算量。在后面计算其他公式时,我们还会见到类似的操作。

另外,我们这里的演示仅用于单张图片,对于 Batch 输入,我们会随机生成 $N$ 个 $t \in [0, T)$ 作为每个样本的 [t] 索引,为了使 xxx[t] 能与 $(N, C, H, W)$ 的 noise 之间能够正常地进行 Broadcasting 与 Element-wise 运算,还需要做一些额外的对 Shape 的处理,届时再进行介绍。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
image = plt.imread("chai.jpg")
image = transform(image)
 
noisy_image = q_sample(image, timesteps-1)
plt.imshow(transform_rev(noisy_image))
plt.show()
 
# Visualization
fig = plt.figure(dpi=100, constrained_layout=True)
 
idx_sampled = [50, 100, 200, 500, 1000]
num_sampled = len(idx_sampled)
 
for i in range(num_sampled):
    plt.subplot(1, num_sampled, i+1)
    plt.axis("off")
    noisy_image = q_sample(image, idx_sampled[i]-1)
    plt.imshow(transform_rev(noisy_image))
    
plt.show()

使用最新实现的 q_sample 可以一步到位求得 $\mathbf{x}_{t}$, 如下图所示:

扩散模型(一):DDPM 基本原理与 MegEngine 实现-MEGChai
其加噪效果应当与逐步加噪等价
扩散模型(一):DDPM 基本原理与 MegEngine 实现-MEGChai
举例:根据 idx_sampled 设定的 $t$ 分别一步到位加噪,与逐步加噪得到的结果等价

逆向扩散过程

如果我们能够知道 $q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)$ , 那样实现逆向过程就很简单 —— 只需采样出一些随机的高斯分布噪声 $\mathbf{x}_{T} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ 作为输入,去逐步进行 “去噪”,最终就能得到从真实分布 $q(\mathbf{x}_{0})$ 的采样。但 $q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)$ 需要知道整个数据分布,所以该想法不太现实。我们不妨尝试使用一个神经网络模型来近似(学习)该后验概率分布。让我们把该去噪模型称作 $p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)$, 其中 $p_{\theta}$ 是神经网络中的参数,通过梯度下降算法进行更新。

我们假设逆向过程也是一个高斯分布(这是一个强假设,需要 $\beta$ 足够小),那么回忆一下,任何高斯分布都由两个参数定义:

  • 由 $\boldsymbol{\mu}_{\theta}$ 参数定义的均值;
  • 由 $\boldsymbol{\Sigma}_{\theta}$ 参数定义的方差。

那么这个过程就可以参数化为如下形式:

$$
p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)=\mathcal{N}\left(\mathbf{x}_{t-1} ; \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right), \boldsymbol{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)\right)
$$

其中均值和方差也是基于噪声等级 $t$ 来决定的。因此我们的神经网络需要根据条件输入 $\mathbf{x}_{t}$ 和 $t$ 来学习(表示)对应的均值和方差,但在 DDPM 的原始论文中,作者选择了固定方差 $\boldsymbol{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)$ 为只与 $t$ 有关的常数,只要求神经网络去学习均值。理由是经过试验,固定与不固定方差最终的预测效果差不多。在后续的研究论文中尝试对此进行改进,其中神经网络除了学习均值之外,还学习了逆向过程的方差。在这里我们选择和原始 DDPM 论文做法一致,仅要求模型学得后验概率分布的均值,固定方差。怎么个固定法呢?一种选择是直接令 $\boldsymbol{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)=\beta \mathbf{I}$ ,其实还有别的选择,接着往下看。

推导:后验概率的均值与方差

注意到,虽然 $q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)$ 无法给出,但当我们加入 $\mathbf{x}_{0}$ 作为条件时,则整个情况变得可处理了:

$$
q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)=\mathcal{N}\left(\mathbf{x}_{t-1} ; \tilde{\boldsymbol{\mu}}\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right), \tilde{\beta}_{t} \mathbf{I}\right)
$$

根据贝叶斯公式,我们有:

$$
\begin{aligned}
q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)
&= \frac {q (\mathbf{x}_{t-1}, \mathbf{x}_{t}, \mathbf{x}_{0} )} { q(\mathbf{x}_{t}, \mathbf{x}_{0}) } \\
&= \frac {q (\mathbf{x}_{t} \mid \mathbf{x}_{t-1}, \mathbf{x}_{0}) q (\mathbf{x}_{t-1} \mid \mathbf{x}_{0}) {\color{green} q (\mathbf{x}_{0}) } } { {\color{green} q(\mathbf{x}_{t}, \mathbf{x}_{0}) } } \\
&=q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}, \mathbf{x}_{0}\right) \frac{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{0}\right)} {{\color{green}q(\mathbf{x}_{t}, \mathbf{x}_{0}) / q (\mathbf{x}_{0}) }} \\
&={\color{blue} q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}, \mathbf{x}_{0}\right)} \frac{ {\color{red} q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{0}\right) }}{{\color{green} q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)}}
\end{aligned}
$$

注意目前这个式子,左边的 ${\color{blue} q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}, \mathbf{x}_{0}\right) }$ 根据马尔可夫性质可以将不造成影响的 $\mathbf{x}_{0}$ 拿掉,变成之前计算过的 $q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) =
\mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t} \mathbf{x}_{t-1}, \beta_t\mathbf{I})$. 回忆一下,对于任意正态分布 $x \sim \mathcal{N}(\mu, \sigma)$, 它的概率密度函数是 $f(x) = \frac{1}{\sqrt{2\pi}\sigma} \exp \left ( - \frac {(x -\mu)^2} {2\sigma^2} \right) \propto \exp -\frac{1}{2} \left( \frac {(x -\mu)^2}{\sigma^2} \right)$. 因此有 $q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}, \mathbf{x}_{0}\right) \propto \exp -\frac{1}{2} \left( \frac {(\mathbf{x}_t -\sqrt{\alpha_{t}} \mathbf{x}_{t-1})^2}{\beta_{t}} \right)$.

同理,对于式子右边的 ${\color{red} q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{0} \right)}$ 和 ${\color{green} q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)}$, 可代入之前计算过的 $q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)=\mathcal{N}\left(\mathbf{x}_{t} ; \sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0},\left(1-\bar{\alpha}_{t}\right) \mathbf{I}\right)$. 分别有 $q \left( \mathbf{x}_{t-1} \mid \mathbf{x}_{0} \right) \propto \exp -\frac{1}{2} \left( \frac{\left(\mathbf{x}_{t-1}-\sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_{0}\right)^{2}}{1-\bar{\alpha}_{t-1}} \right)$ 以及 $\frac {1} {q \left( \mathbf{x}_{t} \mid \mathbf{x}_{0} \right)} \propto \exp \frac{1}{2} \left( \frac{\left(\mathbf{x}_{t}-\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}\right)^{2}}{1-\bar{\alpha}_{t}} \right)$

因此可以得到:

$$
\begin{aligned}
q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)
& \propto \exp \left( -\frac{1}{2} \left(
\frac {(\mathbf{x}_t -\sqrt{\alpha_{t}} \mathbf{x}_{t-1})^2}{\beta_{t}} +
\frac{\left(\mathbf{x}_{t-1}-\sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_{0}\right)^{2}}{1-\bar{\alpha}_{t-1}} -
\frac{\left(\mathbf{x}_{t}-\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}\right)^{2}}{1-\bar{\alpha}_{t}}
\right) \right) \\
&= \exp \left( -\frac{1}{2} \left(
\frac {\mathbf{x}_t^2 -2\sqrt{\alpha_{t}} \mathbf{x}_t {\color{blue} \mathbf{x}_{t-1}} + \alpha_{t} {\color{red} \mathbf{x}_{t-1}^2}} {\beta_{t}} +
\frac {{\color{red} \mathbf{x}_{t-1}^2} - 2 \sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_{0} {\color{blue} \mathbf{x}_{t-1}} +\bar{\alpha}_{t-1} \mathbf{x}_{0}^2} {1-\bar{\alpha}_{t-1}} - \ldots
\right) \right) \\
&= \exp \left(-\frac{1}{2}\left( {\color{red} \left(\frac{\alpha_{t}}{\beta_{t}}+\frac{1}{1-\bar{\alpha}_{t-1}}\right) } \mathbf{x}_{t-1}^{2}- {\color{blue} \left(\frac{2 \sqrt{\alpha_{t}}}{\beta_{t}} \mathbf{x}_{t}+\frac{2 \sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_{t-1}} \mathbf{x}_{0}\right) } \mathbf{x}_{t-1}+C\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right)\right)\right)
\end{aligned}
$$

上式被整理成关于 $ \mathbf{x}_{t-1}$ 的二次方程形式,其中 $C\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right)$ 不包含有任何与 $ \mathbf{x}_{t-1}$ 有关的项,因此细节于此忽略。

🤔 有趣的是:你可能在很多地方看到了错误的推导细节(点击展开)
扩散模型(一):DDPM 基本原理与 MegEngine 实现-MEGChai
Lilian Weng 未做修正时的推导步骤
在 Lilian Weng 介绍 Diffusion Model 的文章中,在很长一段时间内,上图中红框区存在下标错误,这篇文章的公式被许多人引用,却迟迟没有人指出其中的问题。我发邮件向她询问后,她表示:“Those typos happened when I translated stuff from paper to latex. Just fixed.” 至此才没有问题。

已知对于概率密度函数为 $f(x)=ax^2+bx+c=a(x+\frac{b}{2a})^2+c$ 的标准高斯分布有 $\mu=-\frac{b}{2a}$ 和 $\Sigma=\frac{1}{a}$. 对应上面整理得到的公式有 $a =\left(\frac{\alpha_{t}}{\beta_{t}}+\frac{1}{1-\bar{\alpha}_{t-1}}\right) $ 和 $b=-\left(\frac{2 \sqrt{\alpha_{t}}}{\beta_{t}} \mathbf{x}_{t}+\frac{2 \sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_{t-1}} \mathbf{x}_{0}\right)$. 因此可以计算得到 $q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)=\mathcal{N}\left(\mathbf{x}_{t-1} ; \tilde{\boldsymbol{\mu}}\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right), \tilde{\beta}_{t} \mathbf{I}\right)$ 的均值和方差:

$$
\begin{aligned}
\tilde{\beta}_{t} &=1 /\left(\frac{\alpha_{t}}{\beta_{t}}+\frac{1}{1-\bar{\alpha}_{t-1}}\right) \\
&= 1 / \frac { \alpha_{t}- {\color{blue} \alpha_{t} \bar{\alpha}_{t-1} } + {\color{red} \beta_{t} } } {(1-\bar{\alpha}_{t-1} ) \beta_{t}} \\
&= 1 / \frac { \alpha_{t}- {\color{blue} \bar{\alpha}_{t}} + {\color{red} 1 - \alpha_{t}} } {(1-\bar{\alpha}_{t-1} ) \beta_{t}} \\
&= \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}} \cdot \beta_{t}
\end{aligned}
$$

注意到这里的 $\tilde{\beta}_{t}$ 是一个仅仅与 $t$ 有关的常数,因此我们的去噪模型 $p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)$ 中的 $\boldsymbol{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)$ 也可以被固定成与 $\tilde{\beta}_{t} \mathbf{I}$ 一致。

同样地,这些频繁被用到的数值,可提前被计算好:

1
2
3
4
5
6
7
8
# calculations for posterior q(x_{t-1} | x_t, x_0)
sqrt_recip_alphas = np.sqrt(1. / alphas)
alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
posterior_variance = \
    betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
log_posterior_variance = np.log(  # posterior variance is 0 at beginning
    np.append(posterior_variance[1], posterior_variance[1:])
)

计算 alphas_cumprod_prev (对应 $\bar{\alpha}_{t-1}$ ) 时,可由前面已经计算好的 alphas_cumprod (对应 $\bar{\alpha}_{t}$ )取 [:-1] 后经过 append, pad 或 concatenate 等处理(无论用哪个接口,效果是一样的)得到,为了在 $t=0$ 时 $\bar{\alpha}_{t-1}$ 能够有对应值,在开头补上值为 1 的项(因为有 $1>\bar{\alpha}_{1}>\ldots >\bar{\alpha}_{t} >0$ )。

求得方差的化简形式后,方便我们继续计算 $q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)$ 的均值:

$$
\begin{aligned}
\tilde{\boldsymbol{\mu}}_{t}\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right) &=
\left(\frac{\sqrt{\alpha_{t}}}{\beta_{t}} \mathbf{x}_{t}+\frac{\sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_{t-1}} \mathbf{x}_{0}\right) /\left(\frac{\alpha_{t}}{\beta_{t}}+\frac{1}{1-\bar{\alpha}_{t-1}}\right) \\
&= \left(\frac{\sqrt{\alpha_{t}}}{\beta_{t}} \mathbf{x}_{t}+\frac{\sqrt{\bar{\alpha}_{t-1}}}{1-\bar{\alpha}_{t-1}} \mathbf{x}_{0}\right) \cdot \frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}} \cdot \beta_{t} \\
&= \frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_{t}} \mathbf{x}_{t}+\frac{\sqrt{\bar{\alpha}_{t-1}} \beta_{t}}{1-\bar{\alpha}_{t}} {\color{blue} \mathbf{x}_{0} }
\end{aligned}
$$

前文在计算 $q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)$ 时借助参数重整化我们有 $\mathbf{x}_{t} = \sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}$, 可被整理成 ${\color{blue} \mathbf{x}_{0}=\frac{1}{\sqrt{\bar{\alpha}_t}}\left(\mathbf{x}_{t}-\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}_{t}\right) }$, 代入上式化简:

$$
\begin{aligned}
\tilde{\boldsymbol{\mu}}_{t} \left(\mathbf{x}_{t}, \mathbf{z}_{t}\right) &=
\frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_{t}} \mathbf{x}_{t}+\frac{\sqrt{\bar{\alpha}_{t-1}} \beta_{t}}{1-\bar{\alpha}_{t}} {\color{blue} \frac{1}{\sqrt{\bar{\alpha}_t}} \left(\mathbf{x}_{t}-\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}_{t}\right)} \\
& = \left( \frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_{t}} + \frac{ {\color{red} \sqrt{\bar{\alpha}_{t-1}} } \beta_{t}} {1-\bar{\alpha}_{t}} {\color{red} \frac{1}{\sqrt{\bar{\alpha}_t}}} \right) \mathbf{x}_{t} - \frac {\sqrt{\bar{\alpha}_{t-1}} \beta_{t} \sqrt{1-\bar{\alpha}_{t}} } {(1-\bar{\alpha}_{t}) \sqrt{\bar{\alpha}_t}} \mathbf{z}_{t} \\
& = \left( \frac { \sqrt{\alpha_{t}} - \frac {\bar{\alpha}_{t}} {\sqrt{\alpha_{t}}} +
{\color{red} \frac{1}{\sqrt{\alpha_{t}}} } (1-\alpha_{t}) } {1-\bar{\alpha}_{t}} \right) \mathbf{x}_{t} - \frac {\sqrt{\bar{\alpha}_{t-1}} \beta_{t} } {\sqrt{1-\bar{\alpha}_{t}} \sqrt{\bar{\alpha}_t}} \mathbf{z}_{t} \\
&= \left( \frac{ \frac{1}{\sqrt{\alpha_{t}}} (1-\bar{\alpha}_{t}) + \sqrt{\alpha_{t}} - \sqrt{\alpha_{t}}} {1-\bar{\alpha}_{t}} \right) \mathbf{x}_{t} - \frac { \beta_{t} } {\sqrt{1-\bar{\alpha}_{t}} \sqrt{\alpha_t}} \mathbf{z}_{t} \\
&= \frac {1}{\sqrt{\alpha_t}} \mathbf{x}_{t} - \frac {1}{\sqrt{\alpha_t}} \frac { \beta_{t}} {\sqrt{1-\bar{\alpha}_{t}}} \mathbf{z}_{t} \\
&= \frac {1}{\sqrt{\alpha_t}} \left( \mathbf{x}_{t} - \frac { \beta_{t}} {\sqrt{1-\bar{\alpha}_{t}}} \mathbf{z}_{t} \right)
\end{aligned}
$$

如果我们尝试让模型 $\mathbf{z}_{\theta}(\mathbf{x}_{t}, t)$ 去学习预测噪声 $ \mathbf{z}_{t}$, 则可以得到从模型预测的噪声计算出 $p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)$ 均值的公式:

$$
{\color{blue} \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right) = \frac {1}{\sqrt{\alpha_t}} \left( \mathbf{x}_{t} - \frac { \beta_{t}} {\sqrt{1-\bar{\alpha}_{t}}} \mathbf{z}_{\theta}(\mathbf{x}_{t}, t) \right) }
$$

上面提到模型方差已被固定,有了去噪模型的均值和方差后,使用重参数化技巧,我们得到了单步去噪的计算过程:

$$
p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)=\mathcal{N}\left(\mathbf{x}_{t-1} ; \frac {1}{\sqrt{\alpha_t}} \left( \mathbf{x}_{t} - \frac { \beta_{t}} {\sqrt{1-\bar{\alpha}_{t}}} \mathbf{z}_{\theta}(\mathbf{x}_{t}, t) \right), \tilde{\beta}_{t} \mathbf{I} \right)
$$

$$
\mathbf{x}_{t-1} = \frac {1}{\sqrt{\alpha_t}} \left( \mathbf{x}_{t} - \frac { \beta_{t}} {\sqrt{1-\bar{\alpha}_{t}}} \mathbf{z}_{\theta}(\mathbf{x}_{t}, t) \right) + \sqrt{\tilde{\beta}_{t} } \cdot \mathbf{z}
$$

值得注意的一点是,论文官方的实际代码实现中并没有直接使用化简后得到的公式从预测的噪声计算上一时刻的状态,而是选择了先让模型输出 $\mathbf{z}_{\theta}(\mathbf{x}_{t}, t)$ 作为 $\mathbf{z}_{t}$; 接着计算预测的原始数据 $ \mathbf{x}_{0}=\frac{1}{\sqrt{\bar{\alpha}_t}}\left(\mathbf{x}_{t}-\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}_{t}\right) $, 再根据 ${\boldsymbol{\mu}}_{\theta}\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right) = \frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_{t}} \mathbf{x}_{t}+\frac{\sqrt{\bar{\alpha}_{t-1}} \beta_{t}}{1-\bar{\alpha}_{t}} \mathbf{x}_{0}$ 计算得到 $\mathbf{x}_{t-1}$ 的均值,实际上两者的计算结果是理论等价的。对此笔者的解释是,作者在进行对比实验时,分别比对了让预测输出(模型均值)为 $\mathbf{x}_{t-1}$, $\epsilon$ 和直接输出 $\mathbf{x}_{0}$ 三种设置下的效果,为了复用从 $\mathbf{x}_{0}$ 计算得到 $\mathbf{x}_{t-1}$ 的代码(以及 $\mathbf{x}_{0}$ 可能在其它实验中需要被用到),故选择了从 $\mathbf{z}_{\theta}(\mathbf{x}_{t}, t)$ 到 $\mathbf{x}_{0}$ 再到 $\mathbf{x}_{t-1}$ 的计算路线,而没有套用如上论文中所给出算法的 $\mathbf{z}_{\theta}(\mathbf{x}_{t}, t)$ 直接到 $\mathbf{x}_{t-1}$ 的计算公式。

另外由于原始的 $\mathbf{x}_{0}$ 输入经过变换范围到 $[-1, 1]$, 因此对于预测得到的 $\mathbf{x}_{0}$ 通常会采取 clip[-1, 1] 的截断计算确保此分布特征一致,而对中间过程得到的去噪样本不会进行此操作。下面用代码演示一下这两种不同的计算过程——

代码实现:逐步去噪

假定我们已经有一个训练好的理想去噪模型 $\mathbf{z}_{\theta}(\mathbf{x}_{t}, t)$, 并且已知模型输出的是预测的噪声 $\mathbf{z}_{t}$, 那么扩散模型的逆向过程实现如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def p_sample(model, x, t):
    """Sample from p_{theta} (x_{t-1} | x_t)"""
    noise = np.full_like(x, 0) \
        if t == 0 else np.random.normal(0, 1, x.shape)
    mean = sqrt_recip_alphas[t] * \
        (x - betas[t] * model(x, t) /\
        sqrt_one_minus_alphas_cumprod[t])
    var = posterior_variance[t]
    return mean + np.sqrt(var) * noise
 
def p_sample_loop(model, shape):
    """Sample from p_{theta} (x_{0:t-1} | x_t)"""
    x = np.random.normal(0, 1, shape)
    for i in reversed(range(timesteps)):
        x = p_sample(model, x, i)
    return x

单步的去噪 p_sample 计算步骤为:先根据 model(x, t) 得到模型预测的噪声 $\mathbf{z}_{\theta}(\mathbf{x}_{t}, t)$, 再根据上文提到的公式直接计算得到 ${\color{blue} \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right) = \frac {1}{\sqrt{\alpha_t}} \left( \mathbf{x}_{t} - \frac { \beta_{t}} {\sqrt{1-\bar{\alpha}_{t}}} \mathbf{z}_{\theta}(\mathbf{x}_{t}, t) \right) }$, 由于 $\boldsymbol{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right) = \tilde{\beta}_{t} \mathbf{I}$ 是固定的,因此接下来只需要随机生成 $\boldsymbol{\epsilon}$, 并借助重整化公式完成 $p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)$ 采样过程的等价计算。

整个逆向过程 p_sample_loop 即从一个纯噪声 $p_{\theta} (\mathbf{x}_{T})$ 逐步去噪(重建)到 $p_{\theta} (\mathbf{x}_{0})$ 的迭代。理想情况下,我们的模型应当拥有这种从噪声分布扩散到原始数据分布 $q (\mathbf{x}_{0})$ 的能力,可看作是一种无条件的生成式模型。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
# calculations for predict x_0 from noise
sqrt_recip_alphas_cumprod = np.sqrt(1. / alphas_cumprod)
sqrt_recipm1_alphas_cumprod = np.sqrt(1. / alphas_cumprod - 1)
 
# calculations for posterior q(x_{t-1} | x_t, x_0)
posterior_mean_coef1 = (
    betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod)
)
posterior_mean_coef2 = (
    (1. - alphas_cumprod_prev) * np.sqrt(alphas) /
    (1. - alphas_cumprod)
)
 
def p_sample(model, x, t, clip_denoised=False):
    """Sample from p_{theta} (x_{t-1} | x_t)"""
    noise = np.full_like(x, 0) \
        if t == 0 else np.random.normal(0, 1, x.shape)
 
    # predict added noise from trained model, given x_t and t
    model_output = self.model(x, t)  
 
    # predict x_0 from t level noise
    predict_x_start = (
        sqrt_recip_alphas_cumprod[t] * x -
        sqrt_recipm1_alphas_cumprod[t] * model_output
    )
 
    if clip_denoised:
        predict_x_start = np.clip(predict_x_start, -1., 1.)
 
    model_mean = (
        posterior_mean_coef1[t] * predict_x_start +
        posterior_mean_coef2[t] * x
    )
 
    model_var = posterior_variance[t]  # we can take the log ---->
 
    return mean + np.sqrt(var) * noise

前面提到,在 DDPM 论文原作者提供的代码实现中,并没有使用上述的计算方式,而是先根据 model(x, t) 得到模型预测的噪声 $\mathbf{z}_{\theta}(\mathbf{x}_{t}, t) = \mathbf{z}_{t}$,

接着根据预测的噪声来计算预测的原始数据 $\mathbf{x}_{0}=\frac{1}{\sqrt{\bar{\alpha}_t}}\left(\mathbf{x}_{t}-\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}_{t}\right) = \sqrt{ \frac{1}{\bar{\alpha}_t}} \mathbf{x}_{t} - \sqrt{ \frac{1}{\bar{\alpha}_t} -1 } \mathbf{z}_{t} $, 再根据 ${\boldsymbol{\mu}}_{\theta}\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right) = \frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_{t}} \mathbf{x}_{t}+\frac{\sqrt{\bar{\alpha}_{t-1}} \beta_{t}}{1-\bar{\alpha}_{t}} \mathbf{x}_{0}$ 计算得到均值。

同样地,式子中的各个常数系数可以被提前计算好,对应地有 posterior_mean_coef1 代表 $\frac{\sqrt{\bar{\alpha}_{t-1}} \beta_{t}}{1-\bar{\alpha}_{t}}$, posterior_mean_coef2 代表 $\frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_{t}} $. 理清楚各种变量的计算关系后,就不难看懂原始论文中的代码了。至于究竟是直接从 $\mathbf{z}_{\theta}(\mathbf{x}_{t}, t)$ 算出 $\boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right)$, 还是先算出模型预测的 $\mathbf{x}_{0}$ 作为中间结果,笔者认为并不关键,两种做法数学上等价。

对于计算得到的 $\mathbf{x}_{0}$, 可以根据需要选择是否进行值的裁剪,即设置 clip_denoised , 使范围限定在 $[-1, 1]$, 满足我们对于 $\mathbf{x}$ 数据分布的假设。

前面提到过,取对数操作可以保持数值稳定,但要避免第一项为 0:

1
2
3
4
5
6
7
8
if model_var_type == "FIXED_LARGE":
   logvar = np.log(np.append(posterior_variance[1], betas[1:]))
elif model_var_type == "FIXED_SMALL"":
   logvar = np.log(np.maximum(posterior_variance, 1e-20))
 
def p_sample(...):
    # ...
    return model_mean + np.exp(0.5 * logvar[t]) * noise

取对数操作可以让数值范围平滑,使得计算更加稳定,因此这里对模型的方差取对数并不稀奇。注意这里对模型方差的类型 model_var_type 进行了区分, FIXED_LARGE 即固定方差为前向加噪过程中所用的方差 $\beta$, 而 FIXED_SMALL 即固定方差为我们根据后验 $q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)$ 计算得到的方差 $\tilde{\beta}$, DDPM 原论文中认为这两个值分别是固定方差参考值的上下界。在后续的论文中提出可以对方差进行学习,因此还会看到 model_var_type 为 LEARNED 等类似的情况,提供不同的配置项方便实验进行对比。而对于模型的均值类型 model_mean_type 作者在实验时也区分成 PREVIOUS_X, START_X, EPSILON三种,分别代表着模型预测 $\mathbf{x}_{t-1}$, $\mathbf{x}_{0}$ 和 $\epsilon$, 这与本文一开始提到的做法一致,最终作者通过实验发现:让模型预测噪声 $\epsilon$ 最终得到的效果最好。在复现这篇论文的最佳实验结果时,可以直接使用 var_type == FIXED_SMALL与 mean_type == EPSILON 的逻辑,使得代码逻辑更加简洁。

更新:为了避免读者在阅读源码时感到疑惑,这里提前预告一下相关源码可能的形式(单样本情况下的伪代码)——

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# handle with model_output according to the variance type (fixed or learned)
if self.model_var_type == "FIXED_SMALL":
    model_log_var = log_posterior_variance[t]
elif self.model_var_type == "FIXED_LARGE":
    # set the initial (log-)variance to get a better decoder log likelihood.
    model_log_var = np.append(log_posterior_variance[1], F.log(betas[1:]))[t]
else:  # model's output contains learned variance value (the 2nd 3 channels)
    model_output, model_var_values = F.split(model_output, 2, axis=1)
    if self.model_var_type == "LEARNED":  # learned variance directly
        model_log_var = model_var_values
    elif self.model_var_type == "LEARNED_RANGE":  # IDDPM Eq. (15)
        min_log = log_posterior_variance[t]
        max_log = F.log(betas[t])
        # The model_var_values is [-1, 1] and should convert to [0, 1] as coff.
        frac = (model_var_values + 1) / 2
        model_log_var = frac * max_log + (1 - frac) * min_log

我们刚才提到了可以固定方差为 $\beta \mathbf{I}$, 即上述 FIXED_LARGE 的情况(注释中解释了对第一个值进行了修改以取得更好的效果),下面即将介绍的是固定方差与后验概率方差 posterior_variance 一致的情况,即 FIXED_SMALL. 而在后面的论文中,提出了方差是可由模型学习的,即模型的输出一部分是预测值,一部分是对应的方差,需要用 split 做分割。其中 LEARNED 指模型直接输出方差值,而 LEARNED_RANGE 指模型输出一个系数,在两个固定值之间加权求和。学习方差已经是后话了,本文不再赘述,我们接下来需要搞清楚,如何设计目标函数,使得模型能够对期望输出进行预测。

目标分布的似然函数

我们的优化目标是使得经由 $p_{\theta}$ 逆扩散过程得到的数据分布尽可能与 $q(\mathrm{x}_{0})$ 一致,可考虑最小化其负对数似然 $-\log p_{\theta}\left(\mathbf{x}_{0}\right)$, 但它无法被直接计算。

借鉴 VAE 模型中用到的变分下界(Variational Lower Bound, VLB)处理思路,在负对数似然 $-\log p_{\theta}\left(\mathbf{x}_{0}\right)$ 的基础上加上一个相关的 KL 散度,能够构成负对数似然的上界,上界越小,负对数似然也就越小,等同于目标分布的似然最大。我们希望使用 $p_{\theta}$ 去近似 $q$, 因此有:

$$
\begin{aligned}
-\log p_{\theta}\left(\mathbf{x}_{0}\right) &\leq
-\log p_{\theta}\left(\mathbf{x}_{0}\right)+D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right) \| p_{\theta}\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)\right) \\
&=-\log p_{\theta}\left(\mathbf{x}_{0}\right)+\mathbb{E}_{\mathbf{x}_{1: T} \sim q\left(\mathbf{x}_{1: T \mid} \mathbf{x}_{0}\right)}\left[\log \frac{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0: T}\right) / p_{\theta}\left(\mathbf{x}_{0}\right)}\right] \\
&=-\log p_{\theta}\left(\mathbf{x}_{0}\right)+\mathbb{E}_{q}\left[\log \frac{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0: T}\right)}+\log p_{\theta}\left(\mathbf{x}_{0}\right)\right] \\
&=\mathbb{E}{q}\left[\log \frac{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0: T}\right)}\right]
\end{aligned}
$$

上式两边乘上 $-\mathbb{E}_{q\left(\mathbf{x}_{0}\right)}$ 可得到交叉熵形式上界,即:

$$
L_{\mathrm{VLB}}=\mathbb{E}_{q\left(\mathbf{x}_{0: T)}\right)}\left[\log \frac{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0: T}\right)}\right] \geq-\mathbb{E}_{q\left(\mathbf{x}_{0}\right)} \log p_{\theta}\left(\mathbf{x}_{0}\right)
$$

我们希望最小化交叉熵,可以通过最小化交叉熵的上界,即最小化目标函数 $L_{\mathrm{VLB}}$ 来实现,它是可优化的。实际上经过一系列推导,最后得到的需要被优化的目标函数十分简洁,因此对推导过程不感兴趣的可以跳到最后查看简化形式的目标函数。

推导:目标函数重写

为了能够实际地进行计算,需要对上面的 $L_{\mathrm{VLB}}$ 进行重写,希望最终变成已知公式的组合:

$$
\begin{aligned}
L_{\mathrm{VLB}} &=\mathbb{E}_{q\left(\mathbf{x}_{0: T}\right)} \left[ \log \frac{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)} {p_{\theta}\left(\mathbf{x}_{0: T}\right)}\right] \\
&=\mathbb{E}_{q}\left[\log \frac{\prod_{t=1}^{T} q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}\right)}{p_{\theta}\left(\mathbf{x}_{T}\right) \prod_{t=1}^{T} p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)}\right] \\
&= \mathbb{E}_{q}\left[ -\log p_{\theta}\left(\mathbf{x}_{T}\right) + \sum_{t=1}^{T} \log \frac{q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}\right)} { p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)}\right] \\
&= \mathbb{E}_{q}\left[ -\log p_{\theta}\left(\mathbf{x}_{T}\right) + \sum_{t=2}^{T} \log \frac{ {\color{blue} q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}\right) } } { p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)} + \log \frac {q\left(\mathbf{x}_{1} \mid \mathbf{x}_{0}\right)} {p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)}\right]
\end{aligned}
$$

由于马尔可夫性质有 $q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}\right) = q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}, \mathbf{x}_{0} \right)$, 与上文计算后验条件概率 $ q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0} \right)$ 的做法类似,根据贝叶斯公式有:

$$
\begin{aligned}
q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}\right) =
q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}, \mathbf{x}_{0}\right)
&= \frac {q (\mathbf{x}_{t-1}, \mathbf{x}_{t}, \mathbf{x}_{0} )} { q(\mathbf{x}_{t-1}, \mathbf{x}_{0}) } \\
&= \frac {q (\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}) q (\mathbf{x}_{t} \mid \mathbf{x}_{0}) q (\mathbf{x}_{0}) } { q(\mathbf{x}_{t-1}, \mathbf{x}_{0}) } \\
&=q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right) \frac{q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)} {q(\mathbf{x}_{t-1}, \mathbf{x}_{0}) / q (\mathbf{x}_{0})} \\
&=q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right) \frac{q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{0}\right)}
\end{aligned}
$$

代入 $L_{\mathrm{VLB}}$ 式有:

$$
\begin{aligned}
L_{\mathrm{VLB}} &= \mathbb{E}_{q}\left[ -\log p_{\theta}\left(\mathbf{x}_{T}\right) + \sum_{t=2}^{T} \log \left( \frac{ {\color{blue} q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right) } } { p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)} \cdot {\color{blue} \frac{q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{0}\right)} } \right) + \log \frac {q\left(\mathbf{x}_{1} \mid \mathbf{x}_{0}\right)} {p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)}\right] \\
& = \mathbb{E}_{q}\left[ -\log p_{\theta}\left(\mathbf{x}_{T}\right) + \sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)} { p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)} + \sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{0}\right)} + \log \frac {q\left(\mathbf{x}_{1} \mid \mathbf{x}_{0}\right)} {p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)}\right]
\end{aligned}
$$

注意到上式中 $\sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{0}\right)} = \log \frac{q\left(\mathbf{x}_{T} \mid \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{T-1} \mid \mathbf{x}_{0}\right)} + \log \frac{q\left(\mathbf{x}_{T-1} \mid \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{T-2} \mid \mathbf{x}_{0}\right)} + \ldots + \log \frac{q\left(\mathbf{x}_{2} \mid \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{1} \mid \mathbf{x}_{0}\right)} = \log \frac{q\left(\mathbf{x}_{T} \mid \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{1} \mid \mathbf{x}_{0}\right)}$, 因此:

$$
\begin{aligned}
L_{\mathrm{VLB}} &=
\mathbb{E}_{q}\left[ -\log p_{\theta}\left(\mathbf{x}_{T}\right) + \sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)} { p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)} + \log \frac{q\left(\mathbf{x}_{T} \mid \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{1} \mid \mathbf{x}_{0}\right)} + \log \frac {q\left(\mathbf{x}_{1} \mid \mathbf{x}_{0}\right)} {p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)}\right] \\
& = \mathbb{E}_{q}\left[ -\log p_{\theta}\left(\mathbf{x}_{T}\right) + \sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)} { p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)} + \log q\left(\mathbf{x}_{T} \mid \mathbf{x}_{0}\right) -\log p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right) \right] \\
&= \mathbb{E}_{q}\left[ \log \frac{q\left(\mathbf{x}_{T} \mid \mathbf{x}_{0}\right)} {p_{\theta}\left(\mathbf{x}_{T}\right)} + \sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)} { p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)} -\log p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right) \right] \\
& = \mathbb{E}_{q} \left[ \underbrace {D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{T} \mid \mathbf{x}_{0}\right) \| p_{\theta}\left(\mathbf{x}_{T}\right) \right) }_{L_{T}} + \sum_{t=2}^{T} \underbrace{D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right) \| p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)\right)}_{L_{t-1}} \underbrace{ - \log p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)}_{L_{0}} \right]
\end{aligned}
$$

注意到对于 $L_{t-1}$ 式,当 $t=1$ 时有 $D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}, \mathbf{x}_{0}\right) \| p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)\right) = D_{\mathrm{KL}}\left(q\left( 1 \right) \| p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)\right) = - \log p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)$, 因此 $L_{0}$ 式在计算时可并入 $L_{t-1}$ 式情况(至于如何对 $- \log p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)$ 进行计算,DDPM 的代码实现中采用了离散化的分段积分累乘函数 discretized_gaussian_log_likelihood 来进行建模,为了不影响现在的行文逻辑,暂且不提相关细节)。对于 $L_{T}$ 式,$q$ 中不含有可训练的参数,而 $p_{\theta}\left(\mathbf{x}_{T}\right)$ 是一个各向同性的高斯分布,没有未知参数,因此该项无须优化,可以省略。对于将 $L_{0}$ 并入 $L_{t-1}$ 式的处理思路,得到的计算式为:

$$
\mathbb{E}_{q} \left[ \sum_{t=1}^{T} D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right) \| p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)\right) \right]
$$

推导:简化目标函数

我们知道:

$$
D_{\mathrm{KL}} \left(p \| q \right) = \log \frac {\sigma_{2}} {\sigma_{1}} + \frac {\sigma_{1}^2 + (\mu_{1} - \mu_{2})^2}{2\sigma_{2}^2} - \frac {1}{2}
$$

这里折叠了一些推导细节,可点击展开。
回忆一下,对于两个单一变量的高斯分布 $p \sim \mathcal{N}(\mu_{1}, \sigma_{1}^2)$ 和 $q \sim \mathcal{N}(\mu_{2}, \sigma_{2}^2)$, 概率密度公式分别为:

$$
p(x) = \frac{1}{\sqrt{2\pi}\sigma_{1}} \exp \left ( - \frac {(x -\mu_{1})^2} {2\sigma_{1}^2} \right) \quad
q(x) = \frac{1}{\sqrt{2\pi}\sigma_{2}} \exp \left ( - \frac {(x -\mu_{2})^2} {2\sigma_{2}^2} \right)
$$

下面的推导过程中需要用到几个典型结论,摆在这里,对任意的高斯分布有:

$$
\begin{aligned}
\mathbb{E}(x)&=\int_{-\infty}^{\infty} \mathcal{N} \left(x \mid \mu, \sigma^{2}\right) x d x=\mu \\
\mathbb{E}\left(x^{2}\right)&=\int_{-\infty}^{\infty} \mathcal{N} \left(x \mid \mu, \sigma^{2}\right) x^{2} d x=\mu^{2}+\sigma^{2} \\
\operatorname{var}(x)&=\mathbb{E}\left(x^{2}\right)-\mathbb{E}[x]^{2}=\sigma^{2}
\end{aligned}
$$

因此 KL 散度计算过程如下(如果你比较熟悉这些公式的话应该可以直接想起最终结果):

$$
\begin{aligned}
D_{\mathrm{KL}} \left(p \| q \right) &= \mathbb{E}_{p(x)}\left[ \log \frac{p(x)}{q(x)} \right]
= \int [ \log (p(x)) - \log (q(x)) ] p(x) dx \\
&= \int \left[ p(x) \log (p(x)) - p(x) \log (q(x)) \right] dx \\
&= \int p(x) \log \left[ \frac{1}{\sqrt{2\pi}\sigma_{1}} \exp \left ( - \frac {(x -\mu_{1})^2} {2\sigma_{1}^2} \right) \right] dx - \int p(x) \log \left[ \frac{1}{\sqrt{2\pi}\sigma_{2}} \exp \left ( - \frac {(x -\mu_{2})^2} {2\sigma_{2}^2} \right) \right] dx \\
&= \int p(x) \left[ \log \frac{1}{\sqrt{2\pi}\sigma_{1}} + \log \exp \left ( - \frac {(x -\mu_{1})^2} {2\sigma_{1}^2} \right) \right] dx - \int p(x) \left[ \log \frac{1}{\sqrt{2\pi}\sigma_{2}} + \log \exp \left ( - \frac {(x -\mu_{2})^2} {2\sigma_{2}^2} \right) \right] dx \\
&= \left( -\frac{1}{2} \log (2\pi \sigma_{1}^2) + \frac{\int p(x)x^2 dx - \int p(x) 2x \mu_{1} dx + \int p(x) \mu_{1}^2 dx}{2 \sigma_{1}^2} \right) - \left( \frac{1}{2} \log (2\pi \sigma_{2}^2) + \frac{\int p(x)x^2 dx - \int p(x) 2x \mu_{2} dx + \int p(x) \mu_{2}^2 dx}{2 \sigma_{2}^2} \right) \\
&= -\frac{1}{2} \log (2\pi \sigma_{1}^2) + \frac {\mu_{1}^2+\sigma_{1}^2 - 2\mu_{1}\mu_{1} + \mu_{1}^2} {2 \sigma_{1}^2} + \frac{1}{2} \log (2\pi \sigma_{2}^2) - \frac {\mu_{1}^2+\sigma_{1}^2 - 2\mu_{1}\mu_{2} + \mu_{2}^2}{2\sigma_{2}^2} \\
&= -\frac {1}{2} \left[ 1 + \log (2\pi \sigma_{1}^2) \right] + \frac{1}{2} \log (2\pi \sigma_{2}^2) + \frac {\sigma_{1}^2 + (\mu_{1} - \mu_{2})^2}{2\sigma_{2}^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}
$$

由于前面已经固定了二者方差为常数,因此在优化时 $D_{\mathrm{KL}}\left(q \| p_{\theta} \right) $ 中出现的常数项可以被忽略,得到:

$$
\begin{aligned}
L_{t} &= \mathbb{E}_{q} \left[ \frac {1}{2\left\|\boldsymbol{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)\right\|_{2}^{2}} \left\| \tilde{\boldsymbol{\mu}}_{t}\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right)-\boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right)\right\| ^{2} \right] \\
&= \mathbb{E}_{q} \left[ \frac {1}{2\left\|\boldsymbol{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)\right\|_{2}^{2}} \left\| {\color{blue} \frac {1}{\sqrt{\alpha_t}} } \left( {\color{red} \mathbf{x}_{t} } - {\color{blue} \frac { \beta_{t}} {\sqrt{1-\bar{\alpha}_{t}}} } \mathbf{z}_{t} \right) - {\color{blue} \frac {1}{\sqrt{\alpha_t}} } \left( {\color{red} \mathbf{x}_{t} } - {\color{blue} \frac { \beta_{t}} {\sqrt{1-\bar{\alpha}_{t}}} } \mathbf{z}_{\theta}(\mathbf{x}_{t}, t) \right) \right\| ^{2} \right] \\
&= \mathbb{E}_{q} \left[ \frac {\color{blue} {\beta_{t}^{2}} } {2 {\color{blue} \alpha_{t} ( 1-\bar{\alpha}_{t} ) } \left\|\boldsymbol{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)\right\|_{2}^{2} } \left\| \mathbf{z}_{t} - \mathbf{z}_{\theta}(\mathbf{x}_{t}, t) \right\| ^{2} \right]
\end{aligned}
$$

在 DDPM 的原始论文中忽略掉目标函数中的权重项效果会更好,同时将重整化形式的 $\mathbf{x}_{t} = \sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0} +\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}_{t} $ 代入上式:

$$
L_{t}^{\text {simple }}=\mathbb{E}_{\mathbf{x}_{0}, \mathbf{z}_{t}, t}\left[\left\|\mathbf{z}_{t}-\mathbf{z}_{\theta}\left(\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}_{t}, t\right)\right\|^{2}\right]
$$

最终得到简化后的目标函数:

$$
L_{\text {simple}} = L_{t}^{\text {simple }} + C
$$

其中 $C$ 是与 $\theta$ 无关的常量。

代码实现:计算目标函数

在训练模型时,对于模型输出为对噪声预测的情况,使用 MSE 进行计算即可(即简化后的版本):

1
2
3
4
5
6
7
8
9
10
def p_loss(model, x_start, t, noise=None):
    if noise is None:
        noise = np.random.normal(0, 1, x_start.shape)
        
    x_noisy = q_sample(x_start, t, noise)
    predict_noise = model(x_noisy, t)
    
    # MSE
    loss = np.mean((noise - predict_noise) ** 2)
    return loss
  • 从真实数据分布中采样出一张图片 $\mathbf{x}_0$;
  • 给定噪声等级 $t \in \{1, \ldots, T\}$, 对图片进行加噪;
  • 去噪模型根据加噪后的 $\mathbf{x}_t$ 和 $t$ 对噪声进行预测;
  • 计算模型预测的噪声和实际添加的噪声的 MSE.

默认情况下,每个 iter 训练所选取的噪声等级 t 通过均匀随机采样取得,对应 NumPy 中的 numpy.random.randint 接口。即对于单张图片,可能第 1 个 iter 学习 $\mathbf{x}_{0}$ 到 $\mathbf{x}_{233}$ 的加噪与去噪过程,第 2 个 iter 学习 $\mathbf{x}_{0}$ 到 $\mathbf{x}_{500}$ 的加噪与去噪过程... 只要训练的 iter 次数足够多(假定训练了 1000 亿 iters ),就可以认为对 $T=1000$ 的情况,每张图片对 $\mathbf{x}_{0}$ 到 $\mathbf{x}_{t}, t \in \{1, \ldots, T\}$ 的加噪与对应去噪过程都进行 1 亿次训练。对于 Batch 输入,每个样本在单个 iter 中所选取的噪声等级不同,但因为训练了足够多的次数,所以最终的训练效果是均匀且充分的。

当然你也可以依旧选择使用 VLB 来作为优化目标。在本文中为了方便初学者理解,不进行更多描述了。

DDPM 代码实现

用于 DDPM 的 UNet 模型结构

设计什么样的结构可以让模型能输出对噪声的预测呢?从输入 x_noisy 和输出 predict_noise 的 shape 来看,二者具有与最终所需重建的图片一致的维度和形状,比较经典的网络结构是 UNet. 但需要注意的是,如果你只阅读过 UNet 论文,与原始论文中所提到的结构相比,用于 DDPM 的 UNet 模型结构有不少细节上的差异,以 CIFAR10 数据集上进行实验的模型为例子,其结构如下图所示(可以放大):

DDPM-CIFAR

  • 为了保证 High-level 的信息能够在 Low-level 也能用到,原始 UNet 使用了所谓的 skip-connection 操作,即在要进行下采样的时候会将当前的数据存一份,与后续上采样到相同 Level 的数据做 concat 操作,在 Channel 维度上拼接在一起。对于同一个图像大小 Level, 原始 UNet 中的 skip-connection 操作只做一次,而且会对不匹配的宽高比进行处理;DDPM 论文中则是将每一次经过 ResBlock 前的数据也给存了下来,因此同一个 Level 会有多次 concat 操作。请参考上图中进行具体演示的虚线连接;
  • 正因为如此,在 DDPM 所用的 UNet 结构中,如果 downblocks 中每层 ResBlock 的数量为 $n$, 则在 upblocks 中每层 ResBlock 的数量为 $n+1$;
  • 输入到模型的 $t$ 信息以 TimeEmbedding 的形式嵌入到每一个 ResBlock 结构中;
  • DDPM 中根据当前 Level 的图像大小判断是否需要在 downblocks 和 upblocks 的 ResBlock 中加入 AttentionBlock, 比如上图展示的用于 CIFAR10 数据集的模型中,选择在图像分辨率为 $16 \times 16$ 的 Level 使用注意力机制。

具体的细节可以参考开源代码实现: ddpm.py

DDPM 的不足之处

  • DDPM 所用的 UNet 模型结构相当地吃显存,对硬件有了更高的要求(烧钱才能做实验...);
  • 生成(即 p_sample_loop )速度太慢,从噪声生成图片,需要跑 timesteps 次模型前向过程;
  • 最终效果(FID 等指标)依旧不如当时 SOTA 的 GAN 模型实现;
  • 虽然符合直觉,但对整个模型为何有效的解释性略显不足,比较牵强...

如何解决这些问题,则可以引出后续的相关实验和论文了。

参考材料

DDPM 原始论文相关:

  • Deep Unsupervised Learning using Nonequilibrium Thermodynamics: 开山论文,数学符号有所不同
  • Denoising Diffusion Probabilistic Models: 让 Diffusion Model 变得 Practical 的突破性文章
  • MegDiffusion:https://github.com/MegEngine/MegDiffusion (里面补充给出了参考的代码实现)
  • 推荐入门教程:What are Diffusion Models? | Lil'Log ,请阅读最新英文原文,其它转载、二创版本可能没有勘误
  • Diffusion Model 扩散模型理论与完整 PyTorch 代码详细解读(中文视频),含部分推导,略枯燥
  • Denoising Diffusion Probabilistic Models (英文视频),前半部分讲 DDPM, 后半部分讲 Improved DDPM

其他版本的 DDPM Tutorial:

  • The Annotated Diffusion Model: Hugging face 博客,有在 Fashion-MNIST 上的 Toy Example
  • Diffusion models from scratch in PyTorch: 使用的 Stanford Dog 数据集,演示代码很简洁直观
  • Diffusion Models | Paper Explanation | Math Explained: 对数学公式推导进行了适度的可视化

拓展阅读:

  • DDPM - Diffusion Models Beat GANs on Image Synthesis: 应该是最早介绍 DDPM 的 YouTuber
  • Denoising Diffusion-based Generative Modeling: Foundations and Applications:CVPR 2022 Tutorial
  • Diffusion Model 最近在图像生成领域大红大紫,如何看待它的风头开始超过 GAN? - 我想唱high C的回答
  • 生成扩散模型漫谈(一):DDPM = 拆楼 + 建楼 (通俗,数学记号与原论文不同)
  • 生成扩散模型漫谈(二):DDPM = 自回归式VAE (VAE 视角下的推导,Keras 实验,模型有修改)
  • 生成扩散模型漫谈(三):DDPM = 贝叶斯 + 去噪 (贝叶斯视角下的推导,也即本文的推导思路)
  • 欢迎补充 (๑ `▽´๑)
人工智能 扩散模型 深度学习 编程
赞赏 赞(15)
订阅
提醒
guest
guest
0 评论
内嵌评论
查看所有评论
15
  • 15
  • 0
Copyright © 2020-2023 MEGChai.
  • 文章
    • 随笔
    • 笔记
    • 教程
  • 关于
# 生活 # # 心理 # # 编程 # # 音乐 # # 写作 #
Chai
95
文章
4
评论
58
喜欢
wpDiscuz