本周阅读的是"Isolating Sources of Disentanglement in VAEs"。

本文是一篇典型的探索VAE损失函数，促进模型的Disentangled Representation的文章。这类文章的套路是：从损失函数的KL项中q(z∣x){q(z | x)}贴合的分布形式出发，提出一种直觉上看起来合理的损失函数形式，在几个ToyDataset上报告看起来还不错的结果就完事。本文对KL项的分解借鉴了几篇之前的文献，比较可取，同时比较符合直观感受，而且近似计算上下了很大的功夫，对比实验也比较详实。公式看的我比较辛苦。。。。

PaperFrame

本文的结构如下：

  • 简介：本文工作总结
  • 研究背景: 回顾InfoGAN和β{\beta}-VAE，以及 β{\beta}-VAE中提出的度量disentanglement的指标。
  • 主要内容：
    • 证据下界ELBO分解、ELBO在Minibatch上的估计
    • 使用互信息差MIG（Mutual Information Gap）衡量Disentanglement。
  • 实验
  • 结论
  • 附录:
    • [A] 随机样本：主要是CelebA上结果展示
    • [B] MIG相关，包括：(1) I(zk;vk){I(z_k;v_k)}近似计算推导；(2) 使用Normalization的考量
    • [C] ELBO分解相关，包括：(1) KL项分解形式推导；(2) MWS近似计算公式(Equation:estimate-entropy)的推导；(3) MSS近似计算q(z){q(z)}的推导，近似计算Eq(z,n)(log⁡q(z)){\mathbb{E}_{q(z,n)}(\log q(z))}的推导。(4) MSS 和 MWS 估计效果的对比实验。
    • [C] 更多的对比实验：(1) 移除1项结果；(2)Factorial Normalizing Flow，等价于移除3项的结果；(3) 在 β{\beta}-TCVAE估计计算ELBO中，batchsize的取值对最终效果的影响。
    • [D] 比较最好的 β{\beta}-VAE 和 β{\beta}-TCAVE 模型，通过可视化disentangled的隐分量。
    • [E] Higgins Metric和Kim & Mnih Metric 都受batchsize影响的说明性实验。
    • [F] MIG Traversal实验：展示了不同的可视化学习到的disentangled的隐分量和最终MIG的得分的结果。
    • [G] MIG 和 Higgins Metric对比实验，比对两者结果相差很大的模型进行可视化检验。

本文的两个重点是从ELBO的KL项出发深入分析改进Loss项以及新的Disentanglement评价方法。

Content

Decompose ELBO

ELBO，即证据下界(Evident Lower Boundary)是VAE的变分下界的近似计算表示，
在β{\beta}-VAE中写成：

Lβ=1N∑n=1N(Eq[log⁡p(xn∣z)]−βKL⁡(q(z∣xn)∥p(z)))(Equation:ELMO-original){ \begin{aligned} \tag{Equation:ELMO-original} \mathcal{L}_{\beta}=\frac{1}{N} \sum_{n=1}^{N}\left(\mathbb{E}_{q}\left[\log p\left(x_{n} | z\right)\right]-\beta \operatorname{KL}\left(q\left(z | x_{n}\right) \| p(z)\right)\right) \end{aligned} }

等式(Equation:ELMO-original)右项第一项为重建误差，它和Disentanglement无关，
作者基于另一篇文章[1]对ELMO中KL散度项进行了进一步的分解。

首先做以下的记号约定：

q(z∣n)≜q(z∣xn)q(n)≜1Nq(n,z)≜q(z∣n)q(n)q(z)=∑n=1Nq(z∣n)q(n)=1N∑n=1Nq(z∣n){ \begin{aligned} q(z\vert n) &\triangleq q(z \vert x_{n}) &\qquad q(n) &\triangleq \frac{1}{N} \\ q(n,z) &\triangleq q(z \vert n) q(n) &\qquad q(z) &= \sum_{n=1}^{N} q(z \vert n) q(n) = \frac{1}{N} \sum_{n=1}^{N} q(z \vert n) \end{aligned} }

其中q(z){q(z)}称为合计后验分布(aggregate posterior)或者平均编码分布(average encoding distribution)。它表示所有样本推断得到的隐变量分布的平均。

假设p(z){p(z)}各分量独立，p(z)=∏j=1dp(zj){p(z) = \prod_{j = 1}^{d} p(z_{j})}，其中zj{z_j}表示z{z}第j{j}维度的分量，公式(Equation:ELMO-original)中的KL项分解如下：

1N∑n=1NKL⁡(q(z∣xn)∥p(z))=Ep(n)[KL⁡(q(z∣n)∥p(z))]=Ep(n)[Eq(z∣n)[log⁡q(z∣n)−log⁡p(z)+log⁡q(z)−log⁡q(z)+log⁡∏jq(zj)−log⁡∏jq(zj)]]=Eq(z,n)[log⁡q(z∣n)q(z)]+Eq(z)[log⁡q(z)∏jq(zj)]+Eq(z)[∑jlog⁡q(zj)p(zj)]=Eq(z,n)[log⁡q(z∣n)p(n)q(z)p(n)]+Eq(z)[log⁡q(z)∏jq(zj)]+∑jEq(z)[log⁡q(zj)p(zj)]=Eq(z,n)[log⁡q(z∣n)p(n)q(z)p(n)]+Eq(z)[log⁡q(z)∏jq(zj)]+∑jEq(zj)q(z\j∣zj)[log⁡q(zj)p(zj)]=Eq(z,n)[log⁡q(z∣n)p(n)q(z)p(n)]+Eq(z)[log⁡q(z)∏jq(zj)]+∑jEq(zj)[log⁡q(zj)p(zj)]=KL⁡(q(z,n)∥q(z)p(n))⏟1Index-Code MI)+KL⁡(q(z)∥∏jq(zj))⏟2 Total Correlation +∑jKL⁡(q(zj)∥p(zj))⏟3 Dimension-wise KL (Equation:TC-Decompose){ \begin{aligned} \tag{Equation:TC-Decompose} \begin{aligned} & \frac{1}{N} \sum_{n=1}^{N} \operatorname{KL}\left(q\left(z | x_{n}\right) \| p(z)\right)=\mathbb{E}_{p(n)}[\operatorname{KL}(q(z | n) \| p(z))] \\\\ =& \mathbb{E}_{p(n)}\left[\mathbb{E}_{q(z | n)}\left[\log q(z | n)-\log p(z)+\log q(z)-\log q(z)+\log \prod_{j} q\left(z_{j}\right)-\log \prod_{j} q\left(z_{j}\right)\right]\right] \\\\ =& \mathbb{E}_{q(z, n)}\left[\log \frac{q(z | n)}{q(z)}\right]+\mathbb{E}_{q(z)}\left[\log \frac{q(z)}{\prod_{j} q\left(z_{j}\right)}\right]+\mathbb{E}_{q(z)}\left[\sum_{j} \log \frac{q\left(z_{j}\right)}{p\left(z_{j}\right)}\right] \\\\ =& \mathbb{E}_{q(z, n)}\left[\log \frac{q(z | n) p(n)}{q(z) p(n)}\right]+\mathbb{E}_{q(z)}\left[\log \frac{q(z)}{\prod_{j} q\left(z_{j}\right)}\right]+\sum_{j} \mathbb{E}_{q(z)}\left[\log \frac{q\left(z_{j}\right)}{p\left(z_{j}\right)}\right] \\\\ =& \mathbb{E}_{q(z, n)}\left[\log \frac{q(z | n) p(n)}{q(z) p(n)}\right]+\mathbb{E}_{q(z)}\left[\log \frac{q(z)}{\prod_{j} q\left(z_{j}\right)}\right]+\sum_{j} \mathbb{E}_{q\left(z_{j}\right) q\left(z_{\backslash j} | z_{j}\right)}\left[\log \frac{q\left(z_{j}\right)}{p\left(z_{j}\right)}\right] \\\\ =& \mathbb{E}_{q(z, n)}\left[\log \frac{q(z | n) p(n)}{q(z) p(n)}\right]+\mathbb{E}_{q(z)}\left[\log \frac{q(z)}{\prod_{j} q\left(z_{j}\right)}\right]+\sum_{j} \mathbb{E}_{q\left(z_{j}\right)}\left[\log \frac{q\left(z_{j}\right)}{p\left(z_{j}\right)}\right] \\\\ =&\underbrace{\operatorname{KL}(q(z, n) \| q(z) p(n))}_{ 1 \text {Index-Code MI} } )+\underbrace{\operatorname{KL}\left(q(z) \| \prod_{j} q\left(z_{j}\right)\right)}_{ 2 \text { Total Correlation }}+\underbrace{\sum_{j} \operatorname{KL}\left(q\left(z_{j}\right) \| p\left(z_{j}\right)\right)}_{3 \text { Dimension-wise KL }} \end{aligned} \end{aligned} }

将等式(Equation:TC-Decompose)的右端分别记为 1, 2, 3项：

  • 其中1记为索引-编码互信息(index-code mutual information)，该项描述了数据分布和隐变量分布的互信息量的大小，一般认为该互信息量越大，disentanglement效果越好。

  • 2记为总体相关性(Total-Correlation)，它迫使模型学习到的推断的后验分布各分量之间相互独立，作者在实验中说明了该项的作用，并且认为优化该项是 β{\beta}-VAE 起作用的关键，符合直观感觉。

  • 3记为分维度KL散度(Dimension-wise KL Divergence)，它保证推断的各维度的分量不会离先验的各维度分量分布太远，限制合计后验的复杂性。

ELMO在Mini-batch上的近似计算

由于q(z)=1N∑n=1Nq(z∣n){q(z) = \frac{1}{N} \sum_{n=1}^{N} q(z | n)}，每次更新模型参数后，要估计q(z){q(z)}就得对所有样本计算q(z∣n){q(z | n)}，这将使得计算非常低效。在Mini-batch上近似计算q′(z)=1M∑n=1Mq(z∣n){q^{\prime}(z) = \frac{1}{M} \sum_{n=1}^{M} q(z | n)}，可能并不能取得很好的估计。(为什么？可能的原因是随机性太大)，作者受重要性采样的启发，认为某些zj{z_j}的编码的信息只有数据集中的某些样本才有，这些样本对q(z){q(z)}的贡献大，而其余样本可以忽略不计，因此提出了加权采样MWS(Minibatch-Weighted Sampling)的方法近似计算：

(下式的推导见附录,但是作者也说这是一个有偏估计，它的期望是一个下界，但是因为没有引入超参数所以使用了，到底有偏估计对它带来多大的影响呢？这是一个问题。这个式子是直接可计算的，因为在这里q(z∣nj){q(z | n_j)}仍然取的是各向同性的高斯分布，那么q(z∣nj){q(z | n_j)}分布已知，求概率也比较容易。)

Eq(z)[log⁡q(z)]≈1M∑i=1M[log⁡1MN∑j=1Mq(z(ni)∣nj)]z(nj)sample fromq(z∣nj)(Equation:estimate-entropy){ \begin{aligned} \tag{Equation:estimate-entropy} \mathbb{E}_{q(z)}\left[\log q(z)\right] \approx \frac{1}{M} \sum_{i=1}^{M} \left[ \log \frac{1}{MN} \sum_{j=1}^{M} q(z(n_i) | n_j) \right] \qquad z(n_j) \quad \text{sample from} \quad q(z | n_j) \end{aligned} }

而其余的项可以分解化成Eq(z)[log⁡q(z)]{\mathbb{E}_{q(z)}\left[ \log q(z)\right]}项的形式计算。（这里其实我还没有完全搞明白到底怎么计算的。）

进一步，和 β{\beta}-VAE类似，作者提出对公式(Equation:TC-Decompose)的三项加入超参数α,β,γ{\alpha,\beta,\gamma}控制三项优化的力度：

Lβ−TC:=Eq(z∣n)p(n)[log⁡p(n∣z)]−αIq(z;n)−βKL⁡(q(z)∥∏jq(zj))−γ∑jKL⁡(q(zj)∥p(zj))(Equation:beta-TC-Decompose){ \begin{aligned} \tag{Equation:beta-TC-Decompose} \mathcal{L}_{\beta-\mathrm{TC}}:=\mathbb{E}_{q(z | n) p(n)}[\log p(n | z)]-\alpha I_{q}(z ; n)-\beta \operatorname{KL}\left(q(z) \| \prod_{j}q\left(z_{j}\right)\right)-\gamma \sum_{j} \operatorname{KL}\left(q\left(z_{j}\right) \| p\left(z_{j}\right)\right) \end{aligned} }

经验上，作者选取α=γ=1{\alpha = \gamma = 1}，而通过控制β{\beta}调整Disentanglement能力，称为β{\beta}-TCVAE模型。（由于该式影响Disentanglement效果最大的是TC项，作者设计了控制变量的对比实验说明了这一点。）

Disentanglement评价指标：互信息差MIG

作者从两个角度考量衡量Disentanglement的指标：

1. 隐变量分量zj{z_j}和真实的因素vk{v_k}的互信息量高。

2. 坐标对齐(axis-aligned)，意思是：单个隐变量分量仅对应单个变化因子。很显然，如果一个隐分量可以控制多个变换因子，Disentanglement效果肯定比较差。

隐变量分量zj{z_j}和真实变化因子vk{v_k}的联合分布q(zj,vk){q(z_j,v_k)}可以表示如下：

q(zj,vk)=q(zj∣vk)p(vk)=(∑n=1Nq(zj∣n)p(n∣vk))p(vk)=p(v)p(z∣x)q(x∣v){ \begin{aligned} q(z_j, v_k) &= q(z_j | v_k) p(v_k) &= \left( \sum_{n=1}^{N} q(z_j | n)p(n | v_k) \right) p(v_k) &= p(v) p(z | x) q(x | v) \end{aligned} }

记Xvk{\mathcal{X}_{v_k}}表示样本中p(n∣vk){p(n | v_k)}的支撑集，互信息量的计算公式如下：
(推导公式见附录)

In(zj,vk)=Eq(zj,vk)[log⁡∑n∈Xvkq(zj∣n)p(n∣vk)]+H(zj)(Equation:estimate-mi){ \begin{aligned} \tag{Equation:estimate-mi} I_{n}(z_j, v_k) = \mathbb{E}_{q(z_j, v_k)}\left[ \log \sum_{n \in \mathcal{X}_{v_k}} q(z_j | n)p(n | v_k) \right] + H(z_j) \end{aligned} }

为了保证axis-aligned的要求，作者提出使用经过normalized的互信息量，而且对某一个真实变化因子vk{v_k}，使用最大的互信息量和次大的互信息量之差衡量对该因子的Disentangled表示的好坏。（二者之差越大，表示只有一个分量突出表示该因子，更加符合axis-aligned的要求），综合的指标计算公式如下，称为MIG（Mutual Information Gap）：

1K∑k=1K1H(vk)(In(zj(k);vk)−max⁡j≠j(k)In(zj;vk)){ \frac{1}{K} \sum_{k=1}^{K} \frac{1}{H\left(v_{k}\right)}\left(I_{n}\left(z_{j^{(k)}} ; v_{k}\right)-\max _{j \neq j^{(k)}} I_{n}\left(z_{j} ; v_{k}\right)\right) }

结论

作者使用提出的MIG作为主要的评价指标，主要得到的实验结论有：

在第六节（实验小节）得到的结论：

1. 比起β{\beta}-VAE，β{\beta}-TCVAE模型中β{\beta}起到的概率密度估计和disentanglement的trade-off的效果更好，即 β{\beta}-TCVAE的β{\beta}有更大的调整空间。当β{\beta}增大时， β{\beta}-VAE会因为估计不准的问题，导致Disentanglement效果下降，而 β{\beta}-TCVAE 仍能保持好的Disentanglement效果。

2. 分别去除公式(Equation:TC-Decompose)中的1、3项的模型的Disentanglement的效果
变换很小，这表明2项确实是影响Disentanglement效果的关键。

3. β{\beta}-TCVAE和VAE, β{\beta}-VAE, InfoGAN, FactorVAE比较，可以得到更高的平均MIG得分。

4. β{\beta}和MIG的负相关性在 β{\beta}-TCVAE中更加显著，意味着在 β{\beta}-TCVAE中，β{\beta}可以更好地控制Disentanglement效果。

5. 样本中不同的disentangled的因子的采样比例会影响模型的Disentanglement能力， β{\beta}-TCVAE 对不同因子采样不同的数据集学习到的Disentanglement的效果更好。

6. 定性比较中， β{\beta}-TCVAE 比 β{\beta}-VAE 发现了更多的disentangled的因子。

附录中的实验结论：

1. C.1 通常认为互信息量越大，disentanglement效果越好，如果不限制该项，即去掉互信息量大小的约束，可能效果更好，但是实验中移除Index-Code MI项(α=0{\alpha = 0})对Disentanglement能力影响不大，而且α{\alpha}的选取偏经验性，和具体数据集有关。

2. C.2 使用 Factorial Normalizing Flow（等价忽略dimension-wise KL项）也影响不大。

3. C.3 BatchSize取100到1000之间对MIG影响不大。小batch结果也不错。

4. E Higgins Metric和Kim& Mnih Metric受BatchSize大小影响大。小batch表现更好，
batchsize 增大，越来越不能分辨。

总体实验结果：

1. 通常认为互信息量越大，disentanglement效果越好，如果不限制该项，即去掉公式(Equation:TC-Decompose)中1项，可能效果更好，但是实验中移除Index-Code MI项(α=0{\alpha = 0})对Disentanglement结果影响不大，而且去掉3项(γ=0{\gamma = 0})对结果影响也不大。这表明公式(Equation:TC-Decompose)中对Disenanglement表现最关键的是2项，即TC项。另外α{\alpha}的最优取值并不是dataset independent的，这可能和不同数据集的因子的复杂度有关。

2. 比起β{\beta}-VAE，β{\beta}-TCVAE模型中β{\beta}起到的trade-off的效果
更好，β{\beta}-TCVAE的β{\beta}有更大的调整空间。β{\beta}取相同比较大的值时，β{\beta}-VAE会因为重建误差的惩罚太小，学习不到准确的近似后验分布，导致Disentanglement效果下降，而 β{\beta}-TCVAE仍能保持好的Disentanglement效果。

3. 定量实验中，在dSprites和3DFaces数据集上β{\beta}-TCVAE比VAE, β{\beta}-VAE, InfoGAN, FactorVAE取得了更高的平均MIG得分。在定性实验中，在3DChairs数据集上， β{\beta}-TCVAE 比 β{\beta}-VAE 多学习到了两个disentangled的因子。

评价

本文是一篇典型的探索VAE损失函数，促进模型的Disentangled Representation的文章。这类文章的套路是：从损失函数的KL项中q(z∣x){q(z | x)}贴合的分布形式出发，提出一种直觉上看起来合理的损失函数形式，在几个ToyDataset上报告看起来还不错的结果就完事。本文对KL项的分解借鉴了几篇之前的文献，比较可取，同时比较符合直观感受，而且近似计算上下了很大的功夫，对比实验也比较详实。

不足之处有：

1. 在CelebA数据集上仅报告了定性实验的结果，缺少MIG的量化结果，该模型的Disentangled Representation
的效果可能还是局限在几个Toy Dataset上。

2. Compactness是Disentangled Representation另一个要求，即单个变化因子仅由一个
隐分量编码，如果存在几个分量控制一个变化因子，这几个分量就比较entangled。而作者的Disentanglement评价指标
似乎没有涉及到这一点。

3. 公式(Equation:TC-Decompose)中的互信息项去除之后却没有带来Disentangled表示能力的上升，
这一点比较反直觉，而且作者没有对这一点进一步说明。

附录

支撑集合

定义在集合X{\mathcal{X}}上的实值函数f{f}的支撑集合C{C}的定义为：

1. C⊆X{C \subseteq \mathcal{X}} ，即 C{C} 是X{\mathcal{X}}的子集。

2. f{f} 在 C{C} 上非0。

来源：[wiki]

在小节中：Xvk{\mathcal{X}_{v_k}}为p(n∣vk){p(n | v_k)}的支撑集，意味着X{\mathcal{X}}
是使得p(n∣vk){p(n | v_k)}不为0的样本的集合。从更加实际的意义就是：所有包含vk{v_k} 因子（特征）的图像的集合咯？？？？

实现中的计算公式

此处的推导需要进一步的补充

公式(Equation:TC-Decompose)三项Minibatch近似估计计算形式推导

1. 公式(Equation:estimate-entropy)的推导

2. 整体公式(Equation:TC-Decompose)三项通过分解成公式(Equation:estimate-entropy)的
形式近似估计。

3. 使用MWS(Minibatch-Weighted Sampling)对公式(Equation:estimate-entropy)的近似
计算推导。原文附录C.1。

4. 使用MSS(Minibatch-Stratified Sampling)对q(z){q(z)}分布近似计算和另一种近似计算
公式(Equation:estimate-entropy)的方法。原文附录C.2。

5. 推导完成后，怎么实现的傻瓜式说明？

并且作者用实验说明了MSS和MWS估计计算elbo的差距不大。

MIG的Minibatch近似计算

1. 主要是公式(Equation:estimate-mi)近似计算公式的推导。

2. 推导完成后，怎么实现的傻瓜式说明？

Reference

  1. Andrea Asperti. Sparsity in Variational Autoencoders. 2018

  2. James Lucas, George Tucker, Roger Grosse and Mohammad Norouzi. Don’t Blame the ELBO! A Linear VAE Perspective on Posterior Collapse. 2019

  3. Serena Yeung, Anitha Kannan, Yann Dauphin and Li Fei-Fei. Tackling Over-pruning in Variational Autoencoders. 2017

  4. Yuri Burda, Roger Grosse and Ruslan Salakhutdinov. Importance weighted autoencoders. 2016

  5. Junxian He, Daniel Spokoyny, Graham Neubig and Taylor Berg-Kirkpatrick. Lagging inference networks and posterior collapse in variational autoencoders. 2019