图像生成之变分自动编码器（VAE）

简要介绍

“概率图模型 + 神经网络 ”、“ EM算法、变分推断 ”

​ 自动编码器是一种无监督学习方法，将高维的原始数据映射到一个低维特征空间，然后从低维特征学习重建原始的数据。变分自编码器（Variational Autoencoder，简称VAE）是一种生成模型，结合了自编码器和概率图模型的思想。VAE在建模生成模型时是显式地定义了条件概率分布，通过最大似然估计来学习生成模型的参数，使其能够生成与训练数据相似的样本。

符号表示

$$\begin{array}{cc}\text{Symbol}& \text{Mean}\\ {\mathbb{E}}_{x\sim p(x)}[f(x)]& \text{表示对}f(x)\text{算期望，其中}x\text{的分布为}p(x)\\ p_{\theta }(\mathbf{z})& \text{先验分布}\\ p_{\theta }(\mathbf{z}|\mathbf{x})& \text{后验分布}\\ q_{\varphi }(\mathbf{z}|\mathbf{x})& \text{近似的后验概率函数，即概率编码器}\\ p_{\theta }(\mathbf{x}|\mathbf{z})& \text{给定隐变量生成真实数据样本的概率，即概率解码器}\end{array}$$

训练原理

​ 模型学习、优化模型参数$\theta$是通过最大似然估计。我们的目的是学习一个模型以最大化所观察到$x$的似然$p_{\theta }(x)$，贝叶斯公式得：
$$p_{\theta }(z|x)=\frac{p(x,z;\theta )}{p_{\theta }(x)}=\frac{p_{\theta }(z)p_{\theta }(x|z)}{{\int }_z{p}_{\theta }(x,z)dz}$$
由（1）式可以发现可以通过积分和利用联合分布概率算得$p_{\theta }(x)$，但是也容易发现联合分布概率未知，而积分的方法中，隐变量$z$是高维数据，积分要积很多次，需要指数级别的时间去计算，所以要另寻出路了：引入一个函数去近似，即变分推断，$q_{\varphi }(z|x)\approx p_{\theta }(z|x)$。

图1 VAE的概率图模型图示

自编码器（Autoencoder）：

​ 自动编码器是一种神经网络，旨在以无监督的方式学习恒等函数，以重建原始输入，同时压缩过程中的数据，从而发现更有效和压缩的表示。编码器将输入数据映射到潜在空间中的低维表示，而解码器将潜在表示映射回重构数据。自编码器的目标是最小化输入数据与重构数据之间的重构误差，从而学习到数据的压缩表示。

​ 编码器网络本质上实现了降维，就像主成分分析 （PCA）或矩阵分解 （MF） 一样。

图2 自动编码器模型体系结构图示

• 编码器（Encoder）：编码器部分的作用是学习输入数据（比如图片）到隐含空间的映射。在这个过程中，编码器试图将输入数据压缩为一个潜在空间中的点，这个点的坐标由潜在向量（latent vector）表示。这个潜在向量不是一个确定的点，而是一个分布（彩色/灰度图像假设是正态分布，二值图像用伯努利分布），由均值和标准差决定。

• 解码器（Decoder）：解码器部分的作用是学习从隐含空间到输入数据的反向映射。解码器从编码器产生的潜在向量中采样，然后试图将这个采样点恢复（或者说解码）为原始的输入数据。

损失函数：变分推断中的证据下界ELBO

​ VAE的优化目标包括两部分：一部分是重构损失，即解码的数据与原始输入数据的差异；另一部分是KL散度损失，即编码器得到的潜在变量分布与事先假设的分布（比如标准正态分布）的差异。下式推导略，详见参考。
$$\begin{array}{rl}{L}_{\text{VAE}}(\theta ,\varphi )& =-\log p_{\theta }(\mathbf{x})+D_{\text{KL}}(q_{\varphi }(\mathbf{z}|\mathbf{x})\Vert p_{\theta }(\mathbf{z}|\mathbf{x}))\\ & =-{\mathbb{E}}_{\mathbf{z}\sim q_{\varphi }(\mathbf{z}|\mathbf{x})}\log p_{\theta }(\mathbf{x}|\mathbf{z})+D_{\text{KL}}(q_{\varphi }(\mathbf{z}|\mathbf{x})\Vert p_{\theta }(\mathbf{z}))\\ {\theta }^{\ast },{\varphi }^{\ast }& =\arg \underset{\theta ,\varphi }{\min }{L}_{\text{VAE}}\end{array}$$

• 重构损失：如果x是二值图像，这个概率一般用伯努利分布，而伯努利分布的对数似然就是交叉熵损失binary cross entropy，对decoder用sigmoid函数激活；如果x是彩色/灰度图像，这个概率取正态分布，那么高斯分布的对数似然就是平方差MSE。

• 使用KL散度来最小化近似分布时的信息损失量。将KL散度与神经网络相结合，可以学习非常复杂的数据近似分布。KL散度可以写成期望的形式，这允许我们对其进行采样计算。

K L ( p ( x ) ∥ q ( x ) ) = ∫ p ( x ) ln ⁡ p ( x ) q ( x ) d x = E x ∼ p ( x ) [ ln ⁡ p ( x ) q ( x ) ] KL\Big(p(x)\Big\Vert q(x)\Big) = \int p(x)\ln \frac{p(x)}{q(x)} dx=\mathbb{E}_{x\sim p(x)}\left[\ln \frac{p(x)}{q(x)}\right] KL(p(x) ​q(x))=∫p(x)lnq(x)p(x)​dx=Ex∼p(x)​[lnq(x)p(x)​]

重新参数化技巧

​ 假设正态分布，损失函数的计算要使用到$\mathbf{z}\sim q_{\varphi }(\mathbf{z}|\mathbf{x})$，分布如下：

q ( z ∣ x ) = 1 ∏ k = 1 D 2 π σ ( k ) 2 ( z ) exp ⁡ ( − 1 2 ∥ x − μ ( z ) σ ( z ) ∥ 2 ) q(z|x)=\frac{1}{\prod\limits_{k=1}^D \sqrt{2\pi \sigma_{(k)}^2(z)}}\exp\left(-\frac{1}{2}\left\Vert\frac{x-\mu(z)}{\sigma(z)}\right\Vert^2\right) q(z∣x)=k=1∏D​2πσ(k)2​(z) ​1​exp(−21​ ​σ(z)x−μ(z)​ ​2)
​ 这里的均值方差都是靠模型算出来的，要靠这个过程反过来优化均值方差的模型。采样是一个随机过程，因此我们不能反向传播梯度。为了可训练，引入重新参数化：从正态分布$N(\mu ,{\sigma }^2)$中采样$z$时，相当于从$N(0,1)$中采样出来一个$ϵ$，然后再来计算$z=\mu +ϵ\times \sigma$。从$N(\mu ,{\sigma }^2)$采样变成了从$N(0,1)$中采样，然后通过参数变换得到从$N(\mu ,{\sigma }^2)$中采样的结果。
$$\begin{array}{rl}\mathbf{z}& \sim q_{\varphi }(\mathbf{z}|{\mathbf{x}}^{(i)})=\mathcal{N}(\mathbf{z};{\mathbf{\mu }}^{(i)},{\mathbf{\sigma }}^{2(i)}\mathbf{I})\end{array}$$
​ 这样一来，“采样”这个操作就不用参与梯度下降了，改为采样的结果参与，使得整个模型可训练了。

详细代码

import torch
import torch.nn as nn

class VAE(nn.Module):
    """
    变分自动编码器（Variational Autoencoder，VAE）
    """

    def __init__(self, hiddens=[16, 32, 64, 128, 256], latent_dim=128) -> None:
        super().__init__()

        # 编码器
        prev_channels = 3
        modules = []
        img_length = 64
        for cur_channels in hiddens:
            # 卷积层与批量归一化以及ReLU激活
            modules.append(
                nn.Sequential(
                    nn.Conv2d(prev_channels,
                              cur_channels,
                              kernel_size=3,
                              stride=2,
                              padding=1),
                    nn.BatchNorm2d(cur_channels),
                    nn.ReLU())
            )
            prev_channels = cur_channels
            img_length //= 2
        self.encoder = nn.Sequential(*modules)

        # 潜在空间中均值和对数方差的线性层
        self.mean_linear = nn.Linear(prev_channels * img_length * img_length,
                                     latent_dim)
        self.var_linear = nn.Linear(prev_channels * img_length * img_length,
                                    latent_dim)
        self.latent_dim = latent_dim

        # 解码器
        modules = []
        # 从潜在空间投影到初始形状的线性层
        self.decoder_projection = nn.Linear(
            latent_dim, prev_channels * img_length * img_length)
        self.decoder_input_chw = (prev_channels, img_length, img_length)

        # 转置卷积层与批量归一化以及ReLU激活
        for i in range(len(hiddens) - 1, 0, -1):
            modules.append(
                nn.Sequential(
                    nn.ConvTranspose2d(hiddens[i],
                                       hiddens[i - 1],
                                       kernel_size=3,
                                       stride=2,
                                       padding=1,
                                       output_padding=1),
                    nn.BatchNorm2d(hiddens[i - 1]),
                    nn.ReLU())
            )
        # 生成RGB图像的最终层
        modules.append(
            nn.Sequential(
                nn.ConvTranspose2d(hiddens[0],
                                   hiddens[0],
                                   kernel_size=3,
                                   stride=2,
                                   padding=1,
                                   output_padding=1),
                nn.BatchNorm2d(hiddens[0]),
                nn.ReLU(),
                nn.Conv2d(hiddens[0], 3, kernel_size=3, stride=1, padding=1),
                nn.ReLU())
        )
        self.decoder = nn.Sequential(*modules)

    def forward(self, x):
        # 编码器
        encoded = self.encoder(x)
        encoded = torch.flatten(encoded, 1)

        # 潜在空间中的均值和对数方差
        mean = self.mean_linear(encoded)
        logvar = self.var_linear(encoded)

        # 重参数化技巧
        eps = torch.randn_like(logvar)
        std = torch.exp(logvar / 2)
        z = eps * std + mean

        # 解码器
        x = self.decoder_projection(z)
        x = torch.reshape(x, (-1, *self.decoder_input_chw))
        decoded = self.decoder(x)

        return decoded, mean, logvar

    def sample(self, device='cuda'):
        # 从潜在空间生成一个随机样本
        z = torch.randn(1, self.latent_dim).to(device)
        x = self.decoder_projection(z)
        x = torch.reshape(x, (-1, *self.decoder_input_chw))
        decoded = self.decoder(x)
        return decoded

参考

变分自编码器（一）：原来是这么一回事 - 科学空间|Scientific Spaces

Diffusion model (1) 概述：从VAE谈起_哔哩哔哩_bilibili

从自动编码器到贝塔-VAE |利尔日志 (lilianweng.github.io)