08: Auto-Encoding Variational Bayes (VAE)

VAE（变分自动编码器）是一类结合概率图模型与神经网络的生成模型，它通过引入可参数化的近似后验 \(q_\phi(z|x)\) 来摊销推断成本，并用最大化 ELBO 的方式同时学习数据的潜在表示与生成过程：其中重建项确保模型能从潜变量还原数据，KL 项则将潜空间约束为接近先验的连续结构。借助重参数化技巧，VAE 能在端到端训练中高效地学习一个平滑、可采样的潜空间，从而实现表示学习、插值、生成等功能，是现代深度生成模型的重要基础。

# Preliminary

1 VAE

1.1 Experiment

2 Summary

3 Key Concepts

4 Q & A

5 Related resource & Further Reading

介绍了Vision Model 和 Language Model， 以及如何通过Self-Supervised Learning 来学习Vision的 Representation DINO, 接下来我们将会更新一系列的生成式模型（Generative Models）。生成式模型在当下是很热门的话题，比如Stable Diffusion， Nana Banana 这些模型都是基本思想都是基于Latent Variable Models，接下来，我们就来看看Variable Autoencoder这一篇，Latent Variable 模型。

6 Preliminary

在这一节，我们先了解一些前置知识，以便我们可以更好的理解这篇文章。

6.1 Latent Variable Models

在深度学习中，我们经常处理 高维、复杂、带噪声的真实世界数据：图像、语音、文本…… 这些数据的背后，其实常常有一些 未被直接观测到的隐藏结构 —— 这就是 “Latent Variables”（潜变量）。

想象一下，你看一张人脸照片。照片是观测变量（observed variable）。 但是让这张脸“看起来像某个人”的，是一些不能直接看到的因素： • 光照（lighting） • 情绪（expression） • 姿态（pose） • 脸部特征（identity） • 背景（background）

这些因素虽然没有在数据中显式标注，却真实存在，并决定了我们看到的图像。 潜变量模型的目标，就是用数学方式把这些“隐藏因素”建模出来。

Latent Variable Model（LVM）是一类 假设观测数据是由一些隐藏变量生成的概率模型。 比如我们看到的数据\(\mathrm{x} \in \mathbb{R}^{d}\), 那它有相对应的隐藏变量 \(z \in \mathbb{R}^{k}\) 其中 \(k \ll d\) . \(z\) 是我们观察不到的，也就是所谓的 Latent Variable.

在使用Latent Variable Model时，我们有两个关键的任务： 1. Inference: 根据 \(\mathrm{x}\) 我们来推断出 \(\mathrm{z}\) 是什么，比如，给定一张人脸的照片，这个模型需要找出 - 这张脸是男生还是女生 - 表情是开心还是悲伤 - 是什么样的姿势 在数学上，就是求后验分布 \(p(\mathrm{z} | \mathrm{x})\) 2. Generation: 通过 \(\mathrm{z}\)， 我们来生成一个 \(\mathrm{x}\)。 这个就是生成模型。

Latent Variable Model 是现代生成式模型的基石，基本上所有的生成式模型，比如 GAN， DDPM， Flow Model 等，都是以Latent Variable Model为基础，在此条件下，通过不同求 Latent Variable 的方法，来解决这种问题。

6.2 AutoEncoder

AutoEncoder 是一种Self-Supervised Learning 的方法。 它可以让神经网络学会压缩，并且在还原数据，通过这种方法，我们可以学习到的低纬度的Latent Variable \(\mathrm{z}\)。 AutoEncoder 也可以看作是Latent Variable Model的一种学习方法。 AutoEncoder 也通常用在Representation Learning 表征学习。

6.3 KL-Divergence

KL 散度（Kullback–Leibler Divergence）是用来衡量两个概率分布之间差异的一种度量方法。定义如下: \[ D_{KL}(Q \| P) = \mathbb{E}_{x \sim Q}\left[ \log \frac{Q(x)}{P(x)} \right] \]

直观的解释KL-Divergence就是: - 如果数据 \(x\) 是从 \(Q\) 分布中来的，但如果我们用 \(P\) 分布 来解释这些数据，会损失多少信息量

对于高斯分布（Gaussian Distribution），我们KL- Divergence有以下的形式：

$$

D_{KL}(Q | P) = $$

如果是Diagonal Gaussian， 那么KL-Divergence 可以简化为： $$ D_{KL}(q ,|, p)

_{i=1}^{d} $$

\[ D_{KL}(q(\mathbf{z}) \ \| \ \mathcal{N}(0, I)) = \frac{1}{2} \sum_{i=1}^{d} \left( \mu_{q,i}^2 + \sigma_{q,i}^2 - \log \sigma_{q,i}^2 - 1 \right) \]

需要注意的一个点是，KL-Divergence是不对称的， \[ D_{KL}(Q \| P) \neq D_{KL}(P \| Q) \]

为什么我们要强调这一点，是因为： 对于不同位置的Q，P我们所求的是不一样的，简单来说，就是用 \(\|\) 后面的分布，来approximate \(\|\) 前面的分布，具体如下图。

6.4 Variational Inference

Variational Inference（VI，变分推断）是一种用可优化、易计算的分布来近似一个难以求解的后验分布，通过最小化两者之间的 KL 距离，从而实现高效概率推断的方法。 在 Latent Variable 模型中，我们想求： \[ p (\mathrm{z} | \mathrm{x}) = \frac{p(\mathrm{x}, \mathrm{z})}{p(\mathrm{x})} \]

通常 \(p(\mathrm{x})\) 是不可求的，因为: \[ p(\mathrm{x}) = \int p(\mathrm{x}, \mathrm{z}) d\mathrm{z} \] 通常是不可能直接求的。因此我们采取一种 “曲线救国” 的方式： 我们不去计算真正的后验，而是找一个 可计算并且可优化的分布来近似它： \[ q(z | x) \approx p(z | x) \] 实现起来也是很简单的 1. 选一个可计算的分布族 （Variational Family） 2. 让它尽可能的接近真实的后验： \[ \underset{\phi}{\min} D_{KL}(q_{\phi}(z | x ) \| p(z | x)) \]

直观的来说，Variational Inference 就是：不过不断拉升，旋转，一个椭圆型 \(q\), 来使它尽可能的可以和云朵形状的 \(p\) 来重叠

7 VAE

论文试图解决的核心问题是：

How can we perform efficient inference and learning in directed probabilistic models, in the presence of continuous latent variables with intractable posterior distributions, and large datasets Auto-Encoding Variational Bayes, p.1

如何在具有连续潜变量、后验分布不可解（intractable posterior）、以及大规模数据集（large datasets）情况下，在有向概率模型（directed probabilistic model）中，进行有效的推断（inference）和学习（learning）？  也就是说，当我们设定一个生成模型 p_(x,z) （观测 x、潜变量 z）但其后验 p_(z|x) 难以直接计算的时候，如何使用变分贝叶斯（variational Bayes）方法来进行近似推断，并且该方法能扩展到大数据、可用神经网络进行“编码-解码”结构使用。

Variational Autoencoder（VAE）是一类通过引入可参数化的变分分布 \(q_\phi(z|x)\) 来近似难以解析的真实后验 \(p(z|x)\)，并以连续潜变量为核心结构，从而实现对数据潜在空间的建模与可采样生成的深度生成模型。

我们有一些数据 \(\mathrm{x}_{1}, \dots \mathrm{x}_{n}\)， 希望计算 \(\log p_{\theta}(\mathrm{X}) = \sum_{i=1}^{N} \log p_{\theta}(\mathrm{x}_{i})\) ## Evidence Lower BOund(ELBO)

Also known as variational lower bound

我们知道，对于高纬度的数据，我们不能直接求出所谓的 Likelihood。对于一个 data point, 我们有: \[ \begin{split} \log p_\theta(x) & = \log \int p_{\theta}(x, z)\ dz \\ & = \log \int p_{\theta}(x, z)\ \frac{q_\phi(z \mid x)}{q_\phi(z \mid x)}\ dz \\ & \geq \boxed{ \mathbb{E}_{q_{\phi}}[\log p_{\theta}(x, z) - \log q_{\phi}(z | x)] }\\ & = \mathbb{E}_{q_{\phi}}[\log p_{\theta}(x |z) + \log p_{\theta}(z) - \log q_{\phi}(z | x)] \\ &= \mathbb{E}_{q_{\phi}}[\log p_{\theta}(x | z) ] - \mathbb{E}_{q_{\phi}}\left[ \log \frac{q_{\phi}(z | x)}{p_{\theta}(z)} \right]\\ &= \mathbb{E}_{q_{\phi}}[\log p_{\theta}(x | z) ] - D_{KL}[q_{\phi}(z | x) \| p_{\theta}(z)] \\ & \approx \log p_{\theta}(x | z) - D_{KL}[q_{\phi}(z | x) \| p_{\theta}(z)] \end{split} \]

This is known as SGVB estimator > SGVB (Stochastic Gradient Variational Bayes) estimator can be used for ef- ficient approximate posterior inference in almost any model with continuous latent variables and/or parameters, and is straightforward to optimize using standard stochastic gradient ascent techniques. > Auto-Encoding Variational Bayes, p.1

我们可以看到，对于一个data point，我们可以通过这个ELBO来优化这个NN，对于一个 Batch的数据，我们则有: \[ \log p_{\theta}(X) = \sum_{i = 1}^{B} \log p_{\theta}(x_{i} | z) - D_{KL}[q_{\phi}(z | x_{i}) \| p_{\theta}(z)] \]

不过对于这个方法，存在一个问题就是 \(z\) 是从 \(q_{\phi}(z | x_{i})\) 中 sample出来的，对于这种类型的数据，我们不能计算它的Gradient，于是作者提出了 Re-parameterization Trick

7.1 Re-parameterization Trick

如上图可见，\(z\) 是由 \(x, \phi\) 决定的。通过这种方式，那么我们改取怎么样的 \(g_{\phi}(x, \epsilon)\) 呢。文章中给出了3个基本的方法： 1. 方法1 2. 2 3. c

REINFORCE

对于 Re-Parametrization Trick是针对 \(z\) 是 Continuous的情况，如果 \(z\) 是离散（Discrete) 的 那么我们则可以使用 REINFOCE 的方法。

7.2 Amortized Inference

还有一个比较容易被忽略的一点就是，VAE 还运用了Amortized Inference。什么意思呢？

传统的变分推断中，我们需要为每一个观测样本 x 单独优化一个变分分布 q(z|x) 的参数，这意味着每来一个新样本都要重新做一遍变分优化，成本非常高。而在 VAE 中，我们不再为每个样本单独优化后验，而是训练一个 共享的推断网络（Encoder） 来预测 q_(z|x) 的参数。也就是说，模型通过学习一个函数 f_(x) 来一次性“摊销”所有样本的推断成本，使得对任意新样本 x，只需一次前向传播就能得到近似后验，不再需要昂贵的 per-sample 优化。这种方式极大地提升了推断效率，也让变分推断能够在深度学习规模上落地。

也就是用神经网络来一次性学习后验

将上面的几个结合起来，我们就得到的了 Auto-Encoding VB Algorithm。
## 和 AutoEncoder的关系

7.3 Experiements

VAE Loss can be defined as this one:

class VAELoss(nn.Module):
    def __init__(self, rec_loss="bce", kl_beta=1.0):
        super().__init__()

        self.rec_loss = rec_loss.lower()
        self.kl_beta = kl_beta

        self.eval()

    def forward(self, x, x_recon, mu, logvar):
        B = x.shape[0]
        if self.rec_loss == "bce":
            rec = F.binary_cross_entropy(x_recon, x, reduction="sum")
        elif self.rec_loss == "mse":
            rec = F.mse_loss(x_recon, x, reduction="sum")

        # KL divergence: D_KL(q(z|x) || p(z))
        kl = 0.5 * torch.sum(logvar.exp() + mu.pow(2) - logvar - 1)

        total = (rec + self.kl_beta * kl) / B

        return total, {
            "recon": rec.detach().cpu() / B,
            "kl": kl.detach().cpu() / B,
        }

The Loss curve and illustration examples:

8 Summary

总结一下，VAE 这篇文章 (Kingma and Welling 2022) 讲了写什么： VAE 通过ELBO 来优化 Latent Variable Model，通过 Re-parameterization trick 和 Amortized Inference，使得Latent Variable Model 在大数据集上得到显著的效果。VAE 为之后的很多生成式模型奠定了基础。

9 Key Concepts

Concept	Explanation
Latent Variable Model	用潜变量建模数据生成过程
Variational Inference	用可学习的分布近似后验
Evidence Lower Bound (ELBO)	\(\log p(x)\) 的可优化下界, 通常 $$
Amortized Inference	利用一个 Encoder 网络来表示整体的变分推断
Recognition Model / Encoder Network	输出 \(z\) 的分布函数 \(p_\phi(z \| x)\)
Generative Model / Decoder Network	从latent variable \(z\) 来生成 \(x\) \(p_{\theta}(x \| z)\)
AutoEncoder	数据压缩的一种方法，让
Re-parameterization trick	使得梯度可以传递到Recognition Model里面

10 Q & A

Question 1: What is the machine learning framework of VAE-type models? VAE 属于 生成模型（generative model）+ 变分推断（variational inference） 框架。它将生成过程 p_(z) 和 p_(x|z) 与推断网络 q_(z|x) 结合，通过最大化 ELBO 来同时学习 和 。

Question 2: 在 VAE 中，为什么直接对 p_(x) 进行优化不可行？ 因为 p_(x) 里有一个对潜变量 z 的积分（或求和），通常是不可解析的： \[ \log p_\theta(x) = \log \int p_\theta(x, z)\, dz = \log \int p_\theta(x \mid z)\, p(z)\, dz \]

Question 3: reparameterization trick）是如何让 q_(z|x) 的梯度可计算的？

Question 4: 为什么 VAE 有时生成模糊？根本原因是什么？如何改进？

11 Related resource & Further Reading

虽然VAE提出的概念很简单的，但是训练时，还是会存在许多的问题，就比如 Poster Collapse。

(Higgins et al. 2017) \(\beta\)-VAE: https://openreview.net/forum?id=Sy2fzU9gl

(Burgess et al. 2018)

文章中有提到 (Vahdat and Kautz 2021) NVAE: A Deep Hierarchical Variational Autoencoder NVAE 相比传统 VAE 的最大不同在于：它不是只使用一个简单的潜变量 \(\mathrm{z}\) 来一次性生成整张图像，而是采用深度分层的潜变量结构，让图像生成过程从粗到细逐层细化。同时，它为了解决深层 VAE 容易训练不稳定的问题，使用了专门设计的深度解码器结构（ResNet + GroupNorm）和训练技巧来防止 KL 崩塌，使每一层潜变量都能有效工作。简单来说，传统 VAE 是一次性生成图像，而 NVAE 是像画画一样分多层逐步完善图像，并通过改进架构保证训练稳定和生成质量更高。

可以发现，我们这个VAE，是无法根据我们的要求来生成数据，只能Random Generate Image。那如何给他一个 Conditional 呢 (Sohn, Lee, and Yan 2015) Conditional VAE（CVAE） \[ p_\theta(x \mid y, z) \]

\[ q_\phi(z \mid x, y) \]

\[ \mathcal{L}(x, y) = \mathbb{E}_{q_\phi(z\mid x,y)}[\log p_\theta(x \mid y, z)] - D_{KL}(q_\phi(z\mid x,y) \Vert p(z)) \]

References

Burgess, Christopher P., Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. 2018. “Understanding Disentangling in $β$-VAE.” April 10, 2018. https://doi.org/10.48550/arXiv.1804.03599.

Higgins, Irina, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. 2017. “Β-VAE: LEARNING BASIC VISUAL CONCEPTS WITH A CONSTRAINED VARIATIONAL FRAMEWORK.”

Kingma, Diederik P., and Max Welling. 2022. “Auto-Encoding Variational Bayes.” December 10, 2022. https://doi.org/10.48550/arXiv.1312.6114.

Sohn, Kihyuk, Honglak Lee, and Xinchen Yan. 2015. “Learning Structured Output Representation Using Deep Conditional Generative Models.” In Advances in Neural Information Processing Systems. Vol. 28. Curran Associates, Inc. https://proceedings.neurips.cc/paper_files/paper/2015/hash/8d55a249e6baa5c06772297520da2051-Abstract.html.

Vahdat, Arash, and Jan Kautz. 2021. “NVAE: A Deep Hierarchical Variational Autoencoder.” January 8, 2021. https://doi.org/10.48550/arXiv.2007.03898.