变分自编码器系列：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 维。由于我们总是有