Mathematical Establishment#

Assume a data generation process, where $z \sim p(z)$ represent latent variables (certain features like size, color, position) and $x \sim p_{real}(x)$ represent real observed variables (e.g. images, videos, texts). Our aim is to generate an observed variable via latent variables and a generator. Here, we bridge the latent distribution and the observed distribution with a conditional probability function $p_\theta(x|z)$, which is a mathematical description of the generator.

Thus, we establish an man-made observed distribution

and we wish this distribution to match the real observed distribution $p_{real}(x)$

Optimization Metrics#

To measure the similarities between the distributions, we typically use the Kullback-Leibler (KL) divergence, and we try to minimize the divergence

However, this is intractable since we don’t know $p_{real}(x)$ explicitly. To find feasible methods, try decompose the divergence

So instead, we could maximize the log-likelihood of our data under the model

Okay, now, another problem emerged. Computing $p_\theta(x) = \int p(z)p_\theta(x|z)dz$ is also intractable due to hardness of the integral over all possible latent variables. Detailed reasons are listed as follows:

1. High-dimensional integration: The latent space $z$ is typically high-dimensional, making the integral computationally expensive or impossible to evaluate analytically.

2. Complex posterior distribution: Even if we could compute $p_\theta(x)$, we would still need to compute the posterior $p_\theta(z|x) = \frac{p(z)p_\theta(x|z)}{p_\theta(x)}$ for inference, which requires the same intractable integral in the denominator.

3. No closed-form solution: For most practical choices of $p(z)$ and $p_\theta(x|z)$ (e.g., Gaussian distributions with neural network parameterization), the integral has no closed-form solution.

4. Sampling inefficiency: Monte Carlo sampling methods would require an impractical number of samples to get good estimates, especially in high dimensions where most of the probability mass is concentrated in a small region.

Variational Inference Solution#

To address these computational challenges, VAE introduces variational inference by approximating the intractable posterior $p_\theta(z|x)$ with a tractable variational distribution $q_\phi(z|x)$ parameterized by $\phi$ (typically implemented as an encoder neural network).

The key insight is to establish an exact decomposition of the log-likelihood. Starting from the log-likelihood we want to maximize

We introduce the variational distribution $q_\phi(z|x)$ and derive

Therefore, we have the exact decomposition

where the Evidence Lower Bound (ELBO) is

Key insights from this decomposition

1. Since $D_{KL}[q_\phi(z|x) \| p_\theta(z|x)] \geq 0$, we have $\log p_\theta(x) \geq \mathcal{L}(\theta, \phi; x)$, hence ELBO is indeed a lower bound.

2. The equality $\log p_\theta(x) = \mathcal{L}(\theta, \phi; x)$ holds if and only if $q_\phi(z|x) = p_\theta(z|x)$, meaning the variational posterior perfectly matches the true posterior.

3. Maximizing ELBO w.r.t. $\phi$ minimizes $D_{KL}[q_\phi(z|x) \| p_\theta(z|x)]$, making $q_\phi(z|x)$ a better approximation to the true posterior.

The ELBO consists of two terms

1. Reconstruction term: $\mathbb{E}_{z \sim q_\phi(z|x)}[\log p_\theta(x|z)]$ - encourages the decoder to reconstruct the input
2. Regularization term: $D_{KL}[q_\phi(z|x) \| p(z)]$ - keeps the learned latent distribution close to the prior

Since maximizing $\log p_\theta(x)$ is intractable, we instead maximize the tractable ELBO as a surrogate objective.

Final Optimization Objective#

Therefore, the final optimization objective function is:

where the ELBO for each sample is:

This can be equivalently written as:

or in expanded form:

Key points:

1. We optimize both $\theta$ (decoder parameters) and $\phi$ (encoder parameters) simultaneously
2. The KL divergence term also requires expectation over $x$, since $q_\phi(z|x)$ depends on $x$
3. Important: We assume $p(z) = \mathcal{N}(0, I)$ is a fixed standard Gaussian prior. Since this distribution is analytically known, the KL divergence $D_{KL}[q_\phi(z|x) \| p(z)]$ can be computed in closed form without sampling from $z$
4. In practice, both expectations over $p_{\text{real}}(x)$ are approximated using mini-batch sampling from the training dataset

Why no sampling over $z$ in the KL term:

• The KL divergence $D_{KL}[q_\phi(z|x) \| \mathcal{N}(0,I)]$ has a closed-form analytical solution
• For Gaussian $q_\phi(z|x) = \mathcal{N}(\mu_\phi(x), \sigma^2_\phi(x)I)$, we get:

• This eliminates the need for Monte Carlo sampling in the regularization term
• Only the reconstruction term $\mathbb{E}_{z \sim q_\phi(z|x)}[\log p_\theta(x|z)]$ requires sampling (via reparameterization trick)

VAE Architecture#

The VAE consists of two neural networks:

1. Encoder $q_\phi(z|x)$: Maps input $x$ to latent distribution parameters (typically mean and variance for Gaussian)
2. Decoder $p_\theta(x|z)$: Maps latent variable $z$ back to reconstruction of $x$

The training objective becomes:

Motivation#

Standard VAE generates random samples from $p(z)$. Conditional VAE allows controlled generation by conditioning on additional information $c$ (e.g., class labels, attributes).

Mathematical Framework#

The conditional generative model becomes:

Both encoder and decoder are conditioned:

• Conditional encoder: $q_\phi(z|x,c)$
• Conditional decoder: $p_\theta(x|z,c)$

Modified ELBO#

Applications#

• Class-conditional generation: Generate images of specific classes
• Style transfer: Control artistic style while preserving content
• Text-to-image: Generate images from textual descriptions

Motivation#

Standard VAE suffers from posterior collapse - the latent codes may be ignored during generation. VQ-VAE addresses this with discrete latent representations.

Key Innovation: Vector Quantization#

Instead of continuous latent variables, VQ-VAE uses a discrete codebook:

1. Codebook: $\mathcal{C} = \{e_k\}_{k=1}^K$ where $e_k \in \mathbb{R}^d$
2. Quantization: $\text{VQ}(z) = e_k$ where $k = \argmin_j \|z - e_j\|$

Architecture Changes#

Encoder: x → z_e (continuous)
Vector Quantization: z_e → z_q (discrete)
Decoder: z_q → x̂

Training Objective#

where:

• Reconstruction loss: $\|x - x̂\|^2$
• Codebook loss: Updates codebook vectors toward encoder outputs
• Commitment loss: Encourages encoder outputs to commit to codebook entries
• $\text{sg}[\cdot]$ denotes stop-gradient operation

Advantages of VQ-VAE#

1. No posterior collapse: Discrete codes are always used
2. Better reconstruction: Avoids blurry outputs common in VAE
3. Interpretable latents: Discrete codes often correspond to meaningful features
4. Hierarchical modeling: Can be stacked for multi-scale representations

VQ-VAE-2#

Extends VQ-VAE with:

• Hierarchical quantization: Multiple resolution levels
• PixelCNN decoder: Autoregressive modeling of quantized codes
• Better sample quality: Competitive with GANs on image generation

Applications and Impact#

• Image generation: DALL-E uses VQ-VAE for tokenizing images
• Audio modeling: VQ-VAE for speech and music synthesis
• Representation learning: Learning disentangled and interpretable features
• Data compression: Efficient encoding of high-dimensional data