本文我们继续之前的变分自编码器系列,分析一下如何防止 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