2
$\begingroup$

I'm reading about VQ-VAE and I'm not sure why do they say we can view it as a VAE. Can you explicitly show:

  • what is the latent z-space - is it the discrete space where z can take the integers 1..K where K is the number of embedding vectors?
  • what is the true posterior p(z|x)? is it indeed intractable to use for calculating the marginal p(x), if we are dealing just in a scalar that can hold integers 1..K where K is a moderate 128?
  • why do we need to sample from z's distribution in VAE while in VQ-VAE we just take a deterministic distribution and just use the closest embedding to the encoder output?
  • can you show the ELBO formulation, and show how it boils down to only $log p(x|z_q(x))$?
$\endgroup$
2
  • $\begingroup$ The z-space is probably the space where the vectors live, rather than where their indices 1..K live. $\endgroup$
    – root
    Commented Jul 6, 2023 at 14:02
  • $\begingroup$ There are situations/variants where you don't need to sample in the VAE, but simply take the mean. For example when you use a trained VAE for dimensionality reduction. $\endgroup$
    – root
    Commented Jul 6, 2023 at 14:09

1 Answer 1

0
$\begingroup$

As far as I understand, we can do something as follows. One of the variants which the ELBO can be written is:

$$ \mathcal{L}_{\theta,\phi}(\textbf{x}) = \mathbb{E}_{q_\phi(\textbf{z}|\textbf{x})}\left[\log p_\theta(\textbf{x}|\textbf{z})\right] - D_{KL}(q_\phi(\textbf{z}|\textbf{x})\|p_\theta(\textbf{z})). $$

In the paper they assume that the prior $p_\theta(\textbf{z})$ follows a uniform distribution over a $K$-categorical space, and that the posterior $q_\phi(\textbf{z}|\textbf{x})$ can be obtained by

$$q_\phi(\textbf{z} = e_k|\textbf{x}) = \text{one_hot}(\textbf{z}_q(\textbf{x});k) = \begin{cases} 1 & \text{if } \textbf{z}_q(\textbf{x}) = e_k \\ 0 & \text{otherwise} \end{cases},$$

where $e_k$ are the vectors of the $K$-codebook.

With these, the KL divergence of the discrete distributions from the ELBO remains constant, since:

\begin{align*} D_{\text{KL}}(q_\phi(\textbf{z}|\textbf{x})||p(z)) &= \sum_{k=1}^K q_\phi(\textbf{z} = e_k|\textbf{x}) \log \frac{q_\phi(\textbf{z} = e_k|\textbf{x})}{p(\textbf{z} = e_k)} \\ &= \log \frac{1}{1/K} = \log K\\ \end{align*}

Therefore, we can ignore this part for the training.

The other term of the ELBO can be simplified as

$$ \mathbb{E}_{q_\phi(\textbf{z}|\textbf{x})}\left[\log p_\theta(\textbf{x}|\textbf{z})\right] = \sum_{k=1}^K q_\phi(\textbf{z} = e_k|\textbf{x}) \log p_\theta(\textbf{x}|\textbf{z} = e_k) = \log p_\theta(\textbf{x}|\textbf{z}_q(\textbf{x})) $$

Thus, the ELBO is

$$ \mathcal{L}_{\theta,\phi}(\textbf{x}) = \log p_\theta(\textbf{x}|\textbf{z}_q(\textbf{x})) $$

However, I am a bit annoyed since the ELBO is something that needs to be maximized, therefore we look for $$ \arg\min_{\theta,\phi} \mathcal{L}_{\theta,\phi}(\textbf{x}) = \arg\min_{\theta,\phi} \log p_\theta(\textbf{x}|\textbf{z}_q(\textbf{x})). $$ But in the paper the loss function they want to minimize uses $\log p_\theta(\textbf{x}|\textbf{z}_q(\textbf{x}))$ instead of $-\log p_\theta(\textbf{x}|\textbf{z}_q(\textbf{x}))$. Therefore, either I am completely wrong, or they missed a "$-$" sign there. Also, note that the term they minimize in its code implementation is the reconstruction loss, that is, $$ \|D(\textbf{z}_q(\textbf{x})) - \textbf{x}\|_2^2, $$ which can be related with $-\log p_\theta(\textbf{x}|\textbf{z}_q(\textbf{x}))$, see this for more info. So my guess is that there is a missing sign in the paper.

Regarding the other questions:

  1. There are several ways to see the latent space, after the encoder we are in a continuous real valued latent space ($\textbf{z} = E(\textbf{x}) \in \mathbb{R}^D$). Then we apply quantization and project $\textbf{z}$ to the vectors of the codebook, which discretizes $\mathbb{R}^D$. Since we have $K$ indexed vectors in the codebook, we can write each projected $\textbf{z}$ by its corresponding $k$. With this setting is how you see the space $\{1,\dots,k\}$ as the latent space, this is useful e.g. to generate images with transformers. However, in the VQ-VAE forward and backward pass we don't need to go as down and never get to work with the indices, remaining in the discrete latent space $\{e_1, \dots, e_k\}$.

  2. Don't know what the true posterior is.

  3. In VAE we sample from the prior distribution because although Gaussians are simple, the formulation obtained of the ELBO is too complex and required to be approximated by a Monte Carlo method. That is, by sampling on the prior distribution and obtaining an empirical estimate of the ELBO and its gradients. Assuming a uniform prior in a $K$-categorical distribution, leaves the ELBO good enough, not needing from an empirical estimate, i.e. not needing from sampling.

This is what I understood reading about VAEs and VQ-VAEs. I hope my insights can help you.

$\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.