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

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?

Some notes about VQ-VAE:

1. In the paper, they used PixelCNN to learn the prior. PixelCNN is trained on images.
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 image, 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 to the dimensions in the embedding space. With that, each pixel is a vector in the space of the embedding vectors1. 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 to 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 for the prior model and the embedding vectors for the encoder.

1This is just an analogy, pixels values belong to a closed interval.