beta-VAE

beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework

$\beta -VAE$的目标是学习独立的特征，让某种特征对应某个生成因素，而独立于其他因素。For example,a model trained on photos of human faces might capture the gentle, skincolor, hair color, hair length, emotion, whether wearing a pair of glasses andmany other relatively independent factors in separate dimensions. Such a disentangled representation is very beneficial to facial image generation.
数据集 D = { X , V , W } \mathcal{D}=\{X, V, W\} D={X,V,W}，其中$\mathbf{x}\in {\mathbb{R}}^N$代表图片，而$\mathbf{v}\in {\mathbb{R}}^K$为条件独立的隐变量，即$\log p(\mathbf{v}|\mathbf{x})={\sum }_k\log p\left(v_k|\mathbf{x}\right)$，还有条件依赖的隐变量$\mathbf{w}\in {\mathbb{R}}^H$。那么图片由着两个隐变量共同生成，则$p(\mathbf{x}|\mathbf{v},\mathbf{w})=\mathrm{Sim}(\mathbf{v},\mathbf{w})$。
我们现在希望通过一种无监督的方法，仅仅利用数据$X$就能得到$\mathbf{x}$和$\mathbf{z}$的联合分布，其中$\mathbf{z}\in {\mathbb{R}}^M$，$M\ge K$。也就是说有$p(\mathbf{x}|\mathbf{z})\approx p(\mathbf{x}|\mathbf{v},\mathbf{w})=\mathrm{Sim}(\mathbf{v},\mathbf{w})$。那么目标函数为$$\underset{\theta }{\max }{\mathbb{E}}_{p_{\theta }(\mathbf{z})}\left[p_{\theta }(\mathbf{x}|\mathbf{z})\right]$$与VAE一样，引入一个inferred后验分布$q_{\varphi }(\mathbf{z}|\mathbf{x})$，我们的目标是$q_{\varphi }(\mathbf{z}|\mathbf{x})$能够以解耦的方式将$\mathbf{v}$解耦出。为了能够达到目的，我们可以让这个后验接近$p(\mathbf{z})$，这样既能限制隐变量的信息瓶颈，又能达到之前说的解耦的效果，且$p(\mathbf{z})=\mathcal{N}(0,I)$。则目标函数为$$\underset{\varphi ,\theta }{\max }{\mathbb{E}}_{x\sim \mathbf{D}}\left[{\mathbb{E}}_{q_{\varphi }(\mathbf{z}|\mathbf{x})}\left[\log p_{\theta }(\mathbf{x}|\mathbf{z})\right]\right]\text{ subject to }{D}_{KL}\left(q_{\varphi }(\mathbf{z}|\mathbf{x})\Vert p(\mathbf{z})\right)<ϵ$$求解上式，利用拉格朗日KKT条件$$\mathcal{F}(\theta ,\varphi ,\beta ;\mathbf{x},\mathbf{z})={\mathbb{E}}_{q_{\varphi }(\mathbf{z}|\mathbf{x})}\left[\log p_{\theta }(\mathbf{x}|\mathbf{z})\right]-\beta \left(D_{KL}\left(q_{\varphi }(\mathbf{z}|\mathbf{x})\Vert p(\mathbf{z})\right)-ϵ\right)$$又因为$\beta ,ϵ\ge 0$，则$$\mathcal{F}(\theta ,\varphi ,\beta ;\mathbf{x},\mathbf{z})\ge \mathcal{L}(\theta ,\varphi ;\mathbf{x},\mathbf{z},\beta )={\mathbb{E}}_{q_{\varphi }(\mathbf{z}|\mathbf{x})}\left[\log p_{\theta }(\mathbf{x}|\mathbf{z})\right]-\beta D_{KL}\left(q_{\varphi }(\mathbf{z}|\mathbf{x})\Vert p(\mathbf{z})\right)$$可以发现与vanilla-VAE的唯一区别就是在KL项引入了。Therefore a higher $\beta$ encourages more efficient latent encoding and further encouragesthe disentanglement. Meanwhile, a higher $\beta$ may create a trade-off between reconstruction quality and the extent of disentanglement.

Understanding disentangling in $\beta$-VAE

从信息瓶颈的角度看$\beta -VAE$。首先先科普下什么是information bottlenec。假设我们面对分类任务，标注数据对是$\left(x_1,y_1\right),\dots ,\left(x_N,y_N\right)$。我们把这个任务分为两步来理解，第一步是编码，第二步就是分类。第一步是把$x$编码为一个隐变量$z$，然后分类器把$z$识别为类别$y$。
我们试想在$z$处加一个“瓶颈”$\beta$，它像一个沙漏，进入的信息量可能有很多，但是出口就只有$\beta$那么大，所以这个瓶颈的作用是：不允许流过$z$的信息量多于$\beta$。跟沙漏不同的是，沙漏的沙过了瓶颈就完事了，而信息过了信息瓶颈后，还需要完成它要完成的任务（分类、回归等），所以模型迫不得已，只好想办法让最重要的信息通过瓶颈。这就是信息瓶颈的原理！
那么衡量信息量大小的式子刚好有互信息$$I(Z,Y;\mathbf{\theta })=\int dzdyp(z,y|\mathbf{\theta })\log \frac{p(z,y|\mathbf{\theta })}{p(z|\mathbf{\theta })p(y|\mathbf{\theta })}$$显然上式最大的时候刚好为$Z=Y$时，这样我们的学习就没有意义了。因此可以考虑对上式加上一定的约束$$\underset{\mathbf{\theta }}{\max }I(Z,Y;\mathbf{\theta })\text{ s.t. }I(X,Z;\mathbf{\theta })\le I_c$$让$X,Z$的互信息约束在一个范围中，这样就是相当于一个瓶颈，提取十分必要的信息，而略去一些没有意义的信息！求解上式同样适用拉格朗日乘子，最后得到$$R_{IB}(\mathbf{\theta })=I(Z,Y;\mathbf{\theta })-\beta I(Z,X;\mathbf{\theta })$$Intuitively, the first term in RIB encourages Z to be predictive of Y ; thesecond term encourages Z to“forget”X. Essentially it forces Z to act likea minimal sufficient statistic of X for predicting Y .
那么$\beta -VAE$就能用信息瓶颈的理论解释了。首先第一部分为一个重建问题$$\max E_{q(\mathbf{z}|\mathbf{x})}[\log p(\mathbf{x}|\mathbf{z})]$$使得上式最大时，就是在隐含着互信息$I(Z,Y;\mathbf{\theta })$变大，其中$Y=X$。只有$X,Z$的相关程度很高的时候，重建才能很好！更加极端的情况就是AE自动编码机，$Z=encoder(X)$为一个决定性关系，此时重构可以做到很高的精度，而且$X,Z$的互信息也很大！
而$KL$部分$D_{KL}\left(q_{\varphi }(\mathbf{z}|\mathbf{x})\Vert p(\mathbf{z})\right)$，在IB理论中是减少$X,Z$的互信息，当我们减少$KL$散度时，则有$q_{\varphi }(\mathbf{z}|\mathbf{x})\to p(\mathbf{z})$，那就是说明$X,Z$就是独立了，从而互信息也就为0了。那么$KL$部分越大，说明该维度的隐变量含有的信息量就越大。
$\beta -VAE$虽然能够学习到很好的特征，但是其重建能力真的很差，如下图所示。

本文发现逐渐增大信息瓶颈可以让重建能力上升，且能够学习到很棒的特征。因此提出了如下的目标函数：$$\mathcal{L}(\theta ,\varphi ;\mathbf{x},\mathbf{z},C)={\mathbb{E}}_{q_{\varphi }(\mathbf{z}|\mathbf{x})}\left[\log p_{\theta }(\mathbf{x}|\mathbf{z})\right]-\gamma \left|D_{KL}\left(q_{\varphi }(\mathbf{z}|\mathbf{x})\Vert p(\mathbf{z})\right)-C\right|$$其中$C$为一个逐渐增大的数字，一般$\gamma$取一个很大的数字，保证$KL$部分能满足$C$。