Preface

A Paper A Day Keep the Doctor Away. 今天的文章是：

Learning Disentangled Joint Continuous and Discrete Representations. (NIPS2018) [paper] [代码]

本文整体基于AnnealVAE，将隐空间中连续变量和离散变量的空间分开建模，并分开限制信道容量。在离散变量的分布的重参数化上，作者使用了[1]的"Gumbel Softmax"方法，在MNIST、FashionMNIST数据集上学习到了一定的离散分量。

Main Contents

JointVAE

JointVAE的损失基于AnnealVAE，只是将原来的capacity容量的限制分别拆分成离散变量的和连续变量的两个限制之和。

首先使用c{\mathbf{c}}和z{\mathbf{z}}分别表示隐空间中的离散和连续的隐分量，即原VAE的ELBO的隐变量z=(z,c){\mathbf{z} = (\mathbf{z}, \mathbf{c})}，β{\beta}-VAE的ELBO变为：

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

假定连续变量和离散变量条件独立，上式的KL项可以继续拆分成：

DKL(qϕ(z,c∣x)∥p(z,c))=DKL(qϕ(z∣x)∥p(z))+DKL(qϕ(c∣x)∥p(c)){ D_{K L}\left(q_{\phi}(\mathbf{z}, \mathbf{c} | \mathbf{x}) \| p(\mathbf{z}, \mathbf{c})\right)=D_{K L}\left(q_{\phi}(\mathbf{z} | \mathbf{x}) \| p(\mathbf{z})\right)+D_{K L}\left(q_{\phi}(\mathbf{c} | \mathbf{x}) \| p(\mathbf{c})\right) }

整体β{\beta}-VAE的ELBO又变为：

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

但是作者在实验中发现直接优化上式容易导致模型忽略离散隐变量。而如果借鉴AnnealVAE的方法（如下式），又会导致模型偏向将模型容量分配给连续变量。

LAnneal(θ,ϕ)=Eqϕ(z,c∣x)[log⁡pθ(x∣z,c)]−γ∣DKL(qϕ(z∣x)∥p(z))+DKL(qϕ(c∣x)∥p(c))−C∣{ \mathcal{L}_{Anneal}(\theta, \phi)=\mathbb{E}_{q_{\phi}(\mathbf{z}, \mathbf{c} | \mathbf{x})}\left[\log p_{\theta}(\mathbf{x} | \mathbf{z}, \mathbf{c})\right]-\gamma | D_{K L}\left(q_{\phi}(\mathbf{z} | \mathbf{x}) \| p(\mathbf{z})\right) + D_{K L}\left(q_{\phi}(\mathbf{c} | \mathbf{x}) \| p(\mathbf{c})\right) - C | }

最后作者分别对离散隐变量的容量和连续隐变量的容量分别限制，得到：

LAnneal(θ,ϕ)=Eqϕ(z,c∣x)[log⁡pθ(x∣z,c)]−γ∣DKL(qϕ(z∣x)∥p(z))−Cz∣−γ∣DKL(qϕ(c∣x)∥p(c))−Cc∣{ \mathcal{L}_{Anneal}(\theta, \phi)=\mathbb{E}_{q_{\phi}(\mathbf{z}, \mathbf{c} | \mathbf{x})}\left[\log p_{\theta}(\mathbf{x} | \mathbf{z}, \mathbf{c})\right]-\gamma | D_{K L}\left(q_{\phi}(\mathbf{z} | \mathbf{x}) \| p(\mathbf{z})\right)- C_{z} | - \gamma | D_{K L}\left(q_{\phi}(\mathbf{c} | \mathbf{x}) \| p(\mathbf{c})\right) - C_{c} | }

Reparameterization

对连续变量的重参数化和原始VAE的策略一样，而对离散变量的重参数化作者使用了[1]中的Gumbel Softmax。

下面着重描述下我理解的Gumbel Softmax重参数化过程。

对重参数化要求有两点：(1) 采样或近似采样原分布。 (2) 采样的结果对参数可以求导。

一般采用Multinoulli分布对离散随机变量建模： X{\mathbf{X}}服从参数为[α1,⋯ ,αn]{\left[\alpha_{1}, \cdots, \alpha_{n}\right]}的Multinoulli分布，即X=[X1,⋯ ,Xn]∼p(α1,⋯ ,αn){\mathbf{X} = \left[ X_{1}, \cdots, X_{n} \right] \sim p(\alpha_{1}, \cdots, \alpha_{n})}，其中[X1,⋯ ,Xn]{\left[ X_{1}, \cdots, X_{n} \right]}为一个one hot向量，仅有一个维度为1。

第一种方法：从该分布中采样，最简单的方法就是利用它的累积分布函数(CDF)F(x)=P(X≤x){F(x) = P(X \leq x)}和它的反函数。首先取[0,1]{[0,1]}上的均匀分布的随机数p1{p_1}，采样就可以是x∗=max⁡x{x∣F(x)≤p1}{x^{\ast} = \max_{x} \left\{ x | F(x) \leq p_1 \right\}}，最后对x∗{x^{\ast}}取one hot。

第二种方法是[1]中提到的Gumbel Max Trick的方法近似采样：

令π=[α1,⋯ ,αn]{\pi = \left[\alpha_{1}, \cdots, \alpha_{n}\right]}，πi=αi{\pi_i = \alpha_i}，要对X{\mathbf{X}}近似采样，可以利用以下公式：

x∗=one_hot(arg max⁡i(log⁡πi+gi)){ x^{\ast} = \text{one\_hot}\left( \argmax_{i} (\log \pi_i + g_i) \right) }

其中gi(i=1,⋯ ,n){g_i (i = 1, \cdots, n)}为从参数为(0,1){(0,1)}的Gumbel分布中采样的独立同分布的样本。

然而arg max⁡{\argmax}无法求导数，这引出了第三种方法——Gumbel Softmax。

第三种方法是[1]中的Gumbel Softmax，使用Softmax替代上面公式的one_hot(arg max⁡(⋅)){\text{one\_hot}\left( \argmax(\cdot) \right)}，采样公式变为：

x∗=exp⁡((log⁡π+g)/τ)∑inexp⁡((log⁡πi+gi)/τ){ x^{\ast} = \frac{\exp\left( (\log \pi + g) / \tau \right)}{\sum_{i}^{n} \exp \left( (\log \pi_i + g_i) /\tau \right)} }

当然，原文使用的是分量的形式：

xk=exp⁡((log⁡πk+gk)/τ)∑inexp⁡((log⁡πi+gi)/τ)fork=1,⋯ ,n{ x_k = \frac{\exp\left( (\log \pi_k + g_k) / \tau \right)}{\sum_{i}^{n} \exp \left( (\log \pi_i + g_i) /\tau \right)} \qquad for \quad k = 1, \cdots, n }

其中τ{\tau}为annealing temperature。[1]将这样得到的x∗{x^{\ast}}的分布称为Gumbel Softmax Distribution。当τ→0{\tau \to 0}时，x∗{x^{\ast}}和one hot的结果越来越接近（仅有一个维度分量非常大）；而当τ→∞{\tau \to \infty}时，x∗{x^{\ast}}各分量的值变得越来越平均。

可以发现，使用Gumbel Softmax采样的样本是可以对参数π{\pi}求导的，所以它可以用于Multinoulli分布的近似采样。

以上是对单个满足Mulitnoulli分布的离散随机变量的重参数化过程，下面进一步把它扩展到建模的所有的离散分量。

记ci{c_i}为单个满足Mulitnoulli分布的离散随机变量X{\mathbf{X}}。推断后验分布建模成多个由Gumbel Softmax近似的离散分布p(ci∣x){p(c_i | \boldsymbol{x})}的乘积：

qϕ(c∣x)=∏ip(ci∣x){ q_{\phi}(\boldsymbol{c} | \boldsymbol{x}) = \prod_i p(c_i | \boldsymbol{x}) }

而先验p(c){p(\boldsymbol{c})}作者将其设置为均匀Multinoulli分布的乘积：

p(c)=∏ip(ci;α0){ p(\boldsymbol{c}) = \prod_i p(c_i ; \alpha_0) }

其中p(ci;α0){p(c_i; \alpha_0)}是以α0{\alpha_0}为参数的Multinoulli分布，α0=[1n,⋯ ,1n]{\alpha_0 = \left[\frac{1}{n}, \cdots, \frac{1}{n} \right]}，n{n}为ci{c_i}的维数。

Choices about hyperparameter

超参数的选取原则上，作者的观点是：

  • γ{\gamma} 应取的足够的大，保证使得重建效果不会下降太多的容量最大化。
  • Cc{C_c}应该选择最大的离散分量的信道容量，也就是可能的最大的互信息量。
  • Cz{C_z}的选择作者认为比较困难，应选择能保持disentangled的最大容量（Cz{C_z}是最需要调整的参数。）

Conclusions

作者在实验中得到的值得注意的结论有：

  • 在MNIST上，JointVAE成功学习到了“数字的类别”这一离散分量，而且仅通过观察这一隐分量的值，在测试阶段取得了88%的分类正确率，这算是一个比较好的下游应用的展示，而且是无监督的方法。
  • 如果数据集中没有特别突出的离散表示，比如说数字的类型，模型可能无法学习到有意义的离散表示，而且模型不一定能够学习到所有的类别，如FashionMNIST上，模型没有区别shirt和T-shirt。这表明离散表示的建模还有待提升。

References

  1. Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:1611.01144, 2016.