0
$\begingroup$

I'm implementing a VQ-VAE for a LDM for biological time series data. I trained the VQ-VAE, and reconstructions works somewhat reasonable, but I have an understanding problem with how a VQ-VAE works. As I understand it, a VQ-VAE, in a 1D case, has an encoder part that takes inputs with shape ($B$, $L_1$, 1) (= (Batch, series length)) and outputs a variable $z_e$, with shape ($B$, $L_2$, $D$) . The VQ-Layer then takes each of the $L_2$ vectors, and replaces it with the closest vector inside a ($K$, $D$) vector space (called the codebook).

So as I understand it, the discretization comes from the limited size of the vector space. This vector space is learnt during training.

My question then is: How does a VQ-VAE generate new samples?

The prior distribution over the discrete latents p(z) is a categorical distribution, and can be made autoregressive by depending on other z in the feature map. Whilst training the VQ-VAE, the prior is kept constant and uniform. After training, we fit an autoregressive distribution over z, p(z), so that we can generate x via ancestral sampling. We use a PixelCNN over the discrete latents for images, and a WaveNet for raw audio.

which I don't understand at all.

I guess my question is this: In a "normal" VAE, we learn latent distributions instead of latents. We can then sample from these distributions to get actual latents to pass to the decoder. This is how we generate new images that were not part of the training set. How does one sample from a VQ-VAE? Where is the "variational" part in this?

Edit: Had a big misunderstanding of how VQ works, so I changed the question.

$\endgroup$

1 Answer 1

1
$\begingroup$

With the help of reddit I figured out the answer.

We train the VQ-VAE, and then use it to encode our dataset into the latent space. This latent space, of course, will have a particular distribution (or prior). The original paper uses a autoregressive model (the PixelCNN) to learn this distribution. So we can sample by using another model to generate latent spaces that we then pass to the decoder to generate our images/ data.

In the case of a latent diffusion model, we use a diffusion model to learn this prior, instead of an autoregressive one.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.