翻译自英文版本

在这篇博客中，我们将尝试理解扩散家族中的成员——DDPM。鉴于变分自编码器（VAE）和扩散模型共享的常见结构，我们首先讨论一些重要的VAE内容，然后介绍DDPM。这篇博客试图做到自成体系并在数学上清晰，但假定读者有一定的VAE和生成模型的先验知识。

参考文献:

VAE 在 Wikipedia

Lil’Log 什么是扩散模型?

滑铁卢大学 STAT940 第17讲 by Ali Ghodsi

VAE 和 ELBO

在生成模型中，我们的主要目标是观察数据集，学习其底层分布，然后从学习到的分布中生成新样本。在VAE方案中，我们采用类似于自编码器的结构，如下图所示：

graph LR A[输入] --> B[编码器] B --> C[潜在空间] C --> D[解码器] D --> E[输出] N((噪声)) --> C

在推理过程中，我们会舍弃编码器，采样一些噪声来构建潜在变量，然后使用解码器生成样本。潜在变量通常从标准高斯分布中采样。

我知道这可能仍然不清楚，所以让我们深入数学细节。

假设我们有一些数据，我们可以用参数 $\theta$ 来建模数据分布。分布可以表示为 $p_\theta(x) = p(x|\theta)$。问题是我们不知道参数 $\theta$，这正是我们想要学习的。我们能最大化数据的似然吗？可以，但这可能难以优化好，并且我们失去了生成新样本的能力。

VAE引入了潜在变量 $z$ 来建模数据的复杂结构，它与数据 $x$ 具有联合分布。然后 $p_\theta(x)$ 是 $p_\theta(x, z)$ 的边缘分布。我们可以将边缘分布写为： $$ p_\theta(x) = \int p_\theta(x, z) dz $$

根据简单的概率规则，我们进一步写成： $$ p_\theta(x) = \int p_\theta(x|z)p_\theta(z) dz $$

这里的基本假设是，数据分布由潜在变量 $z$ 控制，$z$ 具有分布 $p_\theta(z)$。按照贝叶斯方案，我们在模型中有三样东西：

  1. 先验 $p_\theta(z)$
  2. 似然 $p_\theta(x|z)$
  3. 后验 $p_\theta(z|x)$

通常，我们选择先验为简单的，例如参数化的高斯分布或简单的标准正态分布，这样从中采样是便宜且直接的。似然 $p_\theta(x|z)$ 是给定潜在变量 $z$ 下观察到的数据 $x$ 的概率。它可以容易地实现为一个神经网络，其中输入是一个潜在向量，而输出可以是正态分布的均值和方差。

然而，后验是 不可求 的。为了计算它，我们需要应用贝叶斯定理： $$ p_\theta(z|x)=\frac{p_\theta(x|z)p_\theta(z)}{p_\theta(x)} $$ 显然，我们不知道 $p_\theta(x)$。虽然这是个坏消息，但我们进一步引入一个函数 $q_\phi(z|x)\approx p_\theta(z|x)$ 来用一组不同的参数 $\phi$ 近似后验。注意 这并不是凭空设计出来的，而是源自于一般潜变量模型的证据下界（ELBO）不等式！

graph LR Z((z)) --似然-->X((x)) X((x)) --后验--> Z

那么什么是ELBO呢？ELBO提供了一种在 $p_\theta(z|x)$ 不可求的情况下进行最大似然估计（MLE）的方法。我们可以从 $p_\theta(x)=log p(x|\theta)$ 的MLE推导出ELBO。

$$ \begin{aligned} \log p(x|\theta) &= \log \int p(x, z |\theta) dz \\ &= \log \int p(x, z |\theta) \frac{q(z|x, \phi)}{q(z|x, \phi)} dz\\ &= \log \mathbb{E}_{z \sim q(\cdot|x,\phi)} \left[ \frac{p(x, z | \theta)}{q(z|x, \phi)} \right] \\ &\geq \mathbb{E}_{z \sim q(\cdot|x,\phi)} \left[ \log \frac{p(x, z | \theta)}{q(z|x, \phi)} \right] \quad \text{(詹森不等式)} \end{aligned} $$

ELBO可以进一步分解为 $$ \begin{aligned} & \ \mathbb{E}_{z \sim q(\cdot|x,\phi)} \left[ \log \frac{p(x, z | \theta)}{q(z|x, \phi)} \right]\\ & =\mathbb{E}_{z \sim q(\cdot|x,\phi)} \left[ \log \frac{p(x, z | \theta)}{q(z|x, \phi)} \right]\\ & = \mathbb{E}_{z \sim q(\cdot|x,\phi)} \left[ \log \frac{p(x|z, \theta)p(z|\theta)}{q(z|x, \phi)} \right]\\ & = \mathbb{E}_{z \sim q(\cdot|x,\phi)} \left[ \log p(x|z, \theta) + \log \frac{p(z|\theta)}{q(z|x, \phi)} \right]\\ & = \mathbb{E}_{z \sim q(\cdot|x,\phi)} \left[ \log p(x|z, \theta) \right] - \mathbb{E}_{z \sim q(\cdot|x,\phi)} \left[ \log \frac{q(z|x, \phi)}{p(z|\theta)} \right]\\ & = \mathbb{E}_{z \sim q(\cdot|x,\phi)} \left[ \log p(x|z, \theta) \right] - D_{KL}(q_\phi(z|x)||p_\theta(z)) \end{aligned} $$

这里需要记住的是，两个分布之间的KL散度为： $$ D_{KL}(q||p) = \int q(x) \log \frac{q(x)}{p(x)} dx\\ D_{KL}(q||p) = \mathbb{E}_{x\sim q(\cdot)}[\log \frac{q(x)}{p(x)}] $$

我们可以最大化这个下界以最大化数据的对数似然，但仍然不太方便。因此在实际实现中，我们最大化以下等价项： $$ L_{\theta, \phi}(x) = \mathbb{E}_{z\sim q_\phi(\cdot|x)}[\log p_\theta(x|z)] - D_{KL}(q_\phi(\cdot|x)||p_\theta(\cdot)) $$

给定上述模型，它学习分布并提供生成新样本的方法，我们需要找到一种方法来学习参数。对于VAE，我们希望优化两件事：

  • 我们希望重建是精确的。给定数据 $x_i$，我们可以计算相应的潜在变量 $z_i$，解码器应该能够精确地给我 $x_i$ 给定 $z_i$
  • 我们希望近似 $q_\phi(z|x)\approx p_\theta(z|x)$ 是准确的，否则从训练数据生成的潜在变量将毫无意义。

我们优化的两个项对应我们希望优化的两个目标。由于我们将潜在变量 $p_\theta(z)$ 模型为标准高斯分布，并将 $q_\phi(z|x)$ 模型为高斯分布，因此KL散度有封闭形式表达。注意 $q_\phi(z|x

)$ 的均值和方差通过神经网络得到。因此我们有：

$$ q_\phi(z|x) = N(z; \mu(x), \Sigma(x)) $$

如果我们进一步假设 $p_\theta(z)=N(z; 0, I)$ 具有标准正态分布，我们得到： $$ \begin{aligned} D_{KL}(q_\phi(z|x)||p_\theta(z)) & = \int q_\phi(z|x) \log \frac{q_\phi(z|x)}{p_\theta(z)} dz \\ & = -\frac{1}{2}\sum_{i=1}^{L} (1 + \log(\sigma_i^2)-\mu_i^2-\sigma_i^2) \end{aligned} $$

注意 $L$ 是潜在向量的大小。

第一项 $\mathbb{E}_{z\sim q_\phi(\cdot|x)}[\log p_\theta(x|z)]$ 实现重建误差。我们假设它是给定潜在向量 $z$ 的观察数据 $x$ 的高斯分布：

$$ p_\theta(x|z) = N(x;\mu(z), \sigma^2 I) $$

这里 $\mu(z)$ 由解码器网络输出。重建误差是：

$$ \begin{aligned} \mathbb{E}_{z\sim q_\phi(\cdot|x)}[\log p_\theta(x|z)] &= \mathbb{E}_{z\sim q_\phi(\cdot|x)}[\log N(x; \mu(z), \sigma^2 I)]\\ &= \mathbb{E}_{z\sim q_\phi(\cdot|x)}[-\frac{1}{2\sigma^2} ||x - \mu(z)||^2]\\ &= -\frac{1}{2\sigma^2} \mathbb{E}_{z\sim q_\phi(\cdot|x)} [||x - \mu(z)||^2] \end{aligned} $$

损失函数

最终的目标函数是： $$ L_{\theta, \phi}(x) = - \frac{1}{2\sigma^2} \mathbb{E}_{z\sim q_\phi(\cdot|x)} [||x - \mu(z)||^2] + \frac{1}{2}\sum_{i=1}^{L} (1 + \log(\sigma_i^2)-\mu_i^2-\sigma_i^2) $$

我们通过最大化上述公式来优化模型。

DDPM

DDPM 的训练是基于逐渐去噪过程。我们开始对数据添加噪声（例如正态分布），最终得到一个几乎纯噪声的样本。然后，我们反转这个过程：从纯噪声开始，逐步去除噪声以生成新样本。

前向过程

在前向过程中，我们逐步添加噪声，使数据变得越来越接近正态分布。这个过程通过一个固定的调度器（比如线性）定义。具体来说，对于每个时间步 $t$，我们通过以下方式添加噪声：

$$ q(x_t|x_{t-1}) = N(x_t; \sqrt{1-\beta_t} x_{t-1}, \beta_t I) $$

这里 $\beta_t$ 是在每个时间步添加的噪声强度。这导致以下逐步添加噪声的过程：

$$ q(x_t|x_0) = N(x_t; \sqrt{\bar{\alpha}_t} x_0, (1-\bar{\alpha}_t) I) $$

这里 $\alpha_t = 1-\beta_t$，$\bar{\alpha}_t = \prod_{i=1}^{t}\alpha_i$。

反向过程

反向过程中，我们从纯噪声开始，逐步去除噪声以生成新样本。我们使用一个参数化的模型 $p_\theta(x_{t-1}|x_t)$ 来反转前向过程：

$$ p_\theta(x_{t-1}|x_t) = N(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t)) $$

损失函数

我们希望使模型生成的样本与原始数据尽可能接近。因此，损失函数基于数据的似然定义。具体来说，DDPM 的损失函数是前向过程和反向过程之间的 KL 散度：

$$ L_{DDPM} = \sum_{t=1}^{T} D_{KL}(q(x_{t-1}|x_t, x_0) || p_\theta(x_{t-1}|x_t)) $$

优化这个损失函数使我们能够训练模型，以从纯噪声生成逼真的样本。

总结

本文介绍了变分自编码器（VAE）的基本概念和数学基础，并通过这个框架引入了扩散模型（特别是DDPM）的基本原理。VAE通过最大化证据下界（ELBO）来近似数据分布，而DDPM则通过逐步去噪过程来生成数据。两者都展示了生成模型在学习复杂数据分布和生成新样本方面的强大能力。