变分自编码器系列:VAE + BN = 更好的VAE

©PaperWeekly 原创 · 作者|苏剑林

单位|追一科技

研究方向|NLP、神经网络

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

这里我们简单回顾一下 VAE 模型,并且讨论一下 VAE 在 NLP 中所遇到的困难。关于 VAE 的更详细介绍,请读者参考笔者的旧作 变分自编码器 VAE:原来是这么一回事再谈变分自编码器 VAE:从贝叶斯观点出发等

1.1 VAE的训练流程

VAE 的训练流程大概可以图示为:

▲ VAE训练流程图示

写成公式就是:

其中第一项就是重构项,是通过重参数来实现;第二项则称为 KL 散度项,这是它跟普通自编码器的显式差别,如果没有这一项,那么基本上退化为常规的 AE。更详细的符号含义可以参考 再谈变分自编码器 VAE:从贝叶斯观点出发

1.2 NLP中的VAE

在 NLP 中,句子被编码为离散的整数 ID,所以 q(x|z) 是一个离散型分布,可以用万能的“条件语言模型”来实现,因此理论上 q(x|z) 可以精确地拟合生成分布,问题就出在 q(x|z) 太强了,训练时重参数操作会来噪声,噪声一大,z 的利用就变得困难起来,所以它干脆不要 z 了,退化为无条件语言模型(依然很强),则随之下降到 0,这就出现了  KL 散度消失现象

这种情况下的 VAE 模型并没有什么价值:KL 散度为 0 说明编码器输出的是 0 向量,而解码器则是一个普通的语言模型。而我们使用 VAE 通常来说是看中了它无监督构建编码向量的能力,所以要应用 VAE 的话还是得解决 KL 散度消失问题。

事实上从 2016 开始,有不少工作在做这个问题,相应地也提出了很多方案,比如退火策略、更换先验分布等,读者 Google 一下“KL Vanishing”就可以找到很多文献了,这里不一一溯源。

1.3 BN的巧与秒

本文的方案则是直接针对 KL 散度项入手,简单有效而且没什么超参数。其思想很简单:

KL 散度消失不就是 KL 散度项变成 0 吗?我调整一下编码器输出,让 KL 散度有一个大于零的下界,这样它不就肯定不会消失了吗?

这个简单的思想的直接结果就是:在后面加入 BN 层,如图

往VAE里加入BN

1.4 推导过程简述

为什么会跟 BN 联系起来呢?我们来看 KL 散度项的形式:

上式是采样了 b 个样本进行计算的结果,而编码向量的维度则是 d 维。由于我们总是有,所以,因此:

留意到括号里边的量,其实它就是在 batch 内的二阶矩,如果我们往加入 BN 层,那么大体上可以保证的均值为,方差为(是 BN 里边的可训练参数),这时候:

所以只要控制好(主要是固定为某个常数),就可以让 KL 散度项有个正的下界,因此就不会出现 KL 散度消失现象了。这样一来,KL 散度消失现象跟 BN 就被巧妙地联系起来了,通过 BN 来“杜绝”了 KL 散度消失的可能性。

1.5 为什么不是LN?

善于推导的读者可能会想到,按照上述思路,如果只是为了让 KL 散度项有个正的下界,其实 LN(Layer Normalization)也可以,也就是在式(3)中按 j 那一维归一化。

那为什么用BN而不是LN呢?

这个问题的答案也是 BN 的巧妙之处。直观来理解,KL 散度消失是因为的噪声比较大,解码器无法很好地辨别出 z 中的非噪声成分,所以干脆弃之不用。

而当给加上 BN 后,相当于适当地拉开了不同样本的 z 的距离,使得哪怕 z 带了噪声,区分起来也容易一些,所以这时候解码器乐意用 z 的信息,因此能缓解这个问题;相比之下,LN 是在样本内进的行归一化,没有拉开样本间差距的作用,所以 LN 的效果不会有 BN 那么好。

进一步的结果

事实上,原论文的推导到上面基本上就结束了,剩下的都是实验部分,包括通过实验来确定的值。然而,笔者认为目前为止的结论还有一些美中不足的地方,比如没有提供关于加入 BN 的更深刻理解,倒更像是一个工程的技巧,又比如只是加上了 BN,没有加上,未免有些不对称之感。

经过笔者的推导,发现上面的结论可以进一步完善。

2.1 联系到先验分布

对于 VAE 来说,它希望训练好后的模型的隐变量分布为先验分布,而后验分布则是,所以 VAE 希望下式成立:

两边乘以 z,并对 z 积分,得到:

两边乘以,并对 z 积分,得到:

如果往都加入 BN,那么我们就有:

所以现在我们知道一定是 0,而如果我们也固定,那么我们就有约束关系:

2.2 参考的实现方案

经过这样的推导,我们发现可以往都加入 BN,并且可以固定,但此时需要满足约束(9)。

要注意的是,这部分讨论还仅仅是对 VAE 的一般分析,并没有涉及到 KL 散度消失问题,哪怕这些条件都满足了,也无法保证 KL 项不趋于 0。结合式 (4) 我们可以知道,保证 KL 散度不消失的关键是确保,所以,笔者提出的最终策略是:

其中是一个常数,笔者在自己的实验中取了,而是可训练参数,上式利用了恒等式。

关键代码参考(Keras):

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

更多阅读

# 投 稿 通 道 #

让你的论文被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢? 答案就是:你不认识的人。

总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。 

PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是 最新论文解读 ,也可以是 学习心得技术干货 。我们的目的只有一个,让知识真正流动起来。

:memo:  来稿标准:

• 稿件确系个人 原创作品 ,来稿需注明作者个人信息(姓名+学校/工作单位+学历/职位+研究方向) 

• 如果文章并非首发,请在投稿时提醒并附上所有已发布链接 

• PaperWeekly 默认每篇文章都是首发,均会添加“原创”标志

:mailbox_with_mail:  投稿邮箱:

• 投稿邮箱: hr@paperweekly.site 

• 所有文章配图,请单独在附件中发送 

• 请留下即时联系方式(微信或手机),以便我们在编辑发布时和作者沟通

:mag:

现在,在 「知乎」 也能找到我们了

进入知乎首页搜索 「PaperWeekly」

点击 「关注」 订阅我们的专栏吧

关于PaperWeekly

PaperWeekly 是一个推荐、解读、讨论、报道人工智能前沿论文成果的学术平台。如果你研究或从事 AI 领域,欢迎在公众号后台点击 「交流群」 ,小助手将把你带入 PaperWeekly 的交流群里。

我来评几句
登录后评论

已发表评论数()

相关站点

+订阅
热门文章