VAE

思路

从最大似然估计开始

如果已知数据集$ \mathbf{x}={\mathbf{x}{i}, i = 0,1,2,…,N} $,如果我们能找一个联合概率分布$ p(\mathbf{x}) = p(\mathbf{x{0}},\mathbf{x{1}},\mathbf{x{2}}…,\mathbf{x_{N}}) $,那么我们可以用联合概率分布在分布中采样得到新样本(程序实现上就是直接用np.random.normal之类的函数,至于具体怎么实现就是另一个问题了,一般也不关心)。常用的方法就是对数据集$ \mathbf{x} $进行最大似然估计(Maximum likelihood estimation,MLE)。最大似然估计需要假设数据遵循一个已知分布来求参数,不过往往预先很难知道数据集$ \mathbf{x} $遵循什么分布。

隐变量模型

虽然我们无法知道数据集$ \mathbf{x} $的分布,但是设想我们要画一种类型的画$ \mathbf{x} $,我们首先需要构思这类画共同的布局、颜色等信息,然后再根据这些信息画出来想要的画。像画的布局,颜色,在更抽象的层面描述一幅画,我们将它们记为$ p(\mathbf{z}) $,将通过这些信息画画的过程记作$ P(x \mid \mathbf{z}) $,整个画画的过程可以记作:

$ \begin{equation}p(x)=\int p(x \mid \mathbf{z}) p(\mathbf{z}) \mathrm{d} z\end{equation} $

虽然我们不能假设整个数据集$ x $的分布,但是对于数据集更抽象的特征我们通常可以假设它的分布遵循某种特定分布(神奇)。然后再通过$ p(x \mid \mathbf{z}) $将先验分布映射到目标分布$ p(x) $上。这种方法被称为隐变量模型。背后关键的思想是:任何一个概率分布经过足够复杂的函数变换后可以映射到任意概率分布。

有了这个模型,直接的想法是:

  • 使用一个网络代替$ p(x \mid \mathbf{z}) $,只要有一个先验分布$ p(\mathbf{z}) $,我们在这个分布中采样然后喂给网络,然后让网络产生结果,产生的结果再和数据集中的样本计算Loss即可。但是这个想法存在两个问题:
  • 怎么确定产生的先验分布采样和数据集$ \mathbf{X} $中样本的对应关系。也就是我们采样的输入,到底对应哪一个数据集样本?用哪一个样本计算Loss?
  • $ p(\mathbf{z}) $的采样空间很大,没有利用好数据集信息。

从这两个问题我们注意到我们忽略了数据集可以提供的信息。因此,我们选择从后验概率分布$ p(\mathbf{z}\mid x) $中采样,这样即利用了数据新信息,也减少了采样空间。

近似后验概率分布

由此我们产生了更进一步的想法:

  • 通过$ p(\mathbf{z} \mid x) $采样$ \mathbf{z} $,我们在这个分布中采样然后喂给网络,然后让网络产生结果,产生的结果再和数据集中的样本计算Loss即可。

基于这个想法,我们的问题变成了如何得到$ p(\mathbf{z} \mid x) $?再进一步,我们是否可以用$ p_{\theta}(\mathbf{z} \mid x;\theta) $近似$ p(\mathbf{z} \mid x) $?我们自然期待$ p_{\theta}(\mathbf{z} \mid x;\theta) $尽可能接近$ p(\mathbf{z} \mid x) $。由此我们可以使用KL散度来度量两个概率分布的相似程度(距离),有:

$ \begin{equation}\begin{aligned}& K L\left(q\theta(z \mid x) | p(z \mid x)\right) \& =\int q\theta(z \mid x) \ln \frac{q\theta(z \mid x)}{p(z \mid x)} \mathrm{d} z \& =\mathbb{E}{z \sim q\theta(z \mid x)}\left[\ln \frac{q\theta(z \mid x)}{p(z \mid x)}\right] \& =\mathbb{E}{z \sim q\theta(z \mid x)}\left[\ln q\theta(z \mid x)-\ln p(z \mid x)\right] \& =\mathbb{E}{z \sim q\theta(z \mid x)}\left[\ln q\theta(z \mid x)-\ln \frac{p(x \mid z) p(z)}{p(x)}\right] \& =\mathbb{E}{z \sim q\theta(z \mid x)}\left[\ln q\theta(z \mid x)-\ln p(z)-\ln p(x \mid z)\right]+\ln p(x) \& =K L\left(q\theta(z \mid x)|| p(z)\right)-\mathbb{E}{z \sim q\theta(z \mid x)}[\ln p(x \mid z)]+\ln p(x)\end{aligned}\end{equation} $

整理后可以得到:

$ \begin{equation}\ln p(x)-K L\left(q\theta(z \mid x) | p(z \mid x)\right)=\mathbb{E}{z \sim q\theta(z \mid x)}[\ln p(x \mid z)]-K L\left(q\theta(z \mid x) | p(z)\right)\end{equation} $

如果$ q\theta(z \mid x) $和 $ p(z \mid x) $完全相同,那么$ K L\left(q\theta(z \mid x) | p(z \mid x)\right) = 0 $,则有:

$ \begin{equation}\ln p(x)=\mathbb{E}{z \sim q\theta(z \mid x)}[\ln p(x \mid z)]-K L\left(q_\theta(z \mid x) | p(z)\right)\end{equation} $

左边表示$ x $的最大似然估计,右边被称为对数似然下界(Evidence Lower Bound,ELBO),也就说对 $ p(z \mid x) $的近似等价于$ x $的最大似然估计(KL散度已经为0了)。因此,我们希望$ (4) $式中右边项(ELBO)尽可能大,那么似然函数也最大,$ q_\theta(z \mid x) $和 $ p(z \mid x) $越接近。

最大化ELBO(重构误差+正则项)

ELBO中$ \mathbb{E}{z \sim q\theta(z \mid x)}[\ln p(x \mid z)] $表示把$ x $编码成$ z $然后采样得到$ z $,然后恢复到$ x $的期望,期望越大,编码后恢复的误差越小。这一部分可以看成一个Encoder,一个采样和一个Decoder,为了让ELBO最大,这部分的期望也要最大,也就是经过Encoder和Decoder恢复的采样误差应该越小,所以这部分被称为重构误差。同样为了让ELBO最大,ELBO中的$ K L\left(q_\theta(z \mid x) | p(z)\right) $应该尽可能小(KL散度大于等于0),作为正则项

整理一下,为了从$ p(\mathbf{z} \mid x) $中采样(利用好数据集信息),我们想估计一个采样$ p_{\theta}(\mathbf{z} \mid x;\theta) $近似$ p(\mathbf{z} \mid x) $,为了让近似更接近,我们引入KL散度来度量近似程度。通过推导我们发现为了让近似程度最高,等同于让ELBO最大,而ELBO的最大可以理解为一次重构的误差最小以及正则项最小,重构过程可以拆分为Encoder、采样以及Decoder并让他们的误差最小。加上正则项,就是最终VAE的Loss。

重构误差 & 重参数化技巧

重构误差可以使用常用的度量来确定,比较好计算。但是在重构过程中涉及到一次从编码后分布$ p_{\theta}(\mathbf{z} \mid x;\theta) $的采样,采样过程不可导,没法求Loss,所以需要重参数化技巧。所谓重参数技巧就是指从编码后分布$ p_{\theta}(\mathbf{z} \mid x;\theta) = \mathcal{N}\left(\mu, \sigma^2\right) $中采样,相当于从$ \mathcal{N}(0, I) $中采样一个$ \epsilon $,然后计算$ Z=\mu+\epsilon \times \sigma $。$ Z $是可导的,那么就可以放入Loss的计算了。

正则项

下面我们来处理正则项$ K L\left(q\theta(z \mid x) | p(z)\right) $。一般还是选用正态分布来建模$ q\theta(z \mid x) = z \sim \mathcal{N}(\mu, \sigma^{2}) $,$ p(z) $还是使用标准正态分布建模,$ p(z) = \mathcal{N}(0,I) $。选择$ p(z) $是标准正态分布的原因是为了让方差足够小但是不为$ 0 $(否则方差为0,样本间没有差异,等于失去了生成新样本的能力)。同时KL散度可以计算如下:

$ \begin{equation}\begin{aligned}& K L\left(q_\theta(z \mid x)|| \mathcal{N}(0, I)\right) \= & \int \frac{1}{\sqrt{2 \pi \sigma^2}} e^{-\frac{(x-\mu)^2}{2 \sigma^2}} \ln \frac{\frac{1}{\sqrt{2 \pi \sigma^2}} e^{-\frac{(x-\mu)^2}{2 \sigma^2}}}{\frac{1}{\sqrt{2 \pi}} e^{-\frac{x^2}{2}}} \mathrm{~d} x \= & \int \frac{1}{\sqrt{2 \pi \sigma^2}} e^{-\frac{(x-\mu)^2}{2 \sigma^2}} \ln \frac{1}{\sqrt{\sigma^2}} \times e^{\frac{x^2}{2}-\frac{(x-\mu)^2}{2 \sigma^2}} \mathrm{~d} x \= & \int \frac{1}{\sqrt{2 \pi \sigma^2}} e^{-\frac{(x-\mu)^2}{2 \sigma^2}}\left[-\frac{1}{2} \ln \sigma^2+\frac{1}{2} x^2-\frac{1}{2} \frac{(x-\mu)^2}{\sigma^2}\right] \mathrm{d} x \= & \frac{1}{2} \int \frac{1}{\sqrt{2 \pi \sigma^2}} e^{-\frac{(x-\mu)^2}{2 \sigma^2}}\left[-\ln \sigma^2+x^2-\frac{(x-\mu)^2}{\sigma^2}\right] \mathrm{d} x \= & \frac{1}{2}\left(-\ln \sigma^2+\mathbb{E}\left[x^2\right]-\frac{1}{\sigma^2} \mathbb{E}\left[(x-u)^2\right]\right) \= & \frac{1}{2}\left(-\ln \sigma^2+\sigma^2+\mu^2-1\right)\end{aligned}\end{equation} $

两个Loss加起来就是VAE的Loss:

$ \begin{equation}L o s s=\frac{1}{n} \sum{i=1}^n\left[\ln p\left(x_i \mid z_i\right)+\frac{1}{2} \sum{i=1}^d\left(-\ln \sigma^2+\sigma^2+\mu^2-1\right)\right]\end{equation} $

其中$ \ln p\left(x_i \mid z_i\right) $是重构误差,会使用网络产生的结果与样本的Loss代替,并不直接计算,$ n $是样本数,$ d $是维数。

CVAE

Learning Structured Output Representation using Deep Conditional Generative Models (nips.cc)

与VAE的差别就是加入了Condition(标签)信息,也就是说要使用网络近似的分布变成了$ p(\mathbf{z} \mid x,y) $,$ y $是标签信息。同$ (2) $式有:

$ \begin{equation}\begin{aligned}& K L\left(q\theta(z \mid x,y) | p(z \mid x,y)\right) \& =\int q\theta(z \mid x,y) \ln \frac{q\theta(z \mid x,y)}{p(z \mid x,y)} \mathrm{d} z \& =\mathbb{E}{z \sim q\theta(z \mid x,y)}\left[\ln \frac{q\theta(z \mid x,y)}{p(z \mid x,y)}\right] \& =\mathbb{E}{z \sim q\theta(z \mid x,y)}\left[\ln q\theta(z \mid x,y)-\ln p(z \mid x,y)\right] \& =\mathbb{E}{z \sim q\theta(z \mid x,y)}\left[\ln q\theta(z \mid x,y)-\ln \frac{p(x \mid z,y) p(z)}{p(x)}\right] \& =\mathbb{E}{z \sim q\theta(z \mid x,y)}\left[\ln q\theta(z \mid x,y)-\ln p(z)-\ln p(x \mid z,y)\right]+\ln p(x) \& =K L\left(q\theta(z \mid x,y)|| p(z)\right)-\mathbb{E}{z \sim q\theta(z \mid x,y)}[\ln p(x \mid z,y)]+\ln p(x)\end{aligned}\end{equation} $

如果$ q\theta(z \mid x,y) $和 $ p(z \mid x,y) $完全相同,那么$ K L\left(q\theta(z \mid x,y) | p(z \mid x,y)\right) = 0 $,则有:

$ \begin{equation}\ln p(x)=\mathbb{E}{z \sim q\theta(z \mid x,y)}[\ln p(x \mid z,y)]-K L\left(q_\theta(z \mid x,y) | p(z)\right)\end{equation} $

回忆$ q_\theta(z \mid x,y) $会使用Encoder表示,所以在实现网络时,需要将标签信息$ y $也告诉网络,通常使用的方法是将$ y $进行Embedding,然后和$ x $一起构成$ (x,\text{Embedding}(y)) $送入Encoder;同样,$ p(x \mid \mathbf{z},y) $会使用Decoder近似,这里同样需要将$ \mathbf{z} $和Embedding后的$ y $一起构成$ (\mathbf{z},\text{Embedding}(y)) $送入Decoder中;KL散度部分计算不变(同样还是高斯分布,有没有$ y $都一样),这就是CVAE了。

根据上面的原理实现的VAE和CVAE在这里:https://github.com/wxhaooo/aigcLearn

Reference

VAE的pytorch实现, reproduction loss使用了binary_cross_entropy

CVAE的原始Paper

CVAE使用keras的实现