Diffusion【2】：VAE

系列文章目录

---

---

前言

VAE 可以说是最近大热的生成模型的重要基石。本文侧重对论文内容的翻译和一定程度上的解释。若想完全弄懂 VAE 的具体流程和数学公式推导，建议大家移步苏剑林老师的博客。

原论文链接：Auto-Encoding Variational Bayes
原论文链接：苏神博客

---

1. Abstract

本文主要探讨了在概率模型中处理潜在变量推断的难点。传统方法难以处理以下两种主要情况：

1. 后验分布的不可解性：当模型的后验分布无法通过解析或简单的数值方法计算时，标准方法如期望最大化（EM）或常规变分贝叶斯方法可能无法适用。
2. 大规模数据处理：面对大规模数据集时，传统的批量优化或采样方法（如MCMC）变得计算代价高昂。

为了解决上述问题，作者提出了 随机梯度变分贝叶斯估计器（SGVB） 和 自动编码变分贝叶斯算法（AEVB）：

• SGVB：通过重新参数化的技巧，将随机变量的采样过程转化为可微函数，这样可以使用高效的随机梯度方法进行优化。
• AEVB：结合 SGVB，用一个神经网络训练的识别模型去近似后验分布，从而实现对潜在变量的高效推断和生成模型参数的高效学习。

意义：

1. 效率提升：算法规避了传统推断方法的复杂计算，尤其适用于连续潜在变量的模型。
2. 多功能性：学习到的识别模型不仅可以进行后验推断，还能用于数据重建、降噪、表示学习和可视化。

---

2. Introduction

2.1. Motivation

定向概率模型中的推断与学习是一个难题，尤其是当涉及连续潜在变量时：

• 后验分布往往无法直接求解（不可解）。
• 需要优化的目标通常非常复杂，例如在大数据场景中批量优化计算开销高昂。

传统的变分贝叶斯方法（VB）通过用近似分布替代真实后验分布，并优化这一近似分布来解决问题。但这一方法通常依赖解析计算，其适用性受到限制。

2.2. Contribution

本文的主要贡献：

• 提出 SGVB：通过重新参数化，将不可微的随机采样操作替换为可微的操作，从而可以对目标函数求导并优化。
• 提出 AEVB：利用 SGVB 优化一个称为“识别模型”的神经网络，这个模型能够高效地近似后验分布。
  • 通过识别模型，可以高效地对每个数据点的潜在变量进行采样推断，而无需依赖传统的复杂方法（如 MCMC）。
  • 这种方法既适用于模型参数的优化，也可以生成与原始数据类似的人工数据。

---

3. Methods

3.1. Problem Scenario

考虑一个数据集 X = { x ( i ) } i = 1 N X = \{ x^{(i)} \}_{i=1}^N X={x(i)}i=1N​，包含 $N$ 个独立同分布的样本，每个样本 $x^{(i)}$ 可以是连续或离散的。我们假设数据是通过一个随机过程生成的，该过程涉及一个未观察到的连续随机变量 $z$。生成过程包括两个步骤：

1. 先验采样：从某个先验分布 $p_{{\theta }^{\ast }}(z)$ 中生成一个潜在变量样本 $z^{(i)}$。
2. 条件采样：给定 $z^{(i)}$，从条件分布 $p_{{\theta }^{\ast }}(x|z)$ 中生成观测数据 $x^{(i)}$。

这里，${\theta }^{\ast }$ 表示真实的参数，但我们并不知道它的具体值。我们假设先验分布 $p_{{\theta }^{\ast }}(z)$ 和似然函数 $p_{{\theta }^{\ast }}(x|z)$ 属于参数化的分布族 $p_{\theta }(z)$ 和 $p_{\theta }(x|z)$，并且它们的概率密度函数对于 $\theta$ 和 $z$ 都几乎处处可微。

需要注意的是，这个生成过程的大部分对我们来说是不可见的：真实的参数 ${\theta }^{\ast }$ 和潜在变量 $z^{(i)}$ 都是未知的。

一个关键的地方是，我们并不对边缘概率或后验概率做常见的简化假设。相反，我们感兴趣的是一种通用的算法，即使在以下困难情况下也能高效工作：

1. 不可解性：边缘似然 $p_{\theta }(x)=\int p_{\theta }(z)p_{\theta }(x|z)dz$ 的积分是不可解的（因此我们无法直接计算或对边缘似然求导）；真实的后验分布 $p_{\theta }(z|x)=\frac{p_{\theta }(x|z)p_{\theta }(z)}{p_{\theta }(x)}$ 也是不可解的（因此无法使用 EM 算法）；对于任何合理的平均场变分贝叶斯算法所需的积分也是不可解的。这些不可解性在当似然函数 $p_{\theta }(x|z)$ 稍微复杂一些（例如包含非线性隐藏层的神经网络）时就会出现。
2. 大规模数据集：当数据量非常大时，批量优化计算代价太高；我们希望使用小批量甚至单个数据点来更新参数。基于采样的方法，例如 Monte Carlo EM，在这种情况下通常太慢，因为它需要对每个数据点进行昂贵的采样过程。

在上述场景下，我们希望解决以下三个相关问题：

1. 模型参数 $\theta$ 的高效近似最大似然或最大后验估计。参数 $\theta$ 本身可能是我们感兴趣的对象，例如在分析某个自然过程时。通过估计 $\theta$，我们还可以模拟隐藏的随机过程，生成与真实数据类似的人工数据。
2. 给定观测数据 $x$ 时，对潜在变量 $z$ 进行高效的近似后验推断。这对于编码或数据表示任务非常有用。
3. 对变量 $x$ 进行高效的近似边缘推断。这使我们能够执行各种需要对 $x$ 有先验的推断任务，常见的计算机视觉应用包括图像去噪、修复和超分辨率。

数据生成过程：

• 潜在变量采样：每个数据点 $x^{(i)}$ 都有一个对应的潜在变量 $z^{(i)}$，从先验分布 $p_{\theta }(z)$ 中采样。
• 观测数据生成：给定 $z^{(i)}$，从条件分布 $p_{\theta }(x|z)$ 中生成 $x^{(i)}$。

未知量：

• 参数 ${\theta }^{\ast }$：真实的模型参数，我们希望估计它。
• 潜在变量 $z^{(i)}$：未观察到的随机变量，我们希望对其进行推断。

挑战：

• 后验分布不可解：由于 $p_{\theta }(z|x)$ 无法解析求解，传统的 EM 算法或变分推断方法无法直接应用。
• 大规模数据：需要高效的算法，能够在不牺牲精度的情况下处理大数据量。

为了解决上述问题，我们引入了一个 识别模型（recognition model） $q_{\varphi }(z|x)$：这是对不可解的真实后验 $p_{\theta }(z|x)$ 的近似。需要注意的是，与平均场变分推断中的近似后验不同，这里的 $q_{\varphi }(z|x)$ 不一定是可分解的，其参数 $\varphi$ 也不是通过某种封闭形式的期望计算得出的。相反，我们将介绍一种方法，通过学习来联合优化识别模型的参数 $\varphi$ 和生成模型的参数 $\theta$。

引入识别模型 $q_{\varphi }(z|x)$：

• 目的：近似真实的后验分布 $p_{\theta }(z|x)$。
• 特点：不像传统的平均场方法，不需要可解析的期望计算。$q_{\varphi }(z|x)$ 的形式可以更加灵活。

从编码理论的角度来看，未观察到的变量 $z$ 可以解释为一个潜在表示或编码。因此，在本文中，我们也将 识别模型 $q_{\varphi }(z|x)$ 称为概率编码器（probabilistic encoder），因为给定一个数据点 $x$，它会生成一个关于可能的 $z$ 值的分布（例如高斯分布），这些 $z$ 值可能生成了观测数据 $x$。类似地，我们将 $p_{\theta }(x|z)$ 称为 概率解码器（probabilistic decoder），因为给定一个编码 $z$，它会生成一个关于可能对应的 $x$ 值的分布。

编码器和解码器的解释：

• 概率编码器 $q_{\varphi }(z|x)$：给定输入 $x$，输出关于潜在变量 $z$ 的分布，即编码过程。
• 概率解码器 $p_{\theta }(x|z)$：给定潜在变量 $z$，输出关于观测数据 $x$ 的分布，即解码过程。

3.2. The variational bound

边缘似然是各数据点边缘似然的总和：$\log p_{\theta }(x^{(1)},\cdots ,x^{(N)})={\sum }_{i=1}^N\log p_{\theta }(x^{(i)})$，其中每个数据点的边缘似然可以重写为：

$\log p_{\theta }(x^{(i)})=D_{\mathrm{KL}}(q_{\varphi }(z|x^{(i)})||p_{\theta }(z|x^{(i)}))+\mathcal{L}(\theta ,\varphi ;x^{(i)})$

公式右侧的第一项是近似后验分布 $q_{\varphi }(z|x^{(i)})$ 与真实后验分布 $p_{\theta }(z|x^{(i)})$ 的 KL 散度，由于 KL 散度非负，因此第二项 $\mathcal{L}(\theta ,\varphi ;x^{(i)})$ 被称为数据点 $i$ 的边缘似然的 变分下界，其定义为：

$\mathcal{L}(\theta ,\varphi ;x^{(i)})=\mathbb{E}{q}_{\varphi }(z|x)[-\log q_{\varphi }(z|x)+\log p_{\theta }(x,z)]$

也可以重写为：
$\mathcal{L}(\theta ,\varphi ;x^{(i)})=-D_{\mathrm{KL}}(q_{\varphi }(z|x^{(i)})||p_{\theta }(z))+\mathbb{E}{q}_{\varphi }(z|x^{(i)})[\log p_{\theta }(x^{(i)}|z)]$

我们希望对变分参数 $\varphi$ 和生成参数 $\theta$ 同时优化变分下界 $\mathcal{L}(\theta ,\varphi ;x^{(i)})$

边缘似然分解：

• 数据点的边缘似然 $\log p_{\theta }(x^{(i)})$ 被分解为两部分：
  • KL 散度 $D_{\mathrm{KL}}(q_{\varphi }(z|x^{(i)})||p_{\theta }(z|x^{(i)}))$。
  • 变分下界 $\mathcal{L}(\theta ,\varphi ;x^{(i)})$。

变分下界的意义：

• 由于 KL 散度非负，变分下界 $\mathcal{L}(\theta ,\varphi ;x^{(i)})$ 是边缘似然的下界。
• 优化变分下界可以间接提高边缘似然的值。

下界的分解：

• 第一项是 KL 散度 $D_{\mathrm{KL}}(q_{\varphi }(z|x^{(i)})||p_{\theta }(z))$，用来衡量近似后验 q_\phi(z|x) 与先验 $p_{\theta }(z)$ 的相似度。
• 第二项是重构误差的期望 $\mathbb{E}{q}_{\varphi }(z|x^{(i)})[\log p_{\theta }(x^{(i)}|z)]$，表示生成模型 $p_{\theta }(x|z)$ 重构数据的能力。

3.3. The SGVB estimator and AEVB algorithm

3.3.1. Stochastic Gradient Variational Bayes Estimator

在这一节中，我们提出了一种实用的下界及其梯度估计器，并假设近似后验为 $q_{\varphi }(z|x)$。需要注意的是，这一技术同样适用于非条件情况下的近似后验 $q_{\varphi }(z)$。

在选定的近似后验分布 $q_{\varphi }(z|x)$ 的基础上，在一些轻微约束条件下，我们可以通过一个辅助噪声变量 $\epsilon$ 的可微变换 $g_{\varphi }(\epsilon ,x)$ 对随机变量 $z$ 进行重新参数化：

$z=g_{\varphi }(\epsilon ,x),\epsilon \sim p(\epsilon )$

对于任何关于 $z$ 的函数 $f(z)$，可以通过以下公式用 $\epsilon$ 表达期望值：

$\mathbb{E}{q}_{\varphi }(z|x)[f(z)]=\mathbb{E}p(\epsilon )[f(g_{\varphi }(\epsilon ,x))]$

通过蒙特卡罗采样估计：
$\mathbb{E}{q}_{\varphi }(z|x)[f(z)]\approx \frac{1}{L}{\sum }_{l=1}^{L}f(g_{\varphi }({\epsilon }^{(l)},x))$

其中 ${\epsilon }^{(l)}\sim p(\epsilon )$。

应用于变分下界，可以得到通用的随机梯度变分贝叶斯（SGVB）估计器：

$\tilde{\mathcal{L}}A(\theta ,\varphi ;x^{(i)})\approx \frac{1}{L}\sum {l=1}^L[\log p_{\theta }(x^{(i)},z^{(i,l)})-\log q_{\varphi }(z^{(i,l)}|x^{(i)})]$

对于某些情况（例如 KL 散度可以解析求解），SGVB 估计器可以被简化为更低方差的形式：
$\tilde{\mathcal{L}}B(\theta ,\varphi ;x^{(i)})=-D\mathrm{KL}(q_{\varphi }(z|x^{(i)})||p_{\theta }(z))+\frac{1}{L}{\sum }_{l=1}^L\log p_{\theta }(x^{(i)}|z^{(i,l)})$

其中 $z^{(i,l)}=g_{\varphi }({\epsilon }^{(i,l)},x^{(i)})$，且 ${\epsilon }^{(l)}\sim p(\epsilon )$。

SGVB的核心是利用 重新参数化技巧，将原始需要对 $q_{\varphi }(z|x)$ 的积分替换为对辅助噪声变量 $\epsilon$ 的积分。这种重新参数化使得：

• 随机变量 z 由噪声变量 \varepsilon 和输入 x 通过可微函数 g_\phi 确定。
• 可以直接通过采样 \varepsilon 来估计目标函数的梯度。

相比传统蒙特卡罗方法，SGVB估计器有以下优点：

• 梯度可微，可以使用高效的优化算法。
• 在某些情况下（如KL散度可解析求解），可以进一步降低方差。

3.3.2. Auto-Encoding Variational Bayes

在独立同分布（i.i.d.）数据集 $X$ 的情况下，针对整个数据集的变分下界可以通过小批量（minibatch）估计器进行计算：

$\mathcal{L}(\theta ,\varphi ;X)\approx \tilde{\mathcal{L}}M(\theta ,\varphi ;X_M)=\frac{N}{M}\sum {i=1}^M\tilde{\mathcal{L}}(\theta ,\varphi ;x^{(i)})$

其中， X M = { x ( i ) } i = 1 M X_M = \{x^{(i)}\}_{i=1}^M XM​={x(i)}i=1M​ 是从完整数据集随机抽取的小批量样本

AEVB结合了SGVB估计器和小批量随机梯度方法，用于优化变分下界。其主要创新点包括：

• 通过小批量样本提高了大数据集上的训练效率。
• 将近似后验建模为神经网络（即识别模型），可以高效地进行推断。
• 使用简单的采样（无需复杂的MCMC），直接进行参数优化。

AEVB的意义

• 效率高：针对大规模数据集，能够以小批量方式快速优化模型参数。
• 灵活性强：可以适配各种含连续潜在变量的生成模型。

算法步骤如下：

1. 随机初始化参数 $\theta ,\varphi$。
2. 在每次迭代中：
   • 从完整数据集中随机抽取一个小批量 $X_M$。
   • 从噪声分布 $p(\epsilon )$ 中采样随机变量 $\epsilon$。
   • 计算小批量估计器的梯度 ${\nabla }_{\theta ,\varphi }{\tilde{\mathcal{L}}}_M(\theta ,\varphi ;X_M,\epsilon )$。
   • 使用梯度更新参数（例如通过随机梯度下降或 Adagrad 优化器）。
3. 重复直至收敛。

3.4. The reparameterization trick

3.4.1. Trick

为了解决我们的问题，提出了一种替代方法来从 $q_{\varphi }(z|x)$ 中生成样本。重新参数化的核心思想非常简单。假设 $z$ 是一个连续随机变量，且 $z\sim q_{\varphi }(z|x)$ 是某个条件分布。那么，通常可以将随机变量 $z$ 表示为一个确定性的变量：

$z=g_{\varphi }(\epsilon ,x)$

其中，$\epsilon$ 是一个具有独立边缘分布 $p(\epsilon )$ 的辅助变量，而 $g_{\varphi }(\cdot )$ 是某个参数化的向量值函数。

这种重新参数化对我们的情况非常有用，因为它可以将关于 $q_{\varphi }(z|x)$ 的期望重写为关于 $p(\epsilon )$ 的期望，从而使得期望的蒙特卡罗估计对 $\varphi$ 是可微的。证明如下：

给定确定性映射 $z=g_{\varphi }(\epsilon ,x)$，我们知道 $q_{\varphi }(z|x){\prod }_{i}d{z}_i=p(\epsilon ){\prod }_{i}d{\epsilon }_i$。因此：

$\int q_{\varphi }(z|x)f(z)dz=\int p(\epsilon )f(g_{\varphi }(\epsilon ,x))d\epsilon$

这表明可以构造一个可微的估计器：

$\int q_{\varphi }(z|x)f(z)dz\approx \frac{1}{L}{\sum }_{l=1}^{L}f(g_{\varphi }(x,{\epsilon }^{(l)}))$

其中 ${\epsilon }^{(l)}\sim p(\epsilon )$。

重新参数化的作用

1. 将随机性移到辅助变量 \varepsilon：
   原始的 $z$ 依赖于复杂的后验分布 $q_{\varphi }(z|x)$，直接优化非常困难。重新参数化将 $z$ 表示为由可控参数 $\varphi$、输入 $x$ 和简单分布 $p(\epsilon )$ 生成的确定性变量，从而简化了采样和优化。
2. 使梯度计算成为可能：
   原始的采样过程通常不可微，但通过重新参数化，可以将目标函数关于 $\varphi$ 的梯度转化为辅助变量 $\epsilon$ 的期望，这种期望是可微的。
3. 简化变分推断：
   在变分推断中，我们需要优化变分下界（涉及关于 $q_{\varphi }(z|x)$ 的积分）。重新参数化技巧有效地将积分转化为简单分布 $p(\epsilon )$ 的采样问题，降低了计算复杂度。

3.4.2. Example

以单变量高斯分布为例：假设 $z\sim q(z|x)=\mathcal{N}(\mu ,{\sigma }^2)$。在这种情况下，一个有效的重新参数化是：

$z=\mu +\sigma \epsilon ,\epsilon \sim \mathcal{N}(0,1)$

因此：

$\mathbb{E}\mathcal{N}(z;\mu ,{\sigma }^2)[f(z)]=\mathbb{E}\mathcal{N}(\epsilon ;0,1)[f(\mu +\sigma \epsilon )]\approx \frac{1}{L}{\sum }_{l=1}^{L}f(\mu +\sigma {\epsilon }^{(l)})$

其中 ${\epsilon }^{(l)}\sim \mathcal{N}(0,1)$。

在高斯分布中，重新参数化技巧将采样过程分解为简单的均值和标准差调整：

• 均值 $\mu$ 控制生成的中心位置。
• 标准差 $\sigma$ 控制生成的随机性。

通过引入标准正态分布的噪声 $\epsilon$，采样过程变得简单且可微。

3.4.3. Scope

对于哪些 $q_{\varphi }(z|x)$，我们可以选择这种可微变换 $g_{\varphi }(\cdot )$ 和辅助变量 $\epsilon \sim p(\epsilon )$。以下三种基本方法可行：

1. 可逆累计分布函数（CDF）：
   如果分布的逆CDF是可求的，可以令 $\epsilon \sim U(0,1)$，并令 $g_{\varphi }(\epsilon ,x)$ 为分布 $q_{\varphi }(z|x)$ 的逆CDF。例如：指数分布、洛伦兹分布（Cauchy）、对数正态分布等。
2. 位置-尺度分布：
   对于任何“位置-尺度”分布族，可以选择标准分布（位置=0，尺度=1）作为辅助变量 $\epsilon$，并定义 $g(\cdot )=\text{位置}+\text{尺度}\cdot \epsilon$。例如：高斯分布、拉普拉斯分布、学生 t 分布等。
3. 组合变换：
   一些复杂分布可以通过辅助变量的多种变换表达。例如：
   • 对数正态分布：通过指数化正态分布的变量得到。
   • Gamma分布：通过指数分布变量的加和得到。
   • Dirichlet分布：通过加权的Gamma分布变量得到。

---

总结

Auto-Encoding Variational Bayes (AEVB) 提出了一个高效的框架，将传统变分推断与深度学习结合，为潜变量模型的训练与推断提供了一种自动化的方法。通过 重参数化技巧 和 证据下界（ELBO） 的优化，该方法解决了复杂后验分布的推断问题，并实现了生成建模的高效优化。AEVB 的思想奠定了 变分自编码器（VAE） 的基础，极大地推动了生成模型领域的发展。