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是典型的生成模型，那就从下面公式开始：
log ⁡ p ( x ) = log ⁡ p ( x , z ) p ( z ∣ x ) = log ⁡ p ( x , z ) q ( z ) − log ⁡ p ( z ∣ x ) q ( z )

$$\begin{aligned}\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{aligned}$$ logp(x)​=logp(z∣x)p(x,z)​=logq(z)p(x,z)​−logq(z)p(z∣x)​​
两边关于$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 torch
import torchvision
from torch import nn
from torch import optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.utils import save_image
from torchvision.datasets import MNIST
import os

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


def to_img(x):
    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, 0.5), (0.5, 0.5, 0.5))
])

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


class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()

        self.fc1 = nn.Linear(784, 400)
        self.fc21 = nn.Linear(400, 20)
        self.fc22 = nn.Linear(400, 20)
        self.fc3 = nn.Linear(20, 400)
        self.fc4 = nn.Linear(400, 784)

    def encode(self, x):
        h1 = F.relu(self.fc1(x))
        return self.fc21(h1), self.fc22(h1)

    def reparametrize(self, mu, logvar):
        std = logvar.mul(0.5).exp_()
        if torch.cuda.is_available():
            eps = torch.cuda.FloatTensor(std.size()).normal_()
        else:
            eps = torch.FloatTensor(std.size()).normal_()
        eps = Variable(eps)
        return eps.mul(std).add_(mu)

    def decode(self, z):
        h3 = F.relu(self.fc3(z))
        return F.sigmoid(self.fc4(h3))

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparametrize(mu, logvar)
        return self.decode(z), mu, logvar


model = VAE()
if torch.cuda.is_available():
    model.cuda()

reconstruction_function = nn.MSELoss(size_average=False)


def loss_function(recon_x, x, mu, logvar):
    """
    recon_x: generating images
    x: origin images
    mu: latent mean
    logvar: latent log variance
    """
    BCE = reconstruction_function(recon_x, x)  # mse loss
    # loss = 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    KLD_element = mu.pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar)
    KLD = torch.sum(KLD_element).mul_(-0.5)
    # KL divergence
    return BCE + KLD


optimizer = optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(num_epochs):
    model.train()
    train_loss = 0
    for batch_idx, data in enumerate(dataloader):
        img, _ = data
        img = img.view(img.size(0), -1)
        img = Variable(img)
        if torch.cuda.is_available():
            img = img.cuda()
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(img)
        loss = loss_function(recon_batch, img, mu, logvar)
        loss.backward()
        train_loss += loss.data[0]
        optimizer.step()
        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch,
                batch_idx * len(img),
                len(dataloader.dataset), 100. * batch_idx / len(dataloader),
                loss.data[0] / len(img)))

    print('====> Epoch: {} Average loss: {:.4f}'.format(
        epoch, train_loss / len(dataloader.dataset)))
    if epoch % 10 == 0:
        save = to_img(recon_batch.cpu().data)
        save_image(save, './vae_img/image_{}.png'.format(epoch))

torch.save(model.state_dict(), './vae.pth')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123