VAE变分自编码器原理推导+Python代码实现

1、前言

变分自编码器是近些年较火的一个生成模型，我个人认为其本质上仍然是一个概率图模型，只是在此基础上引入了神经网络。本文将就变分自编码器（VAE）进行简单的原理讲解和数学推导。

2、引入

2.1、高斯混合模型

生成模型，可以简单的理解为生成数据$\boxed{(不止，但我们暂且就这么理解它)}$。假如现在我们有样本数据，而我们发现这些样本符合正态分布且样本具有充分的代表性，因此我们计算出样本的均值和方差，就能得到样本的概率分布。然后从正态分布中抽样，就能得到样本。$\boxed{\mathbf{这种生成样本的过程就是生成过程}}$。

可是，假如我们的数据长这样

很显然，它的数据是由两个不同的正态分布构成。我们可以计算出这些样本的概率分布。但是一种更为常见的方法就是将其当作是两个正态分布。我们引入一个隐变量z。

$\boxed{\mathbf{假设}\mathbf{z}\mathbf{的取值为}0，1}$，如果z为0，我们就从蓝色的概率分布中抽样；否则为1，则从橙色的概率分布中抽样。这就是生成过程。

但是这个隐变量z是什么？它其实就是隐藏特征$\boxed{\mathbf{训练数据}\mathbf{x}\mathbf{的抽象出来的特征}}$，比如，如果x偏小，我们则认为它数据蓝色正太分布，否则为橙色。这个$\boxed{“偏小”}$就是特征，我们把它的取值为0，1（0代表偏小，1代表偏大）。

那这种模型我们如何取训练它呢？如何去找出这个z呢？$\boxed{\mathbf{一种很直观的方法就是重构代价最小}}$，我们希望，给一个训练数据x，由x去预测隐变量z，再由隐变量z预测回x，得到的误差最小。比如假如我们是蓝色正态分布，去提取特征z，得到的z再返回来预测x，结果得到的却是橙色的正态分布，这是不可取的。其模型图如下

这个模型被称为GMM高斯混合模型

2.2、变分自编码器（VAE）

那它和VAE有什么关联呢？其实VAE的模型图跟这个原理差不多。只是有些许改变，$\boxed{\mathbf{隐变量}\mathbf{z}的维度是连续且高维的，不再是简单的离散分布}$，因为假如我们生成的是图片，我们需要提取出来的特征明显是要很多的，传统的GMM无法做到。

在VAE中，$q(z|x)$$\boxed{（也就是用样本x预测变量z）}$，其服从高斯分布，接下来，我们来看模型图

也就是将训练样本x给神经网络，让神经网络计算出均值和协方差矩阵$\mu ,\log {\sigma }^2$$\boxed{\mathbf{取}\log \mathbf{的原因是传统的神经网络输出值总是有正有负}}$。有了这两个值就可以在对应的高斯分布中采样，得到隐变量z。再让z经过神经网络重构回样本，得到新样本。这就是整个VAE的大致过程了。

再次强调，$\boxed{训练过程我们希望每次重构的时候，新样本和训练样本尽可能的相似}$

如果我们这样直接去训练的化，可以吗？可以！但是会有个问题，神经网络会趋向于将协方差变0，而让概率分布$q(z|x)$将不再具备随机性，概率分布空间会坍缩成一个点，每次采样都是均值，这种情况我们俗称过拟合。

为什么协方差会变成0？因为采样具有随机性，也就是存在噪声，噪声是肯定会增加重构的误差的。神经网络为了让误差最小，是肯定让这个随机性越小越好，因为只有这样，才能重构误差最小

但是我们肯定是希望有随机性的，为什么？因为有随机性，我们才可以生成不同的样本啊！

所以，对于概率分布$q(z|x)$，我们不希望它的协方差为0，所以我们需要对其进行约束。在论文中，它对其进行约束，要求它尽量的往$N(0,I)$靠近（$\boxed{其实与先验分布P(z)有关，后续数学推导中可见,假设P(z)\sim N(z|0,I)}$）

所以，有KL散度去衡量两个概率分布的相似性
$$\min KL\left(P(z|x)||P(z)\right)$$
KL散度是大于等于0的值，越小则证明越相似
$$KL(q(x)||p(x))={\int }_{x}q(x)\log \frac{q(x)}{p(x)}dx$$
所以，我们就是两个优化目标
$$\begin{gathered} \boxed{①\mathbf{最小化重构代价}} \\ \boxed{②\mathbf{最小化上述的}\mathbf{KL}\mathbf{散度}} \end{gathered}$$
依照这两个条件，建立目标函数，直接梯度下降$\boxed{(其实还需要重参数化，后面会讲到)}$，刷刷刷地往下降，最终收敛。

下面，我们就对其进行简单的数学推导，并以此推导出目标函数

3、原理推导

3.1、引入目标函数

以VAE的简略图为例

定义隐变量先验分布$P(z)\sim N(0,I)$。很自然的想法，我们直接对x求log极大似然，假设我们有N的样本，记作X。设所需求解的参数为$\theta$，第$i$个样本记为$x^i$，某个样本的第$j$个维度记作$x_j$
$$\begin{array}{rl}\log P(X)=& \log \prod _{i=1}^{N}P(x^i)\\ =& \sum _{i=1}^N\log P(x_i)\end{array}$$
现在，我们先单独看看里面某一个样本的似然，某个样本记为$x$
$$\begin{array}{rl}\log P(x)=& \log \frac{P(x,z)}{P(z|x)}\\ =& \log P(x,z)-\log P(z|x)\end{array}$$
等式左右分别对$q(z|x)$求积分
$$\begin{array}{rl}& 左边：{\int }_z\log P(x)q(z|x)dz=\log P(x){\int }_{z}q(z|x)dz=\log P(x)\end{array}$$
所以
$$\begin{array}{rl}\log P(x)=& {\int }_z(\log P(x,z)-\log P(z|x))q(z|x)dz\\ =& {\int }_z(\log \frac{P(x,z)}{q(z|x)}-\log \frac{P(z|x)}{q(z|x)})q(z|x)dz\\ =& {\int }_z\log \frac{P(x,z)}{q(z|x)}q(z|x)dz-\log \frac{P(z|x)}{q(z|x)}q(z|x)dz\\ =& \underset{①}{{\int }_z\log \frac{P(x,z)}{q(z|x)}q(z|x)dz}-\underset{②}{KL(q(z|x)||P(z|x))}\end{array}\text{\hspace{1em}(a)(b)}$$
（式a）到（式b）用到了$\log$的性质。

对于里面的第②项，因为KL是大于等于0的。在给定x的情况下，$\log P(x)$的值是确定的，所以$\log P(x)$的最大值，其实就等价于最大化第①项（最小化第②项），即
$$\max {\int }_z\log \frac{P(x,z)}{q(z|x)}q(z|x)dz\leftrightarrow \min KL(q(z|x)||P(z|x))$$
好，现在我们只需要最大化
$$\begin{array}{rl}& \max {\int }_z\log \frac{P(x,z)}{q(z|x)}q(z|x)dz\\ =& \max {\int }_z\log \frac{P(x|z)P(z)}{q(z|x)}q(z|x)dz\\ =& \max {\int }_z\left(\log \frac{P(z)}{q(z|x)}+\log P(x|z)\right)q(z|x)dz\\ =& \max {\int }_z\log P(x|z)q(z|x)dz-\int \log \frac{q(z|x)}{P(z)}q(z|x)dz\\ =& \max \left(\underset{①}{{\mathbb{E}}_{q(z|x)}\left[\log P(x|z)\right]}-\underset{②}{KL(q(z|x)||P(z))}\right)\end{array}$$

发现了吗，最大化里面的第一项就期望，不就是从$q(z|x)$采样z，再让$\log P(x|z)$概率最大。这不就是重构代价最小吗；而对于第二项，最大化$-KL$散度，就相当于最小化$KL$散度。这和我们上面提到的两个优化目标是一样的。

3.2、细化目标函数

既然得到了目标函数，那么我们就对$\log P(x|z)$似然和KL散度都求出具体的表达。

$\boxed{先来看KL散度}$

前面我们提到，$q(z|x)$需要逼近$P(z)$，而$P(z)\sim N(0,I)$的多维高斯分布，并且各个维度之间相互独立，所以$q(z|x)$也是如此设定。那么最小化其KL散度，只需要对每一个维度求KL最小即可，单独看某一个维度（为了表达的简洁，正态分布中不再写出随机变量，里面的z都表示和$z_j$是同一个东西，懒得改了，知道就行）
$$\begin{array}{rl}& \min KL(q(z_j|x)||P(z_j))\\ =& \min \int q(z_j|x)\log \frac{q(z_j|x)}{P(z_j)}dz\\ =& \min \int N(\mu ,{\sigma }^2)\log \frac{N(\mu ,{\sigma }^2)}{N(0,1)}dz\\ =& \min {\int }_{z}N(\mu ,{\sigma }^2)\log \frac{\frac{1}{\sqrt{2\pi }\sigma }\exp \{-\frac{(z-\mu {)}^2}{2{\sigma }^2}\}}{\frac{1}{\sqrt{2\pi }}\exp \{-\frac{z^2}{2}\}}dz\\ =& \min \frac{1}{2}{\int }_{z}N(\mu ,{\sigma }^2)\left(z^2-\log {\sigma }^2-\frac{(z-\mu {)}^2}{{\sigma }^2}\right)dz\end{array}$$
可以分为三部分
$$\begin{array}{rl}①：& {\int }_{z}N(\mu ,{\sigma }^2)z^{2}dz=\mathbb{E}[z^2]=D(z)+\mathbb{E}[z{]}^2={\sigma }^2+{\mu }^2\\ ②：& {\int }_{z}N(\mu ,{\sigma }^2)\log {\sigma }^{2}dz=\log {\sigma }^2{\int }_{z}N(\mu ,{\sigma }^2)dz=\log {\sigma }^2\\ ③：& {\int }_{z}N(\mu ,{\sigma }^2)\frac{(z-\mu {)}^2}{{\sigma }^2}dz=\frac{1}{{\sigma }^2}{\int }_{z}N(\mu ,{\sigma }^2)(z-\mu {)}^{2}dz=\frac{1}{{\sigma }^2}\mathbb{E}[(z-\mu {)}^2]=1\end{array}\text{\hspace{1em}(A)(B)(C)}$$
$\boxed{对于（式\mathbf{A}），里面用到了方差的计算公式\mathbf{D}(\mathbf{z})=\mathbf{E}({\mathbf{z}}^2)-\mathbf{E}(\mathbf{z}{)}^2}$

$\boxed{对于（式\mathbf{B}）,\mathbf{是因为}\mathbf{N}(\mu ,{\sigma }^2)\mathbf{积分为}1}$

$\boxed{对于（式\mathbf{C}）}$
$$\begin{array}{rl}& \frac{1}{{\sigma }^2}\mathbb{E}[(z-\mu {)}^2]\\ & \\ =& \frac{1}{{\sigma }^2}\mathbb{E}[z^2-2\mu z+{\mu }^2]\\ =& \frac{1}{{\sigma }^2}\left(\mathbb{E}(z^2)-2\mu \mathbb{E}(z)+{\mu }^2\right)\\ =& \frac{1}{{\sigma }^2}\left({\sigma }^2+{\mu }^2-2{\mu }^2+{\mu }^2\right)\\ =1& \end{array}$$
如果你熟悉高斯分布的高阶矩的话，式A和式C完全就是二阶原点矩和中心距，是直接可以的得出答案的。

$\boxed{\mathbf{所以对于所有维度的}\mathbf{KL}\mathbf{散度},\mathbf{有}(\mathbf{假设隐变量有}\mathbf{J}\mathbf{维})}$
$$\min \frac{1}{2}\sum _{j=1}^J({\sigma }_j^2+{\mu }_j^2-\log {\sigma }_j^2-1)$$

$\boxed{\mathbf{再来看}\log \mathbf{P}(x|z)极大似然，即\max {\mathbb{E}}_{\mathbf{q}(z|x)}[\log \mathbf{P}(x|z)]}$。

值得注意的是，我看很多文章中都说此处就直接采用均方差来计算。这种说法是不准确的，在论文中提到$P(x|z)$是服从一个概率分布的，而不是无端的就计算其差值。它不像是GAN一样，对于其隐藏在内部的概率分布不作约束，VAE是仍然对$P(x|z)$进行约束。

$\boxed{“\log p\theta (x^{(i)}|z^{(i,l)})是伯努利或高斯MLP，这取决于我们建模的数据类型”}$

当然了，其实我们也可以不对其概率分布进行约束，归根究底，其让然是最小重构代价，那么我们的目标函数如果可以充分表达出“最小重构代价”，那么是什么又有何关系呢？

在论文中，其假设$P(x|z)\sim N(\mu ,{\sigma }^{2}I)$，其中$\mu$和${\sigma }^2$都是需要使用神经网络去逼近。

但是，一般地，我们假设$P(x|z)\sim (f(z),cI)$，也就是其均值用神经网络去逼近，对于其协方差矩阵，我们设定为常数c和$I$相乘，所以依然是各个维度之间相互独立。我们来看看它的极大似然估计得什么（假设采样n个样本）

$$\begin{array}{rl}\max {\mathbb{E}}_{q(z|x)}[\log P(x|z)]=& \max \frac{1}{n}\sum _{i=1}^n\log P(x^i|z^i)\\ =& \max \frac{1}{n}\sum _{i=1}^n\log \frac{1}{{2\pi }^{D/2}|cI{|}^{1/2}}\exp \left\{-\frac{1}{2}(x-f(z^i))(cI{)}^{-1}(x-f(z^i){)}^T\right\}\\ =& \max \frac{1}{n}\sum _{i=1}^n\log \frac{1}{{2\pi }^{D/2}|cI{|}^{1/2}}-\frac{1}{2}(x-f(z^i))(cI{)}^{-1}(x-f(z^i){)}^T\\ =& \min \frac{1}{n}\sum _{i=1}^n\frac{(x-f(z^i))(x-f(z^i){)}^T}{2c}\\ =& \min \frac{1}{n}\sum _{i=1}^n\frac{||x-f(z^i)||}{2c}\end{array}$$
看出来了吧，如果让$2c^2=1$，那么它就是一个均方差了。所以归根究底，c就像是神经网络的一个惩罚系数一样。

$\boxed{\mathbf{综合得到最终目标函数表达式}}$
$$\min \left(\frac{1}{2}\sum _{j=1}^J({\sigma }_j^2+{\mu }_j^2-\log {\sigma }_j^2-1)+\frac{1}{n}\sum _{i=1}^n\frac{||x-f(z^i)||}{2c}\right)$$

3.3、重参数化技巧

有了目标函数，理论上我们直接梯度下降就可以了。然而，别忘了，我们是从$q(z|x)$中采样出z来。可是我们却是用的神经网络去计算的均值和方差，得到的高斯分布再去采样，这种情况是不可导的。中间都已经出现了一个断层了。神经网络是一层套一层的计算。而采样计算了一层之后，从这一层中去采样新的值，再计算下一层。因此，采样本身是不可导的。

所以要引入重参数化技巧，即假定$q(z|x)\sim N(\mu ,{\sigma }^2)$。那么可以构造一个概率分布$p(ϵ)\sim N(0,1)$。有
$$z=\mu +ϵ\sigma$$
我们从$P(ϵ)$采样，然后利用上述公式，就相当于得到了从$q(z|x)$采样的采样值。

$\boxed{\mathbf{证明}}$
$$\begin{array}{rl}\mathbb{E}(z)=& \mathbb{E}(\mu )+\mathbb{E}(ϵ)\sigma =\mu \\ Var(z)=& \mathbb{E}[(z-\mu {)}^2]=\mathbb{E}[(ϵ\sigma {)}^2]={\sigma }^2\end{array}$$

4、代码实现

效果一般，不晓得论文里面用了什么手段，效果看起来比这个好。（这个结果甚至还是我加了一层隐藏层的）

import torch
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
from torchvision.transforms import transforms
from  torch import nn
from tqdm import tqdm
import matplotlib.pyplot as plt
class VAE(nn.Module):
    def __init__(self,input_dim,hidden_dim,gaussian_dim):
        super().__init__()
        #编码器
        #隐藏层
        self.fc1=nn.Sequential(
            nn.Linear(in_features=input_dim,out_features=hidden_dim),
            nn.Tanh(),
            nn.Linear(in_features=hidden_dim, out_features=256),
            nn.Tanh(),
        )
        #μ和logσ^2
        self.mu=nn.Linear(in_features=256,out_features=gaussian_dim)
        self.log_sigma=nn.Linear(in_features=256,out_features=gaussian_dim)



        #解码（重构）
        self.fc2=nn.Sequential(
            nn.Linear(in_features=gaussian_dim,out_features=256),
            nn.Tanh(),
            nn.Linear(in_features=256, out_features=512),
            nn.Tanh(),
            nn.Linear(in_features=512,out_features=input_dim),
            nn.Sigmoid() #图片被转为为0，1的值了，故用此函数
        )
    def forward(self,x):
        #隐藏层
        h=self.fc1(x)

        #计算期望和log方差
        mu=self.mu(h)
        log_sigma=self.log_sigma(h)

        #重参数化
        h_sample=self.reparameterization(mu,log_sigma)

        #重构
        reconsitution=self.fc2(h_sample)

        return reconsitution,mu,log_sigma

    def reparameterization(self,mu,log_sigma):
        #重参数化
        sigma=torch.exp(log_sigma*0.5) #计算σ
        e=torch.randn_like(input=sigma,device=device)

        result=mu+e*sigma #依据重参数化技巧可得

        return result
    def predict(self,new_x): #预测
        reconsitution=self.fc2(new_x)

        return reconsitution
def train():

    transformer = transforms.Compose([
        transforms.ToTensor(),
    ]) #归一化
    data = MNIST("./data", transform=transformer,download=True) #载入数据

    dataloader = DataLoader(data, batch_size=128, shuffle=True) #写入加载器

    model = VAE(784, 512, 20).to(device) #初始化模型

    optimer = torch.optim.Adam(model.parameters(), lr=1e-3) #初始化优化器

    loss_fn = nn.MSELoss(reduction="sum") #均方差损失
    epochs = 100 #训练100轮

    for epoch in torch.arange(epochs):
        all_loss = 0
        dataloader_len = len(dataloader.dataset)

        for data in tqdm(dataloader, desc="第{}轮梯度下降".format(epoch)):
            sample, label = data
            sample = sample.to(device)
            sample = sample.reshape(-1, 784) #重塑
            result, mu, log_sigma = model(sample) #预测

            loss_likelihood = loss_fn(sample, result) #计算似然损失

            #计算KL损失
            loss_KL = torch.pow(mu, 2) + torch.exp(log_sigma) - log_sigma - 1

            #总损失
            loss = loss_likelihood + 0.5 * torch.sum(loss_KL)

            #梯度归0并反向传播和更新
            optimer.zero_grad()

            loss.backward()

            optimer.step()
            with torch.no_grad():
                all_loss += loss.item()
        print("函数损失为：{}".format(all_loss / dataloader_len))
        torch.save(model, "./model/VAE.pth")
if __name__ == '__main__':
    #是否有闲置GPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    #训练
    train()

    #载入模型，预测
    model=torch.load("./model/VAE (1).pth",map_location="cpu")
    #预测20个样本
    x=torch.randn(size=(20,20))
    result=model.predict(x).detach().numpy()
    result=result.reshape(-1,28,28)
    #绘图
    for i in range(20):
        plt.subplot(4,5,i+1)
        plt.imshow(result[i])
        plt.gray()
    plt.show()

VAE有很多的变种优化，感兴趣的读者自行查阅。

5、结束

以上，就是VAE的原理和推导过程了。能力有限，过程并不严谨，如有问题，还望指出。阿里嘎多