Auto-Encoding Variational Bayes

Reading guide

Problem and motivation

Directed probabilistic models with continuous latent variables $\mathbf{z}$ are useful for generative modeling, unsupervised representation learning, and marginal inference. However, learning them via maximum likelihood or Bayesian methods faces severe computational hurdles when:

1. Posterior intractability: For non-linear relationships (e.g., deep neural networks), the marginal likelihood integral $p_{\mathbf{\theta }}(\mathbf{x})=\int p_{\mathbf{\theta }}(\mathbf{z})p_{\mathbf{\theta }}(\mathbf{x}|\mathbf{z})d\mathbf{z}$ and the posterior $p_{\mathbf{\theta }}(\mathbf{z}|\mathbf{x})$ lack closed-form analytical solutions, rendering standard Expectation-Maximization (EM) and mean-field variational Bayes inapplicable.
2. Large datasets: Traditional sampling approaches such as Monte Carlo EM require slow, iterative MCMC sampling per datapoint across the dataset, making large-scale mini-batch optimization prohibitively expensive.

The standard naive score-function gradient estimator ${\nabla }_{\mathbf{\varphi }}{\mathbb{E}}_{q_{\mathbf{\varphi }}(\mathbf{z})}[f(\mathbf{z})]$ exhibits excessively high variance, preventing reliable optimization with stochastic gradient ascent.

Key ideas

• Reparameterization Trick: Instead of sampling $\mathbf{z}$ directly from an intractable parameterized distribution $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$, the variable is rewritten as a deterministic transformation of an independent noise variable $\mathbf{ϵ}\sim p(\mathbf{ϵ})$: $\mathbf{z}=g_{\mathbf{\varphi }}(\mathbf{ϵ},\mathbf{x})$. This allows differentiation with respect to $\mathbf{\varphi }$ through standard backpropagation with dramatically reduced variance.
• Amortized Inference (Recognition Model): Instead of optimizing variational parameters per datapoint, an encoder neural network $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ predicts variational parameters (such as the mean and variance of a latent Gaussian) conditioned on $\mathbf{x}$.
• Auto-Encoding Variational Bayes (AEVB): Jointly optimizes generative parameters $\mathbf{\theta }$ (probabilistic decoder) and variational parameters $\mathbf{\varphi }$ (probabilistic encoder) using stochastic gradient ascent on the reparameterized variational lower bound.

Method

The variational objective maximizes the Evidence Lower Bound (ELBO) $\mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})$: $$\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)})\ge \mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})=-D_{KL}(q_{\mathbf{\varphi }}(\mathbf{z}|{\mathbf{x}}^{(i)})\parallel p_{\mathbf{\theta }}(\mathbf{z}))+{\mathbb{E}}_{q_{\mathbf{\varphi }}(\mathbf{z}|{\mathbf{x}}^{(i)})}\left[\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)}|\mathbf{z})\right]$$

Applying the reparameterization trick ${\mathbf{z}}^{(i,l)}=g_{\mathbf{\varphi }}({\mathbf{ϵ}}^{(i,l)},{\mathbf{x}}^{(i)})$ with ${\mathbf{ϵ}}^{(i,l)}\sim p(\mathbf{ϵ})$, the generic SGVB estimator (Estimator A) is: $${\tilde{\mathcal{L}}}^A(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})=\frac{1}{L}\sum _{l=1}^L\left[\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)},{\mathbf{z}}^{(i,l)})-\log q_{\mathbf{\varphi }}({\mathbf{z}}^{(i,l)}|{\mathbf{x}}^{(i)})\right]$$

When the prior $p_{\mathbf{\theta }}(\mathbf{z})=\mathcal{N}(0,\mathbf{I})$ and approximate posterior $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})=\mathcal{N}(\mathbf{\mu },{\mathbf{\sigma }}^2\mathbf{I})$ are Gaussian, the KL divergence can be computed analytically, yielding Estimator B: $$\mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})\simeq \frac{1}{2}\sum _{j=1}^J\left(1+\log (({\sigma }_j^{(i)}{)}^2)-({\mu }_j^{(i)}{)}^2-({\sigma }_j^{(i)}{)}^2\right)+\frac{1}{L}\sum _{l=1}^L\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)}|{\mathbf{z}}^{(i,l)})$$ where ${\mathbf{z}}^{(i,l)}={\mathbf{\mu }}^{(i)}+{\mathbf{\sigma }}^{(i)}\odot {\mathbf{ϵ}}^{(l)}$ and ${\mathbf{ϵ}}^{(l)}\sim \mathcal{N}(0,\mathbf{I})$.

For mini-batches of size $M$, the full lower-bound estimator is: $${\tilde{\mathcal{L}}}^M(\mathbf{\theta },\mathbf{\varphi };{\mathbf{X}}^M)=\frac{N}{M}\sum _{i=1}^M\tilde{\mathcal{L}}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})$$ In practice, setting $L=1$ sample per datapoint is sufficient when paired with mini-batch sizes such as $M=100$.

Results

• Model Settings: Generative models were evaluated on MNIST (Bernoulli MLP decoder, 500 hidden units) and Frey Face (Gaussian MLP decoder, 200 hidden units), optimized via Adagrad with learning rates in $\{0.01,0.02,0.1\}$ and weight decay.
• Lower Bound Optimization: Across various latent dimensionalities ($N_{\mathbf{z}}\in \{2,5,10,20,400\}$ on Frey Face and MNIST), AEVB converged significantly faster and attained a higher lower bound than the wake-sleep algorithm.
• Overfitting Resistance: Increasing the dimensionality of $\mathbf{z}$ did not cause overfitting due to the regularizing effect of the analytical KL divergence term toward the prior.
• Marginal Likelihood Estimation: On 3-dimensional latent space MNIST experiments using an HMC-based marginal likelihood estimator, AEVB converged faster than wake-sleep and far exceeded the feasibility of Monte Carlo EM (which was too slow to run on the full MNIST dataset).
• Runtime: Training took roughly 20–40 minutes per million samples on an Intel Xeon CPU (~40 GFLOPS).

Limitations and open questions

• Continuous Latent Variables: The reparameterization trick directly relies on differentiability with respect to parameters, making it non-trivial to apply directly to discrete latent variable distributions.
• Approximation Gap: Assuming a diagonal Gaussian distribution for $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ limits the expressive capacity when the true posterior is multimodal or strongly coupled.
• Extension to Full Variational Bayes: The empirical evaluation focused on MAP/ML parameter estimation for $\mathbf{\theta }$; variational inference over global parameters $\mathbf{\theta }$ was formulated in the appendix but left for future empirical validation.