UniVAE: A Single-Model, Multi-Scale VAE Based on Transformer

UniAE

I won’t provide a general introduction to VAE (Variational Autoencoder) here, as there are already several articles on this site covering it; feel free to search for them. A VAE can be understood as an AE (Autoencoder) with a regularization term. Generally, the Encoder is responsible for encoding the input into a vector that satisfies a certain distribution, while the Decoder is responsible for reconstructing the input from that encoding vector. Therefore, to implement UniVAE, we must first implement the corresponding UniAE.

In "From Language Models to Seq2Seq: Transformer is All About the Mask", we introduced UniLM (Uni stands for Unified), which uses the Attention Mask shown on the left in the figure below to enable the Transformer to perform Seq2Seq tasks. However, UniLM is not the UniAE we are looking for, because the Decoder part of UniLM relates to the entire encoded sequence of the input, rather than a single vector.

However, we can further adjust the Attention Mask based on UniLM to the pattern shown on the right. In this way, decoding can only rely on the [CLS] vector of the encoding part and the already completed decoding results. This is the UniAE-style Attention Mask we are looking for. Since the input depends only on the [CLS] vector, and the size of the [CLS] vector is fixed, the source information during the generation process is essentially a fixed-size vector, and the input is also encoded into this fixed-size vector. This fulfills the AE function.

Multi-scale

In other words, through the UniAE-style Attention Mask, we can implement a Seq2Seq model similar to UniLM, which is equivalent to the Encoder encoding the input into a fixed-length vector and the Decoder decoding that vector. If this is still not clear, we can also decompose it into an Encoder-Decoder architecture to understand it, as shown below:

The difference from a conventional Seq2Seq architecture is that the weights of the Encoder and Decoder here are shared. As can be seen from the figure above, if we add this kind of Mask to every layer of Attention, the Decoder will depend on the [CLS] vector of every input layer. This means that if there are L layers of Attention, the concatenation of all [CLS] vectors from the input sequences of these L layers constitutes the complete encoding vector of the input text (of course, the first layer can be omitted because the [CLS] of the first layer is its Embedding vector, which is a constant vector for every input). The [CLS] vector of a single layer is not the complete encoding vector.

For the Decoder, every layer of Attention has a [CLS] vector passed in, which actually forms a multi-scale structure. In Computer Vision (CV), the most advanced generative models are basically multi-scale structures, such as StyleGAN, Glow, and NVAE. However, this seems less common in NLP. It is not hard to imagine that in a multi-scale structure, the degree of control that inputs at different levels have over the generation results varies. Variables closer to the input layer control "trivial" parts, while variables closer to the output layer control the key information of the generation results. Ideally, after training a multi-scale model, we can achieve control over different levels of the generation results by editing input variables at different levels.

Reducing Dimensionality

Some readers might think: if the dimension of each layer is d and there are L layers, then concatenating all [CLS] vectors would result in Ld dimensions. For BERT-base, this would be 12 \times 768 = 9216 dimensions. Isn’t this encoding vector dimension too large? Indeed, for an ordinary AE or VAE, an encoding vector of nearly ten thousand dimensions is too large.

The solution is actually very simple. We only need to use a fully connected layer to first reduce the dimensionality of the [CLS] vector of each layer, and then use another fully connected layer to increase the dimensionality back, finally concatenating it with the remaining (L-1) vectors of dimension d, as shown in the figure above. In this way, although the input sequence is still of size L \times d, the [CLS] vector can actually be expressed by a lower-dimensional vector. We only need to concatenate these lower-dimensional vectors from each layer to serve as the total encoding vector.

Decoupling Capability

The previous design and discussion were targeted at ordinary AEs. For a VAE, it involves adding a reparameterization operation to the AE’s encoding vector and adding a KL divergence term to the loss function. Therefore, once UniAE is designed, UniVAE is theoretically designed as well.

However, in practice, there is still room for improvement. Theoretically, a well-trained VAE possesses a certain degree of disentanglement capability, meaning that each dimension of the latent variables is independent and controls a specific aspect of the generation result, allowing for random adjustment. Disentanglement is a very challenging task. If the VAE’s Encoder can produce disentangled encoding vectors, its fitting capability must be relatively strong; in other words, its structure needs to be sufficiently complex.

Looking back at the Encoder of UniAE, its encoding vector is the concatenation of the [CLS] vectors (or corresponding low-dimensional vectors) from each layer. For the earlier layers, their [CLS] vectors are only the outputs of a few Transformer layers. Their encoding capability is very weak and insufficient to produce disentangled vectors. Therefore, using them as latent variables for the VAE is inappropriate.

Thus, when designing UniVAE in practice, we should not use all [CLS] vectors of UniAE as the encoding vector. We should set a starting layer number; the Decoder only uses [CLS] vectors from layers greater than this number, while [CLS] vectors from layers less than or equal to this number are not used. This corresponds to using the Attention Mask on the right in the figure below:

This is equivalent to the following Encoder-Decoder structure:

Other Details

At this point, the key parts of UniVAE have been introduced. Below, I will share some important details from the implementation process.

First is the issue of length leakage. Whether it is UniLM or UniVAE, because the Encoder and Decoder are integrated into a single model, we concatenate the input and output as a single sample for training. Consequently, the starting position of the Decoder part for each sample is different, depending on the length of the input text. This means the input length is also passed into the Decoder as an input condition, which is length leakage.

There are two solutions to this problem: the first is to make all inputs the same length through truncation or padding, which prevents length leakage; the second is even simpler—do nothing. That is, treat the length as a conditional input and control the generation length by controlling the starting position. However, a potential issue is that the length information might not be completely decoupled from the encoding vector, so the same encoding vector paired with different lengths might yield unreasonable results.

Then there is the issue of choosing the number of layers and dimensions. As mentioned earlier, to give the latent variables better decoupling capability, we add an independent Attention Mask to the first k layers of Attention, and the remaining L-k layers use the UniAE-style Attention Mask. How should k be chosen? This is a hyperparameter that needs careful adjustment. A smaller k retains more information, which is beneficial for reconstruction but not for decoupling; conversely, a larger k is more conducive to decoupling but detrimental to reconstruction. In my experiments, k=8 was used.

A similar issue arises in the choice of dimensionality for reduction. A larger dimension is naturally beneficial for reconstruction but also hinders decoupling, and vice versa. This parameter needs to be adjusted specifically based on the complexity of the task itself. The general direction for adjustment is to observe the random sampling effect and the reconstruction effect. If most of the randomly sampled samples are readable and the reconstruction of natural sentences is also good, then the dimension is appropriate; otherwise, it needs to be adjusted accordingly.

Finally, it is worth mentioning that the UniAE design can be used not only for VAEs but also for constructing VQ-VAEs. One only needs to quantize each [CLS] vector, and it becomes a VQ-VAE model that encodes variable-length sentences into fixed-length discrete sequences.

Reference Implementation

Here is a reference implementation of UniVAE:

The code uses the vMF-VAE variant, implemented based on bert4keras, with RoFormer as the base architecture (which can also be replaced with BERT). Below are the results of UniVAE trained on a question dataset.

Summary

This article introduced the UniVAE design I conceived. It follows a logic similar to UniLM, integrating a VAE into a single Transformer model through a specific Attention Mask, and it also possesses multi-scale characteristics. In addition to conventional VAE models, this design can also be applied to models like VQ-VAE.

Reprinted from: https://kexue.fm/archives/8475

For more details on reprinting, please refer to: "Scientific Space FAQ"