Preface

A Paper A day Keep the Doctor Away。 今天的文章是：

Understanding disentangling in β-VAE. (NIPS2017) [paper]

本文从Information BottleNeck的角度分析了beta-VAE的ELBO，第一项为图像重建的q(z∣x){q(z|x)}的BottleNeck，KL项为X和Z的互信息量。优化KL项，使得Z和X的互信息量减小，阻碍了图像重建。除此之外，作者认为beta-VAE会根据KL项限制（BottleNeck的容量）按对重建贡献大小的顺序学习不同的generative factors。由此，作者提出限制BottleNeck的容量从小到大的训练方法修改了beta-VAE的ELBO，使得模型可以按对重建贡献大小顺序逐渐学习不同的generative factors的Disentangled表示。

本文主要围绕的问题有：

  • 用Information BottleNeck的角度解释β{\beta}-VAE为什么可以学习到Disentanlged表示，和β{\beta}起到的tradeoff。
  • 为什么β{\beta}-VAE可以学习到对应真实factor的不同维度的表示?
  • 通过提升latent bottleneck的编码容量，改进β{\beta}-VAE。

Main Contents

Information BottleNeck 角度解释

β{\beta}-VAE的ELBO公式为：

L(θ,ϕ,z,x(i),β)=Eqϕ(z∣x(i))[log⁡pθ(x∣z)]−βDKL(qϕ(z∣x(i))∥p(z)){ \mathcal{L}(\theta,\phi,\mathbf{z},x^{(i)},\beta) = \mathbb{E}_{q_{\phi}(\mathbf{z} | x^{(i)})}\left[ \log p_{\theta}(x | \mathbf{z}) \right] - \beta \mathcal{D}_{\text{KL} }\left( q_{\phi}(\mathbf{z} | x^{(i)}) \| p(\mathbf{z}) \right) }

又x∼p(x){x \sim p(x)}：

L(θ,ϕ,z,x,β)=Ex∼p(x)[Eqϕ(z∣x)[log⁡pθ(x∣z)]]−βEx∼p(x)[DKL(qϕ(z∣x)∥p(z))]{ \mathcal{L}(\theta,\phi,\mathbf{z},\mathbf{x},\beta) = \mathbb{E}_{ \mathbf{x} \sim p(\mathbf{x})} \left[ \mathbb{E}_{q_{\phi}(\mathbf{z} | \mathbf{x})}\left[ \log p_{\theta}(\mathbf{x} | \mathbf{z}) \right] \right] - \beta \mathbb{E}_{ \mathbf{x} \sim p(\mathbf{x})} \left[ \mathcal{D}_{\text{KL} }\left( q_{\phi}(\mathbf{z} | \mathbf{x}) \| p(\mathbf{z}) \right) \right] }

β{\beta}-VAE的ELBO的KL项可以看成X和Z的互信息熵，所以在最大化ELBO同时，会减小KL项，导致X和Z的互信息量减少，互信息量的减少表示Z中可以得到的X的信息减少，decoder从Z中要再次重建X就会变得困难。而当I(X;Z)=0{I(X;Z) = 0}时，Z和X相互独立，从Z从无法得到关于X的信息。增大β{\beta}导致模型整体偏向KL项的优化，互信息量的减少程度更大。一定程度可以解释β{\beta}过大时，β{\beta}-VAE的Disentanglement效果下降的现象。

从另一个角度看，优化过程中减小KL项通过两个方面：一是使得不同样本得到的推断后验q(z∣x(i)){q(z | x^{(i)})}的均值尽量集中；二是增大各推断后验的方差。(为什么？)这个过程使得q(z∣x(i)){q(z | x^{(i)})}的重叠程度（高概率的区域重叠程度）增大，q(z∣x(i)){q(z | x^{(i)})}趋向相同，这导致不同样本在隐空间的表示也变得更加难以区分，进一步导致重建效果的下降。

这个角度解释β{\beta}-VAE和文章[1]中的解释非常类似，[1]也使用了overlap的程度解释β{\beta}-VAE的Disentanglement效果随β{\beta}的增大先上升后降低的现象。

Reasoning Axes-Alignment

作者认为β{\beta}-VAE会寻找对ELBO优化提升的贡献不同的成分。一些成分在重建中起的作用更大，以dSprites数据集为例，如图1所示：

图1 dSprites上不同factor的学习
图1 dSprites上不同factor的学习

从下到上分别为模型学习到position、scale、shape、rotation时生成的图像。这些成分对重建的作用从大到小。（图片来源于原文Figure3）

当β{\beta}很大时，受KL项的信息瓶颈容量的限制，最佳的优化方法是优先编码图像中那些可以带来最显著的重建效果提升的信息。提升信息瓶颈的容量，可以使得模型逐渐学习到更加次要的重建成分。

如此我们可以通过控制Information BottleNeck的方法来逐渐学习到对ELBO优化贡献从大到小的一系列factor，但是如何保证学习新的factor时，不会干扰原有的编码？（模型自动选择新的隐变量的维度表示新的factor）。作者认为β{\beta}-VAE本身的两个性质可以保证这一点：

1. 这是满足data locality pressure（不同的样本在隐空间激活的区域只是一个小的局部）最自然的做法？ （我表示难以理解这一点。）

2. 先验p(z){p(z)}的方差为对角阵，限制模型在新的维度中编码时不会对其他维度干扰。

Improvement

最后作者根据以上思考提出了逐渐提高Information BottleNeck容量的方法，使得模型逐渐学习最优、次优的编码信息，修改的ELBO如下：

L(θ,ϕ,z,x(i),C)=Eqϕ(z∣x(i))[log⁡pθ(x∣z)]−γ∣DKL(qϕ(z∣x(i))∥p(z))−C∣{ \mathcal{L}(\theta,\phi,\mathbf{z},x^{(i)},C) = \mathbb{E}_{q_{\phi}(\mathbf{z} | x^{(i)})}\left[ \log p_{\theta}(x | \mathbf{z}) \right] - \gamma | \mathcal{D}_{\text{KL} }\left( q_{\phi}(\mathbf{z} | x^{(i)}) \| p(\mathbf{z}) \right) - C | }

训练中C{C}逐渐从0增加到一个合理的比较大的值。

在dSprites实验中，作者也观察到了随着C{C}的增大，模型逐渐学习到一些新的成分，如图1所示，模型在C=3.1{C = 3.1}时学习到position因子，在C=7.3{C=7.3}时学习到scale因子，等等。验证了该方法的有效性。

Reference

  1. Disentangling Disentanglement in Variational AutoEncoder.