本文我们继续之前的变分自编码器系列，分析一下如何防止 NLP 中的 VAE 模型出现“KL 散度消失（KL Vanishing）”现象。本文受到参考文献是 ACL 2020 的论文 A Batch Normalized Inference Network Keeps the KL Vanishing Away [1] 的启发，并自行做了进一步的完善。
值得一提的是，本文最后得到的方案还是颇为简洁的——只需往编码输出加入BN（Batch Normalization），然后加个简单的 scale——但确实很有效，因此值得正在研究相关问题的读者一试。同时，相关结论也适用于一般的 VAE 模型（包括 CV 的），如果按照笔者的看法，它甚至可以作为 VAE 模型的“标配”。
VAE简单回顾
这里我们简单回顾一下 VAE 模型，并且讨论一下 VAE 在 NLP 中所遇到的困难。关于 VAE 的更详细介绍，请读者参考笔者的旧作变分自编码器 VAE：原来是这么一回事、再谈变分自编码器 VAE：从贝叶斯观点出发等。
1.1 VAE的训练流程
VAE 的训练流程大概可以图示为：
写成公式就是：
1.2 NLP中的VAE
这种情况下的 VAE 模型并没有什么价值：KL 散度为 0 说明编码器输出的是 0 向量，而解码器则是一个普通的语言模型。而我们使用 VAE 通常来说是看中了它无监督构建编码向量的能力，所以要应用 VAE 的话还是得解决 KL 散度消失问题。
1.3 BN的巧与秒
KL 散度消失不就是 KL 散度项变成 0 吗？我调整一下编码器输出，让 KL 散度有一个大于零的下界，这样它不就肯定不会消失了吗？
1.4 推导过程简述
1.5 为什么不是LN？
善于推导的读者可能会想到，按照上述思路，如果只是为了让 KL 散度项有个正的下界，其实 LN（Layer Normalization）也可以，也就是在式（3）中按 j 那一维归一化。
那为什么用BN而不是LN呢？
进一步的结果
事实上，原论文的推导到上面基本上就结束了，剩下的都是实验部分，包括通过实验来确定 的值。然而，笔者认为目前为止的结论还有一些美中不足的地方，比如没有提供关于加入 BN 的更深刻理解，倒更像是一个工程的技巧，又比如只是 加上了 BN， 没有加上，未免有些不对称之感。
2.1 联系到先验分布
2.2 参考的实现方案
class Scaler(Layer):
"""特殊的scale层
"""
def __init__(self, tau=0.5, **kwargs):
super(Scaler, self).__init__(**kwargs)
self.tau = tau
def build(self, input_shape):
super(Scaler, self).build(input_shape)
self.scale = self.add_weight(
name='scale', shape=(input_shape[-1],), initializer='zeros'
)
def call(self, inputs, mode='positive'):
if mode == 'positive':
scale = self.tau + (1 - self.tau) * K.sigmoid(self.scale)
else:
scale = (1 - self.tau) * K.sigmoid(-self.scale)
return inputs * K.sqrt(scale)
def get_config(self):
config = {'tau': self.tau}
base_config = super(Scaler, self).get_config()
return dict(list(base_config.items()) + list(config.items()))
def sampling(inputs):
"""重参数采样
"""
z_mean, z_std = inputs
noise = K.random_normal(shape=K.shape(z_mean))
return z_mean + z_std * noise
e_outputs # 假设e_outputs是编码器的输出向量
scaler = Scaler()
z_mean = Dense(hidden_dims)(e_outputs)
z_mean = BatchNormalization(scale=False, center=False, epsilon=1e-8)(z_mean)
z_mean = scaler(z_mean, mode='positive')
z_std = Dense(hidden_dims)(e_outputs)
z_std = BatchNormalization(scale=False, center=False, epsilon=1e-8)(z_std)
z_std = scaler(z_std, mode='negative')
z = Lambda(sampling, name='Sampling')([z_mean, z_std])
文章内容小结
本文简单分析了 VAE 在 NLP 中的 KL 散度消失现象，并介绍了通过 BN 层来防止 KL 散度消失、稳定训练流程的方法。这是一种简洁有效的方案，不单单是原论文，笔者私下也做了简单的实验，结果确实也表明了它的有效性，值得各位读者试用。因为其推导具有一般性，所以甚至任意场景（比如 CV）中的 VAE 模型都可以尝试一下。
参考链接
[1] https://arxiv.org/abs/2004.12585
