前段时间在做毕设的时候遇到了变分自编码器(Variational Auto-Encoder,VAE),网上找了许多教程进行了学习,在看了一些博客,代码以及原文后,有了自己的理解,决定自己把这个部分进行梳理一下。
一、VAE概念
VAE的目的是将输入数据编码到隐空间,再将隐空间的数据解码到原数据空间。在实际使用中,为了完成不同的任务,在测试阶段也可以只使用单独的编码部分或解码部分。VAE的设计让模型在训练的过程中也存在一种对抗,与GAN网络相似,因此可以应用在无监督学习的任务中。
二、VAE原理
首先来看一组数据\(x=\{x_1,x_2,x_3,...,x_K\}\),我们要训练一个模型将数据\(x\)转换到隐空间得到\(z\),为便于理解,本文只考虑输入维度与编码维度相等的情况,即\(z=\{z_1,z_2,z_3,...,z_K\}\),模型可以表示为\(q_\phi(z\vert x)\),再将隐空间的数据解码到原始的数据空间得到重构的\(\widehat{x}=\{\widehat{x}_1,\widehat{x}_2,\widehat{x}_3,...,\widehat{x}_K\}\),该模型可以表示为\(p_\theta(x\vert z)\),之前在看博客和原文时对这个模型一直存在着比较大疑问,认为应该是\(p_\theta(\widehat{x}\vert z)\),但是在博客和原文中都没有出现过这种写法,在反复对原文和博客的阅读中,理解了这种写法,在这里\(x\)与\(\widehat{x}\)都同属于原始数据空间,\(x\)也只是原始数据空间中的一些样本,不能表示原始空间的分布,若能直接表示原始数据空间的分布,那就不需要费劲地将原始数据转换到隐空间,再将隐空间的数据转换到原始数据空间。同时这里的模型输出的\(x\)与\(z\)表示的是一种分布,而最终输出的数据是对这两个分布采样得到的结果。
2.1 公式推导
假定所有数据都是独立同分布的,我们要对\(p_\theta(x\vert z)\)做参数估计,可以使用对数最大似然法,即最大化如下的似然函数:
\[\ln p_\theta(x^{(1)},x^{(2)},x^{(3)},...,x^{(N)})=\sum_{i=1}^N\ln p_\theta(x^{(i)})\]然后在整个模型中需要令\(q_\phi(z\vert x)\)逼近\(p_\theta(x\vert z)\),可以使用KL散度计算,即
\[\begin{aligned} D_{KL}(q_\phi(z|x^{(i)})||p_\theta(z|x^{(i)})) &= \int q_\phi(z|x^{(i)}) \ln \frac{q_\phi(z|x^{(i)})}{p_\theta(z|x^{(i)})}dz \\ &=\mathbb{E}_{z \sim q_\phi(z|x^{(i)})}\ln \frac{q_\phi(z|x^{(i)})}{p_\theta(z|x^{(i)})}\\ &=\mathbb{E}_{z \sim q_\phi(z|x^{(i)})}\ln \frac{q_\phi(z|x^{(i)})p_\theta(x^{(i)})}{p_\theta(z,x^{(i)})}\\ &=\mathbb{E}_{z \sim q_\phi(z|x^{(i)})}\ln \frac{q_\phi(z|x^{(i)})}{p_\theta(z,x^{(i)})} +\mathbb{E}_{z \sim q_\phi(z|x^{(i)})}\ln p_\theta(x^{(i)})\\ &=\mathbb{E}_{z \sim q_\phi(z|x^{(i)})}\ln \frac{q_\phi(z|x^{(i)})}{p_\theta(z,x^{(i)})}+\ln p_\theta(x^{(i)}) \end{aligned}\]因此可以得到
\[\begin{aligned} \ln p_\theta(x^{(i)}) &= D_{KL}(q_\phi(z|x^{(i)})||p_\theta(z|x^{(i)}))- \mathbb{E}_{z \sim q_\phi(z|x^{(i)})}\ln \frac{q_\phi(z|x^{(i)})}{p_\theta(z,x^{(i)})} \end{aligned}\]由于KL散度是一个非负值,因此可以得到
\[\begin{aligned} \ln p_\theta(x^{(i)}) &\geq -\mathbb{E}_{z \sim q_\phi(z|x^{(i)})}\ln \frac{q_\phi(z|x^{(i)})}{p_\theta(z,x^{(i)})}\\ &=\mathbb{E}_{z \sim q_\phi(z|x^{(i)})}\ln \frac{p_\theta(z,x^{(i)})}{q_\phi(z|x^{(i)})}\\ &=\int q_\phi(z|x^{(i)})\ln p_\theta(z,x^{(i)})dz - \int q_\phi(z|x^{(i)}) \ln q_\phi(z|x^{(i)})dz\\ \end{aligned}\]从这里就可以得到我们的损失函数
\[\begin{aligned} \mathcal{L(\phi,\theta;x^{(i)})}&=\int q_\phi(z|x^{(i)}) \ln q_\phi(z|x^{(i)})dz -\int q_\phi(z|x^{(i)}) \ln p_\theta(z,x^{(i)})dz\\ &=\int q_\phi(z|x^{(i)}) \ln q_\phi(z|x^{(i)})dz -\int q_\phi(z|x^{(i)}) \ln p_\theta(x^{(i)}|z)p_\theta(z)dz\\ &=\int q_\phi(z|x^{(i)}) \ln \frac{q_\phi(z|x^{(i)})}{p_\theta(z)}dz -\int q_\phi(z|x^{(i)}) \ln p_\theta(x^{(i)}|z)dz\\ &=D_{KL}(q_\phi(z|x^{(i)})||p_\theta(z))-\mathbb{E}_{z \sim q_\phi(z|x^{(i)})}\ln p_\theta(x^{(i)}|z)\\ \end{aligned}\]对于隐空间中的数据分布,我们可以假定为任意的分布,因为我们可以将原始数据转换到一个我们想要的空间,因此在此处我们可以假定为高斯分布,对于假定为高斯分布的优势下面会进行叙述。当\(z\)假定服从标准高斯分布可以得到
\[\begin{aligned} p_\theta (z) &\sim \mathcal{N}(0,1)\\ q_\phi(z|x_i) &\sim \mathcal{N}(\mu(x_i,\phi),\sigma ^2(x_i,\phi)) \end{aligned}\]可以将损失函数\(\mathcal{L}(\phi,\theta;x^{(i)})\)的第一项进一步地化简得到
\[\begin{aligned} D_{KL}(q_\phi(z|x^{(i)})||p_\theta(z))&=\mathbb{E}_{z \sim \mathcal{N}(\mu,\sigma ^2)}\ln \mathcal{N}(\mu,\sigma ^2) - \mathbb{E}_{z \sim \mathcal{N}(\mu,\sigma ^2)}\ln \mathcal{N}(0,1)\\ &=\mathbb{E}_{z \sim \mathcal{N},\sigma ^2)}\ln {\frac{1}{\sqrt{2\pi\sigma^2}}e^{-\frac{(z-\mu)^2}{2\sigma^2}}}-\mathbb{E}_{z\sim \mathcal{N}(\mu,\sigma^2)}\ln {\frac{1}{\sqrt{2\pi}}e^{-\frac{z^2}{2}}}\\ &=-\frac{1}{2}\ln {2\pi}-\frac{1}{2}\ln {\sigma^2}-\frac{1}{2\sigma^2}\mathbb{E}_{z\sim \mathcal{N}(\mu,\sigma^2)}(z-\mu)^2 + \frac{1}{2}\ln2\pi + \frac{1}{2}\mathbb{E}_{z\sim \mathcal{N}(\mu,\sigma^2)}{z^2}\\ &=-\frac{1}{2}(\ln \sigma^2 +1) +\frac{1}{2}(\mu^2 +\sigma^2)\\ &=\frac{1}{2}\sum_{j=1}^{J}\mu_j^2(x^{(i)}) +\sigma_j^2(x^{(i)})-\ln \sigma_j^2(x^{(i)})-1 \end{aligned}\]其中\(\mathbb{E}_{z\sim \mathcal{N}(\mu,\sigma^2)}(z-\mu)^2)\)为方差,\(\sigma^2\),\(\mathbb{E}_{z\sim \mathcal{N}(\mu,\sigma^2)}{z^2}\)为变量\(z\)的二阶矩,结果是\(\mu^2+\sigma^2\)。这里的\(\mu(x)\)与!\(\sigma(x)\)可以通过神经网络直接求得。
对于损失函数\(\mathcal{L}(\phi,\theta;x^{(i)})\)的第二项,首先需要了解蒙特卡洛模拟,此处只需要知道该理论的结论即可。即
\[\begin{aligned} \mathbb{E}_{x\sim p(x)}f(x) &\approx \frac{1}{L}\sum _{i=1}^Lf(x_i)\quad x_i \sim p(x)\\ \end{aligned}\]由此可以得到损失函数\(\mathcal{L}(\phi,\theta;x^{(i)})\))的最终结果为
\[\begin{aligned} \mathcal{L}(\phi,\theta;x^{(i)}) &\approx \widetilde{\mathcal{L}}(\phi,\theta;x^{(i)}) \\ &=\frac{1}{2}\sum_{j=1}^{J}\big[\mu_j^2(x^{(i)})+\sigma_j^2(x^{(i)})-\ln \sigma_j^2(x^{(i)})-1\big]-\frac{1}{L}\sum _{l=1}^L\ln p_\theta(x^{(i)}|z^{(i,l)}) \end{aligned}\]其中\(z^{(i,l)} \sim q_\phi(z\vert x^{(i)})\)
所以总的损失函数\(\mathcal{L}(\phi,\theta;x)\)为
\[\mathcal{L}(\phi,\theta;x)=\sum_{i=1}^{N}\mathcal{\widetilde{L}}(\phi,\theta;x^{(i)})\]因此要求最大话似然函数,即最小化总损失函数\(\mathcal{L}(\phi,\theta;x)\)。
在实际应用中,由于数据量\(N\)非常大,因此需要将数据分为\(M\)个minibatch进行训练,优化后的损失函数为
\[\mathcal{L}(\phi,\theta;x) = \frac{N}{M}\sum_{i=1}^{M}\mathcal{\widetilde L}(\phi,\theta;x^{(i)})\]当\(M\)比较大时,内层估计可以由外层估计来完成,即对于每一个样本来说,内层估计需要从\(q_\phi(z\vert x^{(i)})\)中采样\(L\)个数据来进行估计,当样本量比较大时,内层的估计可以忽略,即\(L=1\),这是因为样本量大到一定程度,每个样本的多次采样对结果的影响不会太大,因此每个样本只需要采样一个数据即可。在该论文的实验中,取\(M=100, L=1\)。
2.2 生成模型的估计
至此,我们可以得到VAE模型的目标函数,但是生成模型还需要选择一个分布,在论文中给出了两种候选方案,伯努利分布和正态分布。
(a) 伯努利分布
伯努利分布的随机变量\(x\)只能取两个值0或1,其概率函数为
\[p(x|\rho)= \begin{cases} \rho^x(1-\rho)^{1-x} & x=0,1\\ 0 & x\neq 0,1 \end{cases}\]对于多元二值向量,生成模型\(p_\theta(x\vert z)\)的概率函数为
\[p_\theta(x^{(i)}|z)=\prod_{k=1}^{K}\rho_k^{x_k^{(i)}}(z)(1-\rho_k(z))^{1-x_k^{(i)}}\]这里的\(\rho^{x^{(i)}}(z)\)就是相当于一个decoder,将隐空间的数据解码到数据空间。
因此损失函数中的最后一项可以化简为
\[\ln p_\theta(x^{(i)}|z)=\sum_{k=1}^{K} \big[x_k^{(i)}\ln\rho_k(z)+({1-x_k^{(i)}})(1-\rho_k(z))\big]\](b) 正态分布
当生成模型为正态分布时,生成模型的概率函数为
\[p_\theta(x^{(i)}|z)=\frac{1}{ \prod_{k=1}^K \sqrt { 2 \pi \widetilde\sigma^2_k(z) } } e^{-\frac{1}{2} \big\lVert{\frac{x^{(i)}-\widetilde\mu(z)}{\widetilde\sigma(z)}}\big\rVert^2}\]这里的\(\widetilde\mu(z)\)与\(\widetilde\sigma(z)\)均可以使用神经网络得到,但是我们往往想要得到一个确定的结果,因此我们会将这里的方差为固\(\widetilde\sigma(z)\)定为一个常数,均值\(\widetilde\mu(z)\)相当于一个decoder。则损失函数的最后一项可以化简为
\[\begin{aligned} \ln p_\theta(x^{(i)}|z)&=-\frac{K}{2}\ln 2\pi-\frac{1}{2}\sum_{k=1}^K\ln \widetilde\sigma_k^2(z) -\frac{1}{2}\big\lVert\frac{x^{(i)}-\widetilde\mu(z)}{\widetilde\sigma(z)}\big\rVert^2\\ &\simeq -\lVert{x^{(i)}-\widetilde\mu(z)}\rVert^2 \end{aligned}\]2.3 重参数技巧
在生成模型\(p_\theta(x\vert z)\)中,\(z\)是通过对识别模型采样\(q_\phi(z\vert x)\)得到的,然而采样的操作是不可导的,这使得直接采样不能使用链式求导法则进行求导,所有的信息反馈到隐空间就断了,无法反馈到识别模型\(q_\phi(z\vert x)\),不能对\(\mu(x)\)和\(\sigma(x)\)的参数进行修正。为解决这个问题,采用了重参数的技巧,即\(z^{(i,l)}\)通过如下公式计算得到
\[z^{(i,l)}=\mu(x^{(i)})+\sigma(x^{(i)})\cdot\epsilon^{(l)} \quad\epsilon^{(l)} \sim \mathcal{N}(0,1)\]通过对变量\(\epsilon\)进行采样,避免了直接采样\(z\)引发的无法采用梯度下降法学习的问题。采用这种方式,对于模型\(\mu(x)\)与\(\sigma(x)\)来说,\(\epsilon\)只是一个常数,可以应用梯度下降法。
总结
在VAE模型中,需要构建三个神经网络,分别是\(\mu(x),\sigma(x)\)与\(\widetilde\mu(z)\),其中构成\(\mu(x),\sigma(x)\)了识别模型,即为\(q_\phi(z\vert x)\)编码过程,构成了生\(\widetilde\mu(z)\)成模型,即为解码\(p_\theta(x\vert z)\)过程。
以生成模型为正态分布为例,其损失函数为
\[\widetilde{\mathcal{L}}(\phi,\theta;x^{(i)}) = \frac{1}{2}\sum_{j=1}^{J}\big[\mu_j^2(x^{(i)})+\sigma_j^2(x^{(i)})-\ln \sigma_j^2(x^{(i)})-1\big]+\lVert{x^{(i)}-\widetilde\mu(z)}\rVert^2\]这里的第一项为KL散度损失,第二项为重构损失。
不考虑KL散度损失,模型就退化为了普通的自动编码器(Auto-Encoder,AE)。在VAE中,重构损失希望生成的数据与原始数据相等,但是KL散度损失会使得隐空间的数据呈标准正态分布,在生成隐空间数据\(z\)时,都会对数据进行采样,这会导致对于同一个输入数据,生成的隐空间数据都会不同,进一步地也就会在生成的数据中引入噪声。因此重构损失与KL散度损失是一种对抗的过程,保证了模型的生成数据接近原始数据,同时还增加了模型的泛化能力。