VQ-VAE

• VQ-VAE is a type of variational autoencoder that uses vector quantisation to obtain a discrete latent representation.
  • It differs from VAEs in two key ways:
    • the encoder network outputs discrete, rather than continuous, codes
    • the prior is learnt rather than static
• In order to learn a discrete latent representation, ideas from vector quantisation (VQ) are incorporated.
  • Using the VQ method allows the model to circumvent issues of posterior collapse - where the latents are ignored when they are paired with a powerful autoregressive decoder - typically observed in the VAE framework.

How it works

• 

• Define a latent embedding space $e\in {\mathbb{R}}^{K\times D}$ ,

  • where $K$ is the size of the discrete latent space (i.e. K-way categorical, size of vocab)
  • $D$ is the embedding size.
  • Thus, $z$ is K-way categorical variable

• First-step Encoder $z_e$

  • Takes an input $x$, and outputs an embedding $z_e(x)\in {\mathbb{R}}^{K\times D}$

• Posterior categorical distribution $q(z|x)$

  • Obtained by calculating nearest neighbour in the latent embedding space vocab. $q(z=k|x)=1[k={\mathrm{argmin}}_j||z_e(x)-e_j|{|}_2]$
  • Outputs either 0 or 1
  • It’s deterministic

• Prior $p(z)$

  • If you choose uniform prior over $z$, then you obtain a constant KL divergence, and equal to $\log K$
  • No posterior collapse

• What is actually fed to the decoder, it’s $z_q(x)$

  • Just the indexing of the embedding table, given the discretized representation of $x$
  • $z_q(x)=e_k$ where $k={\mathrm{argmin}}_j||z_e(x)-e_j|{|}_2$

Learning

• Because of the usage of argmin, and nearest neighbour, there is no real gradient defined between encoder and decoder.
  • They approximate the gradient similar to the straight-through estimator and just copy gradients from decoder input $z_q(x)$ to encoder output $z_e(x)$.
  • Why ?
    • During forward computation the nearest embedding $z_q(x)$ is passed to the decoder, and during the backwards pass the gradient ${\nabla }_{z}L$ is passed unaltered to the encoder.
    • Since the output representation of the encoder and the input to the decoder share the same $D$ dimensional space, the gradients contain useful information for how the encoder has to change its output to lower the reconstruction loss.

Defining the loss

• The loss is defined as $L=\log p_{\theta }(x|z_q(x))+||\mathrm{sg}[z_e(x)-e]|{|}_2^2+\beta ||z_e(x)-\mathrm{sg}[e]|{|}_2^2$

  • $\mathrm{sg}[]$ stands for the stopgradient operator, defined as identity at forward computation time, and has zero partial derivatives at backward time, thus effectively constraining its operand to be a non-updated constant.

  • the first term is the reconstruction loss

    • optimizes decoder and encoder
    • due to the straight-through gradient estimation of mapping from $z_e(x)$ to $z_q(x)$, the embeddings $e_i$ receive no gradients from the reconstruction loss

  • Second term: Vector Quantisation (VQ)

    • In order to learn the embedding space, they use one of the simplest dictionary learning algorithm, Vector Quantisation
      • The VQ objective uses the $l_2$ error to move the embedding vectors $e_i$ towards the encoder outputs $z_e(x)$

  • Third term: commitment loss

    • Forcing the volume of the embedding space outputted by $z_e(x)$ to stay bounded
    • Since the volume of the embedding space is dimensionless, it can grow arbitrarily if the embeddings $e_i$ do not train as fast as the encoder parameters. To make sure the encoder commits to an embedding and its output does not grow, we add the commitment loss