VAE(变分自编码器) 详解

近期看论文要用到VAE，看了很多资料，有这样一种感觉，要么过度过于偏向数学原理，要么只是讲了讲网络结构。本文将两者结合，以简洁易懂的语言结合代码实现来介绍VAE。

1 解决问题

VAE是变分推断（variational inference ）以及自编码器（Auto-encoder）的组合，是一种非监督的生成模型，其概率图模型和深度学习有机结合，近年来比较火热。VAE可以用于但不局限于降维信息检索等任务，我看文献遇到的是一篇做配准的论文，也用到了VAE。

2 从神经网络理解

先从神经网络的角度去看VAE,VAE实际上是在Atuo-encoder（AE）的变种，其基本架构也如Atuo-encoder，包含两部分encoder（编码器）和decoder（解码器）。大概如下：

图片摘自：github
以最开始的自编码器为例，其loss函数一般是输入和输出的MSE,通过调整Encoder输出层的节点数（需要低于输入的维度），我们可以从低维度的数据（code）通过Decoder重建出输入。

自编码器存在这样的问题，倘若模型过完备（中间层维度大于输入），模型会直接复制模型的输入作为输出。

在一般实际使用中，我们往往会添加正则项。
除此之外，还有VAE,可以学习出高容量且过完备（中间层维度大于输入）模型。VAE的网络结构如下：

本图摘自：李宏毅2020深度学习课程
可以看到VAE和AE的区别在于两方面：
1.中间层引入了一个noise；
2.loss函数的改变多了：${\sum }_{i=1}^3(xep({\sigma }_i-(1+{\sigma }_i)+(m_i^2)))$这一项。根据上面的结构，我们基本可以很容易地代码实现。

那么如何直观地理解上面地改变呢？
1.为什么要引入noise？首先直观理解，就算有一个noise也要尽量输入和输出相似，这样的decoder更加鲁棒。另一个直观理解就是，在引入noise之前，我们的decoder的输入和输出地映射可以看作是离散的，但是在加入noise之后，可以看作把不连续地变成连续的了。
2.为什么更改loss？首先假设没有更改，其实最好的情况肯定是方差为0，即$\sigma =0$。这就回到了AE的形式。直观地说，多加地这一项避免了这一点。那么怎么要这样更改loss呢？首先公式地前两项的值域大于等于0，最后一项可以看作一个L2正则项。

上面的理解都是直观的，感性的认识。下面一部分将从变分推断的角度推导出所谓的引入noise其实是重参数化技巧的结果，而loss的改变也是推导所得，本质是一个KL散度。如果读者到目前为止还有求知欲，就可以继续往下看。

3 从变分推断理解VAE

设：x:为观测数据，可以看作是样本
y：为隐变量，包含但不限于模型的参数
首先变分推断的核心思想是：因为一般情况下后验概率$p(z|x)$是不可求解的，所以变分推断采用了一种迂回的策略，即使用$q(z)$去近似$p(z|x)$。如果读者对生成模型不是很理解，可以把$p(z|x)$看作是AE中的编码器，$p(x|z)$看作是AE中的解码器。
VAE是典型的生成模型，那就从下面公式开始：
$$\begin{array}{rl}\log p(x)& =\log \frac{p(x,z)}{p(z|x)}\\ & =\log \frac{p(x,z)}{q(z)}-\log \frac{p(z|x)}{q(z)}\end{array}$$
两边关于$q(z)$同时求期望，则：$$左边={\int }_{z}q(z)\log p(x)dz=\log p(x)$$$$\begin{array}{rl}右边& ={\int }_{z}q(z)\log \frac{p(x,z)}{q(z)}dz-{\int }_{z}q(z|x)\log \frac{p(z|x)}{q(z)}dz\\ & =\mathcal{L}(q)+KL(q(z)||p(z|x))\end{array}$$其中右边化简后的$\mathcal{L}(q)$通常被称为变分下界，$KL(q(z)||p(z|x))$是KL散度，用来衡量两个分布之间相似性。
对$\mathcal{L}(q)$做进一步化简，$$\begin{array}{rl}\mathcal{L}(q)& ={\int }_{z}q(z)\log \frac{p(x,z)}{q(z)}dz\\ & ={\int }_{z}q(z)\log \frac{p(x,z)q(z)}{q(z)p(z)}-{\int }_{z}q(z|x)\log q(z)p(z)\\ & ={\mathbb{E}}_{q(z)}\log p(x|z)-KL(q(z)||p(z))\end{array}$$回到最初的目标，我们要用$q(z)$去近似$p(z|x)$，那么我们优化的目标就是最大化$KL(q(z)||p(z|x))$。即：$$\begin{array}{rl}q(z|x)& =arg\underset{q(z)}{\max }KL(q(z)||p(z|x))\\ & =arg\underset{q(z)}{\max }\log p(x)-\mathcal{L}(q)\\ & =arg\underset{q(z)}{\max }\log p(x)-{\mathbb{E}}_{q(z)}\log p(x|z)+KL(q(z)||p(z))\end{array}$$看到这里，我们可能会想到。$\log p(x)-{\mathbb{E}}_{q(z|x)}\log p(x|z)$不就是真实的X预测的X之间的差异吗？我们用MSE或者MAE代替，后面的KL散度根据预测和假设的结果不是可以直接算出来，就已经解释了上一节的网络结构。但是细想一下感觉又不对，因为这里虽然表示为q(z),但是q(z)要近似的是p(z|x),所以这个z必定和x有关，这样的话上面的想法就不对了。
在此之前我们把q(z)表示为$q(z|x)$，这时通常要使用重参数化技巧对上式进一步变形，现在想要把$q(z|x)$中的x的成分消去。那么上面是重参数化技巧呢？先举个例子吧，一个随机变量a服从概率分布N(0,1),那么对于随机变量b=a+m，服从高斯分布N(m,1)。现在我们采样这个b的时候采用这样的策略：
1.从高斯分布N(0,1)中采样得a。
2.取b = a+m。
其实这就是重采样技巧。对于我们的$q(z|x)$，我们假设$z=g_{\Phi }(x,ϵ)$，然后$ϵ$服从某个分布，记为$p(ϵ)$，一般我们假设其服从标准正态分布。那么我们采样$q(z|x)$就变成了，先根据$p(ϵ)$采样一个${ϵ}^i$,再根据$z=g_{\Phi }(x,ϵ)$计算出z。根据重采样技巧，我们忘掉之前的结果，重新推导KL(q||p)：$$\begin{array}{rl}KL(q(z|x)||p(z|x))& =\log p(x)-{\int }_{z}q(z)\log \frac{p(x,z)q(z|x)}{q(z|x)p(z)}+{\int }_{z}q(z|x)\log q(z|x)p(z)\\ & =\log p(x)-{\int }_{z}q(z)\log p(x|z)dz+KL(q(z|x)||p(z))\\ & =\log p(x)-{\int }_{ϵ}p(ϵ)\log p(x|g_{\Phi }(x,ϵ))dϵ+KL(q(z|x)||p(z))\\ & =\log p(x)-{\mathbb{E}}_{p(ϵ)}\log p(x|g_{\Phi }(x,ϵ))+KL(q(z|x)||p(z))\end{array}$$
其实整个VAE的构建就是根据上面的等式
$g_{\Phi }(x,ϵ)$不知道是什么，那就用一个神经网络代替。$p(x|z)$不知道是什么，也用一个神经网络代替。下面文字叙述一下VAE的前向传播。
1.先从假设的$p(ϵ)$中采样一个$ϵ$，即上一节网络图中的$e$。
2.从假设的encoder中输入x以及$ϵ$,输出隐变量$z$，即上一节网络图中的$c$。
3.将隐变量z输入decoder，输出$\hat{x}$。
而这个前向的过程表示在上面公式里，就是${\mathbb{E}}_{p(ϵ)}\log p(x|g_{\Phi }(x,ϵ))dϵ$，显然优化这个网络，我们要让$KL(q(z|x)||p(z|x))$最小，$\log p(x)-{\mathbb{E}}_{p(ϵ)}\log p(x|g_{\Phi }(x,ϵ))$就用MSE表示，$KL(q(z|x)||p(z))$是可以求出来的。具体的推导也不难，如果假设是高斯分布，即根据多维高斯分布的KL散度，结合我们重采样的q(z|x)，推导出来最后的结果就是第一节中图中的公式，这里就省略推导了。

能看到这里的宝贝都很厉害，毕竟我感觉自己也写的不是很清楚，才疏学浅了。不过最困难的部分也过去了，我们不妨看看VAE的pytorch代码实现，看看自己理解的是不是对的。

4 pytorch代码

本文的代码来自GITHUB

__author__ = 'SherlockLiao'

import os

import torch
import torchvision
from torch import nn
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST
from torchvision.utils import save_image

if not os.path.exists('./mlp_img'):
    os.mkdir('./mlp_img')


def to_img(x):
    x = 0.5 * (x + 1)
    x = x.clamp(0, 1)
    x = x.view(x.size(0), 1, 28, 28)
    return x


num_epochs = 100
batch_size = 128
learning_rate = 1e-3

img_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5])
])

dataset = MNIST('./data', transform=img_transform)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)


class autoencoder(nn.Module):
    def __init__(self):
        super(autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28 * 28, 128),
            nn.ReLU(True),
            nn.Linear(128, 64),
            nn.ReLU(True), nn.Linear(64, 12), nn.ReLU(True), nn.Linear(12, 3))
        self.decoder = nn.Sequential(
            nn.Linear(3, 12),
            nn.ReLU(True),
            nn.Linear(12, 64),
            nn.ReLU(True),
            nn.Linear(64, 128),
            nn.ReLU(True), nn.Linear(128, 28 * 28), nn.Tanh())

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x


model = autoencoder().cuda()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(
    model.parameters(), lr=learning_rate, weight_decay=1e-5)

for epoch in range(num_epochs):
    for data in dataloader:
        img, _ = data
        img = img.view(img.size(0), -1)
        img = Variable(img).cuda()
        # ===================forward=====================
        output = model(img)
        loss = criterion(output, img)
        # ===================backward====================
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    # ===================log========================
    print('epoch [{}/{}], loss:{:.4f}'
          .format(epoch + 1, num_epochs, loss.data[0]))
    if epoch % 10 == 0:
        pic = to_img(output.cpu().data)
        save_image(pic, './mlp_img/image_{}.png'.format(epoch))

torch.save(model.state_dict(), './sim_autoencoder.pth')