目前网络上关于 Diffusion Model 的中文讲解性材料目前还很少，这里简单梳理一下，主要参考如下文章：
- 综述类介绍
- 清华 TSAIL 朱军组博士生写的一篇 知乎回答 (我还不知道他的真实姓名)
- 离散扩散模型
- What are Diffusion Models? —— Lilian Weng
- The Annotated Diffusion Model —— Niels Rogge and Kashif Rasul
- 连续扩散模型
这篇文章是接触 Diffusion Model 领域的第一篇笔记，也略带有些 Step-by-step 教程的性质，里面记录的观点也是基于 20 年 DDPM 原始论文的认知给出的，未必解释得合理；而现今前沿的一些研究可能会以更加严谨、系统而全面的角度来对 Diffusion Model 进行解读，因此本文的将假想读者定位为和我一样的领域新手（尚处于读论文做复现的阶段），且不保证会有系列后续的文章。
另外，此文章处于早期样稿阶段，难免存在疏漏，各位读者劳请于评论区指出谬误之处。
什么是扩散模型？
相较于 GANs, VAEs 与 Normalizing Flows 等生成式模型，扩散模型（Diffusion Model）的思路非常简单：向数据中逐步地加入随机噪声，然后学习其逆向过程，希望最终能够从噪声中重建所需的数据样本 —— 有意思的地方在于，我们可以训练一个神经网络模型来学习逆向扩散过程。
在逆向扩散的每一步，神经网络模型的输入是当前的样本 $\mathbf{x}_{t}$ 和加噪程度 $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$ 指在有限的 $T$ 步内（论文中使用 $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}$ 将逐渐地丢失掉原有的特征，当 $T \rightarrow \infty$ 时，最终得到的 $\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} \cdot \mathbf{z}_{t-1}}$, 而 $\mathbf{z}_{t-1} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ 可认为是隐变量。
注意：上式中完全不含有可训练的参数，也即是说扩散模型的前向扩散过程不需要对网络进行训练。只要给定 $\mathbf{x}_0$ 和 $\beta$, 通过不断迭代，就可以算出 $t$ 为任意时刻所对应的 $q(\mathbf{x}_t)$ 分布。实际上在逆扩散的过程中网络也没有进行训练，我们做的仅仅是使用训练好的网络对噪声进行 “采样”。这也是扩散模型的特点之一：你可以使用不同的 Diffusion 与不同的 Denoise 模型进行组合使用，只需提供一致的维度与超参数等信息（比如 $T$ 值）。那么我们的 Denoise 模型是如何进行参数训练的呢？在设计目标函数的时候将揭开谜底，在此之前让我们先搞清楚并实现最基本的扩散过程。
代码实现：逐步加噪
我们用单张图片+纯 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 |
注意 Uint8 图片的数值范围是 $[0,255]$, 我们在加噪前需要将其缩放到噪声所属的标准正态分布 $[-1, 1]$ 的范围（也可使用 Normalize 借助统计的 mean 和 std 计算，此时数据为 Float 类型），确保逆向过程 $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}$, 只要加噪的次数足够多，最终图片中的原始信息将会所剩无几，变成一个各向同性的高斯分布。
|
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() |
实际执行这段代码体验一下，你会发现逐步加噪过程的 计算代价太大了，且计算量与 $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}$, 如下图所示：
逆向扩散过程
如果我们能够知道 $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}$ 有关的项，因此细节于此忽略。
已知对于概率密度函数为 $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 |
# 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) |
计算 alphas_cumprod_prev （对应 $\bar{\alpha}_{t-1}$ ） 时，可由前面已经计算好的 alphas_cumprod （对应 $\bar{\alpha}_{t}$ ）取 [:-1] 后经过 append, pad 或 concatenate 等处理（无论用哪个接口，效果是一样的）得到，为了在 $t=0$ 时 $\bar{\alpha}_{t-1}$ 能够有对应值，通常在开头补上值为 1 的项。
当然你也可以每次对 $t$ 的值进行特判，再确定是否要通过 alphas_cumprod[t-1] 来取值，只是这样做不如提前处理好方便。
我们可以让神经网络直接输出 $\boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right)$ 作为对 $\tilde{\boldsymbol{\mu}}_{t}\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right)$ 的预测，但论文作者发现让模型输出 $\mathbf{z}_{\theta}(\mathbf{x}_{t}, t) $ 作为对噪声 $\mathbf{z}_{t}$ 的预测，会比直接预测均值的效果更好一些。因此接下来的计算其实是找到如何从预测噪声计算出预测的均值。继续进行推导：
求得方差的化简形式后，方便我们继续计算 $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}
$$
有了上述关系，可以得到从模型预测的噪声计算得到模型均值的公式：
$$
{\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 \boldsymbol{\epsilon}
$$
值得注意的一点是，论文官方的实际代码实现中并没有直接使用化简后得到的公式从预测的噪声计算模型均值，而是选择了先让模型输出 $\mathbf{z}_{t} = \mathbf{z}_{\theta}(\mathbf{x}_{t}, 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}_{0}$ 经过变换范围到 $[-1, 1]$, 因此对于预测的原始数据通常会采取 clip[-1, 1] 的截断计算确保此分布特征一致，而对中间过程得到的去噪样本不会进行此操作。下面用代码演示一下这两种不同的计算过程——
代码实现：逐步去噪
假定我们已经有一个训练好的理想去噪模型 $\mathbf{z}_{\theta}(\mathbf{x}_{t}, 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 t level 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{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}$ 计算得到均值。
同样地，式子中的各个常数系数可以被提前计算好，对应地有 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 == "fixedlarge": logvar = np.log(np.append(posterior_variance[1], betas[1:])) elif model_var_type == 'fixedsmall': logvar = np.log(np.maximum(posterior_variance, 1e-20)) def p_sample(...): # ... return model_mean + np.exp(0.5 * logvar[t]) * noise |
想要得到理想中的模型 $\mathbf{z}_{\theta}(\mathbf{x}_{t}, t)$, 这就需要借助梯度下降法学习参数。
目标分布的似然函数
我们的优化目标是使得经由 $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}$ 式情况（论文中选择了单独从 $\mathcal{N}\left(\mathbf{x}_{0} ; \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{1}, 1\right), \mathbf{\Sigma}_{\theta}\left(\mathbf{x}_{1}, 1\right)\right)$ 进行建模）。对于 $L_{T}$ 式，$q$ 中不含有可训练的参数，而 $p_{\theta}\left(\mathbf{x}_{T}\right)$ 是一个各向同性的高斯分布，没有未知参数，因此该项无须优化，可以省略。得到计算式为：
$$
\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]
$$
推导：简化目标函数
回忆一下，对于两个单一变量的高斯分布 $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}}}\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$ 无关的常量。
代码实现：计算目标函数
|
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.
