VAE与扩散模型 - VAE 的计算

4 . VAE的计算

我们先来整理 VAE 接收到一张图片后，究竟做了什么。
假设我们有一张 \(2\times2\) 的灰度图像。这张图片可以表示成一个向量：\(x=[0,1,1,0]\) 。所以，\(x\) 就是输入图像对应的数值数据。
我们设计了一个非常简单的规定潜变量分布采用高斯分布的 Encoder，于是它的输入是四个数字： \(x=[0,1,1,0]\) 。经过神经网络计算，Encoder 输出的是两个数：均值 \(\mu=3\) 和标准差 \(\sigma=2\)
我们规定 \(q_\phi(z\mid x)=\mathcal N(z;3,4)\) ，说明给定输入图像 \(x=[0,1,1,0]\)，Encoder 认为潜变量 \(z\) 服从一个均值为 3 、方差为 4 的正态分布。这里 \(q_\phi(z\mid x)\) 表示的是在输入图像 \(x\) 已知的情况下，由参数 \(\phi\) 控制的 Encoder 给出的潜变量 \(z\) 的条件概率密度。
既然：
\(Z\mid x\sim\mathcal N(3,4)\)
那就意味着：

• 均值是 3。
• 标准差是 2。
• \(z\) 可以取很多不同的值。
• 靠近 3 的值，其概率密度通常较高。
  采样在 PyTorch 中怎么实现？ VAE 有一个关键技术：重参数化技巧（Reparameterization Trick）。我们可以先从标准正态分布中采样：\(\epsilon\sim\mathcal N(0,1)\) ，再计算：\(z=\mu+\sigma\epsilon\) 。假设本次采样\(\epsilon=-0.5\)，\(z=3-2\times 0.5=2\)。这就取到了一个 \(z\) 的值。重参数化把随机性放到独立的 \(\epsilon\) 中，这样反向传播就可以计算重建损失对均值、标准差以及 Encoder 参数的梯度。
  接下来，我们把 \(z=3\) 输入 Decoder。从代码结构来看，Decoder 可以是普通的全连接神经网络，也可以包含卷积层、上采样层等。如果 Decoder 经过神经网络输出了4个数字：\([0.1,0.8,0.7,0.2]\) 。设 Decoder 输出的每个数字，表示相应像素取值为 1 的概率。
  现在我们已经知道：
• 潜变量：\(z=3\) 。
• Decoder 输出的像素概率：\([0.1,0.8,0.7,0.2]\) 。
• 真实图像：\(x=[0,1,1,0]\) 。
  我们想计算：
  \(p_\theta(x\mid z=3)\)
  它表示：已知潜变量 (z=3)，Decoder 给真实图像 \(x=[0,1,1,0]\) 分配的条件概率是多少？
  解：\(\begin{aligned} p_\theta(x\mid z=3) &=0.9\times0.8\times0.7\times0.8\\ &=0.4032 \end{aligned}\)
  我们希望 Decoder 能够正确重建 \(x=[0,1,1,0]\) 。
  如果模型给真实图像分配的概率较低，那么说明它没有很好地解释当前数据。
  于是使用负对数似然：\(\mathcal L_{\mathrm{recon}}=-\log p_\theta(x\mid z)\) 代入刚才的结果：
  \(\mathcal L_{\mathrm{recon}} =-\log0.4032\approx0.9084\)
  训练时，通过反向传播调整 Decoder 参数 \(\theta\)，让它更倾向于生成真实图像对应的像素。例如训练后，Decoder 可能输出：\([0.01,0.99,0.98,0.02]\)
  这组概率显然比之前更符合真实图像：\([0,1,1,0]\)
  因此，真实图像的条件概率提高，重建损失降低。
  当我们训练完成，我们就可以不用输入图片就能生成一张新的图像。

上文引出了\(\mathcal L=\underbrace{\mathbb E_q[-\log p_\theta(x\mid z)]}_{\text{重建损失}}+\underbrace{\mathbb E_q\left[\log\frac{q_\phi(z\mid x)}{p(z)}\right]}_{\text{KL 散度}}\) 这个重要的式子。
\(p_\theta(x\mid z)\)：Decoder 定义的图像条件分布。
\(q_\phi(z\mid x)\)：近似分布。在输入图像 \(x\) 已知的情况下，由参数 \(\phi\) 控制的 Encoder 给出的潜变量 \(z\) 的条件概率密度。
\(p(z)\)：潜变量 \(z\) 的概率密度。一般人为设定为高斯分布。
我们设\(p_\theta(x\mid z)\)、\(q_\phi(z\mid x)\)、\(p(z)\) 都是多维的高斯分布。首先，来计算 \(\mathbb E_q[-\log p_\theta(x\mid z)]\) 这个式子。这很难算。但我们可以使用蒙特卡洛采样。\(\mathbb E_q[-\log p_\theta(x\mid z)]=\frac{1}{L}\sum^{L}_{n=1}-\log p_\theta(x\mid z_n)\)
然而直接在多维的 \(z\) 上采样是不能进行梯度更新的，这让我们的神经网络学习变得非常困难。还记得上面提到的重参数化技巧（Reparameterization Trick）吗？将 \(z=\mu+\sigma\epsilon\) ，通过采样 \(\epsilon\) 来达到采样 \(z\) 的效果。
多维高斯概率密度为：
\(p_\theta(x\mid z) = \frac{1} {(2\pi)^{D/2}|\Sigma|^{1/2}} \exp\left( -\frac12(x-\mu)^T\Sigma^{-1}(x-\mu) \right)\)
代入重建损失得到：
\(\begin{aligned} -\log p_\theta(x\mid z) ={}&\frac12(x-\mu)^T\Sigma^{-1}(x-\mu)+\frac12\log|\Sigma|+\frac D2\log(2\pi) \end{aligned}\)
为了简化计算，我们把高斯分布里的协方差设为 \(I\) 。\(\mid I\mid=1\)。
\(-\log p_\theta(x\mid z) = \frac12\sum_{j=1}^{D}(x_j-\mu_j)^2 +\frac D2\log(2\pi)\)
我们把 \(L\) 设为1，只采样一次。所以\(\mathcal L_{\mathrm{recon}} \propto \sum_{j=1}^{D}(x_j-\hat \mu_j)^2\)
这居然是一个平方误差 \(MSE\) ?

最后，是KL散度。\(\mathbb E_q\left[\log\frac{q_\phi(z\mid x)}{p(z)}\right]\)
我们可以写成 \(\mathbb D_{KL}(q_\phi(z \mid x)\mid\mid p(z))\)
代入公式 \(D_{\mathrm{KL}}(q\|p) = \log\frac{\sigma_p}{\sigma_q} -\frac12 + \frac{\sigma_q^2+(\mu_q-\mu_p)^2} {2\sigma_p^2}\)
又因为 \(p(z)\) 被设为标准正态分布，\(D_{\mathrm{KL}}(q\|p) = \log\frac{1}{\sigma_q} -\frac12 + \frac{\sigma_q^2+\mu_q^2} {2}\)
整理得 \(D_{\mathrm{KL}}(q\|p) = \frac12 (\mu^2+\sigma^2-1-\log\sigma^2)\)
在多个维度中，可以变成 \(\mathbb D_{KL}\sum_{i=1}^{d}(q_\phi(z_i \mid x)\mid\mid p(z_i))\) ，将计算的 \(D_{KL}\) 全部加起来。
\(D_{\mathrm{KL}}(q\|p) = \frac12 \sum^{d}_{i=1}(\mu^2+\sigma^2-1-\log\sigma^2)\) 这就是KL散度的最终形态。