转载自《UniVAE：基于Transformer的单模型、多尺度的VAE模型》，作者：苏剑林，部分内容有修改。
大家都知道，Transformer 的 $\mathscr{O}(n^2)$ 复杂度是它的“硬伤”之一。不过凡事有弊亦有利，$\mathscr{O}(n^2)$ 的复杂度也为 Transformer 带来很大的折腾空间，我们可以灵活地定制不同的 attention mask，来设计出不同用途的 Transformer 模型来，比如 UniLM、K-BERT 等。
本文介绍笔者构思的一个能用于文本的 UniVAE 模型，它沿用类似 UniLM 的思路，将 VAE 做到了一个 Transformer 模型里边，并且还具备多尺度特性～
UniAE
VAE (Variational Autoencoder) 可以理解为带有正则项的 AE (Autoencoder)，一般情况下，Encoder 负责将输入编码为一个向量，并且满足一定的分布，而 Decoder 则负责将编码向量重构为输入。所以很显然，要实现 UniVAE，首先要实现对应的 UniAE。
在《从语言模型到Seq2Seq：Transformer如戏，全靠Mask》中，我们已经介绍了UniLM（Uni 是 Unified 的缩写），它通过下图左的 Attention Mask 来使得 Transformer 能完成 Seq2Seq 任务。然而 UniLM 并不是我们要寻找的 UniAE，因为 UniLM 的 Decoder 部分关联到的是输入的整个编码序列，而不是单个向量。
 |
 |
不过，我们可以在 UniLM 的基础上，进一步调整 Attention Mask 为上图右的模式，这样一来，解码的时候只能依赖于编码部分的 $\texttt{[CLS]}$ 向量以及当前已完成的解码结果，这就是我们要找的 UniAE 式 Attention Mask 了。因为对于输入来说，它只依赖于 $\texttt{[CLS]}$ 向量，而 $\texttt{[CLS]}$ 向量的大小是固定的，所以相当于说生成过程中的源信息只是一个固定大小的向量，而输入也被编码成这个固定大小的向量，这就是 AE 功能了。
多尺度
也就是说，通过 UniAE 式 Attention Mask，我们可以实现类似 UniLM 的 Seq2Seq 模型，它等效于 Encoder 将输入编码为固定长度的向量，然后 Decoder 对该向量进行解码。如果还觉得不够清晰，我们还可以分拆为 Encoder-Decoder 架构来理解，如下图所示：
跟常规的 Seq2Seq 架构不同的地方在于，这里的 Encoder 和 Decoder 的权重是共享的。从上图还可以看出，如果我们每一层 Attention 都加上这种 Mask，那么 Decoder 将依赖于每一层输入的 $\texttt{[CLS]}$ 向量，这也就意味如果有 $L$ 层 Attention，那么这 $L$ 层 Attention 的输入序列的所有 $\texttt{[CLS]}$ 向量拼接起来，才是输入文本的完整的编码向量（当然，第一层可以去掉，因为第一层的 $\texttt{[CLS]}$ 是其 Embedding 向量，对于每个输入来说它都是常向量），单独某一层的 $\texttt{[CLS]}$ 向量，并不是完整编码向量。
对于 Decoder 来说，每一层 Attention 都有一个 $\texttt{[CLS]}$ 向量传入，这其实就形成了一种多尺度结构。在 CV 中，最先进的生成模型基本上都是多尺度结构了，如 StyleGAN、Glow、NVAE 等，但是 NLP 中似乎还不多见。不难想象，在多尺度结构中，不同层次的输入对生成结果的调控程度也是不同的，越靠近输入层的变量，控制的部分越是“无伤大雅”，而越靠近输出层的变量，则控制着生成结果的关键信息。所以理想情况下，训练好一个多尺度模型后，我们可以通过编辑不同层级的输入变量，来实现对生成结果的不同层次的控制。
降低维度
有些读者可能会想到，要是每层的维度是 $d$，共有 $L$ 层，那么全部 $\texttt{[CLS]}$ 向量拼接起来就是 $Ld$ 维了，对于 BERT base 来说就是 $12×768=9216$ 维了，这编码向量维度是不是太大了？确实如此，对于一个普通的 AE 或者 VAE 来说，近万维的编码向量是太大了。
其实解决方法很简单，我们只需要将每层的 $\texttt{[CLS]}$ 向量用一个全连接层先降维，然后再用另一个全连接层升维，最后拼接到剩下的 $L−1$ 个 $d$ 维向量就行了，如上图所示。这样的话，虽然输入序列还是 $L\times d$ 大小，但事实上 $\texttt{[CLS]}$ 向量可以用一个更低维的向量表达出来，我们只需要把每一层的这个更低维向量拼接起来，作为总的编码向量就行了。
解耦能力
前面的设计和讨论还只是针对普通的 AE 的，对于 VAE 来说，就是往 AE 的编码向量里边加入重参数操作，然后损失函数里边加入 KL 散度项，所以，设计好 UniAE 后，理论上就已经设计好 UniVAE 了。
不过，实际操作的时候，我们还有改进的空间。理论上来说，训练好 VAE 是具有一定的解耦 (Disentanglement) 能力的，也就是说，隐变量的每个维度是独立无关的，它们分别控制生成结果的某一方面，可以随机调节。不难理解，解耦是一件非常有挑战性的事情，所以如果 VAE 的 Encoder 能编码出解耦的编码向量，那么其拟合能力必然也是比较强的，换言之，其结构需要有一定的复杂了。
我们再来看 UniAE 的 Encoder，它的编码向量是每一层的 $\texttt{[CLS]}$ 向量（或者对应的低维向量）的拼接，对于前面的层来说，它们的 $\texttt{[CLS]}$ 向量仅仅是有限几层的 Transformer 的输出，它们的编码能力是很弱的，并不足以编码出解耦的向量，因此将它们作为 VAE 的隐变量是不合适的。
所以，在实际设计 UniVAE 的时候，我们不能使用 UniAE 的所有 $\texttt{[CLS]}$ 向量作为编码向量，应该设置一个起始层数，Decoder 只使用大于这个层数的 $\texttt{[CLS]}$ 向量，而小于等于这个层数的 $\texttt{[CLS]}$ 向量则不使用，此时相对于使用下图右的 Attention Mask：
 |
 |
此时它等效于如下的 Encoder-Decoder 结构：
其他细节
至此，UniVAE 的关键部分已经介绍完毕了，下面分享一下在实现过程中一些比较重要的细节。
首先是长度泄漏问题。不管是 UniLM 还是 UniVAE，因为 Encoder 和 Decoder 整合成了一个模型，所以我们都是将输入输出拼接起来作为单个样本训练的，这样的话每个样本在 Decoder 部分的起始位置就不一样了，取决于输入文本的长度，这就意味着输入长度是也是作为了输入条件传入到了 Decoder 中，这就是长度泄漏。
这个问题有两个解决方案：第一个就是所有输入都通过截断或者填充来变为同一长度，这就不会造成长度泄漏了；第二个就更简单了，干脆啥都不做，即确实把长度当成条件输入，解码时通过控制起始位置来控制生成长度，但这样可能带来的问题是长度信息可能没有跟编码向量完全解耦，因此同一编码向量配上不同的长度可能会得到不合理的结果。
然后是层数和维度的选择问题。前面说了，为了让隐变量具有较好的解耦能力，我们将前 $k$ 层的 Attention 加上独立式 Attention Mask，剩下的 $L−k$ 层则加上 UniAE 式 Attention Mask。那么这个 $k$ 怎么选择呢？这是一个需要仔细调整的超参数，比较小的 $k$ 能保留更多的信息，有利于重构，但不利于解耦；反之较大的 $k$ 则更有利于解耦，但是不利于重构。在笔者的实验中，使用的是 $k=8$。
类似的问题出现在降维的维度选择上，较大的维度自然是有利于重构的，但也不利于解耦，反之则利于解耦而有损重构性能。这个参数需要根据任务本身的复杂度来具体调整，调整的大致方向是观察随机采样效果和重构效果，如果随机采样出来的样本多数可读、自然句子的重构效果也不错，那么说明这个维度适中了，否则则需要相应地调整。
最后，值得一提的是，UniAE 的设计不单单可以用来做 VAE，还可以用于构建 VQ-VAE，只需要对每个 $\texttt{[CLS]}$ 向量做一下量化，就成为了一个将不定长句子编码为定长离散序列的 VQ-VAE 模型了。
参考实现
这里给出一个 UniVAE 参考实现：
代码里使用的是 vMF-VAE 变体，基于 bert4keras 实现，基础架构是 RoFormer，当然也可以换成 BERT。下面演示的是用问句训练的 UniVAE 的效果。
随机采样效果：
我在steam下载的游戏，怎样能在电脑上玩啊？？？
呼市男科医院哪家比较好实惠
我血压高，我妈妈手脚麻木，是怎么回事呀
怎样查询交通违章记录和处罚
为什么我提问的问题有点卡顿
小米2s用的是移动卡还是联通卡
幼儿园怎么发展幼儿教育
英国读研学校排名对于英国留学生来说重要吗
有专业的关于excel表格数据库的培训机构吗？
为什么一到晚上就容易咳嗽，不睡觉就不咳
重构效果：
原句：数字电视机顶盒坏了，可以免费维修吗
重构：数字电视机顶盒坏了可以换吗?
原句：青椒跟什么炒好吃
重构：青椒跟什么炒好吃
原句：王者荣耀carryyou什么意思
重构：王者荣耀carry芈月什么意思
原句：没感冒老是咳嗽要吃什么药好
重构：没感冒老是咳嗽要吃什么药好
原句：沁园（金科西城大院店）怎么样，好不好的默认点评
重构：沁园（金源店）怎么样，好不好的默认点评
随机替换前 32 维隐变量：
原句：牙龈出血要吃什么药？
结果：牙龈出血还出血吃什么消炎药好
牙龈出血吃阿莫西林有效吗
牙龈出血是肝火旺吗？
牙龈出血去医院检查大概要多少钱？
牙龈出血去牙科看什么科室
牙龈出血去深圳哪里看牙科好
原句：广州和深圳哪个更好玩？
结果：广州和深圳哪个城市发展得好? 薪资高?
广州和深圳，哪个发达？深圳到广州的飞机票贵吗？
广州和深圳比哪个好
广州和深圳哪个人均gdp高
广州和深圳房价涨幅
广州和深圳自考一样吗
随机替换后 16 维隐变量：
原句：牙龈出血要吃什么药？
结果：未来21年做什么生意好？
湿疹给身体有什么伤害？
朗逸现在要买什么配置？
马来西亚签证要多少钱？
早上给孩子吃什么水果好？
头晕发热去医院看什么科？
原句：广州和深圳哪个更好玩？
结果：99和98相差多少呢？
微信和支付宝怎么更换手机号
我的指甲和肉很不一样怎么回事？
吃了甲硝唑多久才能喝酒？
桂圆和红枣可以一起泡茶吗？
小米和华为哪个更好点？
可以看到，随机采样和重构的效果都不错的，而通过随机替换不同维度的隐变量，我们可以大致观察到多尺度结构的效果：替换前面部分维度的隐变量，大致上保持了主题词不变；替换后面部分维度的隐变量，大致上保持了句式不变。当然，自然语言的结构性本身就很弱，因此例子中通常也夹杂了一些例外情况。
文章小结
本文介绍了笔者构思的 UniVAE 设计，它沿用类似 UniLM 的思路，通过特定的 Attention Mask 将 VAE 做到了一个 Transformer 模型里边，并且还具备多尺度特性。除了常规的 VAE 模型外，该设计还可以用于 VQ-VAE 等模型。
转载自《UniVAE：基于Transformer的单模型、多尺度的VAE模型》，作者：苏剑林，部分内容有修改。