使用Batch Normalization解决VAE训练中的后验坍塌(posterior collapse)问题

前言

在训练VAE模型时，当我们使用过于过于强大的decoder时，尤其是自回归式的decoder比如LSTM时，存在一个非常大的问题就是，decoder倾向于不从latent variable z中学习，而是独立地重构数据，这个时候，同时伴随着的就是KL(p(z|x)‖q(z))倾向于0。先不考虑说，从损失函数的角度来说，这种情况不是模型的全局最优解这个问题（可以将其理解为一个局部最优解)。单单从VAE模型的意义上来说，这种情况也是我们不愿意看到的。VAE模型的最重要的一点就是其通过无监督的方法构建数据的编码向量（即隐变量z)的能力。而如果出现posterior collapse的情况，就意味着后验概率退化为与先验概率一致，即N(0,1)。此时encoder的输出近似于一个常数向量，不再能充分利用输入x的信息。decoder则变为了一个普通的language model,尽管它依然很强。

因此，不管从哪个方面来看，解决它都是必须要面临的课题，事实上，从2016年开始就有很多文章提出了不同的解决方案，这里重点介绍一下使用Batch Normalization来解决这个问题的思路。这篇文章全名A Batch Normalized Inference Network Keeps the KL Vanishing Away发表于2020年，还算是一篇比较新的文章。下面我们开始。

方法介绍

Expectation of the KL’s Distribution

首先基于隐变量空间为高维高斯分布的假设，对于一个mini-batch的数据来说，我们可以计算KL divergence的表达式如下：

其中b代表的是mini-batch的样本个数，n代表的隐变量z的维度。同时作者还假设对于每个不同的维度，其都遵循某个特定的分布，各个维度可以不同。

假设我们认为上述的样本均值可以近似等于总体期望，那么我们可以将上述的样本均值用期望来代替，又因为我们有如下基本等式

最终我们可以得到KL divergence的期望表达式如下。

上述不等式是因为e^x-x>=1恒成立。那么这么一来，我们就有关于KL divergence的一个lower bound。这个lower bound只与隐变量的维度n和μi的分布有关。