Preface

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

Hyperprior Induced Unsupervised Disentanglement of Latent Representations. [paper] [codes] (AAAI2019)

本文通过引入principled hierarchical bayes prior，修改VAE种隐编码的先验p(z)=N(z;0,Σ){p(z) = \mathcal{N}(z;0,\Sigma)}，将先验的方差作为随机变量，让VAE学习方差的分布。该方法的引入使得我们可以通过控制Σ{\Sigma}的期望分布形式，学习隐编码的变量之间的关系的模型，而不仅仅是要求相互独立。

Main Contents

ELBO Decomposition

作者将先验p(z){p(z)}的形式修改，令p(z){p(z)}的协方差参数也为一个随机变量Σ{\mathbf{\Sigma}}。整个生成过程就变成：

p(x,z,Σ)=p(x∣z)p(z∣Σ)p(Σ){ p(\mathbf{x},\mathbf{z},\boldsymbol{\Sigma}) = p(\mathbf{x} | \mathbf{z}) p(\mathbf{z} | \mathbf{\Sigma}) p(\mathbf{\Sigma}) }

并且使用q(z,Σ∣x){q(\mathbf{z},\mathbf{\Sigma} | \mathbf{x})}近似推断p(z,Σ∣x){p(\mathbf{z},\mathbf{\Sigma} | \mathbf{x})}，则推导ELBO如下：

Epd(x)[DKL(q(x,z∣Σ)∥p(x,z∣Σ))]=Epd(x)[Eqϕ(z,Σ∣x)[log⁡qϕ(z,Σ∣x)−log⁡pθ(z,Σ∣x)]]=Epd(x)[Eqϕ(z,Σ∣x)[log⁡qϕ(z,Σ∣x)−log⁡pθ(x∣Σ,z)+log⁡p(z,Σ)−log⁡pθ(x)]]>=0{ \begin{aligned} \mathbb{E}_{p_d(\mathbf{x})} \left[ \mathcal{D}_{\text{KL} }(q(\mathbf{x}, \mathbf{z} | \mathbf{\Sigma}) \| p(\mathbf{x}, \mathbf{z} | \mathbf{\Sigma}) ) \right] &= \mathbb{E}_{p_d(\mathbf{x})} \left[\mathbb{E}_{q_{\phi}(\mathbf{z}, \mathbf{\Sigma} | \mathbf{x})} \left[ \log q_{\phi}(\mathbf{z}, \mathbf{\Sigma} | \mathbf{x}) - \log p_{\theta}(\mathbf{z}, \mathbf{\Sigma} | \mathbf{x}) \right] \right] \\ &= \mathbb{E}_{p_d(\mathbf{x})} \left[\mathbb{E}_{q_{\phi}(\mathbf{z}, \mathbf{\Sigma} | \mathbf{x})} \left[ \log q_{\phi}(\mathbf{z}, \mathbf{\Sigma} | \mathbf{x}) - \log p_{\theta}(\mathbf{x} | \mathbf{\Sigma} , \mathbf{z})+ \log p(\mathbf{z}, \mathbf{\Sigma}) - \log p_{\theta}(x) \right] \right] \\ & >= 0 \end{aligned} }

得到ELBO的表达式：

Epd(x)[log⁡pθ(x)]≥Epd(x)[Eqϕ(z,Σ∣x)[log⁡pθ(x∣z,Σ)]−DKL(q(z,Σ∣x)∥p(z,Σ))]=ELBO{ \begin{aligned} \mathbb{E}_{p_d(\mathbf{x})} \left[\log p_{\theta}(x) \right] & \geq \mathbb{E}_{p_d(\mathbf{x})} \left[ \mathbb{E}_{q_{\phi}(\mathbf{z}, \mathbf{\Sigma} | \mathbf{x})}\left[ \log p_{\theta}(\mathbf{x} | \mathbf{z}, \mathbf{\Sigma}) \right] - \mathcal{D}_{\text{KL} } (q(\mathbf{z}, \mathbf{\Sigma} | \mathbf{x}) \| p(\mathbf{z}, \mathbf{\Sigma})) \right] \\ & = \text{ELBO} \end{aligned} }

以上其实就是标准ELBO的分解过程，只是将原来的一个隐变量z{\mathbf{z}}变成(z,Σ){(\mathbf{z}, \mathbf{\Sigma})}。

但是我们针对KL项，继续分解，并且作者定义后验近似推断的过程为：q(z,Σ∣x)=q(z∣x)q(Σ∣z){q(\mathbf{z}, \mathbf{\Sigma} | \mathbf{x}) = q(\mathbf{z} | \mathbf{x}) q(\mathbf{\Sigma} | \mathbf{z})}，则有以下的分解：

−Eqd(x)[DKL(q(z,Σ∣x)∥p(z,Σ))]=−Eqd(x)Eqϕ(z,Σ∣x)[log⁡qϕ(z,Σ∣x)−log⁡p(z,Σ)+log⁡q(z,Σ)−log⁡q(z,Σ)]=−Eqϕ(z,Σ,x)[log⁡qϕ(z,Σ∣x)q(z,Σ)+log⁡qϕ(z,Σ)p(z,Σ)]=−Eqϕ(z,x,Σ)[log⁡qϕ(z∣x)qϕ(Σ∣z)qϕ(z)qϕ(Σ∣z)]−Eqϕ(Σ∣z)qϕ(z)[log⁡qϕ(Σ∣z)qϕ(z)p(Σ∣z)p(z)]=Eqϕ(z,x)[log⁡qϕ(z∣x)qϕ(z)]−Eqϕ(Σ∣z)qϕ(z)[log⁡qϕ(Σ∣z)p(Σ∣z)+log⁡qϕ(z)p(z)]=−I(x;z)⏟(index-code MI)−Eqϕ(z)DKL(qϕ(Σ∣z)∥p(Σ∣z))⏟(marginal KL to prior)−DKL(qϕ(z)∥p(z))⏟(covariance penalty){ \begin{aligned} - \mathbb{E}_{q_d(x)} \left[ \mathcal{D}_{\text{KL} } (q(\mathbf{z}, \mathbf{\Sigma} | \mathbf{x}) \| p(\mathbf{z}, \mathbf{\Sigma})) \right] &= - \mathbb{E}_{q_d(x)} \mathbb{E}_{q_{\phi}(\mathbf{z}, \mathbf{\Sigma}| \mathbf{x})}\left[ \log q_{\phi}(\mathbf{z}, \mathbf{\Sigma} | \mathbf{x}) - \log p(\mathbf{z}, \mathbf{\Sigma}) + \log q(\mathbf{z}, \mathbf{\Sigma}) - \log q(\mathbf{z}, \mathbf{\Sigma}) \right] \\ &=- \mathbb{E}_{q_{\phi}(\mathbf{z}, \mathbf{\Sigma},\mathbf{x})}\left[ \log \frac{q_{\phi}(\mathbf{z}, \mathbf{\Sigma} | \mathbf{x})}{q(\mathbf{z}, \mathbf{\Sigma}) } + \log \frac{q_{\phi}(\mathbf{z}, \mathbf{\Sigma})}{p(\mathbf{z}, \mathbf{\Sigma})} \right] \\ &= - \mathbb{E}_{q_{\phi}(\mathbf{z}, \mathbf{x}, \mathbf{\Sigma})} \left[ \log \frac{q_{\phi}(\mathbf{z} | \mathbf{x}) q_{\phi}(\mathbf{\Sigma}|\mathbf{z})}{q_{\phi}(\mathbf{z})q_{\phi}(\mathbf{\Sigma} | \mathbf{z})} \right] - \mathbb{E}_{q_{\phi}(\mathbf{\Sigma} | \mathbf{z}) q_{\phi}(\mathbf{z})} \left[ \log \frac{q_{\phi}(\mathbf{\Sigma} | \mathbf{z}) q_{\phi}(\mathbf{z})}{p(\mathbf{\Sigma} | \mathbf{z}) p(\mathbf{z})} \right] \\ &= \mathbb{E}_{q_{\phi}(\mathbf{z}, \mathbf{x})} \left[ \log \frac{q_{\phi}(\mathbf{z} | \mathbf{x}) }{q_{\phi}(\mathbf{z})} \right] - \mathbb{E}_{q_{\phi}(\mathbf{\Sigma} | \mathbf{z}) q_{\phi}(\mathbf{z})} \left[ \log \frac{q_{\phi}(\mathbf{\Sigma} | \mathbf{z})}{p(\mathbf{\Sigma} | \mathbf{z})} + \log \frac{q_{\phi}(\mathbf{z})}{p(\mathbf{z})} \right] \\ &= - \underbrace{\mathbf{I}(\mathbf{x};\mathbf{z})}_{\text{(index-code MI)} } - \underbrace{\mathbb{E}_{q_{\phi}(\mathbf{z})} \mathcal{D}_{\text{KL} }\left( q_{\phi}(\mathbf{\Sigma} | \mathbf{z}) \| p(\mathbf{\Sigma} | \mathbf{z}) \right)}_{\text{(marginal KL to prior)} } - \underbrace{\mathcal{D}_{\text{KL} }(q_{\phi}(\mathbf{z}) \| p(\mathbf{z}))}_{\text{(covariance penalty)} } \end{aligned} }

上式表明引入之后的KL项可以拆分成index-code ML、marignal prior match、covariance penalty三项，整体可以看成原来的ELBO加上新的covariance penaly这一项。

LELBO=LELBOVAE−Eqϕ(z)DKL(qϕ(Σ∣z)∥p(Σ∣z)){ \mathcal{L}_{ELBO} = \mathcal{L}_{ELBO}^{VAE} - \mathbb{E}_{q_{\phi}(\mathbf{z})} \mathcal{D}_{\text{KL} }\left( q_{\phi}(\mathbf{\Sigma} | \mathbf{z}) \| p(\mathbf{\Sigma} | \mathbf{z}) \right) }

Prior Specification and Inference

ELBO分解之后就是指定Encoder和Decoder中的分布形式：

定义先验中的分布为：

p(Σ)=Wp−1(Σ∣Ψ;v)p(z∣Σ)=N(z;0,Σ)p(x∣z)=N(x;f(z),Σ0){ \begin{aligned} p(\mathbf{\Sigma}) = \mathcal{W}_{p}^{-1} (\mathbf{\Sigma} | \mathbf{\Psi}; v) \\ p(\mathbf{z} | \mathbf{\Sigma}) = \mathcal{N}(\mathbf{z} ; 0, \mathbf{\Sigma}) \\ p(\mathbf{x} | \mathbf{z}) = \mathcal{N}(\mathbf{x}; f(\mathbf{z}), \Sigma_0) \end{aligned} }

p(x∣z){p(\mathbf{x} | \mathbf{z})}的定义不变，而p(z∣Σ){p(\mathbf{z} | \mathbf{\Sigma})}需要Σ{\mathbf{\Sigma}}。

而Σ{\mathbf{\Sigma}}满足Inverse-Wishart分布，它的参数为p×p{p \times p}大小的半正定矩阵Ψ{\mathbf{\Psi}}（Scale）和v>p−1{v > p-1}（自由度）（Σ{\mathbf{\Sigma}}的大小也为p×p{p \times p}）。均值E[Σ]=(v−p−1)−1Ψ{\mathbb{E}[\mathbf{\Sigma}] = (v - p - 1)^{-1} \mathbf{\Psi}}，如果期望的z{\mathbf{z}}协方差为Σ0{\mathbf{\Sigma}_{0}}，则Ψ{\mathbf{\Psi}}应取(v−p−1)Σ0{( v - p - 1)\mathbf{\Sigma}_0}。而且v{v}起到控制统计独立性的作用，v{v}增大，随机采样的协方差就越接近Σ0{\mathbf{\Sigma}_{0}}。

而推断后验的分布定义为：

q(z∣x)=N(z∣u~;Σ~=L~L~T)u~=g1(x)L~=g2(x){ \begin{aligned} q(\mathbf{z} | \mathbf{x} ) = \mathcal{N}(\mathbf{z} | \tilde{u}; \tilde{\mathbf{\Sigma}} = \tilde{\mathbf{L}} \tilde{\mathbf{L}}^{T}) \\ \tilde{u}= g_{1}(\mathbf{x}) \qquad \tilde{\mathbf{L}} = g_2(\mathbf{x}) \\ \end{aligned} }

其中L~{\tilde{\mathbf{L}}}为Σ~{\tilde{\mathbf{\Sigma} } }的Cholesky分解矩阵。

至于q(Σ∣z){q(\mathbf{\Sigma} | \mathbf{z})}，作者使用p(Σ∣z){p(\mathbf{\Sigma} | \mathbf{z})}近似估计它：

q(Σ∣z)≈p(Σ∣z)=p(z∣Σ)p(Σ)∫Σ′p(z∣Σ′)p(Σ′)dΣ′{ q(\mathbf{\Sigma} | \mathbf{z}) \approx p(\mathbf{\Sigma} | \mathbf{z})=\frac{p(\mathbf{z} | \mathbf{\Sigma}) p(\mathbf{\Sigma})}{\int_{\mathbf{\Sigma}^{\prime}} p\left(\mathbf{z} | \mathbf{\Sigma}^{\prime}\right) p\left(\mathbf{\Sigma}^{\prime}\right) d \mathbf{\Sigma}^{\prime}} }

推得：（但是我还并没有理解这里怎么得来的，如下是作者的表述）

图1：近似估计思路
图1：近似估计思路

最终得到：

p(Σ∣z)=Wp−1(Ψ+zizi⊤,ν+1){ p(\mathbf{\Sigma} | \mathbf{z})=\mathcal{W}_{p}^{-1}\left(\mathbf{\Psi}+\mathbf{z}_{i} \mathbf{z}_{i}^{\top}, \nu+1\right) }

作者将ELBO继续推导，换了个表述方式（这里的ELBO的表达式和之前的表述不同，从KL项分解的第二个等式继续向下推导）：

−Eqϕ(z,Σ,x)[log⁡qϕ(z,Σ∣x)q(z,Σ)+log⁡qϕ(z,Σ)p(z,Σ)]=−Eqϕ(z,Σ,x)[log⁡qϕ(z∣x)qϕ(z)+log⁡qϕ(Σ∣z)qϕ(z)p(z∣Σ)p(Σ)]=−Eqϕ(z,Σ,x)[log⁡qϕ(z∣x)p(z∣Σ)+log⁡qϕ(Σ∣z)p(Σ)]{ \begin{aligned} - \mathbb{E}_{q_{\phi}(\mathbf{z}, \mathbf{\Sigma},\mathbf{x})}\left[ \log \frac{q_{\phi}(\mathbf{z}, \mathbf{\Sigma} | \mathbf{x})}{q(\mathbf{z}, \mathbf{\Sigma}) } + \log \frac{q_{\phi}(\mathbf{z}, \mathbf{\Sigma})}{p(\mathbf{z}, \mathbf{\Sigma})} \right] \\ = - \mathbb{E}_{q_{\phi}(\mathbf{z}, \mathbf{\Sigma},\mathbf{x})}\left[ \log \frac{q_{\phi}(\mathbf{z} | \mathbf{x})}{q_{\phi}(\mathbf{z})} + \log \frac{q_{\phi}(\mathbf{\Sigma} | \mathbf{z}) q_{\phi}(\mathbf{z})}{p(\mathbf{z} | \mathbf{\Sigma}) p(\mathbf{\Sigma})} \right] \\ = - \mathbb{E}_{q_{\phi}(\mathbf{z}, \mathbf{\Sigma},\mathbf{x})}\left[ \log \frac{q_{\phi}(\mathbf{z} | \mathbf{x})}{p(\mathbf{z} | \mathbf{\Sigma} )} + \log \frac{q_{\phi}(\mathbf{\Sigma} | \mathbf{z})}{p(\mathbf{\Sigma})} \right] \end{aligned} }

总体上，ELBO的表达式为：

LELBO=Eq(z∣x)[log⁡p(x∣z)]−Ep(Σ∣z)[DKL(N(μ~,Σ~)∥N(0,Σ))]−Eq(z∣x)[DKL(Wp−1(Φ,λ)∥Wp−1(Ψ,ν))]{ \begin{aligned} \mathcal{L}_{\mathrm{ELBO}}=& \mathbb{E}_{q(\mathbf{z} | \mathbf{x})}[\log p(\mathbf{x} | \mathbf{z})] \\ &-\mathbb{E}_{p(\mathbf{\Sigma} | \mathbf{z})}\left[D_{\mathrm{KL}}(\mathcal{N}(\tilde{\boldsymbol{\mu}}, \tilde{\mathbf{\Sigma}}) \| \mathcal{N}(\mathbf{0}, \mathbf{\Sigma}))\right] \\ &-\mathbb{E}_{q(\mathbf{z} | \mathbf{x})}\left[D_{\mathrm{KL}}\left(\mathcal{W}_{p}^{-1}(\mathbf{\Phi}, \lambda) \| \mathcal{W}_{p}^{-1}(\mathbf{\Psi}, \nu)\right)\right] \end{aligned} }

这是最终的优化目标。

其他的细节还有：（1）作者使用Bartlett decomposition从inverse-Wishart分布中采样。（2）对使用比较大的v{v}训练模型，作者inference阶段直接使用p(z)=N(z;0,I)){p(z) = \mathcal{N}(z;0,\mathbf{I}))}采样，也可以生成较好的图像（之前说过v{v}比较大时，各分量趋向于独立）。

以上ELBO的三项的具体计算作者在附录中给出，此处不再赘述。

实验

作者与β{\beta}-VAE，FactorVAE进行了对比实验，选取FactorVAE提出的metric，在dSprites，CorrelatedEllipses数据集（有groud true factors）上报告了Quantitatiave结果，在3DFaces,3DChairs,CelebA上报告了Qualitative结果。

值得注意的实验结论有：

(1) dSprites上更高的v{v}训练的模型的disentangled效果更好，但是CorrelatedEllipses上，更低的v{v}(v=200{v = 200})取得了最好的disentangled效果。

(2) 所有模型在捕捉离散的真实factors的表现都不太好。

思维发散

(1) 为什么作者引入Bayes Hyperprior的方法可以更有效地平衡重建误差和Disentangled约束（如作者在第三节开始所述）。也就是，如何思考引入Bayes Hyperpiror的作用？

(2) BF-VAE(Bayes-Factor-VAE)是本文形式的扩展，具体表现在替换了Bayes先验的形式，不是简单地将协方差Σ{\mathbf{\Sigma}}作为随机变量（因为很可能引起计算上的问题，协方差包含的参数量太大了，O(d*d)，d为分量的个数），而是先将p(z){p(z)}定义成分解的形式，再将各分量的方差组成一个向量，将它作为随机变量（参数量降为O(d)），并且定义的是Gamma分布。