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

UniVAE：基于Transformer的单模型、多尺度的VAE模型 苏剑林 · 2021-06-29 · Machine Learning variational unsupervised VAE attention

As we all know, the $\mathcal{O}(n^2)$ complexity of the Transformer is one of its "Achilles' heels." But every coin has two sides — the $\mathcal{O}(n^2)$ complexity also gives the Transformer a lot of room for maneuvering: we can flexibly customize different attention masks to design Transformer models for different purposes, such as UniLM, K-BERT, and others.

This post introduces a UniVAE model I've been thinking about, designed for text. It follows an approach similar to UniLM, packing the VAE into a single Transformer model, while also having multi-scale properties.

UniAE

I won't give an introduction to VAE (Variational Autoencoder) here — this site already has multiple posts covering it, so feel free to search for those. A VAE can be understood as an AE (Autoencoder) with a regularization term added: normally, 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 the encoded vector. So clearly, to implement UniVAE, we first need to implement a corresponding UniAE.

In From Language Model to Seq2Seq: The Transformer Is All About the Mask, we already introduced UniLM (Uni being short for Unified), which enables the Transformer to perform Seq2Seq tasks through the Attention Mask shown on the left below. However, UniLM is not the UniAE we're looking for, because the Decoder part of UniLM attends to the entire encoded input sequence, rather than to a single vector.

UniLM-style Attention Mask UniAE-style Attention Mask

However, we can further adjust the Attention Mask, building on UniLM, into the pattern shown on the right above. This way, decoding can only depend on the [CLS] vector from the encoding part and on the decoding results produced so far — this is the UniAE-style Attention Mask we're looking for. Since the input only depends on the [CLS] vector, and the size of the [CLS] vector is fixed, this means that the source information during generation is just a fixed-size vector, and the input is likewise encoded into this fixed-size vector — which is exactly the AE functionality we want.

UniAE-style attention dependency diagram

Multi-Scale

In other words, through the UniAE-style Attention Mask, we can implement a Seq2Seq model similar to UniLM, which is equivalent to an Encoder that encodes the input into a fixed-length vector, followed by a Decoder that decodes from that vector. If this still isn't entirely clear, we can also break it apart into an Encoder-Decoder architecture to understand it, as shown below:

Understood as a split Encoder-Decoder structure

The difference from a conventional Seq2Seq architecture is that here the Encoder and Decoder weights are shared. From the figure above, we can also see that if we apply this kind of mask at every attention layer, then the Decoder will depend on the [CLS] vector at every layer of input. This means that if there are $L$ attention layers, then it's the concatenation of all the [CLS] vectors from the input sequences of these $L$ layers that forms the complete encoding vector of the input text (of course, the first layer can be excluded, since the [CLS] at the first layer is just its embedding vector, which is a constant vector regardless of the input). A single layer's [CLS] vector by itself is not the complete encoding vector.

For the Decoder, each attention layer receives a [CLS] vector as input, which effectively forms a kind of multi-scale structure. In computer vision, the most advanced generative models are basically all multi-scale structures now, such as StyleGAN, Glow, NVAE, and so on, but this seems less common in NLP so far. It's not hard to imagine that in a multi-scale structure, different levels of input exert different degrees of control over the generated result: variables closer to the input layer control aspects that are more "inconsequential," while variables closer to the output layer control key information about the generated result. So ideally, once we've trained a multi-scale model, we should be able to control the generated result at different levels by editing the input variables at different levels.

Reducing Dimensionality

Some readers might already be thinking: if the dimension of each layer is $d$, and there are $L$ layers in total, then concatenating all the [CLS] vectors gives us a $Ld$-dimensional vector — for BERT base that would be $12\times 768 = 9216$ dimensions. Isn't that too large for an encoding vector? Indeed it is — for an ordinary AE or VAE, an encoding vector with nearly ten thousand dimensions is way too big.

Dimensionality reduction process diagram

Actually, the fix is simple: we just need to first reduce the dimensionality of each layer's [CLS] vector with a fully-connected layer, then project it back up with another fully-connected layer, and finally concatenate it with the remaining $(L-1)$ vectors of dimension $d$, as shown in the figure above. This way, although the input sequence is still of size $L\times d$, the [CLS] vector can in fact be expressed using a much lower-dimensional vector, and we just need to concatenate these lower-dimensional vectors from every layer to form the overall encoding vector.

Encoder-Decoder diagram after dimensionality reduction

Disentanglement Capability

Everything discussed and designed so far applies only to an ordinary AE. To turn it into a VAE, we simply add a reparameterization step to the AE's encoding vector, and add a KL-divergence term to the loss function. So in theory, once UniAE is designed, UniVAE is already designed too.

However, in practice, there's still room for improvement. In theory, a well-trained VAE has a certain amount of disentanglement capability, meaning that each dimension of the latent variable is independent, each controlling some particular aspect of the generated result, and can be adjusted independently. It's not hard to see that disentanglement is a highly challenging property to achieve, so if a VAE's Encoder is able to produce a disentangled encoding vector, its fitting capacity must necessarily be fairly strong — in other words, its structure needs to have a certain degree of complexity.

Let's go back to the UniAE Encoder: its encoding vector is the concatenation of the [CLS] vectors (or their corresponding low-dimensional vectors) from every layer. For the earlier layers, the [CLS] vector is merely the output of just a few Transformer layers, so its encoding capacity is quite weak — not strong enough to produce a disentangled vector. Therefore, using these early-layer [CLS] vectors as latent variables for the VAE is not appropriate.

So, when actually designing UniVAE, we should not use all of UniAE's [CLS] vectors as the encoding vector. Instead, we should set a starting layer index, so that the Decoder only uses [CLS] vectors from layers beyond that index, while [CLS] vectors from layers at or below that index are not used. This corresponds to using the Attention Mask shown on the right below:

Near the output layers, use the UniAE-style Attention Mask Near the input layers, use the independent-style Attention Mask

This is equivalent to the following Encoder-Decoder structure:

Illustration of the effect when the first two attention layers use the independent-style mask

Other Details

At this point, the key components of UniVAE have all been introduced. Below I'll share some important implementation details.

First, there's the issue of length leakage. Whether it's UniLM or UniVAE, because the Encoder and Decoder are merged into a single model, the input and output need to be concatenated and trained together as a single sample. This means that each sample's starting position in the Decoder portion is different, depending on the length of the input text — which means the input length itself is effectively passed into the Decoder as a condition. This is the length leakage problem.

There are two ways to solve this. The first is to pad or truncate all inputs to the same length, which eliminates length leakage entirely. The second is simpler: just do nothing, i.e., genuinely treat length as a conditioning input, and control the generation length by controlling the starting position during decoding. The downside of this approach is that the length information may not be fully disentangled from the encoding vector, so the same encoding vector paired with a different length might produce unreasonable results.

Next is the question of choosing the number of layers and the dimensionality. As mentioned earlier, in order to give the latent variable good disentanglement capability, we apply the independent-style Attention Mask to the first $k$ layers, and the UniAE-style Attention Mask to the remaining $L-k$ layers. So how should we choose $k$? This is a hyperparameter that needs careful tuning: a smaller $k$ retains more information, which is good for reconstruction but bad for disentanglement; conversely, a larger $k$ is better for disentanglement but worse for reconstruction. In my experiments, I used $k=8$.

A similar issue arises when choosing the dimensionality for the reduction step: a larger dimension is naturally better for reconstruction but worse for disentanglement, while a smaller one favors disentanglement at the cost of reconstruction quality. This parameter needs to be tuned according to the specific complexity of the task, and the general direction for tuning is to observe both the quality of random sampling and reconstruction: if the majority of randomly sampled outputs are readable and natural, and reconstructed sentences also come out well, then the dimensionality is well-suited; otherwise, adjust accordingly.

Finally, it's worth mentioning that the UniAE design isn't only useful for building VAEs — it can also be used to build VQ-VAE models. All you need to do is quantize each [CLS] vector, and you get a VQ-VAE model that encodes variable-length sentences into fixed-length discrete sequences.

Reference Implementation

Here's a reference implementation of UniVAE:

Github: https://github.com/bojone/univae

The code uses the vMF-VAE variant, implemented on top of bert4keras, with a RoFormer base architecture (though it could just as easily be swapped for BERT). Below is a demonstration of the results from training UniVAE on questions.

Random sampling results:

I downloaded a game on Steam — how do I play it on my computer???
Which men's health hospital in Hohhot is good and affordable
My blood pressure is high, and my mom's hands and feet are numb — what's going on
How do I check traffic violation records and penalties
Why does my question keep freezing when I try to post it
Does the Xiaomi 2s use a China Mobile SIM or a China Unicom SIM
How should kindergartens develop early childhood education
Does the ranking of UK graduate schools matter for international students going to the UK
Are there professional training institutions for Excel spreadsheet database skills?
Why is it that I only start coughing at night, and stop once I fall asleep

Reconstruction results:

Original: The digital TV set-top box is broken, can it be repaired for free
Reconstructed: The digital TV set-top box is broken, can it be exchanged?
Original: What goes well fried with green peppers
Reconstructed: What goes well fried with green peppers
Original: What does "carryyou" mean in Honor of Kings
Reconstructed: What does "carry Mi Yue" mean in Honor of Kings
Original: I don't have a cold but keep coughing, what medicine should I take
Reconstructed: I don't have a cold but keep coughing, what medicine should I take
Original: How is Qinyuan (Jinke Xicheng Yuan branch), review by default
Reconstructed: How is Qinyuan (Jinyuan branch), review by default

Randomly replacing the first 32 dimensions of the latent variable:

Original: What medicine should I take for bleeding gums?
Results: Bleeding gums, still bleeding, what anti-inflammatory medicine works well
Is amoxicillin effective for bleeding gums
Are bleeding gums a sign of excess liver fire?
How much does it roughly cost to get bleeding gums checked at the hospital?
Which department should I see for bleeding gums
Where's a good place in Shenzhen to see a dentist for bleeding gums
Original: Which is more fun, Guangzhou or Shenzhen?
Results: Which city, Guangzhou or Shenzhen, is more developed? Higher salaries?
Guangzhou vs. Shenzhen, which is more developed? Is a flight from Shenzhen to Guangzhou expensive?
Which is better, Guangzhou or Shenzhen
Which has a higher per-capita GDP, Guangzhou or Shenzhen
Housing price growth in Guangzhou and Shenzhen
Are the self-study exams the same in Guangzhou and Shenzhen

Randomly replacing the last 16 dimensions of the latent variable:

As we can see, both the random sampling and reconstruction results are quite good, and by randomly replacing the latent variables at different dimensions, we can roughly observe the effect of the multi-scale structure: replacing the earlier dimensions tends to preserve the topic words, while replacing the later dimensions tends to preserve the sentence structure. Of course, natural language itself has fairly weak structural regularity, so the examples usually also mix in some exceptions.

Summary

This post introduced the UniVAE design I came up with, which follows an approach similar to UniLM, using a specific Attention Mask to pack a VAE into a single Transformer model, while also giving it multi-scale properties. Besides conventional VAE models, this design can also be applied to models such as VQ-VAE.