VAE模型及pytorch实现

VAE（变分自编码器）是一种生成模型，结合了自编码器和概率图模型的思想。它通过学习数据的潜在分布，可以生成新的数据样本。VAE通过将输入数据映射到潜在空间中的分布，并在训练过程中最大化数据与潜在变量之间的条件概率来实现。其关键思想在于编码器将输入数据编码成潜在分布的参数，解码器则从这个分布中采样生成新的数据。这种生成方式不仅能够生成新的数据，还能够在潜在空间中进行插值和操作，提供了强大的特征学习和数据生成能力。

AE论文：Auto-Encoding Variational Bayes

VAE论文：Semi-supervised Learning with Deep Generative Models

VAE模型推导部分

假设$P(z)$是一个正态分布，$x|z\sim N(\mu (z),\sigma (z))$是x从z分布中进行采样得到的。
$$P(x)={\int }_{z}P(z)P(x|z)dz$$
为了最大化$P(x)$，我们采用极大似然估计
$$L=\sum _{x}logP(x)\mathrm{Maximizing}\mathrm{the}\mathrm{likelihood}\mathrm{of}\mathrm{the}\mathrm{observed}\mathrm{x}$$
对$logP(x)$进一步进行变形
$$\begin{array}{rl}logP(x)=& {\int }_{z}q(z|x)logP(x)dz\mathrm{q}(z|x)\mathrm{can}\mathrm{be}\mathrm{any}\mathrm{distribution}\\ & ={\int }_{z}q(z|x)log\left(\frac{P(z,x)}{P(z|x)}\right)dz={\int }_{z}q(z|x)log\left(\frac{P(z,x)}{q(z|x)}\frac{q(z|x)}{P(z|x)}\right)dz\\ & ={\int }_{z}q(z|x)log\left(\frac{P(z,x)}{q(z|x)}\right)dz+{\int }_{z}q(z|x)log\left(\frac{q(z|x)}{P(z|x)}\right)dz\\ & \ge {\int }_{z}q(z|x)log\left(\frac{P(x|z)P(z)}{q(z|x)}\right)dz\end{array}$$
因为$KL\left(q(z|x)||P(z|x)\right)={\int }_{z}q(z|x)log\left(\frac{q(z|x)}{P(z|x)}\right)dz$是大于0的数，所以，上述式子大于等于前面那一项。

​ 对于给定的$P(x|z)$,让KL尽可能小，就是让$L_b$最大。同时，当$KL$尽可能小，也就是说明$q(z|x)$和$p(z|x)$这两个分布的相似度越高。

​ 接下来我们就对$L_b$进行最大化变形处理，变形后左侧为，右侧为
$$\begin{array}{rl}{L}_b& ={\int }_{z}q(z|x)log\left(\frac{P(z,x)}{q(z|x)}\right)dz={\int }_{z}q(z|x)log\left(\frac{P(x|z)P(z)}{q(z|x)}\right)dz\\ & ={\int }_{z}q(z|x)\log (\frac{P(z)}{q(z|x)})dz+{\int }_{z}q(z|x)logP(x|z)dz\\ & =KL(q(z|x)||P(z))+E_{q(z|x)}[logP(x|z)]\end{array}$$

​ 如下所示，我们需要做的就是最小化$KL(q(z|x)||P(z))$并最大化$E_{q(z|x)}[logP(x|z)]$。对于最小化KL，我们可以理解为输入一个$x$,然后通过神经网络调参输出${\mu }_{(}x),\sigma (x)$，也就是让这个数值尽可能和$P(z)$这个分布接近。这部分相当于Encoder部分。

​ 在Encoder部分结束后，对于第2项，从已知的$z$，也就是数据的隐式特征表示，去采样出$x$，相当于模型的Decoder部分，输出一个均值使之尽可能接近原始的$x$，因为对于这种条件概率，均值最大的时候就是$x$

最小化KL散度推导

为了最小化$q(z|x)$和$P(z)$的KL散度，首先，我们先对正态分布的KL散度计算进行推导。参考链接高斯分布的KL散度-CSDN博客
$$\begin{array}{rl}\mathrm{KL}\left(\mathcal{N}\left({\mu }_1,{\sigma }_1^2\right)\Vert \mathcal{N}\left({\mu }_2,{\sigma }_2^2\right)\right)& ={\int }_{\mathrm{x}}\frac{1}{\sqrt{2\pi }{\sigma }_1}{\mathrm{e}}^{-\frac{{\left(x-{\mu }_1\right)}^2}{2{\sigma }_1^2}}\log \frac{\frac{1}{\sqrt{2\pi }{\sigma }_1}{e}^{-\frac{{\left(x-{\mu }_1\right)}^2}{2{\sigma }_1^2}}}{\frac{1}{\sqrt{2\pi }{\sigma }_2}{e}^{-\frac{{\left(x-{\mu }_2\right)}^2}{2{\sigma }_2^2}}}dx\\ & ={\int }_x\frac{1}{\sqrt{2\pi }{\sigma }_1}{e}^{-\frac{{\left(x-{\mu }_1\right)}^2}{2{\sigma }_1^2}}\left[\log \frac{{\sigma }_2}{{\sigma }_1}-\frac{{\left(x-{\mu }_1\right)}^2}{2{\sigma }_1^2}+\frac{{\left(x-{\mu }_2\right)}^2}{2{\sigma }_2^2}\right]dx\end{array}$$

1. 对于第1项，由于${\sigma }_1,{\sigma }_2$与x无关，则可以直接提取到积分外面，该积分即为正态分布的全概率公式，也就是为1
   $$\log \frac{{\sigma }_2}{{\sigma }_1}{\int }_{\mathrm{x}}\frac{1}{\sqrt{2\pi }{\sigma }_1}{\mathrm{e}}^{-\frac{{\left(x-{\mu }_1\right)}^2}{2{\sigma }_1^2}}\mathrm{dx}=\log \frac{{\sigma }_2}{{\sigma }_1}$$

2. 对于第2项，则是由方差定义式$D(x)={\int }_x(x-\mu {)}^{2}f(x)dx$,可知这个积分的结果为${\sigma }_1^2$
   $$-\frac{1}{2{\sigma }_1^2}{\int }_{\mathrm{x}}{\left(\mathrm{x}-{\mu }_1\right)}^2\frac{1}{\sqrt{2\pi }{\sigma }_1}{\mathrm{e}}^{-\frac{{\left(\mathrm{x}-{\mu }_1\right)}^2}{2{\sigma }_1^2}}\mathrm{dx}=-\frac{1}{2{\sigma }_1^2}{\sigma }_1^2=-\frac{1}{2}$$

3. 对于第3项，首先将其展开，对于$x^2$,由均方值公式，$E(x^2)=D(x)+E(x{)}^2$,后面两项则分别是通过均值公式以及全概率公式进行计算。
   $$\begin{array}{rl}\frac{1}{2{\sigma }_2^2}{\int }_{\mathrm{x}}{\left(\mathrm{x}-{\mu }_2\right)}^2\frac{1}{\sqrt{2\pi }{\sigma }_1}{\mathrm{e}}^{-\frac{{\left(\mathrm{x}-{\mu }_1\right)}^2}{2{\sigma }_1^2}}\mathrm{dx}& =\frac{1}{2{\sigma }_2^2}{\int }_{\mathrm{x}}\left({\mathrm{x}}^2-2{\mu }_2\mathrm{x}+{\mu }_2^2\right)\frac{1}{\sqrt{2\pi }{\sigma }_1}{\mathrm{e}}^{-\frac{{\left(\mathrm{x}-{\mu }_1\right)}^2}{2{\sigma }_1^2}}\mathrm{dx}\\ & =\frac{{\sigma }_1^2+{\mu }_1^2-2{\mu }_1{\mu }_2+{\mu }_2^2}{2{\sigma }_2^2}=\frac{{\sigma }_1^2+{\left({\mu }_1-{\mu }_2\right)}^2}{2{\sigma }_2^2}\end{array}$$

对上述式子进行汇总：
$$\begin{array}{rl}\mathrm{KL}\left(\mathcal{N}\left({\mu }_1,{\sigma }_1^2\right)\Vert \mathcal{N}\left({\mu }_2,{\sigma }_2^2\right)\right)& =\log \frac{{\sigma }_2}{{\sigma }_1}-\frac{1}{2}+\frac{{\sigma }_1^2+({\mu }_1-{\mu }_2{)}^2}{2{\sigma }_2^2}\\ & =\frac{1}{2}({\sigma }_1^2+{\mu }_1^2-{\log }^{{\sigma }_1^2}-1)\end{array}$$

代码部分

损失函数

通过上述推导，我们知道了需要最小化散度，然后最大化那个均值。所以可以得到如下的损失函数。

    def loss_fn(recon_x, x, mean, log_var):
        BCE = torch.nn.functional.binary_cross_entropy(
            recon_x.view(-1, 28*28), x.view(-1, 28*28), reduction='sum')
        KLD = -0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp())

        return (BCE + KLD) / x.size(0)

Encoder部分

class Encoder(nn.Module):
    def __init__(self, layer_sizes, latent_size):
        super(Encoder, self).__init__()
        self.MLP = nn.Sequential()
        for i, (in_size, out_size) in enumerate(zip(layer_sizes[:-1], layer_sizes[1:])):
            self.MLP.add_module(name="L{:d}".format(i), module=nn.Linear(in_size, out_size))
            self.MLP.add_module(name="A{:d}".format(i), module=nn.ReLU())

        # 首先对图像特征进行一些变换处理，然后将其展开成一维向量，然后通过全连接层得到均值和方差
        self.linear_means = nn.Linear(layer_sizes[-1], latent_size)
        self.linear_log_var = nn.Linear(layer_sizes[-1], latent_size)

    def forward(self, x):
        x = self.MLP(x)

        means = self.linear_means(x)
        log_vars = self.linear_log_var(x)

        return means, log_vars

Decoder部分

class Decoder(nn.Module):
    def __init__(self, layer_sizes, latent_size):
        super(Decoder, self).__init__()
        self.MLP = nn.Sequential()
        input_size = latent_size
        
        for i, (in_size, out_size) in enumerate(zip([input_size] + layer_sizes[:-1], layer_sizes)):
            self.MLP.add_module(
                name="L{:d}".format(i), module=nn.Linear(in_size, out_size))
            if i + 1 < len(layer_sizes):
                self.MLP.add_module(name="A{:d}".format(i), module=nn.ReLU())
            else:
                self.MLP.add_module(name="sigmoid", module=nn.Sigmoid())

    def forward(self, z):
        #对输入的z进行全接连操作，最后输出一个重构的x
        x = self.MLP(z)
        return x

VAE整体架构

class VAE(nn.Module):
    def __init__(self, encoder_layer_sizes, latent_size, decoder_layer_sizes):
        super(VAE, self).__init__()
        self.latent_size = latent_size
        self.encoder = Encoder(encoder_layer_sizes, latent_size)
        self.decoder = Decoder(decoder_layer_sizes, latent_size)

    def forward(self, x):
        if x.dim() > 2:
            x = x.view(-1, 28 * 28)
        means, log_var = self.encoder(x)
        z = self.reparameterize(means, log_var)
        recon_x = self.decoder(z)
        return recon_x, means, log_var, z

    def reparameterize(self, mu, log_var):
        """
        用于对encoder部分输出的均值方差进行重参数化，采样得到隐式表示部分z
        :param mu:
        :param log_var:
        :return:
        """
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        return mu + eps * std

    def inference(self, z):
        recon_x = self.decoder(z)
        return recon_x

VAE问题

vae只是记住图片，而不是生成图片

再产生图片时，只是通过像素差异进行评估，则对于关键点像素和可忽略像素之间的图片，两者在vae看来是一致的，但是不是理想的产生图片，因此出现了GAN

参考资料

VAE 模型基本原理简单介绍_vae模型-CSDN博客

高斯分布的KL散度-CSDN博客

ML Lecture 18: Unsupervised Learning - Deep Generative Model (Part II)