[解读] 浅析变分自编码器 (VAE)

变分自编码器 (VAE)

VAE 开篇论文:

(1312) Auto-Encoding Variational Bayes

(1401)Stochastic Backpropagation and Approximate Inference in Deep Generative Models

VAE 教程 (1606) Tutorial on Variational Autoencoders

VAE 综述 (1906) An Introduction to Variational Autoencoders

变分推理 Variational Inference

博文解读: https://blog.youkuaiyun.com/jackytintin/article/details/53641885

https://blog.youkuaiyun.com/weixin_40255337/article/details/83088786

研究背景

在机器学习领域, 我们对学习数据的概率模型非常感兴趣, 概率模型对于未知预测和各种形式的辅助和自动决策有很大帮助.

假设观察变量 $\mathbf{x}$ 服从于一个未知的真实分布 $p^{\star }(\mathbf{x})$, 我们希望学习分布 $p_{\theta }(\mathbf{x})$ 的参数 $\theta$ 来逼近真实分布 $p^{\star }(\mathbf{x})$.

这个逼近方法有很多, 本文考虑深度隐变量模型 (DLVM).

首先使用 $\mathbf{z}$ 来代表隐变量, 有向图模型表示为一种建立于 $\mathbf{x}$ 和 $\mathbf{z}$ 上的联合分布 $p_{\theta }(\mathbf{x},\mathbf{z})$, 计算它的边际分布即可得到
$$p_{\theta }(\mathbf{x})=\int p_{\theta }(\mathbf{x},\mathbf{z})d\mathbf{z}$$
如果 $\mathbf{z}$ 是连续型变量, 则 $p_{\theta }(\mathbf{x})$ 可以看成一个无限的混合模型, 由此可见它可以表示任意的真实分布. 如果使用神经网络来参数化 $p_{\theta }(\mathbf{x},\mathbf{z})$, 则这个网络称为 DLVM.

DLVM 的一个重要优点是即使有向图中的每个因子(先验或条件分布) 是相对简单的分布, $p_{\theta }(\mathbf{x})$ 可以是非常复杂的. 因此我们使用 DLVM 来逼近真实分布 $p^{\star }(\mathbf{x})$.

DLVM 作为一种最简单最常见的方法, 它的分解具有以下结构:
$$p_{\theta }(\mathbf{x},\mathbf{z})=p_{\theta }(\mathbf{z})p_{\theta }(\mathbf{x}|\mathbf{z})$$
其中 $p_{\theta }(\mathbf{z})$ 和 $ p_{\theta}(\mathbf{x}|\mathbf{z})$ 至少一个是被指定的, $p_{\theta }(\mathbf{z})$ 是 $\mathbf{z}$ 的先验分布.

DLVM 可以通过最大似然来学习, 但由于计算 $p_{\theta }(\mathbf{x},\mathbf{z})$ 的积分没有分析解或者有效的估计, 所以无法通过求微分的方式来优化它.

计算 $p_{\theta }(x)$ 的难易程度与 $p_{\theta }(\mathbf{z}|\mathbf{x})$ 难易程度是相关联的, 换句话说, 如果 $p_{\theta }(x)$ 好算,则后者也好算, 注意到联合分布 $p_{\theta }(\mathbf{x},\mathbf{z})$ 是容易计算的, 这些密度分布有以下关系:
$$p_{\mathbf{\theta }}(\mathbf{z}|\mathbf{x})=\frac{p_{\mathbf{\theta }}(\mathbf{x},\mathbf{z})}{p_{\mathbf{\theta }}(\mathbf{x})}$$
近似推理方法能够估计后验分布 $p_{\mathbf{\theta }}(\mathbf{z}|\mathbf{x})$ 和边际分布 $p_{\mathbf{\theta }}(\mathbf{x})$. 然而传统的计算方法是比较复杂的, 例如这需要逐个样本的循环, 或者产生不好的后验近似, 我们的目的是避免这种高代价的处理过程.

变分自编码器

前面说到 DLVM 模型的训练问题, 变分自编码器则是一种非常高效的计算框架. 首先引入一个参数化的推理模型 $q_{\varphi }(\mathbf{z}|\mathbf{x})$, 也被称为编码器或识别模型. 称 $\varphi$ 为变分参数. 我们优化这个参数来使得
$$q_{\varphi }(\mathbf{z}|\mathbf{x})\approx p_{\mathbf{\theta }}(\mathbf{z}|\mathbf{x})$$
对任意一个推理模型 $q$ 和参数 $\varphi$ 有
KaTeX parse error: \cr valid only within a tabular/array environment
其中第二项是 $p$ 和 $q$ 之间的 KL 散度, 这是一个非负值, 当分布相同时值为0. 第一项是一个变分下界, 也被称为 evidence lower bound (ELBO), 记为
$${\mathcal{L}}_{\mathbf{\theta },\varphi }(\mathbf{x})={\mathbb{E}}_{q_{\varphi }(\mathbf{z}|\mathbf{x})}\left[\log p_{\mathbf{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\varphi }(\mathbf{z}|\mathbf{x})\right]$$

已知 KL 散度值是非负的, 由
$$\begin{array}{rl}{\mathcal{L}}_{\mathbf{\theta },\varphi }(\mathbf{x})& =\log p_{\mathbf{\theta }}(\mathbf{x})-D_{KL}\left(q_{\varphi }(\mathbf{z}|\mathbf{x})\Vert p_{\mathbf{\theta }}(\mathbf{z}|\mathbf{x})\right)\\ & \le \log p_{\mathbf{\theta }}(\mathbf{x})\end{array}$$
可知 ELBO 是 $\log p_{\mathbf{\theta }}(\mathbf{x})$ 的下界. 从上式可以看出, 最大化 ELBO (优化参数为 $\theta ,\varphi$) 将会带来两个结果: 首先这能够同时最大化似然 $\log p_{\mathbf{\theta }}(\mathbf{x})$, 使得模型变得更好, 其次能够最小化 KL 散度.

现在我们的目标是最大化 ELBO, 它有一个重要的特性是可以使用随机梯度下降方法来联合优化所有参数.

ELBO 关于 $\theta$ 的无偏梯度为
$$\begin{array}{rl}{\nabla }_{\mathbf{\theta }}{\mathcal{L}}_{\mathbf{\theta },\mathbf{\varphi }}(\mathbf{x})& ={\nabla }_{\mathbf{\theta }}{\mathbb{E}}_{q_{\varphi }(\mathbf{z}|\mathbf{x})}\left[\log p_{\mathbf{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\varphi }(\mathbf{z}|\mathbf{x})\right]\\ & ={\mathbb{E}}_{q_{\varphi }(\mathbf{z}|\mathbf{x})}\left[{\nabla }_{\mathbf{\theta }}\left(\log p_{\mathbf{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\varphi }(\mathbf{z}|\mathbf{x})\right)\right]\\ & \simeq {\nabla }_{\mathbf{\theta }}\left(\log p_{\mathbf{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\varphi }(\mathbf{z}|\mathbf{x})\right)\\ & ={\nabla }_{\mathbf{\theta }}\left(\log p_{\mathbf{\theta }}(\mathbf{x},\mathbf{z})\right)\end{array}$$
其中第四行是第二行的简单的蒙特卡罗估计, 最后两行的 $\mathbf{z}$ 随机采样于 $q_{\varphi }(\mathbf{z}|\mathbf{x})$.

然而关于 $\varphi$ 的无偏梯度比较难以获得, 这是因为与 $q_{\varphi }(\mathbf{z}|\mathbf{x})$ 有关, 微分算子无法穿过期望算子:
$$\begin{array}{rl}{\nabla }_{\varphi }{\mathcal{L}}_{\mathbf{\theta },\mathbf{\varphi }}(\mathbf{x})& ={\nabla }_{\mathbf{\varphi }}{\mathbb{E}}_{q_{\varphi }(\mathbf{z}|\mathbf{x})}\left[\log p_{\mathbf{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})\right]\\ & \ne {\mathbb{E}}_{q_{\varphi }(\mathbf{z}|\mathbf{x})}\left[{\nabla }_{\mathbf{\varphi }}\left(\log p_{\mathbf{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})\right)\right]\end{array}$$

重参数技巧

对于连续型的隐变量 $\mathbf{z}$, 我们可以用重参数的技巧来计算无偏梯度 ${\nabla }_{\varphi }{\mathcal{L}}_{\mathbf{\theta },\mathbf{\varphi }}(\mathbf{x})$. 把随机变量变量 $\mathbf{z}$ 表示为另一个随机变量 $\mathbf{ϵ}\sim p(\mathbf{ϵ})$ 的可微并且可逆的变换, 即
$$\mathbf{z}=\mathbf{g}(\mathbf{ϵ},\mathbf{\varphi },\mathbf{x})$$
其中三个随机变量都是互相独立的. 给定这个变换后, ELBO 可以重写为
$$\begin{array}{rl}{\mathcal{L}}_{\mathbf{\theta },\mathbf{\varphi }}(\mathbf{x})& ={\mathbb{E}}_{q_{\varphi }(\mathbf{z}|\mathbf{x})}\left[\log p_{\mathbf{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})\right]\\ & ={\mathbb{E}}_{p(\mathbf{ϵ})}\left[\log p_{\mathbf{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})\right]\end{array}$$
然后便可以采用简单的蒙特卡洛模拟来估计出关于单点的 ELBO, 整理如下
$$\begin{array}{rl}ϵ& \sim p(ϵ)\\ \mathbf{z}& =\mathbf{g}(\mathbf{\varphi },\mathbf{x},\mathbf{ϵ})\\ {\tilde{\mathcal{L}}}_{\mathbf{\theta },\varphi }(\mathbf{x})& =\log p_{\mathbf{\theta }}(\mathbf{x},\mathbf{z})-\log q_{\varphi }(\mathbf{z}|\mathbf{x})\end{array}$$

算法伪代码如下:

只要选择一个好的 $g()$ 函数, 那么就是关于 $\log q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ 的计算就简单了. 当 $g()$ 是一个可逆函数时, $\mathbf{ϵ}$ 和 $\mathbf{z}$ 的密度的关系为
$$\log q_{\varphi }(\mathbf{z}|\mathbf{x})=\log p(ϵ)-\log d_{\varphi }(\mathbf{x},ϵ)$$
其中第二项为 $\log d_{\varphi }(\mathrm{x},ϵ)=\log \left|\det \left(\frac{\partial \mathrm{z}}{\partial ϵ}\right)\right|$.

然后构造一个灵活的 $g()$ 来使得 $\log d_{\varphi }(\mathrm{x},ϵ)$ 计算更简单, 并且得到高度灵活的推理模型 $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$. 常用的选择是一个简单的 factorized Gaussian encoder $q_{\varphi }(\mathbf{z}|\mathbf{x})=\mathcal{N}\left(\mathbf{z};\mathbf{\mu },\mathrm{diag}\left({\mathbf{\sigma }}^2\right)\right)$:
$$\begin{array}{rl}(\mathbf{\mu },\log \mathbf{\sigma })& ={\text{ EncoderNeuralNet }}_{\mathbf{\varphi }}(\mathbf{x})\\ q_{\varphi }(\mathbf{z}|\mathbf{x})& =\prod _i{q}_{\varphi }\left(z_i|\mathbf{x}\right)=\prod _i\mathcal{N}\left(z_i;{\mu }_i,{\sigma }_i^2\right)\end{array}$$
重参数后可以写为
$$\begin{array}{rl}ϵ& \sim \mathcal{N}(0,\mathbf{I})\\ (\mathbf{\mu },\log \mathbf{\sigma })& ={\text{ EncoderNeuralNet }}_{\varphi }(\mathbf{x})\\ \mathbf{z}& =\mathbf{\mu }+\mathbf{\sigma }\odot \mathbf{ϵ}\end{array}$$
并且有 $\log d_{\varphi }(\mathbf{x},ϵ)=\log \left|\det \left(\frac{\partial \mathbf{z}}{\partial ϵ}\right)\right|={\sum }_i\log {\sigma }_i$. 从而得到
$$\begin{array}{rl}\log q_{\varphi }(\mathbf{z}|\mathbf{x})& =\log p(\mathbf{ϵ})-\log d_{\varphi }(\mathbf{x},\mathbf{ϵ})\\ & =\sum _i\log \mathcal{N}\left({ϵ}_i;0,1\right)-\log {\sigma }_i\end{array}$$