VAE原理详细直观化解析

VAE原理解析

Auto-Encoder

上图为普通的AutoEncoder，输入数据经过Encoder的编码，被压缩到一个称为“潜在空间”的较低维度表示中，这个潜在空间包含了模型学习到的数据的特征，我们将这个低维度表示抽象为潜在空间的某个点，然后Decoder会解析这个点去尽可能的去生成和原始数据接近的输出。但这样会有一个问题：

输入数据x经过Encoder之后输出的是一个确定的潜在变量z，这个z代表着对数据的高度抽象的特征描述。例如途中的圆形、三角形、正方形，经过Encoder之后抽象为潜在空间中的三个点，Decoder通过解析这三个点生成和原始图形接近的输出，可是如果我们在这三个点之外随机取一个点（图中为紫色点），这个点经过Decoder的解析会生成一个没有意义的数据（例如图中紫色的线条），这个数据与原始输入数据之间没有任何关联。

也就是说在潜在空间中，除了这三个点是有意义的之外，其它的点没有任何意义。

这样的潜在空间是离散的，或者说是不平滑的。

VAE

那如果潜在空间变得平滑，也就是变得连续之后，有什么好处呢？

如图，在连续的潜在空间中，相似的输入数据对应的低维表示在这个空间中是接近的，也就是说潜在空间中每一个点都有意义，例如右图中红色的点既和三角形对应的点接近又离圆形对应的点接近，所以由这个红色的点生成的数据会是一个圆角的三角形。

也就是说当我们在这个平滑的潜在空间中改变数值时，生成的数据也会相应地变化，而且这种变化是渐进、连续的。这使得在潜在空间中的相近的点对应于相似的生成数据，同时也意味着我们可以在这个空间中执行插值，即在两个不同的数据点之间进行平滑的过渡，生成新的数据点，而这些新生成的数据点在原始数据中可能是合理的。

VAE是如何使潜在空间变得平滑的？

VAE实际上是通过引入概率的方式来使潜在空间变得平滑。

VAE的目标是学习数据真实的潜在表示，即p(z|x); 同时，将每个输入数据的潜在分布近似为一个标准正态分布（均值为0，方差为1）使潜在空间变得连续。

也就是说对于VAE来说，我们想要的不是一个确定的潜在表示$z$，而是一个概率分布$p(z|x)$，这个概率分布表示在给定输入 $x$的条件下，潜在空间中潜在变量 $z$ 的分布，这个概率分布被称为真实的后验概率分布，可是由于这个真实的后验概率分布难以计算，所以我们希望通过Encoder去生成一个近似的后验概率分布$q(z|x)$，我们希望这个生成的$q(z|x)$去尽可能接近真实的概率分布$p(z|x)$，我们通过KL散度去衡量这两个概率分布之间的差值，我们希望这个差值尽可能地小。$q(z|x)$和$p(z|x)$的差距代表着模型学习数据潜在表示的能力即对输入数据特征抽取的能力。

ps: $p(z|x)$代表着输入数据真实的特征表示z的分布是什么样的，而我们设计各种各样的模型就是为了尽可能地计算或者接近这个特征表示的分布，从而更好地根据这个特征表示去重构数据。

通过让$q(z|x)$近似于真实的后验分布$p(z|x)$，模型可以更好地捕捉数据的结构和特征，使得潜在空间中的点表示的输入数据的特征更加准确。这种潜在表示的学习有助于模型学习到数据的抽象表征，进而用于数据生成和重构。

KL Divergence（KL散度）

KL散度计算公式：$D(q(z|x)||p(z|x))=E_{q(z|x)}[log\frac{q(z|x)}{p(z|x)}]=E_{q(z|x)}[logq(z|x)-logp(z|x)]$ (1)

首先，我们通过以下公式转换KL公式：

$logp(z|x)=log\frac{p(z,x)}{p(x)}$

$D(q(z|x) || p(z|x)) = E_{q(z|x)}[log q(z|x) - log \frac{p(z,x)}{p(x)}] $

$=E_{q(z|x)}[logq(z|x)-logp(z,x)]+logp(x)$

也就是：
$logp(x)=E_{q(z|x)}[logp(z,x)-logq(z|x)]+D(q(z|x)||p(z|x))$ (2)

$p(x)$ 可以被看作是在整个数据集上的概率密度函数，表示了在整个数据集中观察到任何数据点 $x$ 的概率。这个$p(x)$同样是难以计算的，但是，给定一个数据$x$，$p(x)$ 是确定的，也就是说$p(x)$ 是一个定值。

我们想让$D(q(z|x)||p(z|x))$尽可能地小，也就是让第一项越来越大：

$ELBO=E_{q(z|x)}[logp(z,x)-logq(z|x)]$ (3)

我们称之为Evidence lower bound(ELBO)。

同样的$logp(z,x)=logp(z)p(x|z)=logp(x|z)+logp(z)$

代入ELBO可得：

$ELBO=E_{q(z|x)}[logp(x|z)+logp(z)-logq(z|x)]$

$=E_{q(z|x)}[logp(x|z)]-E_{q(z|x)}[log\frac{q(z|x)}{p(z)}]$

我们发现第二项为q(z|x)和p(z)的KL散度，也就是说：

$ELBO=E_{q(z|x)}[logp(x|z)]-D(q(z|x)||p(z))$ (4)

要想最大化ELBO，就要最大化第一项，最小化第二项

第一项我们称为Reconstruction Likelihood(重构似然)，它表示$p(x|z)$的期望，表示给定潜在变量$z$后，生成数据$x$的概率。当我们从$q(z|x)$中采样出$z$，并通过解码器生成数据$x$时，我们希望生成的$x$能够尽可能地接近原始的$x$，也就是说，我们希望$p(x|z)$尽可能地大。也就是说==这一项也代表着输出数据与输入数据之间的损失。==

第二项为Regularization term(正则化项），这里指的是q(z|x)和p(z)的KL散度，也就是Encoder输出的分布和先验分布**p(z)**的相似度，我们最小化这一项就是希望让$q(z|x)$和$p(z)$尽可能地逼近。我们通常将$p(z)$假定为一个单位高斯分布$N(0,I)$即标准正态分布。通过这种正则化，VAE 被迫学习产生更加平滑和连续的潜在表示。

（ps：**标准正态分布是一个连续的、平滑的分布，通过鼓励潜在空间中的点遵循正态分布，VAE 能够推动相似的输入数据在潜在空间中彼此靠近，从而保持潜在空间的连续性和平滑性。**这意味着在潜在空间中邻近的点对应于输入空间中相似的数据点。）

By forcing the distributions to be close, we avoid “holes” in the latent space: we can move smoothly from one distribution to another without generating non-sense reconstructions.

==总结：==我们想让**近似的后验概率分布$q(z|x)$去尽可能接近真实的概率分布$p(z|x)$，最终要做的就是让$q(z|x)$接近先验分布$p(z)$（正则化），并且让模型最终的输出尽可能地接近原始的输入$x$（重构似然）。**基于这个思想，我们稍微转换一下公式(4)便可以得到整个模型的损失函数：

$Loss=-ELBO$ （具体的损失函数在实际的模型训练中会做出一些调整）

Reparameterization（重参数化）

重参数化（Reparameterization）是 Variational Autoencoder（VAE）中的一个关键技术。在 VAE 中，Encoder生成的潜在表示（latent representation）是一个概率分布，然后再从这个概率分布中进行随机采样，通过Decoder生成数据。这个随机采样操作本身是不可导的，也就不能进行反向传播（backpropagation）和梯度下降等优化方法来更新参数。

为了解决这个问题，重参数化技术被引入。它将随机采样操作从网络结构中分离出来，通过随机采样一个噪声分布，然后通过可导的操作将这个采样值转化为所需的潜在表示。这样就可以计算梯度并使用反向传播来训练整个网络。

具体来说，对于VAE模型，编码器网络会输出两个参数，一个表示均值（$\mu$），另一个表示标准差（$\sigma$）。然后我们从标准正态分布中随机采样一个噪声值 $\epsilon$。然后便可通过公式 $z=\mu +\sigma \ast \epsilon$ 将采样操作转化为了一个可微的函数，这样也就使得从潜在分布中采样的过程成为可导的操作，允许梯度通过这个操作传播回编码器，进而进行模型参数的优化。

本篇笔记参考Autoencoders — Neurocomputing (julien-vitay.net)

Pytorch搭建一个简单的VAE模型

导包：

import torch
import torch.nn as nn
import torch.optim as optim

搭建模型：

class Encoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, latent_dim):
        super(Encoder, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.fc2_mean = nn.Linear(hidden_dim, latent_dim)
        self.fc2_logvar = nn.Linear(hidden_dim, latent_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        mean = self.fc2_mean(x)
        log_var = self.fc2_logvar(x)
        return mean, log_var

class Decoder(nn.Module):
    def __init__(self, latent_dim, hidden_dim, output_dim):
        super(Decoder, self).__init__()
        self.fc1 = nn.Linear(latent_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, z):
        z = torch.relu(self.fc1(z))
        reconstructed = torch.sigmoid(self.fc2(z))  # Adjust activation based on the data type
        return reconstructed

定义损失函数和优化器，训练模型：

# Initialize the models
encoder = Encoder(input_dim, hidden_dim, latent_dim)
decoder = Decoder(latent_dim, hidden_dim, output_dim)
model = VAE(encoder, decoder)

# Define the loss function
def loss_function(recon_x, x, mu, logvar):
    # Reconstruction loss
    recon_loss = nn.functional.binary_cross_entropy(recon_x, x, reduction='sum')
    # KL Divergence
    kl_divergence = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + kl_divergence

# Initialize the optimizer
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# Training loop
for epoch in range(num_epochs):
    model.train()
    for batch_idx, data in enumerate(train_loader):
        optimizer.zero_grad()
        recon_batch, mu, logvar = model(data)
        loss = loss_function(recon_batch, data, mu, logvar)
        loss.backward()
        optimizer.step()