直面联合分布

出发点依然没变，这里再重述一下。首先我们有一批数据样本 { x 1 , … , x n } \{x_1, \dots, x_n\} {x1​,…,xn​} ，其整体用 $x$ 来描述，我们希望借助隐变量 $z$ 描述 $x$ 的分布 $\tilde{p}(x)$ ：

$$q(x)=\int q(x|z)q(z)dz,q(x,z)=q(x|z)q(z)\text{\hspace{1em}(7)}$$

这里 $q(z)$ 是先验分布（标准正态分布），目的是希望 $q(x)$ 能逼近 $\tilde{p}(x)$ 。这样（理论上）我们既描述了 $\tilde{p}(x)$ ，又得到了生成模型 $q(x|z)$ ，一举两得。

• $\tilde{p}(x)$ 就是你手里那批真实数据在物理世界中真实存在的客观概率密度
  • 如果是训练的样本是 0-9 的手写数字集合，那么就代表这种手写数字集本质上内在的客观规律。比如可能是某个高维空间中的 10 个小簇

但事实上，直接来对 $p(x,z)$ 进行近似是最为干脆的。具体来说，定义 $p(x,z)=\tilde{p}(x)p(z|x)$ ，我们设想用一个联合概率分布 $q(x,z)$ 来逼近 $p(x,z)$ ，那么我们用KL散度来看它们的距离：

$$KL(p(x,z)\parallel q(x,z))=\iint p(x,z)\ln \frac{p(x,z)}{q(x,z)}dzdx\text{\hspace{1em}(8)}$$

• $p(x,z)$ 的物理意义可以高度概括为：真实世界中，“观测到的表象”与其 “隐蔽的本质特征” 的完美因果配对。
  • 如果某个 $(x,z)$ 的组合在 $p(x,z)$ 下的概率很大，意味着：这张图在现实里很常见，并且它的底层特征确实就是这一组特定的 $z$
  • 如果概率极小，意味着要么这张图本身就不存在（$\tilde{p}(x)$ 极小），要么这张图的特征根本不是你猜的这个 $z$
• 公式 (8) 算的就是 $KL(p(x,z)\parallel q(x,z))$。它的物理意义非常直接：模型在解码侧学习到的 运作规律 $q(x,z)$，在数学距离上逼近 模型在解码端学习到的规律 $p(x,z)$。

KL散度是我们的终极目标，因为我们希望两个分布越接近越好，所以KL散度越小越好。当然，由于现在 $p(x,z)$ 也有参数，所以不单单是 $q(x,z)$ 来逼近 $p(x,z)$ ， $p(x,z)$ 也会主动来逼近 $q(x,z)$ ，两者是相互接近。

我们有 $$p(x,z)=\tilde{p}(x)p(z|x)$$

• $p(z|x)$：在给定一张真实图片 $x$ 的情况下，去推断它的底层特征 $z$。因为真实的推断过程太复杂我们算不出来，所以我们用一个神经网络（编码器）来拟合它。
  $$q(x,z)=q(z)q(x|z)$$
• $q(x|z)$： 这是给定底层特征 $z$ 后，把图像画出来的过程。同样，我们用一个神经网络（解码器）来拟合它。

于是我们有

$$\begin{array}{rl}KL(p(x,z)\parallel q(x,z))& =\int \tilde{p}(x)\left[\int p(z|x)\ln \frac{\tilde{p}(x)p(z|x)}{q(x,z)}dz\right]dx\\ & ={\mathbb{E}}_{x\sim \tilde{p}(x)}\left[\int p(z|x)\ln \frac{\tilde{p}(x)p(z|x)}{q(x,z)}dz\right]\end{array}\text{\hspace{1em}(9)}$$

这样一来利用 (4) 式，把各个 $x_i$ 代入就可以进行计算了，这个式子还可以进一步简化，因为 $\ln \frac{\tilde{p}(x)p(z|x)}{q(x,z)}=\ln \tilde{p}(x)+\ln \frac{p(z|x)}{q(x,z)}$ ，而

$$\begin{array}{rl}{\mathbb{E}}_{x\sim \tilde{p}(x)}\left[\int p(z|x)\ln \tilde{p}(x)dz\right]& ={\mathbb{E}}_{x\sim \tilde{p}(x)}\left[\ln \tilde{p}(x)\int p(z|x)dz\right]\\ & ={\mathbb{E}}_{x\sim \tilde{p}(x)}[\ln \tilde{p}(x)]\end{array}\text{\hspace{1em}(10)}$$

注意这里的 $\tilde{p}(x)$ 是根据样本 $x_1,x_2,\dots ,x_n$ 确定的关于 $x$ 的先验分布，尽管我们不一定能准确写出它的形式，但它是确定的、存在的，因此这一项只是一个常数，所以可以写出

$$\mathcal{L}=KL(p(x,z)\parallel q(x,z))-\text{常数}={\mathbb{E}}_{x\sim \tilde{p}(x)}\left[\int p(z|x)\ln \frac{p(z|x)}{q(x,z)}dz\right]\text{\hspace{1em}(11)}$$

目前最小化 $KL(p(x,z)\parallel q(x,z))$ 也就等价于最小化 $\mathcal{L}$ 。注意减去的常数为 ${\mathbb{E}}_{x\sim \tilde{p}(x)}[\ln \tilde{p}(x)]$ ，所以 $\mathcal{L}$ 拥有下界 $-{\mathbb{E}}_{x\sim \tilde{p}(x)}[\ln \tilde{p}(x)]$ ~注意到 $\tilde{p}(x)$ 不一定是概率，在连续情形时 $\tilde{p}(x)$ 是概率密度，它可以大于 1 也可以小于 1 ，所以 $-{\mathbb{E}}_{x\sim \tilde{p}(x)}[\ln \tilde{p}(x)]$ 不一定是非负，即loss可能是负数。

到这里，我们回顾初衷——为了得到生成模型，所以我们把 $q(x,z)$ 写成 $q(x|z)q(z)$，于是就有

$$\begin{array}{rl}\mathcal{L}=& {\mathbb{E}}_{x\sim \tilde{p}(x)}\left[\int p(z|x)\ln \frac{p(z|x)}{q(x|z)q(z)}dz\right]\\ =& {\mathbb{E}}_{x\sim \tilde{p}(x)}\left[-\int p(z|x)\ln q(x|z)dz+\int p(z|x)\ln \frac{p(z|x)}{q(z)}dz\right]\end{array}\text{\hspace{1em}(12)}$$

再简明一点，那就是

$$\begin{array}{rl}\mathcal{L}=& {\mathbb{E}}_{x\sim \tilde{p}(x)}\left[{\mathbb{E}}_{z\sim p(z|x)}[-\ln q(x|z)]+{\mathbb{E}}_{z\sim p(z|x)}\left[\ln \frac{p(z|x)}{q(z)}\right]\right]\\ =& {\mathbb{E}}_{x\sim \tilde{p}(x)}\left[{\mathbb{E}}_{z\sim p(z|x)}[-\ln q(x|z)]+KL(p(z|x)\Vert q(z))\right]\end{array}\text{\hspace{1em}(13)}$$

括号内的不就是 VAE 的损失函数，只不过我们换了个符号而已。我们就是要想办法找到适当的 $q(x|z)$（解码器） 和 $p(z|x)$（编码器），使得 $\mathcal{L}$ 最小化。

• 我们有一批数据样本 { x 1 , … , x n } \{x_1, \dots, x_n\} {x1​,…,xn​}。在现实世界中，这批数据并不是纯随机的杂音，它们内在蕴含着某种极度复杂但客观存在的统计规律。这个客观存在的、产生这批真实样本的底层概率密度函数，就是 $\tilde{p}(x)$。
• ${\mathbb{E}}_{x\sim \tilde{p}(x)}$：在我的理解中，这是对所有的数据样本进行采样。也就是遍历一遍。
• ${\mathbb{E}}_{z\sim p(z|x)}[-\ln q(x|z)]$：重建误差，经过压缩和解压缩之后，恢复的效果如何。
• $KL(p(z|x)\Vert q(z))$：我们希望后验分布 $p(z|x)$ 靠近我们的先验分布 $q(z)$

• 极端情况 A：强行使 KL 散度最小化（趋于 0）

  • 数学表现： 当 $KL(p(z|x)||q(z))=0$ 时，意味着对于任意的输入数据 $x$，编码器输出的后验分布 $p(z|x)$ 都严格等同于先验分布 $q(z)$（即标准正态分布 $\mathcal{N}(0,\mathbf{I})$）。
  • 信息论后果： 这表示隐变量 $z$ 的分布不再依赖于具体的输入 $x$。换言之，$z$ 与 $x$ 之间的互信息量降为 0，$z$ 丢失了所有关于 $x$ 的特征辨识度。
  • 对另一项的影响： 既然 $z$ 中不包含任何 $x$ 的有效信息，解码器 $q(x|z)$ 就无法基于 $z$ 还原出 $x$。因此，重构误差 $-\ln q(x|z)$ 必然会变得极大。

• 极端情况 B：强行使重构误差最小化（趋于 0）

  • 数学表现： 为了让 $-\ln q(x|z)$ 尽可能小，解码器需要获得极其精确且稳定的输入特征。这就要求编码器 $p(z|x)$ 输出的方差必须极小，使其退化为一个几乎没有随机性的确定性映射（Dirac delta function），从而消除采样过程带来的不确定性。
  • 信息论后果： 隐变量 $z$ 完全记住了 $x$ 的确定性特征，网络退化为传统的非随机自编码器。
  • 对另一项的影响： 一个方差趋于 0 的分布，与方差为 1 的标准正态先验分布 $q(z)$ 存在巨大的差异。这种极度的分布偏离，必然导致 $KL(p(z|x)||q(z))$ 的值急剧上升，甚至趋于无穷大。

因此，必须把 $\mathcal{L}$ 作为一个整体目标来优化。总 $\mathcal{L}$ 的持续变小才代表模型真正收敛。VAE 天生具备这样一个全局统一的收敛指标，这也是它相比于早期 GAN 模型（直到 WGAN 出现前）的一大优势。

客观假设

实际文章做了三个大胆的假设，

• 隐变量的先验分布 $q(z)$：假设它是标准正态分布 $\mathcal{N}(0,\mathbf{I})$。这是为了让隐空间有一个统一、规则的“锚点”，防止隐变量乱跑；为了后面算 KL 散度计算有简单的公式
• 后验分布 $p(z|x)$ 是正态分布：假设它是各分量独立的多元正态分布 $\mathcal{N}(\mathbf{\mu }(x),{\mathbf{\sigma }}^2(x))$
• 解码器的生成分布 $q(x|z)$：
  • 连续数据（如彩色自然图像）：我们确实假设它是正态分布。
  • 离散/二值数据（如黑白的 MNIST 手写体）： 我们就不能假设它是正态分布了，而是假设它服从伯努利分布（Bernoulli Distribution）

编码器

首先，为了便于采样，我们假设 $z\sim N(0,I)$，即标准的多元正态分布。然后，我们 假设 $p(z|x)$ 是 各分量独立的 正态分布，其均值和方差由 $x$ 来决定，这个“决定”，就是一个神经网络：

$$p(z|x)=\frac{1}{{\prod }_{k=1}^d\sqrt{2\pi {\sigma }_{(k)}^2(x)}}\exp \left(-\frac{1}{2}{\Vert \frac{z-\mu (x)}{\sigma (x)}\Vert }^2\right)\text{\hspace{1em}(15)}$$

这里的 $\mu (x),{\sigma }^2(x)$ 是输入为 $x$ 、输出分别为均值和方差的神经网络，其中 $\mu (x)$ 就起到了类似 encoder 的作用。既然假定了高斯分布，那么 (13) 式中的 KL 散度这一项就可以先算出来：

$$KL(p(z|x)\parallel q(z))=\frac{1}{2}\sum _{k=1}^d\left({\mu }_{(k)}^2(x)+{\sigma }_{(k)}^2(x)-\ln {\sigma }_{(k)}^2(x)-1\right)\text{\hspace{1em}(16)}$$

也就是我们所说的 KL loss，这在上一篇文章已经给出。

解码器

接下来我们考虑解码器

$$q(x|z)=\frac{1}{{\prod }_{k=1}^D\sqrt{2\pi {\tilde{\sigma }}_{(k)}^2(z)}}\exp \left(-\frac{1}{2}{\Vert \frac{x-\tilde{\mu }(z)}{\tilde{\sigma }(z)}\Vert }^2\right)\text{\hspace{1em}(20)}$$

这里的 $\tilde{\mu }(z),{\tilde{\sigma }}^2(z)$ 是输入为 $z$ 、输出分别为均值和方差的神经网络， $\tilde{\mu }(z)$ 就起到了 decoder 的作用。于是

$$-\ln q(x|z)=\frac{1}{2}{\Vert \frac{x-\tilde{\mu }(z)}{\tilde{\sigma }(z)}\Vert }^2+\frac{D}{2}\ln 2\pi +\frac{1}{2}\sum _{k=1}^D\ln {\tilde{\sigma }}_{(k)}^2(z)\text{\hspace{1em}(21)}$$

很多时候我们会固定方差为一个常数 ${\tilde{\sigma }}^2$ ，这时候

$$-\ln q(x|z)\sim \frac{1}{2{\tilde{\sigma }}^2}{\Vert x-\tilde{\mu }(z)\Vert }^2\text{\hspace{1em}(22)}$$

这就出现了 MSE 损失函数。对于一般数据，我们用 MSE 作为损失函数，这对应于 $q(x|z)$ 为固定方差的正态分布。

重参数化的技巧

假设 $L$ 是我们要优化的损失函数（例如重构误差），$\varphi$ 是编码器（神经网络）的权重参数。编码器根据输入 $x$ 计算出分布的参数：均值 $\mu$ 和标准差 $\sigma$。即：$\mu =f_{\varphi }(x)$，$\sigma =g_{\varphi }(x)$。我们要计算的目标是损失函数 $L$ 对网络参数 $\varphi$ 的梯度：$\frac{\partial L}{\partial \varphi }$。

1. 重参数化之前：在原始的直接采样中，计算图的路径是：$$\varphi ⟶(\mu ,\sigma )\stackrel{\text{随机采样}}{⟶}z⟶L$$根据多元微积分的链式法则，我们要将 $\frac{\partial L}{\partial \varphi }$ 完全展开：$$\frac{\partial L}{\partial \varphi }=\frac{\partial L}{\partial z}\cdot \frac{\partial z}{\partial \varphi }$$$$=\frac{\partial L}{\partial z}\cdot \left(\frac{\partial z}{\partial \mu }\cdot \frac{\partial \mu }{\partial \varphi }+\frac{\partial z}{\partial \sigma }\cdot \frac{\partial \sigma }{\partial \varphi }\right)$$我们来逐项检查这个公式中的偏导数：

• $\frac{\partial L}{\partial z}$：可导。这是解码器的反向传播，完全由确定的矩阵运算组成。
• $\frac{\partial \mu }{\partial \varphi }$ 和 $\frac{\partial \sigma }{\partial \varphi }$：可导。这是编码器的内部反向传播。
• $\frac{\partial z}{\partial \mu }$ 和 $\frac{\partial z}{\partial \sigma }$：不可导（未定义）。

在 $z\sim \mathcal{N}(\mu ,{\sigma }^2)$ 的操作中，$z$ 是抛掷由 $\mu$ 和 $\sigma$ 决定的“概率骰子”得到的结果。如果 $\mu$ 改变了 $\Delta \mu$，“骰子的重心”变了，但由于每次抛骰子都是独立的随机事件，你无法写出一个确定性的代数表达式来描述 $\Delta z$ 和 $\Delta \mu$ 之间的比例关系。

2. 重参数化之后
   在引入重参数化技巧后，我们改变了 $z$ 的生成方式：$$z=\mu +\sigma \cdot ϵ(\text{其中 }ϵ\sim \mathcal{N}(0,1))$$此时的计算图路径变成了：$$\varphi ⟶(\mu ,\sigma )\stackrel{\text{与 }ϵ\text{ 进行确定性运算}}{⟶}z⟶L$$我们再次根据链式法则展开 $\frac{\partial L}{\partial \varphi }$：$$\frac{\partial L}{\partial \varphi }=\frac{\partial L}{\partial z}\cdot \left(\frac{\partial z}{\partial \mu }\cdot \frac{\partial \mu }{\partial \varphi }+\frac{\partial z}{\partial \sigma }\cdot \frac{\partial \sigma }{\partial \varphi }\right)$$现在，我们重新审视曾经导致断裂的那两项。因为 $z$ 变成了关于 $\mu$ 和 $\sigma$ 的显式、确定性代数方程，我们可以直接对它们求偏导：

   • 对 $\mu$ 求偏导：$$\frac{\partial z}{\partial \mu }=\frac{\partial }{\partial \mu }(\mu +\sigma \cdot ϵ)=1$$
   • 对 $\sigma$ 求偏导：$$\frac{\partial z}{\partial \sigma }=\frac{\partial }{\partial \sigma }(\mu +\sigma \cdot ϵ)=ϵ$$

   联通的证明：将上述存在且极其简单的偏导数结果代回链式法则公式中，我们得到了一个完美的、完全可计算的解析表达式：$$\frac{\partial L}{\partial \varphi }=\frac{\partial L}{\partial z}\cdot \left(1\cdot \frac{\partial \mu }{\partial \varphi }+ϵ\cdot \frac{\partial \sigma }{\partial \varphi }\right)$$