VAEs

• For many modalities, we can think of the data we observe as represented or generated by an associated unseen latent variable, which we can denote by random variable $z$.

• SOURCE: # Understanding Diffusion Models: A Unified Perspective

Evidence Lower Bound

• Mathematically, we can imagine the latent variables and the data we observe as modeled by a joint distribution $p(x,z)$.
• One approach of generative modeling, termed “likelihood-based”, is to learn a model to maximize the likelihood $p(x)$ of all observed $x$. Two ways to go about manipulate this joint distribution to recover the likelihood:
  • Equation 1: marginalization: $p(x)=\int p(x,z)dz$
  • Equation 2: chain rule: $p(x)=\frac{p(x,z)}{p(z|x)}$
  • Computing any of them is kinda hard
• However, using the two equations, we can lower-bound the likelihood, which would give us a good enough proxy objective to maximize. ⇒ ELBO
  • $logp(x)\ge {\mathbb{E}}_{q_{\varphi }(z|x)}[log\frac{p(x,z)}{q_{\varphi }(z|x)}]$
  • $q_{\varphi }(z|x)$ is an approximate variational distribution of $p(z|x)$ with parameters $\varphi$ to optimize
  • ELBO can be derived from writing $logp(x)$ and
    • applying Eq.1, multiply by $1=\frac{q_{\varphi }(z|x)}{q_{\varphi }(z|x)}$, apply Jensen Inequality (not very informative)
    • multiply by $1=\int q_{\varphi }(z|x)dz$, applying Eq.2, multiply by $1=\frac{q_{\varphi }(z|x)}{q_{\varphi }(z|x)}$, split the expectation using the log, and arriving at
      • $logp(x)={\mathbb{E}}_{q_{\varphi }(z|x)}[log\frac{p(x,z)}{q_{\varphi }(z|x)}]+D_{KL}(q_{\varphi }(z|x)||p(z|x))$
      • KL divergence always $\ge 0$
• Conclusion: ELBO proxy is as good as the KL divergence between the approximate posterior $q_{\varphi }(z|x)$ and the true posterior $p(z|x)$
• ⇒ Maximizing the ELBO will minimize the KL divergence

Variational Autoencoders

• In the default formation of the VAE, we directly maximize the ELBO.
  • It’s called variational because we optimize for the best $q_{\varphi }(z|x)$ amongst a family of potential posterior distributions parameterized by $\varphi$
  • It’s called autoencoder because it follow the usual auto-encoder architecture
• In practice, there’s thus two sets of parameters optimized $\varphi$ for the encoder and $\theta$ for the decoder.
• Let’s see how maximizing the ELBO makes sense in this context:
  • 
  • $p_{\theta }(x|z)$ is a deterministic function (decoder) to convert a given latent vector $z$ into an observation $x$. This explicitly assumes a somewhat deterministic mapping between x and z.
  • $q_{\varphi }(z|x)$ can be seen as an intermediate bottlenecking distribution (encoder)

How to train

• A defining feature of the VAE is how the ELBO is optimized jointly over parameters $\varphi$ and $\theta$.
• The encoder of the VAE is commonly chosen to model a multivariate Gaussian with diagonal covariance, and the prior is often selected to be a standard multivariate Gaussian
  • $q_{\varphi }(z|x)=\mathcal{N}(z;{\mu }_{\varphi }(x),{\sigma }_{\varphi }^2(x))$
  • $p(z)=\mathcal{N}(z;0,I)$
  • We learn the bottleneck mean and covariance

Maximizing ELBO

• KL divergence can be computed analytically
• Reconstruction term through Monte Carlo
• Objective: $\underset{\varphi ,\theta }{\mathrm{argmax}}{\sum }_{l=1}^L[log{p}_{\theta }(x|z^{(l)})]-D_{KL}(q_{\varphi }(z|x)||p(z))$
• IMPORTANT where the latents $z^{(l)}$ are sampled from $q_{\varphi }(z|x)$ for every sample $x$ in the dataset i.e. making sure that going x → z → x is a good reconstruction
• the KL divergence ensures that the posterior maps to a well-behaved distribution, ensuring that we’ll be able to sample from $p(z)$ to create samples later on. (obviously also ensures that we’re maximizing the ELBO)

Reparameterization trick

• The latents $z^{(l)}$ that are obtained from going through the encoder and then sampling need to be differentiable because they are passed on the decoder for the reconstruction
• To ensure this, each $z$ is computed as a determinstic function of the input $x$ and auxiliary noise $ϵ$
  • $q_{\varphi }(z|x)=\mathcal{N}(z;{\mu }_{\varphi }(x),{\sigma }_{\varphi }^2(x))$ (in theory)
  • $z={\mu }_{\varphi }(x),+{\sigma }_{\varphi }^2(x)\odot ϵ$ (in practice) with $ϵ\sim \mathcal{N}(0,I)$

Encoder

• We can use $q_{\varphi }(z|x)$ as a meaningful and useful representation (for latent diffusion!!)

Posterior collapse

What is it

• Posterior collapse occurs when the approximate posterior $q(z\mid x)$ collapses to the prior $p(z)$ irrespective of the input $x$.
• This means that the latent variable $z$ carries very little information about the input $x$
  • effectively rendering the latent space meaningless.
• In this scenario, the encoder outputs a distribution that is very close to the prior
  • making the KL divergence term very small,
  • but at the cost of losing meaningful representations of the input data.

Causes

• Overpowering Decoder: If the decoder is too powerful, it can learn to reconstruct the input data directly from the prior distribution, making the latent variable $z$ redundant.

• High KL Weight: During training, the KL divergence term can dominate the loss function, pushing the posterior $q(z\mid x)$ to align closely with the prior $p(z)$.

How to address it

• Adjusting the KL Weight: Introducing an annealing schedule where the weight of the KL divergence term is gradually increased during training can help prevent early posterior collapse.
• VAE Variants: Using alternative VAE architectures such as β-VAE, which introduces an adjustable weight on the KL divergence term, or hierarchical VAEs, which have more complex latent structures, can help avoid posterior collapse.
• Structured Priors: Using more complex priors than a simple Gaussian can help in maintaining a meaningful latent space.