Variational AutoEncoders

Variational AutoEncoders

VAE作为前几年兴起的一个技术,在数据生成方面有极为广阔的应用,并且已经证实了在生成复杂数据方面非常有效。

1. Latent Variables

Latent Variable (潜在变量)是一个我个人认为很神奇的东西,它支撑起了很多以概率角度建立的模型。Latent Variable是一个我们看不见的隐藏的随机变量,往往代表着事物背后的本质。比如一个手写的数字,在计算机看来,你将一个手写的数字的图片传给它,它只能看到一个个的像素点,那现在你叫它分辨一张图片到底是数字几或者生成一张手写数字的图片,它怎么做到呢?这时就需要Latent Variable:

我们将图片传给神经网络的时候,神经网络自己对这张图片进行特征提取,这是神经网络自己内部的行为,我们并不知道它能提取到什么样的特征,我们只能告诉它你以目前这种方式提取的特征所获得的分类/生成图片结果,是正确/错误的,并告诉它有多正确/多错误。它根据我们的指示所提取到的特征,是以神经网络自己内部的形式对图片进行编码(encode)的结果,而编出的码,就是神经网络对Latent Variable进行的模拟。Latent Variable代表的是图片背后的信息,比如数字几,笔画粗细,数字的角度等等。但是需要注意的一点是,以上仅仅为举例,现实中人类往往并不能直接通过这些编码获得有用的信息,只有神经网络自己才看得懂这是啥。

在进行了一定程度上的学习之后,我们希望训练出的神经网络能够自己生成一张图片。这个时候神经网络会对之前训练出的对于latent variable的分布的模拟进行取样,根据取样,来生成一张图片。更具体一点,理想一点,假设latent variable代表的是数字0-9,神经网络会首先在0-9之间根据概率分布取出一个数字,根据这个数字,神经网络会根据之前训练的结果,每个数字大概能长成什么样,来生成相对应的图片。这个过程,被称作解码(decode)。

2. Generative Models

需要明确的一点是VAE(Variational AutoEncoders)是Generative Model(生成模型)。在classification(分类)问题中,Generative Model 是对联合概率密度P(X,Y)P(\mathbf{X},Y)进行模拟,其中X\mathbf{X}是数据点, Y是label,二者都被看作是随机的,前者是随机向量,后者是一个随机变量。

而在VAE的问题中,Generative Model主要应对的是对P(X)P(\mathbf{X})进行模拟:
P(X)=P(Xz;θ)P(z)dz P(\mathbf{X}) = \int P(\mathbf{X}|z;\theta)P(z)dz
上式中的zz是latent variable,是我们通过训练所模拟的一个分布;θ\theta是随机变量XzX|z的概率分布的参数

这个问题的目标是建立一个模型,然后不断的优化θ\theta, 使得在这个模型下,已知数据集X\mathbf{X}发生的可能性最大,这就是Maximum Likelihood(最大似然)。如果我们训练的模型能够使得训练集中的数据的可能性最大,那么就有理由相信,它能生成和我们数据集的数据相似的数据,这里的相似指的是服从同一概率分布。

如果P(Xz)P(\mathbf{X}|z)是任意概率分布,那么对它进行训练显然是不现实的,所以在VAE中,假设P(Xz)P(\mathbf{X}|z)是一个正态分布N(f(z;θ),σ2I)\mathcal{N}(f(z;\theta),\sigma^2 I ),其中f(z;θ)f(z;\theta)是一个关于latent variable的函数,通常由1层或多层神经网络构成。

3. Variational AutoEncoders

3.1 Latent Variables in AutoEncoders

首先需要解决的问题是,我们如何选择Latent Variable需要模拟的分布使得它能够捕捉到数据的特性呢?这里我们假设我们的latent variable是服从标准正态分布N(0,σ2I)\mathcal{N}(0,\sigma^2 I )的,但这样的确失去了普遍性,凭什么你说是标准正态分布,就是标准正态分布了,这也太局限了吧?但其实这没有关系,因为任何一个分布d维的概率分布函数,都可以通过将d个1维标准正态分布传入一个足够复杂的函数来得到(https://informs-sim.org/wsc86papers/1986_0039.pdf) 也就是说,就算我们模拟的latent variable不是真正的能够捕捉到输入数据的特性的latent variable,我们依然可以通过后面添加神经网络来将其映射到真正的latent variable,因为神经网络是 universal function approximator, 意即任何函数都可以通过一个具有有限节点个数、至少一个hidden layer(隐藏层)的神经网络进行逼近。所以理论上,利用假设latent variable是标准正态分布变量配合神经网络,是可行的。

3.2 An Intuitive Way

解决完latent variable的问题之后,剩下的就是想办法设计一个Objective Function(目标函数)来最大化P(X)P(\mathbf{X})。如果我们能够做到通过一个可行的数学表达式的方法计算P(X)P(\mathbf{X}),那么就可以通过Stochastic Gradient Descent来进行优化。

那么一个比较简单的方法是通过计算平均值来估计期望。

我们需要注意到实际上这个式子:
P(X)=P(Xz;θ)P(z)dz P(\mathbf{X}) = \int P(\mathbf{X}|z;\theta)P(z)dz
其实是
EzP(z)(Xz) \mathbb{E}_{z \sim P(z)}(X|z)
那么通过对zz进行取样:{z1,z2...zn}\{z_1,z_2...z_n\},然后计算:
P(X)1niP(Xzi) P(X)\approx \frac{1}{n}\sum_i P(X|z_i)
但其实这个方法不太行,因为它需要非常多的数据点来进行逼近,这样计算的代价就太高了。具体为什么需要很多数据点,有兴趣请移步:arXiv:1606.05908 (2016)

所以我们需要另辟蹊径。

3.3 Setting Up the Objective Function

上面这个方法之所以不行,是因为需要取样太多的点了,在标准正态分布下取样到的点,大多数其实P(Xz)P(X|z)都是非常接近0的,导致取样到了对最终的P(X)P(\mathbf{X})也没有什么影响,可以说我们取样的效率是极低的。所以我们自然而然的就想到了,要在现有的概率分布下,得到一个条件分布Q(zX)Q(z|X),使得在已知输入的XX的条件下,得到一个能够进行高效率取样的分布。这样就能够更加有效地通过取样来计算Ez[P(Xz)]\mathbb{E}_{z}[P(X|z)]

那么怎么得到这个Q(zX)Q(z|X)呢?
我们希望能够让训练得到的Q(zX)Q(z|X)可以尽可能地逼近真正的P(zX)P(z|X)。在信息论中常用的一种方法是KL-Divergence,它能够度量两个概率分布之间的差距(注意,不是距离,KL-Divergence不具有距离的特性)。

KL-Divergence的定义式:
D[Q(zX)P(zX)]=EzQ[log(Q(zX))log(P(zX)] \mathcal{D}[Q(z|X)||P(z|X)] = \mathbb{E}_{z\sim Q}[log(Q(z|X))-log(P(z|X)]
但是这里面并没有我们关心的变量,我们关心的是关于XX的概率,但现在它都被作为条件,那么需要通过贝叶斯定理把它换到前面来:
D[Q(zX)P(zX)]=EzQ[log(Q(zX))]EzQ[log(P(zX)]=EzQ[log(Q(zX))]EzQ[log(P(Xz)P(z)P(X))]=EzQ[log(Q(zX))]EzQ[log(P(Xz))+log(P(z))]+log(P(X)) \begin{aligned} \mathcal{D}[Q(z|X)||P(z|X)] &= \mathbb{E}_{z\sim Q}[log(Q(z|X))]-\mathbb{E}_{z\sim Q}[log(P(z|X)]\\ &=\mathbb{E}_{z\sim Q}[log(Q(z|X))]-\mathbb{E}_{z\sim Q}[log(\frac{P(X|z)P(z)}{P(X)})]\\ &=\mathbb{E}_{z\sim Q}[log(Q(z|X))]-\mathbb{E}_{z\sim Q}[log(P(X|z))+log(P(z))]+log(P(X))\\ \end{aligned}
到这里为止,我们的式子里非常顺利的出现了我们关心的变量:P(X),  P(Xz)P(X), \; P(X|z)
回忆一下,我们的最终目标是maxP(X)\max P(X),而中间需要minD[Q(zX)P(zX)]\min \mathcal{D}[Q(z|X)||P(z|X)],那么把二者结合起来:
log(P(X))D[Q(zX)P(zX)]=log(P(X))(EzQ[log(Q(zX))]EzQ[log(P(Xz))+log(P(z))]+log(P(X)))=EzQ[log(P(Xz)]EzQ[log(Q(zX)log(P(z))]=EzQ[log(P(Xz)]D[Q(zX)P(z)] \begin{aligned} log(P(X))-\mathcal{D}[Q(z|X)||P(z|X)] &=log(P(X))-(\mathbb{E}_{z\sim Q}[log(Q(z|X))]-\mathbb{E}_{z\sim Q}[log(P(X|z))+log(P(z))]+log(P(X)))\\ &=\mathbb{E}_{z\sim Q}[log(P(X|z)]-\mathbb{E}_{z\sim Q}[log(Q(z|X)-log(P(z))]\\ &=\mathbb{E}_{z\sim Q}[log(P(X|z)]-\mathcal{D}[Q(z|X)||P(z)] \end{aligned}
很漂亮了。我们需要最大化左侧,我们根本不可能通过Stochastic Gradient Descent来优化左侧,因为我们根本不知道P(zX)P(z|X),但是通过数学推导我们能够得到右侧的式子:右侧的变量里面不是我们关心的,就是我们已经知道的。

这个伟大的式子将会贯穿在VAE中:
log(P(X))D[Q(zX)P(zX)]=EzQ[log(P(Xz)]D[Q(zX)P(z)]log(P(X))-\mathcal{D}[Q(z|X)||P(z|X)] = \mathbb{E}_{z\sim Q}[log(P(X|z)]-\mathcal{D}[Q(z|X)||P(z)]

3.4 Optimizing the Objective Function

首先如果Q(zX)Q(z|X)也不能是一个任意的分布,我们总要通过参数规定它是什么类型的分布,假设Q(zX)Q(z|X)是高斯分布:N(μ(X),Σ(X))\mathcal{N(}\mu(X),\Sigma(X)),我们通过神经网络,其实就是在模拟均值和方差。Variational AutoEncoders
所以,我们得到了:

D[Q(zX)P(z)]=D[N(μ,Σ)N(0,I)] \mathcal{D}[Q(z|X)||P(z)] = \mathcal{D}[\mathcal{N}(\mu,\Sigma)||\mathcal{N}(0,I)]

3.4.1 KL-Divergence of Two Multi-Variate Gaussians

那么第一个问题来了,如何求多元高斯分布的KL-Divergence呢?
D[N(μ,Σ)N(0,I)]=E[log(det(Σ)0.5(2π)kexp{12(xμ)TΣ1(xμ)})log((2π)kexp{12xTx}]=12EzQ[log(det(Σ))tr((xμ)TΣ1(xμ))+(xμ+μ)T(xμ+μ)]=12(EzQ[log(det(Σ))]tr(EzQ[(xμ)(xμ)T]Σ1)+EzQ[(xμ+μ)T(xμ+μ)])=12(log(det(Σ))k+tr(Σ)+μTμ) \begin{aligned} \mathcal{D}[\mathcal{N}(\mu,\Sigma)||\mathcal{N}(0,I)] &= \mathbb{E}[log(det(\Sigma)^{-0.5}(\sqrt{2\pi})^{-k}exp\{-\frac{1}{2}(x-\mu)^T\Sigma ^{-1}(x-\mu ){}\})-log((\sqrt{2\pi})^{-k}exp\{-\frac{1}{2}x^Tx\}]\\ &=\frac{1}{2}\mathbb{E}_{z\sim Q}[-log(det(\Sigma))-tr((x-\mu)^T\Sigma^{-1}(x-\mu))+(x-\mu +\mu)^T(x-\mu+\mu)]\\ &=\frac{1}{2}(\mathbb{E}_{z\sim Q}[-log(det(\Sigma))]-tr(\mathbb{E}_{z\sim Q}[(x-\mu)(x-\mu)^T]\Sigma^{-1})+\mathbb{E}_{z\sim Q}[(x-\mu +\mu)^T(x-\mu+\mu)])\\ &=\frac{1}{2}(-log(det(\Sigma))-k+tr(\Sigma)+\mu^T\mu) \end{aligned}

3.4.2 Gradient

计算出了这两个多元高斯分布的KL-Divergence,要想使用stochastic gradient descent 对这个目标函数进行优化,还有第二个问题要解决,即如何计算EzQ[log(P(Xz)]\mathbb{E}_{z\sim Q}[log(P(X|z)]
我们采用类似于stochastic gradient descent 的思想,我们随机取样一个zz,计算出log(P(Xz)log(P(X|z)的值,用它来作为这个期望值,然后采用stochastic gradient descent来进行优化,我们可以取任意多的XX,zz的值,然后计算出他们的gradient(梯度)最后取平均值,这样就能够真正计算出这个函数的gradient。

3.4.3 BackPropagation and Reparameterization Trick

但是要想通过神经网络进行优化,一个不可避免的问题就是,如何通过Backpropagation(反向传播)来传送gradient。我们在计算EzQ[log(P(Xz)]\mathbb{E}_{z\sim Q}[log(P(X|z)]的时候偷了个懒,通过取样来计算gradient,这就导致一个非常致命的问题——gradient没办法通过一个随机取样的操作传到更前面的参数去,因为没有任何依赖关系,我们没办法对随机取样这个操作进行求导。这里就有一个非常酷的想法:Reparameterization Trick

我们知道,任何一个正态分布,都可以写成标准正态分布的一个函数,所以我们对Q(zX)Q(z|X)进行改写:
Q(zX)=μ+Σ12N(0,I) Q(z|X) = \mu+\Sigma^{\frac{1}{2}} \mathcal{N}(0,I)
我们需要的是关于μ,Σ\mu,\Sigma的gradient,我们根本不关心对于标准正态分布的gradient,但是就是因为我们通过随机取样才断了gradient这条路,那为什么不把随机取样的这部分给单独分离出去,使得最后的取样等于一个固定的量和一个随机的量进行代数运算呢?这个固定的量是关于μ,Σ\mu ,\Sigma的函数,而每次取样的随机性就来源于这个随机的量,这样既可以通过随机取样估计期望,又可以让gradient安全地被传过来。我这么说很抽象,那么结合图像将会更加直观:
Variational AutoEncoders
Marvelous!!!
很高兴的通知你,到目前为止,我们已经完全可以通过神经网络来训练这个模型了!

4. Interpretation of the Objective Function

4.1 Information Theory Perspective

首先,回忆一下我们的objective function:

log(P(X))D[Q(zX)P(zX)]=EzQ[log(P(Xz)]D[Q(zX)P(z)]log(P(X))-\mathcal{D}[Q(z|X)||P(z|X)] = \mathbb{E}_{z\sim Q}[log(P(X|z)]-\mathcal{D}[Q(z|X)||P(z)]
进行一些变换:
log(P(X))=(EzQ[log(P(Xz)])+(D[Q(zX)P(z)])D[Q(zX)P(zX)]-log(P(X)) = (-\mathbb{E}_{z\sim Q}[log(P(X|z)])+(\mathcal{D}[Q(z|X)||P(z)])-\mathcal{D}[Q(z|X)||P(z|X)]
在信息论中,熵(Entropy) = E[log(p(x))]-\mathbb{E}[log(p(x))]可以被看作是包含信息的总数,如果一个随机变量是必然事件,即概率为1,那么它就不包含不确定性,即不包含信息,如果一个随机变量的概率是0,那么它就有最大的不确定性,所以包含最多信息。log(P(X))-log(P(X))可以被看作是一个数据点(往往是图片)需要被重构所需要的信息。而右侧的式子将重建XX看作了三个步骤:

  1. XX进行编码(encode),用来构建latent variable zz。在信息论中,KL-Divergence可以被看作是二者包含信息的总量的差距。我们可以认为,在P(z)P(z)这个分布下随机取出一个zz,它包含XX的信息的概率为0. 具体参照3.2。所以,用Q(zX)Q(z|X)P(z)P(z)之间的KL-Divergence来度量需要多少信息才能通过给出的XX构建zz
  2. 根据zz来重建XXEzQ[log(P(Xz)]\mathbb{E}_{z\sim Q}[log(P(X|z)] 可以被看作是根据已有的zz来重建XX所需要的信息。
  3. 因为我们无法直接模拟P(z)P(z)来计算期望值,只能间接地通过模拟P(zX)P(z|X)来计算期望值,这样其实是使用更多的信息来模拟P(zX)P(z|X),所以还需减掉这样多余的信息D[Q(zX)P(zX)]\mathcal{D}[Q(z|X)||P(z|X)]

4.2 Regularization Perspective

看到objective function等式左边的那两项的形式,不可避免的想要问自己这样一个问题:是否可以将log(P(X))-log(P(X))看作是一个objective function,然后D[Q(zX)P(zX)]\mathcal{D}[Q(z|X)||P(z|X)]看作是一个regularization term(正则项)?那么问题来了,既然是正则项,必然有hyperparameter(超参数),那么这里有没有呢?
第一个能想到的就是,我们是否可以修改P(z)P(z)的方差,让默认的变成一个具有不是1的方差的正态分布?但其实这不能改变这个模型,因为如3.1所说,即使不是标准正态分布,我们也可以通过后续的hidden layer来将标准正态分布映射到真正的latent variable的分布中
第二个能想到的是,P(Xz)P(\mathbf{X}|z)是一个正态分布N(f(z;θ),σ2I)\mathcal{N}(f(z;\theta),\sigma^2 I ),我可不可以通过修改这个σ\sigma来修改这个模型?如果我们仅仅局限于认为这是个高斯分布的话,的确,这个σ\sigma代表了我们期望最后得到的模型的准确率,因为对于一个正态分布,方差越大,均值点的概率越高,就代表了我们对于取最大概率点均值点的XX是真正服从原始数据所服从的概率分布的信心越强。但是如果P(Xz)P(\mathbf{X}|z)不是高斯分布了呢?那么就根本不存在这个“hyperparameter”了呢。

5. Reference

[1]. Tutorial on variational autoencoders. arXiv 2016
[2]. MIT 6.191 Lecture 4
[3]. UC Berkeley EECS 126 Lecture 10
[4]. UC Berkeley CS 189 Note 14
[5]. Luc Devroye. Sample-based non-uniform random variate generation. Springer-Verlag, New York, 1986.