What is the input for the prior model of VQ-VAE?

Question

I'm trying to implement the VQ-VAE model. In there, a continuous variable $x$ is encoded in an array $z$ of discrete latent variables $z_i$ that are mapped each to an embedding vector $e_i$. These vectors can be used to generate an $\hat{x}$ that approximates $x$.

In order to obtain a reasonable generative model $p_\theta(x)=\int p_\theta(x|z)p(z)$, one needs to learn the prior distribution of the code $z$. However, it is not clear in this paper, or its second version, what should be the input of the network that learns the prior. Is it $z=[z_i]$ or $e=[e_i]$? The paper seems to indicate that it is $z$, but if that's the case, I don't understand how I should encode $z$ properly. For example, a sample of $z$ might be an $n\times n$ matrix with discrete values between $0$ and $511$. It is not reasonable to me to use a one-hot encoding, nor to simply use the discrete numbers as if they were continuous, given that there is no defined order for them. On the other hand, using $e$ doesn't have this problem since it represents a matrix with continuous entries, but then the required network would be much bigger.

So, what should be the input for the prior model? $z$ or $e$? If it is $z$, how should I represent it? If it is $e$, how should I implement the network?

Answer 1

Note: This answer just uses images and pixels as an analogy for simplicity. It's not formal.

Some notes about VQ-VAE:

1. In the paper, they used PixelCNN to learn the prior. Remember, PixelCNN is a good choice for images of discrete values (e.g. 256 value, 3 channels).
2. The discrete latent variables are just the indices of the embedding vectors. For example, you can put your embedding vectors in an array.

For a single input, the number of the output channels of the encoder before quantization equals the dimensionality of the embedding vectors. For an analogy, the output of the encoder is like an image but with a number of channels that equals the dimensions in the embedding space. With that, each pixel is a vector in the space of the embedding vectors. The quantization is done on each pixel by mapping each pixel to the nearest embedding vector $e$.

Now, each pixel can be represented by a single integer number which equals the index of its nearest embedding (what you call the discrete latent variables). Thus, you have two representations for the quantized output, one is a large tensor with many channels (embedding vectors), and another is a simple "discrete image" with one channel (discrete latent variables).

The one-channel "discrete images" are used to train PixelCNN. This is great because they are small-size images. while the large tensor with the embedding vectors is used for the decoder as it holds the information needed for reconstruction.

So, the discrete latent variables are used later to learn the prior and the embedding vectors are used as inputs for the decoder.