1. β-VAE

β-VAE的出发点是对特征空间进行解耦(disentanglement)，即使得隐变量空间$Z$的每一个维度作为一个factor，每一个factor表示独立的特征，而不影响其他factor表示的特征。

如一个在人脸数据集上训练的VAE，训练后隐空间中的每一个factor可以表示性别、肤色、表情…而不互相影响。

与VAE类似，模型希望最大化生成真实数据的概率，同时使得隐变量的后验分布$q(z|x)$与先验分布$p(z)$的距离小于常数$\epsilon>0$：

\[\begin{aligned} \mathop{\max}_{\theta,\phi}& \mathbb{E}_{x \text{~} D} [\mathbb{E}_{q_{\phi}(z|x)} [\log p_{\theta}(x | z)]] \\ \text{s.t. } &D_{KL}(q_{\phi}(z|x)||p_{\theta}(z))<\epsilon \end{aligned}\]

引入拉格朗日乘子$\beta$，问题转换成最大化拉格朗日函数：

\[\begin{aligned} \mathcal{F}(\theta,\phi,\beta) &= \mathbb{E}_{q_{\phi}(z|x)} [\log p_{\theta}(x | z)] - \beta (D_{KL}(q_{\phi}(z|x)||p_{\theta}(z)) - \epsilon) \\ &= \mathbb{E}_{q_{\phi}(z|x)} [\log p_{\theta}(x | z)] - \beta D_{KL}(q_{\phi}(z|x)||p_{\theta}(z)) + \beta\epsilon \\ &≥ \mathbb{E}_{q_{\phi}(z|x)} [\log p_{\theta}(x | z)] - \beta D_{KL}(q_{\phi}(z|x)||p_{\theta}(z)) \end{aligned}\]

因此β-VAE的损失函数定义为：

\[L(\theta,\phi,\beta) = -\mathbb{E}_{q_{\phi}(z|x)} [\log p_{\theta}(x | z)] + \beta D_{KL}(q_{\phi}(z|x)||p_{\theta}(z))\]

当$\beta = 1$时，模型和VAE相同。当$\beta > 1$时，引入了信息瓶颈，限制了模型的重构能力，但增加了模型的解耦能力。

2. 解耦评估得分 disentanglement metric score

如果模型的解耦程度较好，则其学习的隐变量不同维度之间应该具有独立性，则可以使用非常简单的线性分类器实现稳健的分类。作者设计了如下评估方法：

随机选择一个解耦因子$y$（如尺寸）；
采样$L$对图像$x_{1,l}$和$x_{2,l}$，将它们的因子$y$设置为相等$[x_{1,l}]_y=[x_{2,l}]_y$，其他因子随机；
使用编码器$q(z|x)$构造图像对的隐变量$z_{1,l}$和$z_{2,l}$；
计算所有图像对的隐变量差异$z_{diff}^l=|z_{1,l}-z_{2,l}|$并取平均$z_{diff}^b=\frac{1}{L} \sum_{l=1}^{L} z_{diff}^l$；
使用线性分类器$p(y|z_{diff}^b)$预测解耦因子$y$；
分类准确率即可作为最终的解耦评估得分。

作者汇报了不同模型在一个二维形状数据集上的解耦表现，并给出了不同隐变量长度和不同$\beta$对结构结果的影响。

3. β-VAE的pytorch实现

β-VAE的完整pytorch实现可参考PyTorch-VAE，与标准的VAE主要区别在损失函数上，即引入$\beta$：

recons_loss =F.mse_loss(recons, input)
kld_loss = torch.mean(-0.5 * torch.sum(1 + log_var - mu ** 2 - log_var.exp(), dim = 1), dim = 0)
loss = recons_loss + self.beta * kld_loss