Principles of Variational Autoencoder (VAE) and the Reparameterization Trick

Description
The Variational Autoencoder (VAE) is a generative model that combines autoencoder architecture with probabilistic graphical models, learning the underlying distribution of data to generate new samples. Unlike standard autoencoders, VAE's latent space is continuous and structured, allowing data generation via sampling. A core challenge is enabling gradient backpropagation through stochastic sampling, which is resolved by the reparameterization trick.

Detailed Explanation

1. Basic Framework and Objective
   • VAE consists of an encoder and a decoder. The encoder maps input data \(x\) to the posterior distribution \(q_\phi(z|x)\) of latent variables \(z\) (typically assumed to be Gaussian). The decoder reconstructs data \(p_\theta(x|z)\) from \(z\).
   • The objective function is the Evidence Lower Bound (ELBO). Maximizing ELBO is equivalent to minimizing reconstruction error and regularizing the latent space:

\[ \text{ELBO} = \mathbb{E}_{q_\phi(z|x)}[\log p_\theta(x|z)] - D_{\text{KL}}(q_\phi(z|x) \| p(z)) \]

 where the first term is reconstruction loss and the second term is KL divergence (constraining the latent distribution to be close to the prior $ p(z) $, usually a standard normal distribution).

2. Necessity of the Reparameterization Trick

   • Directly sampling \(z\) from \(q_\phi(z|x)\) (e.g., Gaussian distribution \(\mathcal{N}(\mu, \sigma^2)\)) makes the sampling operation non-differentiable, preventing gradient-based optimization of parameters \(\phi\).
   • The reparameterization trick separates randomness: let \(z = \mu + \sigma \odot \epsilon\), where \(\epsilon \sim \mathcal{N}(0,1)\). Here, the randomness of \(z\) comes solely from \(\epsilon\), while \(\mu\) and \(\sigma\) are outputs of the encoder, allowing gradients to backpropagate through \(\mu\) and \(\sigma\).

3. Detailed Training Process

   • Step 1: The encoder takes input \(x\) and outputs parameters of the latent distribution: \(\mu\) and \(\log \sigma^2\) (using log variance for numerical stability).
   • Step 2: Sample \(\epsilon \sim \mathcal{N}(0,1)\), then compute \(z = \mu + \sigma \odot \epsilon\).
   • Step 3: The decoder maps \(z\) to reconstructed data \(\hat{x}\).
   • Step 4: Compute the loss function:
     • Reconstruction loss: measures discrepancy between \(\hat{x}\) and \(x\) (e.g., cross-entropy or mean squared error).
     • KL divergence: closed-form expression is \(-\frac{1}{2} \sum(1 + \log \sigma^2 - \mu^2 - \sigma^2)\), pushing \(q_\phi(z|x)\) toward the standard normal distribution.
   • Step 5: Jointly optimize encoder parameters \(\phi\) and decoder parameters \(\theta\) via gradient descent.

4. Mathematical Principles of Reparameterization

   • Original sampling \(z \sim \mathcal{N}(\mu, \sigma^2)\) blocks gradients because the sampling operation has no derivative.
   • After reparameterization, the gradient path becomes:

\[ \frac{\partial z}{\partial \mu} = 1, \quad \frac{\partial z}{\partial \sigma} = \epsilon, \]

 making $ \nabla_\phi \mathbb{E}_{q_\phi}[\log p_\theta(x|z)] = \mathbb{E}_{\epsilon}[\nabla_\phi \log p_\theta(x|z=\mu+\sigma \epsilon)] $ computable.

5. Differences Between VAE and GAN
   • VAE explicitly learns data distribution, focusing on reconstruction quality, but generated samples may be blurry; GAN generates sharper samples via adversarial training but suffers from unstable training.
   • VAE's latent space is interpretable and supports operations like interpolation.

Summary
VAE introduces a probabilistic framework and the reparameterization trick to solve the gradient propagation problem in latent variable sampling for generative models. Its structured latent space provides a powerful tool for data generation and representation learning.